diff --git a/src/documents/data_models.py b/src/documents/data_models.py index 230af0684..e461ef9af 100644 --- a/src/documents/data_models.py +++ b/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.all() + for custom_field in doc.custom_fields.select_related("field").all() } groups_with_perms = get_groups_with_perms( diff --git a/src/documents/serialisers.py b/src/documents/serialisers.py index a3a92ef41..8d049f100 100644 --- a/src/documents/serialisers.py +++ b/src/documents/serialisers.py @@ -3,6 +3,7 @@ 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 @@ -956,8 +957,106 @@ 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 = serializers.PrimaryKeyRelatedField(queryset=CustomField.objects.all()) + field = _CachingCustomFieldPrimaryKeyField(queryset=CustomField.objects.all()) value = ReadWriteSerializerMethodField(allow_null=True) def create(self, validated_data): @@ -1058,6 +1157,7 @@ class CustomFieldInstanceSerializer(serializers.ModelSerializer[CustomFieldInsta class Meta: model = CustomFieldInstance + list_serializer_class = CustomFieldInstanceListSerializer fields = [ "value", "field", diff --git a/src/documents/tests/test_api_custom_fields.py b/src/documents/tests/test_api_custom_fields.py index 8ad69dd0d..f5e208582 100644 --- a/src/documents/tests/test_api_custom_fields.py +++ b/src/documents/tests/test_api_custom_fields.py @@ -5,7 +5,9 @@ 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 @@ -13,6 +15,9 @@ 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 @@ -530,6 +535,136 @@ 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: diff --git a/src/documents/tests/test_data_models.py b/src/documents/tests/test_data_models.py new file mode 100644 index 000000000..516d231e1 --- /dev/null +++ b/src/documents/tests/test_data_models.py @@ -0,0 +1,58 @@ +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}", + )