How to Make Always-Installed Django Database Instrumentation2020-07-23
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.
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
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.append(slow_query_warner) connection_created.connect(install_slow_query_warner) for connection in connections.all(): install_slow_query_warner(connection=connection)
The best place to add this would be in one of our apps’
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.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
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.
Thanks to Shai Berger for implementing database instrumentation back in Ticket #28595. I hope this post helps you write your own wrapper functions,
📙👉Speed Up Your Django Tests👈📙
One summary email a week, no spam, I pinky promise.
- The Fast Way to Test Django transaction.on_commit() Callbacks
- Disable Instrumentation Packages during Tests
- Setting Python’s Decimal Context for All Threads
© 2020 All rights reserved.