Solve N+1 performance issues forever on Django with Pytest

Have you ever improved the performance of your Django apps by tackling N+1 performance issues, but after some time developing new features on your project, the performance problems came back?

Solving N+1 queries problem in Django is often quite straightforward using prefetch_related and select_related (see the documentation for more insights on how to do it). But testing that there are no N+1 queries to the database is less intuitive. Actually there is no simple solution to do it natively or with a library in Django.

To make it simple, let’s say you have 2 models, a Car and a Color. Each car has a color, thus there is a foreign key from Car on Color.

And let’s say you have a view (basically an API in Django), that returns the list of all the cars in your database with their associated color Car.objects.all(). If you don’t prefetch the Color table, you will have one query to get all the cars in the Car table, and N queries on the Color table to get the associated color for each car ⇒ 1+N queries.

As usual, you write a test that asserts that the response status of your view is HTTP_200_OK and that the data you receive is correct. But then, you would also like to somehow test the performance of the route. To be more precise, you want to test that you don’t have N+1 queries. How should you do it?

In your tests/utils folder, or wherever you think it should be, you can use the following wrapper:

class QueryLogger:
    def __init__(self):
        self.queries = []

    def __call__(self, execute, sql, params, many, context):
        try:
            return execute(sql, params, many, context)
        finally:
            self.queries.append(sql)

@contextmanager
def assert_no_duplicated_queries():
    query_logger = QueryLogger()
    try:
        with connection.execute_wrapper(query_logger):
            yield
    finally:
        duplicated_queries = [
            (query, count)
            for query, count in collections.Counter(query_logger.queries).items()
            if count > 1
        ]
        number_of_duplicated_queries = len(duplicated_queries)
        assert number_of_duplicated_queries == 0, (
            [f"{query} is duplicated {count} times \n " for (query, count) in duplicated_queries]
        )

Then in your class TestViewSetPerformance you can easily use this wrapper to assert that there are no duplicated queries to the database when calling your view:

def test_performance_car_views(self, client):
        with assert_no_duplicated_queries():
            response = client.get(reverse("cars-list"))
        assert response.status_code == 200

Without a prefetch in your view, you should have an assertion error with the duplicated queries and the number of times it is duplicated.

⚠️ The test will only be effective if there is enough data for the N+1 to be measurable. In the exemple above, with only one car in the database, the test will pass wether there are N+1 queries or not.

And that’s it! Now you can be sure that, even if your simple Car/Color models, or your view evolve, you won’t introduce N+1 queries performance issues anymore 🎉