Не так чтобы часто, но с той самой неприятной регулярностью, когда уже забыл как это делал в прошлый раз, бывает нужно посчитать сколько запросов к БД генерирует тот или иной блок кода для django.
При этом мало что лучше закрепляется в памяти, чем очередная неудачная статья на хабре собственного сочинения. Штош, попробуем совместить полезное с неприятным.

Для чего вообще считать запросы?

Можно назвать несколько теоретических поводов:

  1. Предотвращение деградации производительности: при изменениях в коде можно не заметить, что количество запросов увеличилось в разы из-за того что где-то потерялся волшебный prefetch_related или select_related.

  2. Документирование оптимизации: тест явно фиксирует ожидаемое количество SQL-запросов.

  3. Контроль отсутствия примесей: можно посмотреть какие именно запросы и в каком порядке выполняются, чтобы убедиться, что в логику не примешался код не имеющий прямого отношения к выполняемой задаче.

Из практики вспоминается как в DaData нужно было писать и поддерживать автотесты проверяющие, что только нужные запросы делаются в узких местах приложения - и мы даже описывали комментариями почему без этого запроса не получается.
На недавнем code-review коллега заметил у меня в коде подозрительный цикл и, чтобы убедиться в отсутвии N*K запросов, проще всего оказалось написать автотест на этот блок кода - так как воспроизвести всю схему данных и посмотреть на запросы в логах было слишком затратно по времени.

Таким образом, смысл в этом определённо есть. Осталось закрепить то, как это сделать на unittest и pytest.

Как посчитать запросы в unittest?

Django предоставляет контекстный менеджер assertNumQueries, доступный в django.test.TestCase.
Он позволяет проверить, что внутри блока выполняется ровно заданное количество SQL-запросов.

Пример базового использования:

from django.test import TestCase

from myapp.models import Book

class BookTests(TestCase):
    def test_list_books_queries(self):
        # Предварительно создаём данные
        for i in range(5):
            Book.objects.create(title=f"Book {i}")

        # Проверяем количество запросов при выборке всех объектов
        with self.assertNumQueries(1):
            books = list(Book.objects.all())
            self.assertEqual(len(books), 5)

А что с pytest?

Если вы используете pytest для написания автотестов с django, то можете использовать контекстный менеджер CaptureQueriesContext.
Код при использовании расширения pytest-django может выглядеть как-то так:

import pytest
from django.db import connection
from django.test.utils import CaptureQueriesContext

from myapp.models import Book

@pytest.mark.django_db
def test_books_queries():
    Book.objects.create(title="Book 1")
    Book.objects.create(title="Book 2")


    with CaptureQueriesContext(connection) as ctx:
        list(Book.objects.all())

    assert len(ctx.captured_queries) == 1

Этот же трюк можно использовать и в unittest. Внутри ctx.captured_queries есть и sql соответствующих запросов.

Сохраните это себе в закладки.

А еще, если вам понравилась статья, то я с удовольствием писал бы для вас техническую документацию или статьи или, что более вероятно, код на python. Словом, я в поисках работы - напишите, пожалуйста, в личку и я отправлю вам резюме.