From 9d2416c435266d45b3b127879c6c5f4dec63eba3 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Mon, 24 Aug 2026 15:15:00 -0700 Subject: [PATCH] Perf: batch id resolution for TagsField and friends TagsField/CorrespondentField/DocumentTypeField/StoragePathField were plain PrimaryKeyRelatedField subclasses with no batching. When used with many=True (only tags today: DocumentSerializer.tags, WorkflowActionSerializer.assign_tags), DRF's ManyRelatedField resolves each submitted id with its own query -- one query per tag on every PATCH/PUT that sets tags. Added BatchResolvingPrimaryKeyRelatedField as the shared base for all four field classes and overrode many_init so the many=True form (_BatchingManyRelatedField) resolves the whole id list with one pk__in query, falling back to the child relation's normal per-item validation for anything not found in that batch. Only TagsField uses many=True today, but the fix isn't tag-specific -- if a future PR puts many=True on one of the others, it inherits the same batching instead of reintroducing this as a new bug. Independent review caught a real regression: Django's IntegerFieldOverflow guard (out-of-range int -> EmptyResultSet) only covers exact/gt/gte/lt/lte lookups, not `in`, so an absurdly large tag id reached the batched pk__in= query as-is and raised an unhandled OverflowError (SQLite) / DataError (Postgres) instead of the normal 400 the original per-item `exact` lookup produced. Guarded the batch query and fall through to per-item resolution (which goes through the protected `exact` lookup) on failure. Verified via CaptureQueriesContext against a real API PATCH: 20 tags dropped from 54 to 35 queries per request (exactly the 19 saved by collapsing 20 individual lookups into one batched query). Full documents/workflows/bulk-edit/retagger/custom-fields suites green (443 passed). --- src/documents/serialisers.py | 88 +++++++++++++++++++++-- src/documents/tests/test_api_documents.py | 79 ++++++++++++++++++++ 2 files changed, 163 insertions(+), 4 deletions(-) diff --git a/src/documents/serialisers.py b/src/documents/serialisers.py index b2f14b505..e77a5249c 100644 --- a/src/documents/serialisers.py +++ b/src/documents/serialisers.py @@ -24,6 +24,7 @@ from django.core.validators import MaxValueValidator from django.core.validators import MinValueValidator from django.core.validators import RegexValidator from django.core.validators import integer_validator +from django.db import DataError from django.db.models import Count from django.db.models import Q from django.db.models.functions import Lower @@ -43,6 +44,7 @@ from guardian.shortcuts import get_users_with_perms from guardian.utils import get_group_obj_perms_model from guardian.utils import get_user_obj_perms_model from rest_framework import fields +from rest_framework import relations from rest_framework import serializers from rest_framework.exceptions import PermissionDenied from rest_framework.fields import SerializerMethodField @@ -742,22 +744,100 @@ class TagSerializer(MatchingModelSerializer, OwnedObjectSerializer): return super().validate(attrs) -class CorrespondentField(serializers.PrimaryKeyRelatedField[Correspondent]): +class _BatchingManyRelatedField(serializers.ManyRelatedField): + """ + `ManyRelatedField.to_internal_value` resolves each id in the submitted + list with its own `child_relation.to_internal_value()` call -- one query + per item on every PATCH/PUT that sets a `many=True` relation field. + Batch-resolve them instead, falling back to the child relation's normal + (query-per-item) validation for anything that isn't a plausible int pk, + so bad input still gets the usual DRF validation error rather than being + silently dropped. + """ + + @staticmethod + def _normalize_pk(item) -> int | None: + # Excludes bool: DRF's own PrimaryKeyRelatedField rejects it too + # (True == 1 would otherwise silently match pk 1). + if isinstance(item, bool): + return None + try: + return int(item) + except (TypeError, ValueError): + return None + + def to_internal_value(self, data): + if isinstance(data, str) or not hasattr(data, "__iter__"): + self.fail("not_a_list", input_type=type(data).__name__) + if not self.allow_empty and len(data) == 0: + self.fail("empty") + + item_pks = [(item, self._normalize_pk(item)) for item in data] + candidate_pks = {pk for _, pk in item_pks if pk is not None} + + # Django's IntegerFieldOverflow guard (-> EmptyResultSet, i.e. no + # match) only covers exact/gt/gte/lt/lte lookups, not `in` -- an + # out-of-range int in `pk__in=` reaches the DB driver as-is and + # raises OverflowError (SQLite) / DataError (Postgres) instead of + # cleanly matching nothing. The per-item `exact`-lookup fallback + # below IS covered, so on that failure just skip the batch and let + # every item resolve individually -- each still costs one query, + # but reports the normal validation error instead of a raw 500. + try: + resolved_by_pk = { + obj.pk: obj + for obj in self.child_relation.get_queryset().filter( + pk__in=candidate_pks, + ) + } + except (OverflowError, DataError): + resolved_by_pk = {} + + result = [] + for item, pk in item_pks: + obj = resolved_by_pk.get(pk) if pk is not None else None + result.append( + obj if obj is not None else self.child_relation.to_internal_value(item), + ) + return result + + +class BatchResolvingPrimaryKeyRelatedField(serializers.PrimaryKeyRelatedField): + """ + A PrimaryKeyRelatedField whose `many=True` form (a DRF ManyRelatedField) + resolves all submitted ids with one batched query instead of one query + per id. Subclasses only need to implement `get_queryset()` as usual -- + only `TagsField` is used with `many=True` today, but this is the base + for all four so the fix isn't tag-specific: if a future PR puts + `many=True` on correspondent/document_type/storage_path, it inherits the + same batching instead of reintroducing this as a new bug to rediscover. + """ + + @classmethod + def many_init(cls, *args, **kwargs): + list_kwargs = {"child_relation": cls(*args, **kwargs)} + for key, value in kwargs.items(): + if key in relations.MANY_RELATION_KWARGS: + list_kwargs[key] = value + return _BatchingManyRelatedField(**list_kwargs) + + +class CorrespondentField(BatchResolvingPrimaryKeyRelatedField[Correspondent]): def get_queryset(self): return Correspondent.objects.all() -class TagsField(serializers.PrimaryKeyRelatedField[Tag]): +class TagsField(BatchResolvingPrimaryKeyRelatedField[Tag]): def get_queryset(self): return Tag.objects.all() -class DocumentTypeField(serializers.PrimaryKeyRelatedField[DocumentType]): +class DocumentTypeField(BatchResolvingPrimaryKeyRelatedField[DocumentType]): def get_queryset(self): return DocumentType.objects.all() -class StoragePathField(serializers.PrimaryKeyRelatedField[StoragePath]): +class StoragePathField(BatchResolvingPrimaryKeyRelatedField[StoragePath]): def get_queryset(self): return StoragePath.objects.all() diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index b67b87772..2664f6314 100644 --- a/src/documents/tests/test_api_documents.py +++ b/src/documents/tests/test_api_documents.py @@ -1,5 +1,6 @@ import datetime import json +import re import shutil import tempfile import uuid @@ -21,7 +22,9 @@ from django.core import mail from django.core.cache import cache from django.core.files.uploadedfile import SimpleUploadedFile from django.db import DataError +from django.db import connection from django.test import override_settings +from django.test.utils import CaptureQueriesContext from django.utils import timezone from guardian.shortcuts import assign_perm from rest_framework import status @@ -252,6 +255,82 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): doc.refresh_from_db() self.assertEqual(doc.created, date(2023, 6, 28)) + def test_document_update_tags_batches_tag_lookup(self) -> None: + """ + GIVEN: + - A document is being updated with several tags at once + WHEN: + - API PATCH request is made setting the document's tags + THEN: + - The referenced Tag objects are resolved with a single batched + query, not one query per tag + """ + doc = Document.objects.create( + title="none", + checksum="123", + mime_type="application/pdf", + ) + tags = [TagFactory() for _ in range(8)] + + with CaptureQueriesContext(connection) as ctx: + response = self.client.patch( + f"/api/documents/{doc.pk}/", + {"tags": [t.id for t in tags]}, + format="json", + ) + self.assertEqual(response.status_code, status.HTTP_200_OK) + + # Match `"documents_tag"."id" = ` (a single-row WHERE lookup) + # but not the same substring appearing as a JOIN's ON condition + # (`"documents_tag"."id" = "documents_document_tags"."tag_id"`), + # which is a legitimate, unrelated response-serialization query. + single_tag_lookup_re = re.compile(r'"documents_tag"\."id" = \d') + single_tag_lookups = [ + q for q in ctx.captured_queries if single_tag_lookup_re.search(q["sql"]) + ] + self.assertEqual( + len(single_tag_lookups), + 0, + "Expected tags to be resolved with a batched query, not " + f"per-tag lookups, got: {single_tag_lookups}", + ) + + doc.refresh_from_db() + self.assertCountEqual( + doc.tags.values_list("id", flat=True), + [t.id for t in tags], + ) + + def test_document_update_tags_rejects_out_of_range_id(self) -> None: + """ + GIVEN: + - A document is being updated with a tag id too large for the + database's integer column + WHEN: + - API PATCH request is made setting the document's tags + THEN: + - A normal 400 validation error is returned, not an unhandled + OverflowError/DataError escaping as a 500 + + Django's IntegerFieldOverflow guard converts an out-of-range int + into a clean "no match" for exact/gt/gte/lt/lte lookups, but not for + `in` -- the batched tag resolution uses `pk__in=`, so this has to be + guarded explicitly rather than relying on Django to do it. + """ + doc = Document.objects.create( + title="none", + checksum="123", + mime_type="application/pdf", + ) + + response = self.client.patch( + f"/api/documents/{doc.pk}/", + {"tags": [99999999999999999999999999999]}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + def test_document_update_legacy_created_format(self) -> None: """ GIVEN: