diff --git a/src/documents/apps.py b/src/documents/apps.py index 28761a0df2..4a3aa9be9e 100644 --- a/src/documents/apps.py +++ b/src/documents/apps.py @@ -28,9 +28,9 @@ class DocumentsConfig(AppConfig): document_consumption_finished.connect(set_tags) document_consumption_finished.connect(set_storage_path) document_consumption_finished.connect(run_workflows_added) + document_consumption_finished.connect(run_zone_ocr_extraction) document_consumption_finished.connect(add_to_index) document_consumption_finished.connect(add_or_update_document_in_llm_index) - document_consumption_finished.connect(run_zone_ocr_extraction) document_updated.connect(run_workflows_updated) document_updated.connect(send_websocket_document_updated) diff --git a/src/documents/tasks.py b/src/documents/tasks.py index b4df2fff31..c7d1f4910d 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -376,22 +376,22 @@ def extract_zones(document_id: int, template_id: int | None = None) -> None: matches it when no template is given. """ document = Document.objects.get(id=document_id) - if template_id is None: - run_zone_extraction(document, None) - return + templates = None + if template_id is not None: + try: + template = OcrTemplate.objects.prefetch_related( + "zones", + "zones__custom_field", + ).get(id=template_id) + except OcrTemplate.DoesNotExist: + logger.warning( + "OCR template %s no longer exists, not running zone OCR", + template_id, + ) + return + templates = [template] - try: - template = OcrTemplate.objects.prefetch_related( - "zones", - "zones__custom_field", - ).get(id=template_id) - except OcrTemplate.DoesNotExist: - logger.warning( - "OCR template %s no longer exists, not running zone OCR", - template_id, - ) - return - run_zone_extraction(document, None, templates=[template]) + run_zone_extraction(document, templates=templates, reindex=True) @shared_task diff --git a/src/documents/tests/test_api_ocr_templates.py b/src/documents/tests/test_api_ocr_templates.py index b58758c5a2..926ec0499c 100644 --- a/src/documents/tests/test_api_ocr_templates.py +++ b/src/documents/tests/test_api_ocr_templates.py @@ -5,6 +5,7 @@ import subprocess from pathlib import Path from unittest.mock import patch +import pytest from django.contrib.auth.models import Permission from django.contrib.auth.models import User from django.db import IntegrityError @@ -17,6 +18,9 @@ from documents.models import CustomFieldInstance from documents.models import MatchingModel from documents.models import OcrTemplate from documents.models import OcrTemplateZone +from documents.search import get_backend +from documents.signals import document_consumption_finished +from documents.tasks import extract_zones from paperless_testing.dirs import DirectoriesMixin from paperless_testing.factories import DocumentFactory from paperless_testing.permissions import grant_global @@ -25,6 +29,75 @@ from paperless_testing.permissions import grant_object CONTENT = "Northwind Power & Light statement" +@pytest.mark.django_db +@pytest.mark.parametrize("entry_point", ["ingestion", "single", "bulk"]) +def test_zone_ocr_indexes_extracted_metadata( + entry_point, + admin_client, + paperless_dirs, + django_capture_on_commit_callbacks, + mocker, +): + document = DocumentFactory( + title="Before OCR", + content=CONTENT, + mime_type="image/png", + ) + with Image.new("RGB", (100, 100), "white") as image: + image.save(document.source_path, format="PNG") + template = OcrTemplate.objects.create( + name="Invoice template", + match="Northwind", + matching_algorithm=MatchingModel.MATCH_ANY, + source_width=100, + source_height=100, + ) + field = CustomField.objects.create(name="Invoice code", data_type="string") + OcrTemplateZone.objects.create( + template=template, + name="Title", + target="title", + x=0, + y=0, + width=50, + height=50, + ) + OcrTemplateZone.objects.create( + template=template, + name="Invoice code", + custom_field=field, + x=0, + y=50, + width=50, + height=50, + order=1, + ) + mocker.patch( + "documents.zone_ocr.subprocess.run", + side_effect=[mocker.Mock(stdout="After OCR"), mocker.Mock(stdout="INV42")], + ) + + backend = get_backend() + if entry_point != "ingestion": + backend.add_or_update(document) + assert backend.search_ids("title:before", user=None) == [document.pk] + + with django_capture_on_commit_callbacks(execute=True): + if entry_point == "ingestion": + document_consumption_finished.send(sender=None, document=document) + elif entry_point == "single": + response = admin_client.post(f"/api/documents/{document.pk}/run-zone-ocr/") + assert response.status_code == status.HTTP_200_OK + else: + extract_zones(document.pk, template.pk) + + document.refresh_from_db() + assert document.title == "After OCR" + assert backend.search_ids("title:before", user=None) == [] + assert backend.search_ids("title:after", user=None) == [document.pk] + assert backend.search_ids("custom_fields.value:INV42", user=None) == [document.pk] + + class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase): ENDPOINT = "/api/ocr_templates/" @@ -787,7 +860,7 @@ class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.json()["status"], "completed") - extract.assert_called_once_with(document, None, templates=[template]) + extract.assert_called_once_with(document, templates=[template], reindex=True) @patch("documents.zone_ocr.run_zone_extraction") def test_run_zone_ocr_rejects_a_template_the_user_cannot_view(self, extract): diff --git a/src/documents/tests/test_bulk_edit.py b/src/documents/tests/test_bulk_edit.py index 4bf0c530f9..2230e54393 100644 --- a/src/documents/tests/test_bulk_edit.py +++ b/src/documents/tests/test_bulk_edit.py @@ -1941,7 +1941,7 @@ class TestBulkEditRunZoneOcr: extract_zones(doc.id) - extract.assert_called_once_with(doc, None) + extract.assert_called_once_with(doc, templates=None, reindex=True) def test_task_runs_the_chosen_template(self, mocker: MockerFixture) -> None: """ @@ -1962,7 +1962,7 @@ class TestBulkEditRunZoneOcr: extract_zones(doc.id, template.pk) - extract.assert_called_once_with(doc, None, templates=[template]) + extract.assert_called_once_with(doc, templates=[template], reindex=True) def test_task_skips_a_deleted_template(self, mocker: MockerFixture) -> None: """ diff --git a/src/documents/views.py b/src/documents/views.py index df1947978e..aab8b6e011 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -2306,6 +2306,7 @@ class DocumentViewSet( return HttpResponseForbidden("Insufficient permissions") from documents.zone_ocr import get_templates_for_document + from documents.zone_ocr import run_zone_extraction templates = ( [template] if template is not None else get_templates_for_document(document) @@ -2322,9 +2323,7 @@ class DocumentViewSet( status=status.HTTP_404_NOT_FOUND, ) - from documents.zone_ocr import run_zone_extraction - - results = run_zone_extraction(document, None, templates=templates) + results = run_zone_extraction(document, templates=templates, reindex=True) return Response( {"status": "completed", "results": [asdict(result) for result in results]}, ) diff --git a/src/documents/zone_ocr.py b/src/documents/zone_ocr.py index 8ca1e49c0d..d05561f0e9 100644 --- a/src/documents/zone_ocr.py +++ b/src/documents/zone_ocr.py @@ -11,6 +11,7 @@ from datetime import date from datetime import datetime from decimal import Decimal from decimal import InvalidOperation +from functools import partial from pathlib import Path import dateparser @@ -29,6 +30,7 @@ from documents.models import OcrTemplate from documents.models import OcrTemplateZone from documents.permissions import permitted_object_ids from documents.regex import safe_regex_fullmatch +from documents.signals.handlers import add_to_index from documents.versioning import get_root_document logger = logging.getLogger("paperless.zone_ocr") @@ -90,12 +92,16 @@ def get_templates_for_document(document: Document) -> list[OcrTemplate]: def run_zone_extraction( document: Document, - original_file: Path | None, + original_file: Path | None = None, templates: list[OcrTemplate] | None = None, + *, + reindex: bool = False, ) -> list[ZoneOcrResult]: """ Run zone-based OCR extraction for the given templates, or for every template that matches the document when none are given. + + Manual runs can request a search refresh after successful writes commit. """ if templates is None: templates = get_templates_for_document(document) @@ -142,9 +148,18 @@ def run_zone_extraction( _zone_result(template, zone, reason="template_failed") for zone in zones ) + if reindex and any(result.reason is None for result in results): + # queue a reindex of the document after the transaction commits + transaction.on_commit(partial(_reindex_document, document.pk)) return results +def _reindex_document(document_id: int) -> None: + document = Document.objects.filter(pk=document_id).first() + if document is not None: + add_to_index(sender=None, document=document) + + def _resolve_doc_path( document: Document, original_file: Path | None,