Warn on too many queries in context

warn_too_many_queries.py

import logging
from contextlib import contextmanager
from inspect import getframeinfo, stack

from django.db import connections

logger = logging.getLogger(__name__)


@contextmanager
def warn_too_many_queries(*, threshold, using="default"):
    """
    Log a warning when too many queries were executed within a context.
    
    For example, the following code will raise a warning if there's at
    least 3 books::
    
        with warn_too_many_queries(threshold=3):
            for book in Book.objects.all():
                print(f"{book.title} by {book.author.name}")
            
        # Too many queries at myapp/views.py:12
    """
    num_queries = len(connections[using].queries)
    yield
    new_num_queries = len(connections[using].queries)
    diff = new_num_queries - num_queries
    if diff > threshold:
        caller = getframeinfo(stack()[1][0])
        logger.warning(f"Too many queries at %s:%s", caller.filename, caller.lineno)