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).
This commit is contained in:
stumpylog
2026-08-24 15:15:00 -07:00
parent 01c40aa0bd
commit 9d2416c435
2 changed files with 163 additions and 4 deletions
+84 -4
View File
@@ -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()
+79
View File
@@ -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: