From fda50bb2d348eb87e105fb65da899b7ca04fe581 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Sun, 30 Aug 2026 14:49:24 -0700 Subject: [PATCH] fix(search): resolve index-write permissions and effective content in bulk Add WriteBatch.add_or_update_ids() and use it in bulk_update_documents and trash restore, cutting index writes from ~8 queries per document to a constant handful per batch --- src/documents/search/_backend.py | 40 +++++ src/documents/tasks.py | 8 +- src/documents/tests/search/test_backend.py | 190 +++++++++++++++++++++ src/documents/views.py | 3 +- 4 files changed, 236 insertions(+), 5 deletions(-) diff --git a/src/documents/search/_backend.py b/src/documents/search/_backend.py index b783cfd2d..e5bae979e 100644 --- a/src/documents/search/_backend.py +++ b/src/documents/search/_backend.py @@ -284,6 +284,46 @@ class WriteBatch: tantivy.Query.term_query(self._backend._schema, "id", doc_id), ) + def add_or_update_ids(self, ids: Sequence[int]) -> None: + """ + Add or update multiple documents in the batch by primary key. + + Unlike calling ``add_or_update()`` once per document, this resolves + viewer permissions and effective (versioned) content in bulk against + the ids as a whole, instead of once per document -- see + ``_DocumentViewerStream`` and ``annotate_effective_content``. Use + this whenever more than one document is being written in the same + batch. + + An id with no matching document (e.g. deleted between the caller + collecting ids and the batch running) is silently skipped, matching + ``add_or_update()``'s existing single-document deferred-task behavior + rather than erroring or leaving a stale index entry. + + Args: + ids: Primary keys of Document instances to index + """ + from documents.models import Document + from documents.versioning import annotate_effective_content + + ids = list(ids) + if not ids: + return + + queryset = annotate_effective_content( + Document.objects.filter(pk__in=ids) + .select_related("correspondent", "document_type", "storage_path", "owner") + .prefetch_related("tags", "notes__user", "custom_fields__field"), + ) + for document, grant in _DocumentViewerStream(queryset, chunk_size=1000): + self.remove(document.pk) + doc = self._backend._build_tantivy_doc( + document, + viewer_ids=grant.viewer_ids, + viewer_group_ids=grant.viewer_group_ids, + ) + self._writer.add_document(doc) + class TantivyBackend: """ diff --git a/src/documents/tasks.py b/src/documents/tasks.py index ca53468eb..33f7d4f90 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -312,7 +312,10 @@ def bulk_update_documents(document_ids) -> None: from documents.search import get_backend document_ids = list(document_ids) - # Annotated so indexing below doesn't query the versions of each document + # Annotated so the signal handlers below (e.g. matching) don't query the + # versions of each document. Indexing re-queries and re-annotates its own + # copy via add_or_update_ids() below, after these signals (and any + # workflow they trigger) have had a chance to mutate the documents. documents = annotate_effective_content( Document.objects.filter(id__in=document_ids), ) @@ -328,8 +331,7 @@ def bulk_update_documents(document_ids) -> None: post_save.send(Document, instance=doc, created=False) with get_backend().batch_update() as batch: - for doc in documents: - batch.add_or_update(doc) + batch.add_or_update_ids(document_ids) ai_config = AIConfig() if ai_config.llm_index_enabled: diff --git a/src/documents/tests/search/test_backend.py b/src/documents/tests/search/test_backend.py index 21fed8f02..52e1e593c 100644 --- a/src/documents/tests/search/test_backend.py +++ b/src/documents/tests/search/test_backend.py @@ -4,6 +4,8 @@ from pathlib import Path import pytest from django.contrib.auth.models import Group from django.contrib.auth.models import User +from django.db import connection +from django.test.utils import CaptureQueriesContext from guardian.shortcuts import assign_perm from pytest_mock import MockerFixture @@ -102,6 +104,194 @@ class TestWriteBatch: assert len(backend.search_ids("indexable", user=None)) == 1 +class TestAddOrUpdateIds: + """Test WriteBatch.add_or_update_ids(), the bulk id-based upsert path. + + Unlike add_or_update() called once per document, this resolves viewer + permissions and effective (versioned) content in bulk against the ids as + a whole, so it must produce identical indexed output to the per-document + path while issuing a constant number of queries regardless of batch size. + """ + + def test_indexes_all_documents_in_the_batch( + self, + backend: TantivyBackend, + ) -> None: + docs = [ + Document.objects.create( + title="doc", + content=f"unique{i}", + checksum=f"BULK{i}", + pk=i, + ) + for i in range(1, 4) + ] + + with backend.batch_update() as batch: + batch.add_or_update_ids([d.pk for d in docs]) + + for doc in docs: + assert backend.search_ids(f"unique{doc.pk}", user=None) == [doc.pk] + + def test_empty_id_list_is_a_noop(self, backend: TantivyBackend) -> None: + with backend.batch_update() as batch: + batch.add_or_update_ids([]) + + assert backend.search_ids("anything", user=None) == [] + + def test_missing_id_is_skipped_not_errored( + self, + backend: TantivyBackend, + ) -> None: + doc = Document.objects.create( + title="doc", + content="present", + checksum="EXIST1", + pk=1, + ) + missing_pk = 999 + + with backend.batch_update() as batch: + batch.add_or_update_ids([doc.pk, missing_pk]) + + assert backend.search_ids("present", user=None) == [doc.pk] + + def test_query_count_is_constant_regardless_of_batch_size( + self, + backend: TantivyBackend, + ) -> None: + small_docs = [ + Document.objects.create(title="doc", checksum=f"SMALL{i}", pk=i) + for i in range(1, 3) + ] + with CaptureQueriesContext(connection) as ctx_small: + with backend.batch_update() as batch: + batch.add_or_update_ids([d.pk for d in small_docs]) + num_queries_small = len(ctx_small.captured_queries) + + large_docs = [ + Document.objects.create(title="doc", checksum=f"LARGE{i}", pk=i) + for i in range(100, 150) + ] + with CaptureQueriesContext(connection) as ctx_large: + with backend.batch_update() as batch: + batch.add_or_update_ids([d.pk for d in large_docs]) + num_queries_large = len(ctx_large.captured_queries) + + assert num_queries_small == num_queries_large + + def test_resolves_direct_user_grant_in_bulk( + self, + backend: TantivyBackend, + ) -> None: + owner = UserFactory() + user = UserFactory() + doc = Document.objects.create( + title="doc", + checksum="PERM1", + pk=1, + owner=owner, + ) + assign_perm("view_document", user, doc) + + with backend.batch_update() as batch: + batch.add_or_update_ids([doc.pk]) + + assert backend.search_ids("doc", user=user) == [doc.pk] + other = UserFactory() + assert backend.search_ids("doc", user=other) == [] + + def test_resolves_group_grant_in_bulk(self, backend: TantivyBackend) -> None: + owner = UserFactory() + group = Group.objects.create(name="reviewers") + user = UserFactory() + user.groups.add(group) + doc = Document.objects.create( + title="doc", + checksum="GPERM1", + pk=1, + owner=owner, + ) + assign_perm("view_document", group, doc) + + with backend.batch_update() as batch: + batch.add_or_update_ids([doc.pk]) + + assert backend.search_ids("doc", user=user) == [doc.pk] + other = UserFactory() + assert backend.search_ids("doc", user=other) == [] + + def test_indexes_notes_and_custom_fields(self, backend: TantivyBackend) -> None: + note_author = UserFactory(username="noter") + field = CustomField.objects.create( + name="Invoice Number", + data_type=CustomField.FieldDataType.STRING, + ) + doc = Document.objects.create(title="doc", checksum="RICH1", pk=1) + Note.objects.create(document=doc, note="Reviewed", user=note_author) + CustomFieldInstance.objects.create( + document=doc, + field=field, + value_text="INV-42", + ) + + with backend.batch_update() as batch: + batch.add_or_update_ids([doc.pk]) + + assert backend.search_ids("notes.user:noter", user=None) == [doc.pk] + assert backend.search_ids("custom_fields.value:INV-42", user=None) == [ + doc.pk, + ] + + def test_uses_effective_content_for_versioned_documents( + self, + backend: TantivyBackend, + ) -> None: + root = Document.objects.create( + title="Statement", + content="stale text", + checksum="ROOT1", + pk=1, + ) + Document.objects.create( + title="Statement", + content="latest version text", + checksum="VER1", + pk=2, + root_document=root, + version_index=1, + ) + + with backend.batch_update() as batch: + batch.add_or_update_ids([root.pk]) + + assert backend.search_ids("latest", user=None) == [root.pk] + assert backend.search_ids("stale", user=None) == [] + + def test_reindexes_documents_already_in_the_index( + self, + backend: TantivyBackend, + ) -> None: + """add_or_update_ids must upsert, matching add_or_update's behaviour.""" + doc = Document.objects.create( + title="doc", + content="original", + checksum="UP1", + pk=1, + ) + backend.add_or_update(doc) + assert backend.search_ids("original", user=None) == [doc.pk] + + doc.content = "updated" + doc.save() + + with backend.batch_update() as batch: + batch.add_or_update_ids([doc.pk]) + + assert backend.search_ids("original", user=None) == [] + assert backend.search_ids("updated", user=None) == [doc.pk] + + class TestSearch: """Test search query parsing and matching via search_ids.""" diff --git a/src/documents/views.py b/src/documents/views.py index c6f0aeb15..bfda5a4bf 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -5444,8 +5444,7 @@ class TrashView(ListModelMixin, PassUserMixin): from documents.search import get_backend with get_backend().batch_update() as batch: - for doc in restored: - batch.add_or_update(doc) + batch.add_or_update_ids([doc.pk for doc in restored]) elif action == "empty": if doc_ids is None: doc_ids = [doc.id for doc in docs]