Merge remote-tracking branch 'origin/perf/batch-modify-custom-fields' into tmp/perf-integration

src/documents/bulk_edit.py

@@ -305,33 +305,49 @@ def modify_custom_fields(
         else [(field, None) for field in add_custom_fields]
     )
 
-    custom_fields = CustomField.objects.filter(
-        id__in=[int(field) for field, _ in add_custom_fields],
-    ).distinct()
+    custom_fields_by_id: dict[int, CustomField] = {
+        cf.id: cf
+        for cf in CustomField.objects.filter(
+            id__in=[int(field) for field, _ in add_custom_fields],
+        )
+    }
+    # Deferred, not `.only()`: these objects get cached onto the FK
+    # descriptor of newly-created CustomFieldInstance rows below, and
+    # downstream post_save receivers (e.g. the filename-generation signal)
+    # touch other Document fields -- `.only("pk")` would just turn that into
+    # a deferred-field reload per document, trading one N+1 for another.
+    # `content` is the one field guaranteed to be both large (full OCR text)
+    # and unused by anything this function or its receivers touch.
+    docs_by_id: dict[int, Document] = {
+        doc.id: doc
+        for doc in Document.objects.filter(id__in=affected_docs).defer("content")
+    }
     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 = {}
-            custom_field = custom_fields.get(id=field_id)
-            if custom_field:
-                value_field = CustomFieldInstance.TYPE_TO_DATA_STORE_NAME_MAP[
-                    custom_field.data_type
-                ]
-                defaults[value_field] = value
-                if (
-                    custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
-                    and value
-                    and doc_id in value
-                ):
-                    # Prevent self-linking
-                    continue
+            defaults = {value_field: value}
+            if (
+                custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK
+                and value
+                and doc_id in value
+            ):
+                # Prevent self-linking
+                continue
+            # Pass the already-resolved objects, not bare ids: this caches
+            # them on the FK descriptor of any newly-created instance, so a
+            # later `.field`/`.document` access (e.g. auditlog's post_save
+            # receiver calling `str(instance)`, which touches `.field.name`)
+            # doesn't trigger its own per-instance re-fetch.
             CustomFieldInstance.objects.update_or_create(
-                document_id=doc_id,
-                field_id=field_id,
+                document=docs_by_id[doc_id],
+                field=custom_field,
                 defaults=defaults,
             )
             if custom_field.data_type == CustomField.FieldDataType.DOCUMENTLINK:
-                doc = Document.objects.get(id=doc_id)
-                reflect_doclinks(doc, custom_field, value)
+                reflect_doclinks(docs_by_id[doc_id], custom_field, value)
 
     # For doc link fields that are being removed, remove symmetrical links
     for doclink_being_removed_instance in CustomFieldInstance.objects.filter(
@@ -339,12 +355,10 @@ def modify_custom_fields(
         field__id__in=remove_custom_fields,
         field__data_type=CustomField.FieldDataType.DOCUMENTLINK,
         value_document_ids__isnull=False,
-    ):
+    ).select_related("field"):
         for target_doc_id in doclink_being_removed_instance.value:
             remove_doclink(
-                document=Document.objects.get(
-                    id=doclink_being_removed_instance.document.id,
-                ),
+                document=docs_by_id[doclink_being_removed_instance.document_id],
                 field=doclink_being_removed_instance.field,
                 target_doc_id=target_doc_id,
             )

src/documents/tests/test_bulk_edit.py

@@ -348,6 +348,100 @@ 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_doclink_self_link(self) -> None:
         """
         GIVEN: