Merge remote-tracking branch 'origin/perf/batch-tags-field-lookup' into tmp/perf-integration

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()
 

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" = <literal>` (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: