Fix: select_related in remove_doclink() to avoid signal-triggered reload

Same pattern as the update_or_create() fix: target_doc_field_instance was
fetched without select_related, so its .document/.field weren't cached
when .save() fired the post_save signal -- auditlog's receiver touching
.document re-fetched it, once per (source, target) pair being unlinked
with no batching across calls. Also benefits the single-document PATCH
path in serialisers.py, which calls the same helper.

Broadened the removal test's query assertion now that both sides are fixed.

src/documents/bulk_edit.py

@@ -305,46 +305,74 @@ def modify_custom_fields(
         else [(field, None) for field in add_custom_fields]
     )
 
-    custom_fields = CustomField.objects.filter(
+    custom_fields_by_id: dict[int, CustomField] = {
-        id__in=[int(field) for field, _ in add_custom_fields],
+        cf.id: cf
-    ).distinct()
+        for cf in CustomField.objects.filter(
+            id__in=[int(field) for field, _ in add_custom_fields],
+        )
+    }
+    # Deferred, not `.only()`: signal receivers touch other Document fields,
+    # and `.only("pk")` would just turn that into a per-document reload.
+    # `content` is the one field both large and unused here. Skipped
+    # entirely for a remove-only call -- the removal pass below resolves
+    # its own documents.
+    docs_by_id: dict[int, Document] = (
+        {
+            doc.id: doc
+            for doc in Document.objects.filter(id__in=affected_docs).defer("content")
+        }
+        if add_custom_fields
+        else {}
+    )
     for field_id, value in add_custom_fields:
+        custom_field = custom_fields_by_id[field_id]
+        value_field = CustomFieldInstance.TYPE_TO_DATA_STORE_NAME_MAP[
+            custom_field.data_type
+        ]
         for doc_id in affected_docs:
-            defaults = {}
+            defaults = {value_field: value}
-            custom_field = custom_fields.get(id=field_id)
+            if (
-            if custom_field:
+                custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
-                value_field = CustomFieldInstance.TYPE_TO_DATA_STORE_NAME_MAP[
+                and value
-                    custom_field.data_type
+                and doc_id in value
-                ]
+            ):
-                defaults[value_field] = value
+                # Prevent self-linking
-                if (
+                continue
-                    custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
+            # Not update_or_create(): it fetches an existing row via plain
-                    and value
+            # `.get()` before calling .save(), so a signal receiver touching
-                    and doc_id in value
+            # `.field`/`.document` (e.g. auditlog) on that save re-fetches
-                ):
+            # per instance regardless of what's passed in as lookup kwargs.
-                    # Prevent self-linking
+            # Assigning the cached objects ourselves before .save() avoids
-                    continue
+            # that for both the create and update case.
-            CustomFieldInstance.objects.update_or_create(
+            try:
-                document_id=doc_id,
+                instance = CustomFieldInstance.objects.get(
-                field_id=field_id,
+                    document=docs_by_id[doc_id],
-                defaults=defaults,
+                    field=custom_field,
-            )
+                )
+            except CustomFieldInstance.DoesNotExist:
+                instance = CustomFieldInstance(
+                    document=docs_by_id[doc_id],
+                    field=custom_field,
+                )
+            instance.document = docs_by_id[doc_id]
+            instance.field = custom_field
+            for attr, val in defaults.items():
+                setattr(instance, attr, val)
+            instance.save()
             if custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK:
-                doc = Document.objects.get(id=doc_id)
+                reflect_doclinks(docs_by_id[doc_id], custom_field, value)
-                reflect_doclinks(doc, custom_field, value)
 
-    # For doc link fields that are being removed, remove symmetrical links
+    # For doc link fields being removed, remove symmetrical links.
+    # select_related here avoids resolving every affected document up front.
     for doclink_being_removed_instance in CustomFieldInstance.objects.filter(
         document_id__in=affected_docs,
         field__id__in=remove_custom_fields,
         field__data_type=CustomField.FieldDataType.DOCUMENTLINK,
         value_document_ids__isnull=False,
-    ):
+    ).select_related("field", "document"):
         for target_doc_id in doclink_being_removed_instance.value:
             remove_doclink(
-                document=Document.objects.get(
+                document=doclink_being_removed_instance.document,
-                    id=doclink_being_removed_instance.document.id,
-                ),
                 field=doclink_being_removed_instance.field,
                 target_doc_id=target_doc_id,
             )
@@ -1177,10 +1205,13 @@ def remove_doclink(
     """
     Removes a 'symmetrical' link to `document` from the target document's existing custom field instance
     """
-    target_doc_field_instance = CustomFieldInstance.objects.filter(
+    # select_related: a signal receiver (auditlog) touches .document/.field
-        document_id=target_doc_id,
+    # on save() below -- without this, that's a per-call reload query.
-        field=field,
+    target_doc_field_instance = (
-    ).first()
+        CustomFieldInstance.objects.filter(document_id=target_doc_id, field=field)
+        .select_related("document", "field")
+        .first()
+    )
     if (
         target_doc_field_instance is not None
         and document.id in target_doc_field_instance.value

src/documents/data_models.py

@@ -129,7 +129,7 @@ class DocumentMetadataOverrides:
         )
         overrides.custom_fields = {
             custom_field.field.id: custom_field.value
-            for custom_field in doc.custom_fields.select_related("field").all()
+            for custom_field in doc.custom_fields.all()
         }
 
         groups_with_perms = get_groups_with_perms(

src/documents/serialisers.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import logging
 import math
 import re
-from collections.abc import Iterable
 from datetime import datetime
 from datetime import timedelta
 from decimal import Decimal
@@ -877,106 +876,8 @@ def validate_documentlink_targets(user, doc_ids):
         )
 
 
-# drf-writable-nested revalidates a document's custom_fields more than once
-# per request: once as the ordinary nested list, then again per-item while
-# matching existing vs. new CustomFieldInstance rows during save() -- and
-# that second pass builds a brand new serializer (and field) instance per
-# item (see its update_or_create_reverse_relations / _get_serializer_for_field),
-# so a cache on the field instance alone only helps the first pass. It does,
-# however, explicitly pass `context=self.context` to every one of those
-# fresh serializers -- the *same* dict object the outer DocumentSerializer
-# is using, not a copy. That context dict is already request-scoped (DRF
-# builds it fresh per request via get_serializer_context()), so stashing the
-# resolved CustomField objects there -- rather than in some new global/
-# thread-local cache -- lets every later pass reuse them for free while
-# staying entirely within DRF's existing, already-request-scoped machinery.
-_CUSTOM_FIELD_CONTEXT_CACHE_KEY = "_custom_field_lookup_cache"
-
-
-class _CachingCustomFieldPrimaryKeyField(serializers.PrimaryKeyRelatedField):
-    """
-    Resolves CustomField ids with as few queries as possible: a per-instance
-    cache for repeat lookups on this exact field instance, backed by a
-    shared cache on the serializer context (see _CUSTOM_FIELD_CONTEXT_CACHE_KEY
-    above) so later, separately-instantiated fields for the same request
-    reuse what was already resolved instead of re-querying.
-    """
-
-    def __init__(self, **kwargs: Any) -> None:
-        super().__init__(**kwargs)
-        self._cache: dict[int, CustomField] = {}
-
-    def _shared_cache(self) -> dict[int, CustomField]:
-        return self.context.setdefault(_CUSTOM_FIELD_CONTEXT_CACHE_KEY, {})
-
-    @staticmethod
-    def _normalize_pk(data: Any) -> int | None:
-        """
-        Returns `data` coerced to the int a valid CustomField pk would be,
-        or None if `data` isn't a plausible pk (wrong type, unhashable,
-        non-numeric, or a bool -- DRF itself rejects bools as pks since
-        `True == 1` would otherwise silently match). None tells callers to
-        leave `data` alone and let `super().to_internal_value()` report the
-        normal validation error instead of touching the cache/queryset with
-        it directly.
-        """
-        if isinstance(data, bool):
-            return None
-        try:
-            return int(data)
-        except (TypeError, ValueError):
-            return None
-
-    def prefetch(self, ids: Iterable[Any]) -> None:
-        shared_cache = self._shared_cache()
-        candidates = {pk for i in ids if (pk := self._normalize_pk(i)) is not None}
-        missing = {
-            i for i in candidates if i not in self._cache and i not in shared_cache
-        }
-        if missing:
-            for obj in self.get_queryset().filter(pk__in=missing):
-                shared_cache[obj.pk] = obj
-        for i in candidates:
-            obj = shared_cache.get(i)
-            if obj is not None:
-                self._cache[i] = obj
-
-    def to_internal_value(self, data: Any) -> CustomField:
-        pk = self._normalize_pk(data)
-        if pk is None:
-            return super().to_internal_value(data)
-        if pk in self._cache:
-            return self._cache[pk]
-        shared_cache = self._shared_cache()
-        if pk in shared_cache:
-            obj = shared_cache[pk]
-            self._cache[pk] = obj
-            return obj
-        obj: CustomField = super().to_internal_value(data)
-        self._cache[obj.pk] = obj
-        shared_cache[obj.pk] = obj
-        return obj
-
-
-class CustomFieldInstanceListSerializer(serializers.ListSerializer):
-    def to_internal_value(self, data: Any) -> list[Any]:
-        if isinstance(data, list):
-            field_ids = []
-            for item in data:
-                if not isinstance(item, dict) or "field" not in item:
-                    continue
-                try:
-                    hash(item["field"])
-                except TypeError:
-                    continue
-                field_ids.append(item["field"])
-            if field_ids:
-                self.child.fields["field"].prefetch(field_ids)
-        return super().to_internal_value(data)
-
-
 class CustomFieldInstanceSerializer(serializers.ModelSerializer[CustomFieldInstance]):
-    field = _CachingCustomFieldPrimaryKeyField(queryset=CustomField.objects.all())
+    field = serializers.PrimaryKeyRelatedField(queryset=CustomField.objects.all())
     value = ReadWriteSerializerMethodField(allow_null=True)
 
     def create(self, validated_data):
@@ -1077,7 +978,6 @@ class CustomFieldInstanceSerializer(serializers.ModelSerializer[CustomFieldInsta
 
     class Meta:
         model = CustomFieldInstance
-        list_serializer_class = CustomFieldInstanceListSerializer
         fields = [
             "value",
             "field",

src/documents/tests/test_api_custom_fields.py

@@ -5,9 +5,7 @@ from unittest.mock import ANY
 
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
-from django.db import connection
 from django.test import override_settings
-from django.test.utils import CaptureQueriesContext
 from guardian.shortcuts import assign_perm
 from rest_framework import status
 from rest_framework.test import APITestCase
@@ -15,9 +13,6 @@ from rest_framework.test import APITestCase
 from documents.models import CustomField
 from documents.models import CustomFieldInstance
 from documents.models import Document
-from documents.serialisers import CustomFieldInstanceSerializer
-from documents.serialisers import DocumentSerializer
-from documents.tests.factories import DocumentFactory
 from documents.tests.utils import DirectoriesMixin
 
 
@@ -535,136 +530,6 @@ class TestCustomFieldsAPI(DirectoriesMixin, APITestCase):
         doc.refresh_from_db()
         self.assertEqual(len(doc.custom_fields.all()), 10)
 
-    def test_document_serializer_custom_fields_validation_batches_field_lookup(
-        self,
-    ) -> None:
-        """
-        GIVEN:
-            - A document is being validated with several custom field values
-              at once (as happens on every PATCH/PUT/POST)
-        WHEN:
-            - The serializer is validated
-        THEN:
-            - The referenced CustomField objects are resolved with a single
-              query, not one query per custom field
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-        custom_fields = [
-            CustomField.objects.create(
-                name=f"Test Custom Field {i}",
-                data_type=CustomField.FieldDataType.STRING,
-            )
-            for i in range(5)
-        ]
-
-        serializer = DocumentSerializer(
-            doc,
-            data={
-                "custom_fields": [
-                    {"field": custom_field.id, "value": "test value"}
-                    for custom_field in custom_fields
-                ],
-            },
-            partial=True,
-        )
-
-        with CaptureQueriesContext(connection) as ctx:
-            self.assertTrue(serializer.is_valid(), serializer.errors)
-
-        custom_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            len(custom_field_lookups),
-            1,
-            "Expected a single batched query to resolve the custom fields, "
-            f"got {len(custom_field_lookups)}: {custom_field_lookups}",
-        )
-
-    def test_custom_field_lookup_reuses_shared_context_cache(self) -> None:
-        """
-        GIVEN:
-            - A CustomField has already been resolved once, by a serializer
-              sharing a given `context` dict
-        WHEN:
-            - A second, separately-instantiated CustomFieldInstanceSerializer
-              validates the same field id, sharing that same context
-              (this is what drf-writable-nested does: it rebuilds a fresh
-              serializer -- and fresh field instances -- per item while
-              matching existing vs. new instances during save())
-        THEN:
-            - No additional query is issued to resolve the CustomField
-        """
-        custom_field = CustomField.objects.create(
-            name="Test Custom Field",
-            data_type=CustomField.FieldDataType.STRING,
-        )
-
-        context: dict = {}
-        first_pass = CustomFieldInstanceSerializer(
-            data={"field": custom_field.id, "value": "a"},
-            context=context,
-        )
-        self.assertTrue(first_pass.is_valid(), first_pass.errors)
-
-        second_pass = CustomFieldInstanceSerializer(
-            data={"field": custom_field.id, "value": "b"},
-            context=context,
-        )
-        with CaptureQueriesContext(connection) as ctx:
-            self.assertTrue(second_pass.is_valid(), second_pass.errors)
-
-        custom_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            len(custom_field_lookups),
-            0,
-            "Expected the second, separately-instantiated serializer to reuse "
-            f"the already-resolved CustomField, got: {custom_field_lookups}",
-        )
-
-    def test_custom_field_validation_rejects_malformed_field_value(self) -> None:
-        """
-        GIVEN:
-            - A document is being validated with a malformed custom_fields
-              entry whose "field" value is neither a valid CustomField id
-              nor a type DRF's own PrimaryKeyRelatedField can safely reject
-              on its own (unhashable, or a non-numeric scalar)
-        WHEN:
-            - The serializer is validated
-        THEN:
-            - A normal validation error is raised, not an unhandled
-              TypeError/ValueError escaping past DRF's validation layer
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-
-        bad_field_values = {
-            "unhashable-list": [],
-            "unhashable-dict": {},
-            "non-numeric-scalar": "abc",
-        }
-        for case_id, bad_field_value in bad_field_values.items():
-            with self.subTest(case_id):
-                serializer = DocumentSerializer(
-                    doc,
-                    data={
-                        "custom_fields": [
-                            {"field": bad_field_value, "value": "test value"},
-                        ],
-                    },
-                    partial=True,
-                )
-
-                self.assertFalse(serializer.is_valid())
-                self.assertIn("custom_fields", serializer.errors)
-
     def test_change_custom_field_instance_value(self) -> None:
         """
         GIVEN:

src/documents/tests/test_bulk_edit.py

@@ -6,7 +6,9 @@ from unittest import mock
 import pikepdf
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import User
+from django.db import connection
 from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
 from guardian.shortcuts import assign_perm
 from guardian.shortcuts import get_groups_with_perms
 from guardian.shortcuts import get_users_with_perms
@@ -344,6 +346,202 @@ class TestBulkEdit(DirectoriesMixin, TestCase):
         assert _cf_3 is not None
         self.assertNotIn(self.doc3.id, _cf_3.value)
 
+    def test_modify_custom_fields_batches_field_lookup(self) -> None:
+        """
+        GIVEN:
+            - Several documents are being bulk-edited to add several custom
+              fields at once
+        WHEN:
+            - modify_custom_fields runs
+        THEN:
+            - Each CustomField is resolved with one batched query total, not
+              once per (field, document) pair
+        """
+        docs = [
+            Document.objects.create(checksum=f"batch-{i}", title=f"batch-{i}")
+            for i in range(6)
+        ]
+        fields = [
+            CustomField.objects.create(
+                name=f"Batch Field {i}",
+                data_type=CustomField.FieldDataType.STRING,
+            )
+            for i in range(4)
+        ]
+
+        with CaptureQueriesContext(connection) as ctx:
+            bulk_edit.modify_custom_fields(
+                [doc.id for doc in docs],
+                add_custom_fields=[field.id for field in fields],
+                remove_custom_fields=[],
+            )
+
+        field_lookups = [
+            q
+            for q in ctx.captured_queries
+            if 'FROM "documents_customfield"' in q["sql"]
+        ]
+        self.assertEqual(
+            len(field_lookups),
+            1,
+            "Expected a single batched query to resolve the custom fields, "
+            f"got {len(field_lookups)}: {field_lookups}",
+        )
+
+        for doc in docs:
+            self.assertEqual(doc.custom_fields.count(), len(fields))
+
+    def test_modify_custom_fields_batches_document_lookup_for_documentlink(
+        self,
+    ) -> None:
+        """
+        GIVEN:
+            - Several documents are being bulk-edited to add a DOCUMENTLINK
+              custom field at once
+        WHEN:
+            - modify_custom_fields runs
+        THEN:
+            - The Document rows needed to reflect the symmetrical links are
+              resolved with one batched query total, not once per document
+        """
+        docs = [
+            Document.objects.create(checksum=f"link-{i}", title=f"link-{i}")
+            for i in range(6)
+        ]
+        target = Document.objects.create(checksum="link-target", title="link-target")
+        doclink_field = CustomField.objects.create(
+            name="Related",
+            data_type=CustomField.FieldDataType.DOCUMENTLINK,
+        )
+
+        with CaptureQueriesContext(connection) as ctx:
+            bulk_edit.modify_custom_fields(
+                [doc.id for doc in docs],
+                add_custom_fields={doclink_field.id: [target.id]},
+                remove_custom_fields=[],
+            )
+
+        single_document_lookups = [
+            q
+            for q in ctx.captured_queries
+            if 'FROM "documents_document"' in q["sql"]
+            and '"documents_document"."id" = ' in q["sql"]
+        ]
+        self.assertEqual(
+            len(single_document_lookups),
+            0,
+            "Expected document rows to come from a batched query, not "
+            f"per-document lookups, got: {single_document_lookups}",
+        )
+
+        for doc in docs:
+            self.assertEqual(
+                doc.custom_fields.get(field=doclink_field).value,
+                [target.id],
+            )
+
+    def test_modify_custom_fields_update_caches_document_and_field(self) -> None:
+        """
+        GIVEN:
+            - Several documents already have an instance of a custom field
+        WHEN:
+            - modify_custom_fields runs again for the same field, updating
+              the existing instances rather than creating new ones
+        THEN:
+            - No per-instance `.document`/`.field` reload query is issued
+              (e.g. by auditlog's post_save receiver touching them)
+        """
+        docs = [
+            Document.objects.create(checksum=f"update-{i}", title=f"update-{i}")
+            for i in range(6)
+        ]
+        field = CustomField.objects.create(
+            name="Update Field",
+            data_type=CustomField.FieldDataType.STRING,
+        )
+        bulk_edit.modify_custom_fields(
+            [doc.id for doc in docs],
+            add_custom_fields=[field.id],
+            remove_custom_fields=[],
+        )
+
+        with CaptureQueriesContext(connection) as ctx:
+            bulk_edit.modify_custom_fields(
+                [doc.id for doc in docs],
+                add_custom_fields={field.id: "updated value"},
+                remove_custom_fields=[],
+            )
+
+        single_row_reloads = [
+            q
+            for q in ctx.captured_queries
+            if ('FROM "documents_document"' in q["sql"] and '."id" = ' in q["sql"])
+            or ('FROM "documents_customfield"' in q["sql"] and '."id" = ' in q["sql"])
+        ]
+        self.assertEqual(
+            single_row_reloads,
+            [],
+            "Expected no per-instance document/field reload queries when "
+            f"updating existing custom field instances, got: {single_row_reloads}",
+        )
+
+        for doc in docs:
+            self.assertEqual(
+                doc.custom_fields.get(field=field).value,
+                "updated value",
+            )
+
+    def test_modify_custom_fields_removes_symmetrical_doclinks_batched(self) -> None:
+        """
+        GIVEN:
+            - Several source documents link to a shared target via a doc
+              link field
+        WHEN:
+            - The field is removed from all of them in one call
+        THEN:
+            - The symmetrical links are removed from the target
+            - No per-document lookup query is issued, on either side
+        """
+        target = Document.objects.create(checksum="rm-target", title="rm-target")
+        docs = [
+            Document.objects.create(checksum=f"rm-{i}", title=f"rm-{i}")
+            for i in range(6)
+        ]
+        field = CustomField.objects.create(
+            name="Related",
+            data_type=CustomField.FieldDataType.DOCUMENTLINK,
+        )
+        bulk_edit.modify_custom_fields(
+            [doc.id for doc in docs],
+            add_custom_fields={field.id: [target.id]},
+            remove_custom_fields=[],
+        )
+        self.assertEqual(
+            target.custom_fields.get(field=field).value,
+            [d.id for d in docs],
+        )
+
+        with CaptureQueriesContext(connection) as ctx:
+            bulk_edit.modify_custom_fields(
+                [doc.id for doc in docs],
+                add_custom_fields=[],
+                remove_custom_fields=[field.id],
+            )
+
+        single_document_lookups = [
+            q
+            for q in ctx.captured_queries
+            if 'FROM "documents_document"' in q["sql"]
+            and '"documents_document"."id" = ' in q["sql"]
+        ]
+        self.assertEqual(
+            single_document_lookups,
+            [],
+            "Expected batched document resolution, not per-document lookups, "
+            f"got: {single_document_lookups}",
+        )
+        self.assertEqual(target.custom_fields.get(field=field).value, [])
+
     def test_modify_custom_fields_doclink_self_link(self) -> None:
         """
         GIVEN:

src/documents/tests/test_data_models.py

@@ -1,58 +0,0 @@
-from django.db import connection
-from django.test import TestCase
-from django.test.utils import CaptureQueriesContext
-
-from documents.data_models import DocumentMetadataOverrides
-from documents.models import CustomField
-from documents.models import CustomFieldInstance
-from documents.tests.factories import DocumentFactory
-from documents.tests.utils import DirectoriesMixin
-
-
-class TestDocumentMetadataOverridesFromDocument(DirectoriesMixin, TestCase):
-    def test_from_document_batches_custom_field_lookup_after_refresh_from_db(
-        self,
-    ) -> None:
-        """
-        GIVEN:
-            - A document has several custom field values
-            - The document instance has just been refreshed from the database,
-              which drops any prefetched related objects (as
-              send_websocket_document_updated does before building overrides)
-        WHEN:
-            - DocumentMetadataOverrides.from_document() reads the document's
-              custom field values
-        THEN:
-            - The referenced CustomField objects are resolved with a single
-              query, not one query per custom field
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-        for i in range(5):
-            CustomFieldInstance.objects.create(
-                document=doc,
-                field=CustomField.objects.create(
-                    name=f"Test Custom Field {i}",
-                    data_type=CustomField.FieldDataType.STRING,
-                ),
-                value_text="value",
-            )
-
-        doc.refresh_from_db()
-
-        with CaptureQueriesContext(connection) as ctx:
-            overrides = DocumentMetadataOverrides.from_document(doc)
-
-        self.assertEqual(len(overrides.custom_fields), 5)
-        unbatched_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            unbatched_field_lookups,
-            [],
-            "Expected CustomField data to come from the CustomFieldInstance "
-            "join, not a separate per-instance lookup, "
-            f"got: {unbatched_field_lookups}",
-        )