Resolve reprocess with stumyplogs deduplicated query

Co-Authored-By: Trenton H <797416+stumpylog@users.noreply.github.com>

src/documents/bulk_edit.py

@@ -13,8 +13,14 @@ from celery import group
 from celery import shared_task
 from django.conf import settings
 from django.db import transaction
+from django.db.models import Case
+from django.db.models import F
 from django.db.models import Max
+from django.db.models import OuterRef
 from django.db.models import Q
+from django.db.models import Subquery
+from django.db.models import When
+from django.db.models.functions import Coalesce
 from django.utils import timezone
 
 from documents.data_models import ConsumableDocument
@@ -36,6 +42,7 @@ from documents.tasks import remove_document_from_index
 from documents.tasks import update_document_content_maybe_archive_file
 from documents.versioning import get_latest_version_for_root
 from documents.versioning import get_root_document
+from documents.versioning import versions_newest_first
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -408,10 +415,28 @@ def reprocess(doc_ids: list[int], *, remote_ocr: bool = False) -> Literal["OK"]:
 
     Consumption workflows do not run here, so ``remote_ocr`` is how the user
     asks for the remote engine when it is not configured to handle everything.
+
+    A root document with versions reprocesses its latest version, which is the
+    file whose content, archive and thumbnail are shown for it.
     """
-    for document_id in doc_ids:
+    latest_version = versions_newest_first(
+        Document.objects.filter(root_document=OuterRef("pk")),
+    ).values("id")[:1]
+    source_ids = (
+        Document.objects.filter(id__in=doc_ids)
+        .annotate(
+            source_id=Case(
+                When(root_document__isnull=False, then=F("id")),
+                default=Coalesce(Subquery(latest_version), F("id")),
+            ),
+        )
+        .order_by()
+        .values_list("source_id", flat=True)
+        .distinct()
+    )
+    for source_id in source_ids:
         update_document_content_maybe_archive_file.apply_async(
-            kwargs={"document_id": document_id, "remote_ocr": remote_ocr},
+            kwargs={"document_id": source_id, "remote_ocr": remote_ocr},
             headers={"trigger_source": PaperlessTask.TriggerSource.MANUAL},
         )
 

src/documents/tasks.py

@@ -490,18 +490,23 @@ def update_document_content_maybe_archive_file(
                     shutil.move(thumbnail, document.thumbnail_path)
 
             document.refresh_from_db()
+            root_document = (
+                document.root_document if document.root_document_id else document
+            )
             logger.info(
-                f"Updating index for document {document_id} ({document.archive_checksum})",
+                f"Updating index for document {root_document.pk} ({document.archive_checksum})",
             )
             from documents.search import get_backend
 
-            get_backend().add_or_update(document)
+            get_backend().add_or_update(root_document)
 
             ai_config = AIConfig()
             if ai_config.llm_index_enabled:
-                llm_index_add_or_update_document(document)
+                llm_index_add_or_update_document(root_document)
 
             clear_document_caches(document.pk)
+            if root_document.pk != document.pk:
+                clear_document_caches(root_document.pk)
 
         except Exception:
             logger.exception(

src/documents/tests/test_bulk_edit.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from unittest import mock
 
 import pikepdf
+import pytest
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
@@ -12,6 +13,7 @@ from django.test import TestCase
 from django.test.utils import CaptureQueriesContext
 from guardian.shortcuts import get_groups_with_perms
 from guardian.shortcuts import get_users_with_perms
+from pytest_mock import MockerFixture
 
 from documents import bulk_edit
 from documents.models import Correspondent
@@ -23,6 +25,7 @@ from documents.models import StoragePath
 from documents.models import Tag
 from documents.permissions import set_permissions_for_objects
 from paperless_testing.dirs import DirectoriesMixin
+from paperless_testing.factories import DocumentFactory
 from paperless_testing.permissions import grant_object
 
 
@@ -1970,18 +1973,22 @@ class TestPDFActions(DirectoriesMixin, TestCase):
         self.assertIn("Error removing password from document", cm.output[0])
 
 
-class TestBulkEditReprocess(DirectoriesMixin, TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-
-        self.doc = Document.objects.create(
-            title="test",
-            checksum="A",
-            mime_type="application/pdf",
+@pytest.mark.django_db
+class TestBulkEditReprocess:
+    @pytest.fixture
+    def mock_task(self, mocker: MockerFixture) -> mock.MagicMock:
+        return mocker.patch(
+            "documents.bulk_edit.update_document_content_maybe_archive_file",
         )
 
-    @mock.patch("documents.bulk_edit.update_document_content_maybe_archive_file")
-    def test_reprocess_defaults_to_local(self, mock_task: mock.Mock) -> None:
+    @staticmethod
+    def _queued_ids(mock_task: mock.MagicMock) -> list[int]:
+        return [
+            call.kwargs["kwargs"]["document_id"]
+            for call in mock_task.apply_async.call_args_list
+        ]
+
+    def test_reprocess_defaults_to_local(self, mock_task: mock.MagicMock) -> None:
         """
         GIVEN:
             - A reprocess request that says nothing about remote OCR
@@ -1990,18 +1997,17 @@ class TestBulkEditReprocess(DirectoriesMixin, TestCase):
         THEN:
             - The task is queued without asking for the remote engine
         """
-        result = bulk_edit.reprocess([self.doc.id])
+        doc = DocumentFactory()
+
+        assert bulk_edit.reprocess([doc.id]) == "OK"
 
-        self.assertEqual(result, "OK")
         mock_task.apply_async.assert_called_once()
-        _, kwargs = mock_task.apply_async.call_args
-        self.assertEqual(
-            kwargs["kwargs"],
-            {"document_id": self.doc.id, "remote_ocr": False},
-        )
+        assert mock_task.apply_async.call_args.kwargs["kwargs"] == {
+            "document_id": doc.id,
+            "remote_ocr": False,
+        }
 
-    @mock.patch("documents.bulk_edit.update_document_content_maybe_archive_file")
-    def test_reprocess_passes_remote_ocr(self, mock_task: mock.Mock) -> None:
+    def test_reprocess_passes_remote_ocr(self, mock_task: mock.MagicMock) -> None:
         """
         GIVEN:
             - A reprocess request that explicitly asks for remote OCR
@@ -2010,14 +2016,72 @@ class TestBulkEditReprocess(DirectoriesMixin, TestCase):
         THEN:
             - The request is forwarded to the task for every document
         """
-        other = Document.objects.create(
-            title="test2",
-            checksum="B",
-            mime_type="application/pdf",
+        docs = DocumentFactory.create_batch(2)
+
+        bulk_edit.reprocess([doc.id for doc in docs], remote_ocr=True)
+
+        assert mock_task.apply_async.call_count == 2
+        for call in mock_task.apply_async.call_args_list:
+            assert call.kwargs["kwargs"]["remote_ocr"]
+
+    def test_reprocess_root_uses_latest_version(
+        self,
+        mock_task: mock.MagicMock,
+    ) -> None:
+        """
+        GIVEN:
+            - A root document with two versions
+        WHEN:
+            - reprocess is called with the root document
+        THEN:
+            - The latest version is reprocessed, not the root's original file
+        """
+        root = DocumentFactory()
+        DocumentFactory(root_document=root, version_index=1)
+        latest = DocumentFactory(root_document=root, version_index=2)
+
+        bulk_edit.reprocess([root.id])
+
+        assert self._queued_ids(mock_task) == [latest.id]
+
+    def test_reprocess_explicit_version(self, mock_task: mock.MagicMock) -> None:
+        """
+        GIVEN:
+            - A root document with two versions
+        WHEN:
+            - reprocess is called with the older version
+        THEN:
+            - That version is reprocessed
+        """
+        root = DocumentFactory()
+        older = DocumentFactory(root_document=root, version_index=1)
+        DocumentFactory(root_document=root, version_index=2)
+
+        bulk_edit.reprocess([older.id])
+
+        assert self._queued_ids(mock_task) == [older.id]
+
+    def test_reprocess_root_and_latest_version_dispatches_once(
+        self,
+        mock_task: mock.MagicMock,
+    ) -> None:
+        """
+        GIVEN:
+            - A root document with two versions, the latest created on a
+              different date than the root
+        WHEN:
+            - reprocess is called with both the root and its latest version
+        THEN:
+            - The latest version is reprocessed only once
+        """
+        root = DocumentFactory(created=date(2024, 1, 1))
+        DocumentFactory(root_document=root, version_index=1)
+        latest = DocumentFactory(
+            root_document=root,
+            version_index=2,
+            created=date(2025, 1, 1),
         )
 
-        bulk_edit.reprocess([self.doc.id, other.id], remote_ocr=True)
+        bulk_edit.reprocess([root.id, latest.id])
 
-        self.assertEqual(mock_task.apply_async.call_count, 2)
-        for call in mock_task.apply_async.call_args_list:
-            self.assertTrue(call.kwargs["kwargs"]["remote_ocr"])
+        assert self._queued_ids(mock_task) == [latest.id]

src/documents/tests/test_tasks.py

@@ -20,6 +20,7 @@ from documents.sanity_checker import SanityCheckMessages
 from documents.tests.helpers import dummy_preprocess
 from paperless_testing.assertions import FileSystemAssertsMixin
 from paperless_testing.dirs import DirectoriesMixin
+from paperless_testing.factories import DocumentFactory
 
 
 @pytest.mark.django_db
@@ -287,6 +288,82 @@ class TestUpdateContent(DirectoriesMixin, TestCase):
         tasks.update_document_content_maybe_archive_file(doc.pk)
         self.assertNotEqual(Document.objects.get(pk=doc.pk).content, "test")
 
+    def _create_root_with_version(self) -> tuple[Document, Document]:
+        sample1 = self.dirs.scratch_dir / "sample.pdf"
+        shutil.copy(
+            Path(__file__).parent
+            / "samples"
+            / "documents"
+            / "originals"
+            / "0000001.pdf",
+            sample1,
+        )
+        root = DocumentFactory(content="root content", mime_type="application/pdf")
+        version = DocumentFactory(
+            content="my document",
+            filename=sample1,
+            mime_type="application/pdf",
+            root_document=root,
+            version_index=1,
+        )
+        return root, version
+
+    @mock.patch("documents.tasks.clear_document_caches")
+    @mock.patch("documents.search.get_backend")
+    def test_update_content_version_indexes_root(
+        self,
+        mock_get_backend: mock.Mock,
+        mock_clear_caches: mock.Mock,
+    ) -> None:
+        """
+        GIVEN:
+            - A root document with a version
+        WHEN:
+            - Update content task is called for the version
+        THEN:
+            - The version's content is updated
+            - The root document is indexed rather than the version
+            - Caches are cleared for both
+        """
+        root, version = self._create_root_with_version()
+
+        tasks.update_document_content_maybe_archive_file(version.pk)
+
+        self.assertNotEqual(
+            Document.objects.get(pk=version.pk).content,
+            "my document",
+        )
+        self.assertEqual(Document.objects.get(pk=root.pk).content, "root content")
+        indexed = mock_get_backend.return_value.add_or_update.call_args.args[0]
+        self.assertEqual(indexed.pk, root.pk)
+        mock_clear_caches.assert_has_calls(
+            [mock.call(version.pk), mock.call(root.pk)],
+        )
+
+    @override_settings(AI_ENABLED=True, LLM_EMBEDDING_BACKEND="huggingface")
+    @mock.patch("documents.tasks.llm_index_add_or_update_document")
+    @mock.patch("documents.search.get_backend")
+    def test_update_content_version_updates_llm_index_for_root(
+        self,
+        mock_get_backend: mock.Mock,
+        mock_llm_index: mock.Mock,
+    ) -> None:
+        """
+        GIVEN:
+            - A root document with a version
+            - The LLM index is enabled
+        WHEN:
+            - Update content task is called for the version
+        THEN:
+            - The LLM index is updated for the root document, not the version
+        """
+        root, version = self._create_root_with_version()
+
+        tasks.update_document_content_maybe_archive_file(version.pk)
+
+        mock_llm_index.assert_called_once()
+        self.assertEqual(mock_llm_index.call_args.args[0].pk, root.pk)
+
 
 class TestUpdateContentRemoteOCR(DirectoriesMixin, TestCase):
     """