Как стажёр оптимизировал запросы и нашел баг в Django

Привет! Меня зовут Иван, я бэкенд-разработчик-стажёр в KTS. 

Недавно я нашел баг в Django, создал тикет с исправлением и его приняли.

В статье расскажу подробнее — над чем работал, в чём была ошибка и почему её сложно встретить. А также ещё про один баг, который по классике оказался фичей ?.

Что будет в статье:

• Предыстория и немного о работе Django

• Оптимизация запросов

• Ошибка Django

• Причины бага

• Тикет и Pull Request

• Наша оптимизация

• Баг, который оказался фичей

• Заключение

Скрытый текст

У меня высшее медицинское образование и 4 года работы хирургом. 

Со школы меня интересовали математика и компьютерные технологии, но этот интерес долго оставался лишь в виде хобби — например, я делал графические патчи на телефоны. На уроках информатики преподавали лишь Word и PowerPoint, а крепкой математической базы или знакомых, связанных с IT, у меня не было. Так что я не воспринимал IT как свою профессию. 

В свободное время я читал статьи по разработке. Особенно нравилась тематика нейронных сетей и параллели с работой человеческого мозга. На 4-м курсе медакадемии понял, что созрел для изучения языка программирования для общего развития. Для этой цели давно присмотрел Python, и благодаря книге Марка Лутца «Изучаем Python» началось моё погружение в мир программирования. Не описать словами кайф, который я ощутил от первых работающих строк кода! Так из года в год, медленно, но верно, повышались мои знания и навыки по Python. Ещё немного игрался с Unity.

Я люблю медицину, но всё время работы меня не покидало желание заниматься тем, что нравится ещё больше — программированием. В связи с большим количеством негативных моментов в работе хирурга я наконец решился сменить род деятельности и плотно занялся самообучением. Начал со школьных курсов программирования, потом изучал алгоритмы и решал задачи для подготовки к школьным олимпиадам, читал о веб-разработке.

В январе 2022 года по счастливой случайности узнал о наборе в бесплатную школу начинающих бэкенд-разработчиков и буквально в последний день набора успел подать заявку. За месяц обучения познакомился с aiohttp, открыл для себя мощь баз данных, SQL и ORM. Стало намного проще писать асинхронный код. Местами было очень трудно и больно, но оно того стоило. Во время написания курсового проекта изучил работу с git. 

После окончания курса меня пригласили на собеседование, в результате чего удалось попасть на стажировку. 

Пользуясь случаем, хочу поблагодарить за помощь и наставления в работе коллег из KTS, особенно своего наставника Сашу ?.

Предыстория и немного о работе Django

Знакомиться с Django я начал буквально с первых дней стажировки. Ранее был опыт с aiohttp + SQLAlchemy/GINO. 

На мой взгляд, одной из самых крутых особенностей Django является Django ORM, который преобразует запрос в базу данных из Python-кода в язык SQL. Это значительно облегчает процесс разработки, но есть одно «но». При получении данных через отношения эти запросы не оптимизируются, и возникает так называемая проблема N+1 запросов. Чем больше запросов в БД на бэкенде, тем медленнее работает приложение. Для решения необходима оптимизация, чтобы доставать нужные данные за меньшее количество запросов.

Основными методами оптимизации запросов в Django являются:

• select_related

• prefetch_related (включая класс Prefetch)

• only

Для создания API имеется 2 основных варианта: REST API и GraphQL. В рабочем проекте мы используем GraphQL. 

REST API имеет довольно строгую архитектуру, поэтому запросы в БД можно оптимизировать вручную (захардкодить). GraphQL использует более гибкий подход: клиент сам определяет, какие поля необходимо получить, поэтому оптимизировать что-либо вручную довольно сложно. Можно сразу делать select_related и prefetch_related по всем полям, но тогда сложно предусмотреть все варианты запросов, будут лишние «джойны» и из БД будет извлекаться избыточное количество данных. Именно поэтому мы решили сделать оптимизацию запросов.

Оптимизация запросов

Задача автоматической оптимизации запросов показалась мне очень интересной, поэтому я взялся за её решение. Идея состояла в том, чтобы доставать запрашиваемые поля из info (содержит всю информацию о запросе), рекурсивно проходить по внешним ключам моделей и на основании типа отношения применять нужный метод оптимизации. Вся необходимая информация о модели находится в атрибуте _meta. Так как информация о модели статична, ее можно кэшировать при помощи functools.

В info все названия полей указаны в формате CamelCase, поэтому нам нужно привести их к формату snake_case. Это удобно делать при помощи регулярных выражений. В итоге waitingAsyncDjango преобразуется в waiting_async_django:

camel_case_pattern = re.compile(r"(?<!^)(?=[A-Z])")

def string_to_snake_case(value: str) -> str:
	return camel_case_pattern.sub("_", value).lower()

Рекурсивно узнаём, какие именно поля запрашиваются в GraphQL. Для этого потребуется функция для преобразования абстрактного синтаксического дерева в словарь ast_to_dict, данные берём из info:

from graphql.utils.ast_to_dict import ast_to_dict

def collect_gql_fields(
	node: dict,
	fragments: dict,
	snake_case: bool = True,
) -> dict[str, dict]:
	field = {}
	if node.get("selection_set"):
		for leaf in node["selection_set"]["selections"]:
			if leaf["kind"] == "Field":
				key = leaf["name"]["value"]
				if snake_case:
					key = string_to_snake_case(key)
					field.update({key: collect_gql_fields(leaf, fragments)})
			elif leaf["kind"] == "FragmentSpread":
				field.update(
					collect_gql_fields(fragments[leaf["name"]["value"]], fragments)
				)
	return field

def get_gql_fields(
	info: ResolveInfo,
	snake_case: bool = True,
) -> dict[str, dict]:
	fragments = {}
	node = ast_to_dict(info.field_asts[0])
	for name, value in info.fragments.items():
		fragments[name] = ast_to_dict(value)
	return collect_gql_fields(node, fragments, snake_case=snake_case)

Для работы с only находим названия всех полей модели:

@lru_cache
def _get_model_all_field_names(model: Type[models.Model]) -> list[str]:
	return [field.name for field in model._meta.fields]

Для select_related находим названия полей с типом отношений many-to-one и one-to-one, включая реверсивные типы этих отношений:

@lru_cache
def _get_model_fields_to_select(
	model: Type[models.Model],
) -> dict[str, Type[models.Model]]:
	fields = {
		field.name: field.related_model
		for field in model._meta.fields
		if field.many_to_one or field.one_to_one
	}
	related_fields = {
		field.name: field.related_model
		for field in model._meta.fields_map.values()
		if field.many_to_one or field.one_to_one
	}
	return fields | related_fields

Для prefetch_related и Prefetch находим названия полей с типом отношений one-to-many и many-to-many, включая реверсивные типы этих отношений:

@lru_cache
def _get_model_fields_to_prefetch(
	model: Type[models.Model],
) -> dict[str, Type[models.Model]]:
	one_to_many = {
		field.name: field.related_model
		for field in model._meta.fields
		if field.one_to_many or field.many_to_many
	}
	many_to_many = {
		field.name: field.related_model 
		for field in model._meta.many_to_many
	}
	related_fields = {
		field.name: field.related_model
		for field in model._meta.fields_map.values()
		if field.one_to_many or field.many_to_many
	}
	return one_to_many | many_to_many | related_fields

Для удобства работы с данными определяем dataclass:

@dataclass
class ModelFields:
	select_fields: list[str]
	prefetch_fields: list[str | Prefetch]
	only_fields: list[str]

Узнав запрашиваемые поля и типы отношений, можно приступать к рекурсивному заполнению списка полей для каждого метода оптимизации:

def collect_model_fields(
	model: Type[models.Model],
	gql_fields: dict,
	prefix: str = "",
) -> ModelFields:
	all_field_names = _get_model_all_field_names(model)
	to_select = _get_model_fields_to_select(model)
	to_prefetch = _get_model_fields_to_prefetch(model)

	select_fields, prefetch_fields, only_fields = [], [], []

	# проходимся по всем запрашиваемым полям на текущем уровне
	for field_name in gql_fields.keys():
		# чем глубже уровень вложенности запроса, тем длиннее префикс
		location = prefix + field_name
		next_prefix = location + "__"
    
		if field_name in to_select:
			# текущее поле - внешний ключ с типом отношения many-to-one или one-to-one
			
			# рекурсивно достаем запрашиваемые поля из другой модели
			next_fields = collect_model_fields(
				to_select[field_name],
				gql_fields[field_name],
				prefix=next_prefix,
			)
      
			# и добавляем их к уже известным полям
			if next_fields.only_fields:
				# для текущей локации необходимо делать select_related
				# только в том случае, если из нее берутся поля для .only(),
				# так как если берутся поля ТОЛЬКО для prefetch,
				# будет выброшена ошибка.
				select_fields.append(location)
				only_fields.extend(next_fields.only_fields)
        
			select_fields.extend(next_fields.select_fields)
			prefetch_fields.extend(next_fields.prefetch_fields)
      
		elif field_name in to_prefetch:
			# текущее поле - внешний ключ с типом отношения one-to-many или many-to-many
			# рекурсивно достаем запрашиваемые поля из другой модели
			next_fields = collect_model_fields(
				to_prefetch[field_name],
				gql_fields[field_name],
				prefix=next_prefix,
			)
      
			# и добавляем их к уже известным полям
			prefetch_fields.extend(next_fields.prefetch_fields)
			if next_fields.select_fields:
				# если у другой модели запрашиваются поля по внешнему ключу
				# с типом отношения many-to-one или one-to-one,
				# то для бОльшей оптимизации можно использовать класс Prefetch
				related_select_fields = [
					i.removeprefix(next_prefix)
					for i in next_fields.select_fields
				]
				queryset = to_prefetch[field_name].objects.select_related(
					*related_select_fields
				)
				prefetch_fields.append(Prefetch(location, queryset=queryset))
			else:
				prefetch_fields.append(location)
      
		elif field_name in all_field_names:
			# обычное поле модели
			only_fields.append(location)
      
		else:
			# property и прочее.
			pass

	return ModelFields(
		select_fields=select_fields,
		prefetch_fields=prefetch_fields,
		only_fields=only_fields,
	)


def get_optimization_fields(
	model: Type[models.Model],
	info: ResolveInfo,
) -> ModelFields:
	# узнаем запрашиваемые поля
	gql_fields = get_gql_fields(info, snake_case=True)

	# находим поля для каждого способа оптимизации
	return collect_model_fields(model, gql_fields)

Для удобства мы определили менеджер моделей и абстрактный класс моделей:

class OptimizedModelManager(models.Manager):
	def db_prepared(self, info: ResolveInfo) -> models.QuerySet:
		fields = get_optimization_fields(self.model, info)
		return (
			self.get_queryset()
			.select_related(*fields.select_fields)
			.prefetch_related(*fields.prefetch_fields)
			.only(*fields.only_fields)
		)


class OptimizedModel(models.Model):
	class Meta:
		abstract = True

	objects = OptimizedModelManager()

Теперь можно с легкостью оптимизировать почти любые GraphQL-запросы:

# наследуемся от абстрактного класса с реализацией оптимизации запросов
class Book(OptimizedModel):
	"""Обычная модель книги."""

class BookGQLType(DjangoObjectType):
	class Meta:
		model = Book

class Queries(ObjectType):
	books = List(NonNull(BookGQLType), required=True)

	def resolve_books(self, info: ResolveInfo) -> list[Book]:
		# практически любой запрос будет хорошо оптимизирован
		return Book.objects.db_prepared(info)

В дальнейшем этот код оброс новым функционалом, например оптимизацией property моделей.

Ошибка Django

Метод only позволяет задавать список полей, которые нужно получить из базы данных. Получая только нужные поля, можно немного увеличить производительность приложения. В only можно указывать поля и других моделей, которые будут заджойнены.

Важно то, что Django автоматически достает pk (primary key) моделей, даже если это не указано в only.

На основании базовой модели Django позволяет создавать любое количество таких же моделей. Они будут различаться поведением, но при этом в БД будет существовать только одна базовая модель. Это достигается за счет proxy-моделей.

И вот при оптимизации запросов к proxy-модели иногда возникала такая ошибка:

File "django\db\models\query.py", line 302, in len
	self._fetch_all()
File "django\db\models\query.py", line 1507, in _fetch_all
	self._result_cache = list(self._iterable_class(self))
File "django\db\models\query.py", line 71, in iter
	related_populators = get_related_populators(klass_info, select, db)
File "django\db\models\query.py", line 2268, in get_related_populators
	rel_cls = RelatedPopulator(rel_klass_info, select, db)
File "django\db\models\query.py", line 2243, in init
	self.pk_idx = self.init_list.index(self.model_cls._meta.pk.attname)
ValueError: 'id' is not in list

Исследование показало, что ошибка возникала при попытке получить proxy-модель в select_related с применением only. Ошибка возникала только с only.

Причины бага

Судя по всему, у прокси-модели почему-то не доставался pk, хотя это должно происходить автоматически. 

Я принудительно добавил pk (в нашем случае это id) proxy-модели в only, и ошибка исчезла. Стало очевидно, что проблема где-то в Django.

Тогда я начал искать, где Django автоматически добавляет pk в список доставаемых полей. Поиски привели в django/db/models/sql/query.py, класс Query, функция deffered_to_data. 

Во время дебага обнаружил, что pk достается у прокси-модели, а запрашиваемые поля — у базовой.

Этот код показался подозрительным:

opts = cur_model._meta
if not is_reverse_o2o(source):
	must_include[old_model].add(source)
add_to_dict(must_include, cur_model, opts.pk)

Ошибка исчезла, когда pk стали доставать именно у базовой модели: 

opts = cur_model._meta.concrete_model._meta
if not is_reverse_o2o(source):
	must_include[old_model].add(source)
add_to_dict(must_include, cur_model, opts.pk)

Тикет и Pull Request

Убедившись, что в разделе Django Issues не было тикетов с подобным багом, я создал новый. Разработчики приняли его и предложили подготовить патч. Необходимо было создать pull request с фиксом бага и тестом. К счастью, в Django для этого есть отличная инструкция.

Так я сделал свой первый pull request в open-source проект.

В нашей реализации этот баг тоже пофиксили: сделали принудительное добавление pk в only при получении полей proxy-модели.

Наша оптимизация

Интересно, что мы вряд ли нашли бы эту ошибку без кастомной оптимизации запросов. Мы искали готовые реализации и библиотеки, и даже нашли одну — https://github.com/tfoxy/graphene-django-optimizer. Но она оказалась недостаточно гибкой для использования в некоторых сценариях и иногда делала даже больше запросов в БД, чем наша реализация.

Больше ничего похожего мы не нашли, но если знаете, поделитесь в комментариях ?

Баг, который оказался фичей

Во время работы я нашел еще один баг, который значительно увеличивает количество запросов в БД. Он возникает, если делать Prefetch через реверсивное отношение.

class Author(models.Model):
	name = models.CharField(max_length=32)

class Book(models.Model):
	title = models.CharField(max_length=32)
	author = models.ForeignKey("Author", on_delete=models.CASCADE)

	# для модели Publisher это реверсивный внешний ключ
	publisher = models.ForeignKey(
		"Publisher",
		on_delete=models.CASCADE,
		related_name="books",
	)

class Publisher(models.Model):
	address = models.CharField(max_length=32)

При таком GraphQL запросе…

{
	publishers {
		address
		books {
			title
			author {
				name
			}
		}
	}
}

… оптимизация должна быть следующей:

def resolve_publishers(self, info: ResolveInfo) -> list[Publisher]:
	queryset = Book.objects.select_related("author").only("title", "author__name")
	return Publisher.objects.prefetch_related(Prefetch("books", queryset)).only("address")

Но в этом случае наблюдается огромное количество запросов в БД.

Причина проблемы в этой строке:

Book.objects.select_related("author").only("title", "author__name")

Стоит добавить в only реверсивный pk (publisher_id), и запросов в БД становится в 10 раз меньше (в зависимости от количества записей в БД):

Book.objects.select_related("author").only("title", "author__name", "publisher_id")

Найдя причину бага и способ его исправления, я создал еще один тикет. Но...

Разработчики ответили, что функция prefetch_related_objects не должна неявно добавлять какие-либо столбцы, так как это может быть неожиданным и вводящим в заблуждение. Они добавили, что те, кто используют only и defer, несут полную ответственность за передаваемые в них поля. 

Тикет не подлежал фиксу и был отмечен как New feature.

В нашей реализации мы просто убрали использование only в Prefetch.

Заключение

Скрытый текст

В августе 2022 года ожидается релиз Django 4.1, где будет представлен асинхронный интерфейс для Django ORM. Под капотом запросы пока что синхронные, но разработчики продолжают работать над внедрением полноценной асинхронности. Думаю, когда это произойдет, для меня выбор будет очевидным. 

?? 8 августа — старт нашей бесплатной школы начинающих фронтендеров и бэкендеров. 

Если вы начинаете учиться разработке и вам интересны подобные штуки, приходите!