diff --git a/src/documents/permissions.py b/src/documents/permissions.py index ea6b1bc672..d297c43744 100644 --- a/src/documents/permissions.py +++ b/src/documents/permissions.py @@ -576,12 +576,9 @@ def annotate_document_count_by_ids( or CustomFieldInstance) related_object_field: field on the relation pointing back to queryset pk document_ids: the document ids to count against -- a concrete list/set, - or a simple (already resolved) queryset of ids. Callers - that need this filtered by a complex condition (e.g. a - permission check) should resolve it to a concrete list - first if the same ids will be reused across multiple - calls, rather than passing the complex queryset itself - into each -- see ``_get_selection_data_for_queryset``. + or a simple (already resolved) queryset of ids. Keep + the list small: it is embedded in the statement as + literal parameters. target_field: field on the relation pointing to Document id """ diff --git a/src/documents/tests/test_api_bulk_edit.py b/src/documents/tests/test_api_bulk_edit.py index 43dfc1064f..7f0cec816a 100644 --- a/src/documents/tests/test_api_bulk_edit.py +++ b/src/documents/tests/test_api_bulk_edit.py @@ -1141,8 +1141,8 @@ class TestBulkEditAPI(DirectoriesMixin, APITestCase): self.assertCountEqual( response.data["selected_document_types"], [ - {"id": self.c1.id, "document_count": 1}, - {"id": self.c2.id, "document_count": 0}, + {"id": self.dt1.id, "document_count": 1}, + {"id": self.dt2.id, "document_count": 0}, ], ) diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 9a502b2cbc..07616f2d98 100644 --- a/src/documents/tests/test_api_documents.py +++ b/src/documents/tests/test_api_documents.py @@ -49,6 +49,7 @@ from documents.models import WorkflowTrigger from documents.search import get_backend from documents.signals.handlers import run_workflows from documents.tests.utils import ConsumeTaskMixin +from documents.views import DocumentViewSet from paperless_testing.dirs import DirectoriesMixin from paperless_testing.factories import DocumentFactory from paperless_testing.factories import TagFactory @@ -1755,6 +1756,106 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): } self.assertEqual(counts_by_field[custom_field.id], 1) + def test_selection_data_counts_with_shared_by_filter(self) -> None: + """ + GIVEN: + - Several documents owned by one user with the same correspondent + and document type, only one of which is shared + WHEN: + - Requesting the document list filtered by shared_by__id with + include_selection_data=true + THEN: + - Counts reflect only the shared document, not every document the + owner has with that correspondent or document type, regardless of + the batch size + """ + owner = User.objects.create_user(username="owner") + other = User.objects.create_user(username="other") + correspondent = Correspondent.objects.create(name="c1") + doc_type = DocumentType.objects.create(name="dt1") + tag = Tag.objects.create(name="tag") + docs = [ + Document.objects.create( + checksum=str(i), + owner=owner, + correspondent=correspondent, + document_type=doc_type, + ) + for i in range(3) + ] + for doc in docs: + doc.tags.add(tag) + grant_object(other, docs[0], "view_document") + + for batch_size in (10_000, 1): + with ( + self.subTest(batch_size=batch_size), + mock.patch.object( + DocumentViewSet, + "selection_data_batch_size", + batch_size, + ), + ): + response = self.client.get( + f"/api/documents/?shared_by__id={owner.id}" + "&include_selection_data=true", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.data["count"], 1) + selection_data = response.data["selection_data"] + for key, obj in ( + ("selected_correspondents", correspondent), + ("selected_document_types", doc_type), + ("selected_tags", tag), + ): + counts = { + item["id"]: item["document_count"] + for item in selection_data[key] + } + self.assertEqual(counts[obj.id], 1, key) + + def test_selection_data_counts_documents_once_when_filter_joins(self) -> None: + """ + GIVEN: + - A document matching a multi-valued filter through two tags + WHEN: + - Requesting the document list with include_selection_data=true + THEN: + - The document is counted once per correspondent, not once per + matching tag, regardless of the batch size + """ + correspondent = Correspondent.objects.create(name="c1") + tag_a = Tag.objects.create(name="a") + tag_b = Tag.objects.create(name="b") + both = Document.objects.create(checksum="1", correspondent=correspondent) + both.tags.add(tag_a, tag_b) + only_a = Document.objects.create(checksum="2", correspondent=correspondent) + only_a.tags.add(tag_a) + + for batch_size in (10_000, 1): + with ( + self.subTest(batch_size=batch_size), + mock.patch.object( + DocumentViewSet, + "selection_data_batch_size", + batch_size, + ), + ): + response = self.client.get( + f"/api/documents/?tags__id__in={tag_a.id},{tag_b.id}" + "&include_selection_data=true", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + counts = { + item["id"]: item["document_count"] + for item in response.data["selection_data"][ + "selected_correspondents" + ] + } + self.assertEqual(counts[correspondent.id], 2) + def test_statistics(self) -> None: doc1 = Document.objects.create( title="none1", diff --git a/src/documents/tests/test_api_search.py b/src/documents/tests/test_api_search.py index 9c7ff158df..ebcd20280b 100644 --- a/src/documents/tests/test_api_search.py +++ b/src/documents/tests/test_api_search.py @@ -25,6 +25,7 @@ from documents.models import Tag from documents.models import Workflow from documents.search import get_backend from documents.search import reset_backend +from documents.views import DocumentViewSet from paperless_mail.models import MailAccount from paperless_mail.models import MailRule from paperless_testing.dirs import DirectoriesMixin @@ -345,38 +346,39 @@ class TestDocumentSearchApi(DirectoriesMixin, APITestCase): storage_path = StoragePath.objects.create(name="sp1") tag = Tag.objects.create(name="tag") - matching_doc = Document.objects.create( - title="bank statement", - content="bank content", - checksum="A", - correspondent=correspondent, - document_type=doc_type, - storage_path=storage_path, - ) - matching_doc.tags.add(tag) + for checksum, title in ( + ("A", "bank statement"), + ("B", "bank account"), + ("C", "credit card"), + ): + doc = Document.objects.create( + title=title, + checksum=checksum, + correspondent=correspondent, + document_type=doc_type, + storage_path=storage_path, + ) + doc.tags.add(tag) + get_backend().add_or_update(doc) - get_backend().add_or_update(matching_doc) - - response = self.client.get( - "/api/documents/?query=bank&include_selection_data=true", - ) + with mock.patch.object(DocumentViewSet, "selection_data_batch_size", 1): + response = self.client.get( + "/api/documents/?query=bank&page_size=1&include_selection_data=true", + ) self.assertEqual(response.status_code, status.HTTP_200_OK) - self.assertIn("selection_data", response.data) - - selected_correspondent = next( - item - for item in response.data["selection_data"]["selected_correspondents"] - if item["id"] == correspondent.id - ) - selected_tag = next( - item - for item in response.data["selection_data"]["selected_tags"] - if item["id"] == tag.id - ) - - self.assertEqual(selected_correspondent["document_count"], 1) - self.assertEqual(selected_tag["document_count"], 1) + self.assertEqual(response.data["count"], 2) + self.assertEqual(len(response.data["results"]), 1) + for key, pk in ( + ("selected_correspondents", correspondent.pk), + ("selected_document_types", doc_type.pk), + ("selected_storage_paths", storage_path.pk), + ("selected_tags", tag.pk), + ): + self.assertEqual( + response.data["selection_data"][key], + [{"id": pk, "document_count": 2}], + ) def test_search_custom_field_ordering(self) -> None: custom_field = CustomField.objects.create( diff --git a/src/documents/views.py b/src/documents/views.py index 0b289e8da0..1548753400 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -5,11 +5,13 @@ import platform import re import tempfile import zipfile +from collections import Counter from collections import deque from collections.abc import Callable from datetime import datetime from datetime import timedelta from http import HTTPStatus +from itertools import islice from pathlib import Path from time import mktime from time import sleep @@ -168,7 +170,6 @@ from documents.permissions import PaperlessNotePermissions from documents.permissions import PaperlessObjectPermissions from documents.permissions import TrashPermissions from documents.permissions import ViewDocumentsPermissions -from documents.permissions import annotate_document_count_by_ids from documents.permissions import annotate_document_count_for_related_queryset from documents.permissions import get_document_count_filter_for_user from documents.permissions import get_permission_grantees @@ -1014,71 +1015,80 @@ class DocumentViewSet( "custom_field_", ) - def _get_selection_data_for_queryset(self, queryset): - # Resolve once instead of once per model below. `queryset` can carry an - # arbitrarily expensive WHERE clause (user filters plus the permission - # filter); re-embedding it as a subquery inside 5 separate Count(...) - # calls forces the database to re-evaluate that whole thing 5 times, and - # -- for FK relations especially -- can defeat semi-join planning - # entirely at scale. A concrete id list is cheap to reuse. - # order_by() drops the default/user ordering -- irrelevant for a plain - # id list, but left in place it forces a sort over the full filtered - # set before the ids can even be collected. - document_ids = list(queryset.order_by().values_list("pk", flat=True)) + # Bound the IN clauses used for search results and through-table lookups. + selection_data_batch_size = 10_000 - correspondents = Correspondent.objects.annotate( - document_count=Count( - "documents", - filter=Q(documents__id__in=document_ids), - distinct=True, - ), - ) - document_types = DocumentType.objects.annotate( - document_count=Count( - "documents", - filter=Q(documents__id__in=document_ids), - distinct=True, - ), - ) - storage_paths = StoragePath.objects.annotate( - document_count=Count( - "documents", - filter=Q(documents__id__in=document_ids), - distinct=True, - ), - ) - # Tag and CustomField reach Document through an M2M/through-model table; - # a plain Count(filter=...) there is a much more expensive plan than the - # FK relations above once the bridge table is large -- see - # annotate_document_count_by_ids() for why. - tags = annotate_document_count_by_ids( - Tag.objects.all(), - through_model=Document.tags.through, - related_object_field="tag_id", - document_ids=document_ids, - ) - custom_fields = annotate_document_count_by_ids( - CustomField.objects.all(), - through_model=CustomFieldInstance, - related_object_field="field_id", - document_ids=document_ids, + def _get_selection_data_for_queryset( + self, + queryset: QuerySet[Document], + *, + search_ids: list[int] | None = None, + ) -> dict[str, list[dict[str, int]]]: + # Keep permission and filter predicates at document granularity. + # Reusing the filtered queryset as a subquery for each count is slow. + rows = ( + queryset.order_by() + .prefetch_related(None) + .values_list( + "pk", + "correspondent_id", + "document_type_id", + "storage_path_id", + ) ) + batch_size = self.selection_data_batch_size + + def document_rows(): + if search_ids is None: + yield from rows.iterator() + else: + for start in range(0, len(search_ids), batch_size): + yield from rows.filter( + pk__in=search_ids[start : start + batch_size], + ).iterator() + + correspondents: Counter[int | None] = Counter() + document_types: Counter[int | None] = Counter() + storage_paths: Counter[int | None] = Counter() + document_ids: set[int] = set() + for pk, correspondent, document_type, storage_path in document_rows(): + # Filters that join multi-valued relations can repeat a document. + if pk in document_ids: + continue + document_ids.add(pk) + correspondents[correspondent] += 1 + document_types[document_type] += 1 + storage_paths[storage_path] += 1 + + tags: Counter[int | None] = Counter() + custom_fields: Counter[int | None] = Counter() + ids = iter(document_ids) + while batch := list(islice(ids, batch_size)): + for model, field, counts in ( + (Document.tags.through, "tag_id", tags), + (CustomFieldInstance, "field_id", custom_fields), + ): + counts.update( + dict( + model.objects.filter(document_id__in=batch) + .order_by() + .values_list(field) + .annotate(document_count=Count("document_id")), + ), + ) + + def selection(model: type[Model], counts: Counter[int | None]): + return [ + {"id": pk, "document_count": counts[pk]} + for pk in model._default_manager.values_list("pk", flat=True) + ] + return { - "selected_correspondents": [ - {"id": t.id, "document_count": t.document_count} for t in correspondents - ], - "selected_tags": [ - {"id": t.id, "document_count": t.document_count} for t in tags - ], - "selected_document_types": [ - {"id": t.id, "document_count": t.document_count} for t in document_types - ], - "selected_storage_paths": [ - {"id": t.id, "document_count": t.document_count} for t in storage_paths - ], - "selected_custom_fields": [ - {"id": t.id, "document_count": t.document_count} for t in custom_fields - ], + "selected_correspondents": selection(Correspondent, correspondents), + "selected_tags": selection(Tag, tags), + "selected_document_types": selection(DocumentType, document_types), + "selected_storage_paths": selection(StoragePath, storage_paths), + "selected_custom_fields": selection(CustomField, custom_fields), } @classmethod @@ -2725,13 +2735,10 @@ class UnifiedSearchViewSet(DocumentViewSet): if get_boolean( str(request.query_params.get("include_selection_data", "false")), ): - # NOTE: pk__in=ordered_ids generates a large SQL IN clause - # for big result sets. Acceptable today but may need a temp - # table or chunked approach if selection_data becomes slow - # at scale (tens of thousands of matching documents). response.data["selection_data"] = ( self._get_selection_data_for_queryset( - filtered_qs.filter(pk__in=result.ordered_ids), + filtered_qs, + search_ids=result.ordered_ids, ) ) return response