How to Make Always-Installed Django Database Instrumentation

Since version 2.0, Django has provided a hook for installing database instrumentation wrapper functions. These functions are like middleware for database queries, allowing you to inspect and rewrite SQL before it is sent to the database. There are many use cases, for example:

• Blocking database queries in certain blocks of code, as in the feature documentation.
• Tracking all queries, as in Scout APM’s python agent.
• Rewriting queries to add extra SQL features, as used in Django-MySQL’s query hints extensions.

For this post, let’s use an example wrapper function that logs a warning for slow queries. We can implement it like this:

import logging
import time

logger = logging.getLogger(__name__)
THRESHOLD_SECONDS = 1.0


def slow_query_warner(execute, sql, params, many, context):
    start = time.perf_counter()
    result = execute(sql, params, many, context)
    duration = time.perf_counter() - start
    if duration > THRESHOLD_SECONDS:
        logger.warning("Slow query, took %.2f seconds: %s", duration, sql)
    return result

Then for any block of code for which we want to log about slow queries, we can use:

with connection.execute_wrapper(warn_about_slow_queries):
    ...

A small limitation of the API is that it is only available as a context manager. This means having a wrapper function that’s installed for every single query is a little bit of a challenge.

We can use a middleware to install our wrapper:

from django.db import connection


class SlowQueryLoggingMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        with connection.execute_wrapper(slow_query_warner):
            return self.get_response(request)

If we add SlowQueryLoggingMiddleware to the top of our MIDDLEWARE setting, we can inspect all queries that happen within HTTP requests. But our instrumentation is still rather incomplete. We’ll miss all queries made outside of the request-response cycle, such as during management commands or background task functions.

To always install a wrapper function, we have to step outside of the documented database instrumentation API. We can directly manipulate the connection’s execute_wrappers list to insert our wrapper function. The execute_wrappers list is not documented, so it’s not a public API and could change at any time in Django. But in practice execute_wrappers hasn’t changed since Django 2.0 added database instrumentation.

We can install our wrapper from within a handler for the connection_created signal. Django sends this signal each time it connects (or re-connects) to a database. Because Django will probably have already connected to the databases by the time we attach our handler, we also need to directly run our handler once for each connection.

Here’s how this all looks together:

from django.db import connections
from django.db.backends.signals import connection_created


def install_slow_query_warner(connection, **kwargs):
    """
    Install slow_query_warner on the given database connection.
    Rather than use the documented API of the `execute_wrapper()` context
    manager, directly insert the hook.
    """
    if slow_query_warner not in connection.execute_wrappers:
        connection.execute_wrappers.insert(0, slow_query_warner)


connection_created.connect(install_slow_query_warner)
for connection in connections.all():
    install_slow_query_warner(connection=connection)

Update (2022-10-03): Previously the above code used connection.execute_wrappers.append(slow_query_warner), but this is prone to a bug with threads.

The best place to add this would be in one of our apps’ AppConfig.ready() methods.

Naturally, we’d best add some tests for our wrapper function. They’re extra important since we’re using an undocumented API. For our example, we can test that slow queries trigger a log message:

from unittest import mock

from django.db import connection
from django.test import TestCase

import __main__


class SlowQueryWarnerTests(TestCase):
    @mock.patch.object(__main__, "THRESHOLD_SECONDS", 0.01)
    def test_slow_query(self):
        with self.assertLogs("__main__") as logged, connection.cursor() as cursor:
            cursor.execute("SELECT SLEEP(0.02)")

        self.assertEqual(len(logged.records), 1)
        self.assertTrue(logged.records[0].msg.startswith("Slow query, took "))

These tests use the TestCase.assertLogs() method to grab the log messages. If you’re using pytest, you can also use the caplog fixture.

Additionally, to keep the tests fast (my favourite topic), we use the decorator form of mock.patch.object() to temporarily reduce THRESHOLD_SECONDS. Without this we’d need to make a query that lasts at least one second, which is enough time to run hundreds of other tests.

And by the way, we import the application as __main__ here because the example code has been added to a single file application. For normal projects with more than one file we would import from the appropriate modules.

😸😸😸 Check out my new book on using GitHub effectively, Boost Your GitHub DX! 😸😸😸

One summary email a week, no spam, I pinky promise.

Related posts:

• The Fast Way to Test Django transaction.on_commit() Callbacks
• Disable Instrumentation Packages during Tests
• Setting Python’s Decimal Context for All Threads

Tags: django