Performance: stop embedding every document ID in selection_data queries (#14597)

* Fix: count selection_data from one pass over the filtered documents

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

* Chore: use document type ids in the bulk edit selection_data test

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

* Batch search selection data queries

Co-Authored-By: Codex <267193182+codex@users.noreply.github.com>

---------

Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-authored-by: shamoon <4887959+shamoon@users.noreply.github.com>
Co-authored-by: Codex <267193182+codex@users.noreply.github.com>

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},
             ],
         )
 

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",

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(