From 01f5c81bc2163c95d1d856ee0d6987099adc1ceb Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:08:37 -0700 Subject: [PATCH] Fix: count selection_data from one pass over the filtered documents Co-Authored-By: Claude Sonnet 5.5 --- src/documents/permissions.py | 9 +- src/documents/tests/test_api_documents.py | 101 ++++++++++++++++ src/documents/views.py | 138 ++++++++++++---------- 3 files changed, 179 insertions(+), 69 deletions(-) diff --git a/src/documents/permissions.py b/src/documents/permissions.py index ea6b1bc67..d297c4374 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_documents.py b/src/documents/tests/test_api_documents.py index eeb1ab59a..0007b1205 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, however the + link-table lookups are batched + """ + 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(link_batch_size=batch_size), + mock.patch.object( + DocumentViewSet, + "selection_data_link_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, however the link-table lookups are batched + """ + 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(link_batch_size=batch_size), + mock.patch.object( + DocumentViewSet, + "selection_data_link_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/views.py b/src/documents/views.py index 5e17c67cb..b529ba14c 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -5,6 +5,7 @@ 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 @@ -168,7 +169,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 @@ -1022,71 +1022,83 @@ 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)) + # Matching document ids are looked up in the link tables in batches of this + # size: one literal IN (...) over every match is a multi-MB statement per + # relation and, on SQLite, exceeds the bound-parameter limit. + selection_data_link_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): + # The filtered queryset is evaluated once, as the page's own filter is, + # and never wrapped in a subquery: re-deriving a join or a HAVING + # (shared_by__id) per statement is very slow on some databases, and + # regrouping the queryset by another column would miscount. + # order_by() drops the default/user ordering, which only costs a sort. + correspondents: Counter = Counter() + document_types: Counter = Counter() + storage_paths: Counter = Counter() + document_ids: list[int] = [] + seen: set[int] = set() + rows = queryset.order_by().values_list( + "pk", + "correspondent_id", + "document_type_id", + "storage_path_id", ) + for pk, correspondent, document_type, storage_path in rows: + # A filter can join, repeating a document row. + if pk in seen: + continue + seen.add(pk) + document_ids.append(pk) + # Documents without a value count under None, which selection() + # never looks up. + correspondents[correspondent] += 1 + document_types[document_type] += 1 + storage_paths[storage_path] += 1 + + def link_table_counts(link_model, field_name: str) -> Counter: + result: Counter = Counter() + batch_size = self.selection_data_link_batch_size + for start in range(0, len(document_ids), batch_size): + result.update( + dict( + link_model.objects.filter( + document_id__in=document_ids[start : start + batch_size], + ) + .values_list(field_name) + .annotate(document_count=Count("document_id")), + ), + ) + return result + + def selection(model, counts: Counter) -> list[dict[str, int]]: + # Every object is listed, with 0 for those with no matching document. + return [ + {"id": pk, "document_count": counts.get(pk, 0)} + for pk in model.objects.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, + link_table_counts(Document.tags.through, "tag_id"), + ), + "selected_document_types": selection( + DocumentType, + document_types, + ), + "selected_storage_paths": selection( + StoragePath, + storage_paths, + ), + "selected_custom_fields": selection( + CustomField, + link_table_counts(CustomFieldInstance, "field_id"), + ), } @classmethod