Capture stack trace when queries are executed

30 March, 2021

Asserting query counts in Django tests with CaptureQueriesContext is a great way to maintain/optimize code, but it doesn't provide enough context to find what changed when tests fail during a refactor. I found it helpful to capture some lines from the stack trace when each query is being executed and have them accessible during test runs that use CaptureQueriesContext to be logged along with the query, when count assertions fail.

You can achieve this by replacing CursorDebugWrapper in django.db.backends.utils with a new class. The new class can inherit from CursorDebugWrapper and override execute and executemany methods. The replacement methods would wrap the call (of super()) with a try/finally block, and in the finally block, call a new private method that can collect the stack trace.

Real world example

The following example assumes you have an environment variable ENABLE_QUERY_TRACEBACK_CAPTURE=1 set during test runs to override CursorDebugWrapper.

Somewhere in the project, say a testutils.py module:

import traceback

from typing import List
from dataclasses import dataclass

from django.db.backends.utils import CursorDebugWrapper
from django.conf import settings


@dataclass(frozen=True)
class Traceline:
    path: str
    lineno: int
    func: str
    line: str


@dataclass(frozen=True)
class QueryTrace:
    sql: str
    traceback: List[Traceline]


class CustomCursorDebugWrapper(CursorDebugWrapper):
    def __init__(self, cursor, db):
        super().__init__(cursor, db)

        if not hasattr(self.db, "captured_query_traceback"):
            self.db.captured_query_traceback = []

    def capture_traceback(self, sql):
        stack = traceback.extract_stack()

        self.db.captured_query_traceback.append(
            QueryTrace(
                sql=sql,
                traceback=[
                    Traceline(
                        path=path,
                        lineno=lineno,
                        func=func,
                        line=line,
                    )
                    for (path, lineno, func, line) in reversed(stack)
                ],
            )
        )

    def execute(self, sql, params=None):
        try:
            return super().execute(sql, params)
        finally:
            sql = self.db.ops.last_executed_query(self.cursor, sql, params)
            self.capture_traceback(sql)

    def executemany(self, sql, param_list):
        try:
            return super().executemany(sql, param_list)
        finally:
            self.capture_traceback(sql)

In settings.py:

import os

if os.environ.get("ENABLE_QUERY_TRACEBACK_CAPTURE"):
    import django.db.backends.utils as backend_utils

    from testutils import CustomCursorDebugWrapper

    backend_utils.CursorDebugWrapper = CustomCursorDebugWrapper