Prepare and atomically apply OCR template field updates

Co-authored-by: OpenAI Codex <codex@openai.com>

src/documents/tests/test_zone_ocr.py

@@ -1,6 +1,7 @@
 """Tests for the zone-based OCR extraction engine."""
 
 import tempfile
+from datetime import date
 from pathlib import Path
 from unittest.mock import MagicMock
 from unittest.mock import patch
@@ -18,6 +19,7 @@ from documents.models import MatchingModel
 from documents.models import OcrTemplate
 from documents.models import OcrTemplateZone
 from documents.zone_ocr import _apply_transform
+from documents.zone_ocr import _combine_field_value
 from documents.zone_ocr import _convert_value
 from documents.zone_ocr import _crop_zone
 from documents.zone_ocr import _detect_mime
@@ -161,25 +163,150 @@ class TestZoneExtraction(SimpleTestCase):
 
 
 class TestSavedZoneExtraction(DirectoriesMixin, TestCase):
-    @patch("documents.zone_ocr.subprocess.run")
-    def test_saved_extraction_matches_preview_and_rejects_invalid_value(self, run):
-        page = Path(self.dirs.scratch_dir) / "page.png"
+    def setUp(self):
+        super().setUp()
+        self.page = self.dirs.scratch_dir / "page.png"
         with Image.new("RGB", (200, 100), "white") as image:
-            image.save(page)
-        document = DocumentFactory(title="Original")
-        template = OcrTemplate.objects.create(
+            image.save(self.page)
+        self.document = DocumentFactory(title="Original")
+        self.template = OcrTemplate.objects.create(
             name="Account",
             source_width=100,
             source_height=50,
         )
-        zone = OcrTemplateZone.objects.create(
-            template=template,
-            name="Title",
-            target="title",
+
+    def zone(self, name, target="title", **kwargs):
+        return OcrTemplateZone.objects.create(
+            template=self.template,
+            name=name,
+            target=target,
             x=0,
             y=0,
             width=20,
             height=10,
+            **kwargs,
+        )
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_combines_zones_and_writes_each_field_once(self, run):
+        zones = [self.zone("First"), self.zone("Second")]
+        self.template.combine_formats = {"title": "{Second} / {First}"}
+        run.side_effect = [MagicMock(stdout="one"), MagicMock(stdout="two")]
+        with patch.object(self.document, "save", wraps=self.document.save) as save:
+            _process_template(self.document, self.page, self.template, zones)
+        self.document.refresh_from_db()
+        self.assertEqual(self.document.title, "two / one")
+        save.assert_called_once_with(update_fields=["title"])
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_converts_custom_and_builtin_fields_before_saving(self, run):
+        field = CustomField.objects.create(name="Paid", data_type="boolean")
+        zones = [
+            self.zone("Title"),
+            self.zone("ASN", "asn"),
+            self.zone("Created", "created"),
+            self.zone("Paid", "custom_field", custom_field=field),
+        ]
+        run.side_effect = [
+            MagicMock(stdout=value)
+            for value in [
+                "New title",
+                "ASN 42",
+                "2026-04-13",
+                "no",
+            ]
+        ]
+        _process_template(self.document, self.page, self.template, zones)
+        self.document.refresh_from_db()
+        self.assertEqual(self.document.title, "New title")
+        self.assertEqual(self.document.archive_serial_number, 42)
+        self.assertEqual(self.document.created, date(2026, 4, 13))
+        self.assertFalse(self.document.custom_fields.get(field=field).value)
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_invalid_field_values_preserve_existing_values(self, run):
+        field = CustomField.objects.create(name="Number", data_type="integer")
+        CustomFieldInstance.objects.create(
+            document=self.document,
+            field=field,
+            value_int=12,
+        )
+        zones = [
+            self.zone("Title"),
+            self.zone("ASN", "asn"),
+            self.zone("Created", "created"),
+            self.zone("Number", "custom_field", custom_field=field),
+            self.zone("Unassigned", "custom_field"),
+        ]
+        original_created = self.document.created
+        run.side_effect = [
+            MagicMock(stdout=value)
+            for value in [
+                "Valid",
+                "9999999999999999999999999999999",
+                "invalid",
+                "abc",
+                "text",
+            ]
+        ]
+        _process_template(self.document, self.page, self.template, zones)
+        self.document.refresh_from_db()
+        self.assertEqual(self.document.title, "Valid")
+        self.assertIsNone(self.document.archive_serial_number)
+        self.assertEqual(self.document.created, original_created)
+        self.assertEqual(self.document.custom_fields.get(field=field).value, 12)
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_write_failure_rolls_back_template_and_restores_document(self, run):
+        field = CustomField.objects.create(name="Number", data_type="integer")
+        zones = [
+            self.zone("Title"),
+            self.zone("Number", "custom_field", custom_field=field),
+        ]
+        run.side_effect = [MagicMock(stdout="New title"), MagicMock(stdout="42")]
+        save = self.document.save
+
+        def fail_after_save(**kwargs):
+            save(**kwargs)
+            raise RuntimeError("save failed")
+
+        with patch.object(self.document, "save", side_effect=fail_after_save):
+            with self.assertRaisesMessage(RuntimeError, "save failed"):
+                _process_template(self.document, self.page, self.template, zones)
+        self.assertEqual(self.document.title, "Original")
+        self.document.refresh_from_db()
+        self.assertEqual(self.document.title, "Original")
+        self.assertFalse(self.document.custom_fields.exists())
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_last_successful_template_wins(self, run):
+        self.zone("First")
+        second = OcrTemplate.objects.create(
+            name="Second",
+            source_width=100,
+            source_height=50,
+        )
+        self.template = second
+        self.zone("Second")
+        third = OcrTemplate.objects.create(
+            name="Third",
+            source_width=100,
+            source_height=50,
+        )
+        self.template = third
+        self.zone("Third", validation_regex="accepted")
+        templates = list(OcrTemplate.objects.order_by("pk"))
+        run.side_effect = [
+            MagicMock(stdout=value) for value in ["First", "Second", "rejected"]
+        ]
+        run_zone_extraction(self.document, self.page, templates=templates)
+        self.document.refresh_from_db()
+        self.assertEqual(self.document.title, "Second")
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_saved_extraction_matches_preview_and_rejects_invalid_value(self, run):
+        zone = self.zone(
+            "Title",
             transform="uppercase",
             validation_regex="ACCOUNT [A-Z]+",
         )
@@ -189,12 +316,30 @@ class TestSavedZoneExtraction(DirectoriesMixin, TestCase):
         ]:
             with self.subTest(raw_text=raw_text):
                 run.return_value.stdout = raw_text
-                preview = extract_zone_preview(page, zone, 100, 50, 1)
+                preview = extract_zone_preview(self.page, zone, 100, 50, 1)
                 self.assertEqual(preview["raw_text"], raw_text)
                 self.assertEqual(preview["value"], raw_text.upper())
-                _process_template(document, page, template, [zone])
-                document.refresh_from_db()
-                self.assertEqual(document.title, expected_title)
+                _process_template(self.document, self.page, self.template, [zone])
+                self.document.refresh_from_db()
+                self.assertEqual(self.document.title, expected_title)
+
+
+class TestCombineFieldValue(SimpleTestCase):
+    def test_default_format_preserves_zone_order_and_duplicate_names(self):
+        zones = [
+            (OcrTemplateZone(name="Same"), value) for value in ["one", None, "two"]
+        ]
+        self.assertEqual(_combine_field_value("", zones), "one two")
+
+    def test_format_preserves_existing_missing_token_cleanup(self):
+        zones = [
+            (OcrTemplateZone(name="First"), "one"),
+            (OcrTemplateZone(name="Second"), None),
+        ]
+        self.assertEqual(
+            _combine_field_value("{ First } / {Second} / {Unknown}", zones),
+            "one",
+        )
 
 
 class TestTemplateVisibility(DirectoriesMixin, TestCase):

src/documents/zone_ocr.py

@@ -20,6 +20,8 @@ from datetime import datetime
 from pathlib import Path
 
 from django.conf import settings
+from django.core.exceptions import ValidationError
+from django.db import transaction
 from PIL import Image
 
 from documents.matching import matches
@@ -41,6 +43,13 @@ class ZoneExtraction:
     value: str | None = None
 
 
+@dataclass(frozen=True)
+class FieldUpdate:
+    field_name: str
+    value: object
+    custom_field: CustomField | None = None
+
+
 def get_templates_for_document(document: Document) -> list[OcrTemplate]:
     """
     The enabled templates that match the document's content and that its owner
@@ -140,14 +149,7 @@ def _process_template(
     template: OcrTemplate,
     zones: list[OcrTemplateZone],
 ) -> None:
-    """Process all zones in a template against a document.
-
-    Each zone is OCR'd independently, then zones are grouped by their target
-    field and each field is written exactly once. When several zones share a
-    field, their values are combined via the template's per-field format string
-    (or joined in order if none is set) — this avoids the zones overwriting each
-    other's value.
-    """
+    """Extract and combine zone values before applying a template's updates."""
     pages_needed: set[int] = {
         _resolve_page_idx(zone.page, document.page_count) for zone in zones
     }
@@ -162,30 +164,19 @@ def _process_template(
             document.page_count,
         )
 
-        # Pass 1: OCR every zone into a value (or None if it failed/was rejected).
-        zone_values: dict[int, str | None] = {}
+        grouped: dict[str, list[tuple[OcrTemplateZone, str | None]]] = {}
         for zone in zones:
             page_idx = _resolve_page_idx(zone.page, document.page_count)
-
-            if page_idx not in page_images:
-                logger.warning(
-                    "Zone OCR: page %d not available for zone '%s'",
-                    page_idx,
-                    zone.name,
+            extracted = None
+            if page_idx in page_images:
+                result = _extract_zone(
+                    page_images[page_idx],
+                    zone,
+                    zone.zone_source_width or template.source_width,
+                    zone.zone_source_height or template.source_height,
+                    tmp_path,
                 )
-                continue
-
-            src_w = zone.zone_source_width or template.source_width
-            src_h = zone.zone_source_height or template.source_height
-
-            result = _extract_zone(
-                page_images[page_idx],
-                zone,
-                src_w,
-                src_h,
-                tmp_path,
-            )
-            extracted = result.value
+                extracted = result.value
 
             if (
                 extracted is not None
@@ -200,38 +191,28 @@ def _process_template(
                 )
                 extracted = None
 
-            zone_values[id(zone)] = extracted
-
-    # Pass 2: group zones by target field and write each field once.
-    grouped: dict[str, list[OcrTemplateZone]] = {}
-    for zone in zones:
-        grouped.setdefault(_field_key(zone), []).append(zone)
+            grouped.setdefault(_field_key(zone), []).append((zone, extracted))
 
+    updates = []
     combine_formats = template.combine_formats or {}
     for key, field_zones in grouped.items():
         value = _combine_field_value(
             combine_formats.get(key, ""),
             field_zones,
-            zone_values,
         )
         if not value:
             continue
-
-        target_zone = field_zones[0]
-        _write_zone_value(document, target_zone, value)
-        logger.info(
-            "Zone OCR: %s = %r (from %d zone(s))",
-            _zone_target_label(target_zone),
-            value[:100] if len(value) > 100 else value,
-            len(field_zones),
-        )
+        update = _prepare_field_update(field_zones[0][0], value)
+        if update is not None:
+            updates.append(update)
+    _apply_field_updates(document, updates)
 
 
 def _field_key(zone: OcrTemplateZone) -> str:
     """Identify a zone's target field. Custom fields key by id, built-in targets
     by their name. Matches the key used in OcrTemplate.combine_formats and on the
     frontend field select."""
-    target = getattr(zone, "target", None) or "custom_field"
+    target = zone.target
     if target == "custom_field" and zone.custom_field_id:
         return str(zone.custom_field_id)
     return target
@@ -239,8 +220,7 @@ def _field_key(zone: OcrTemplateZone) -> str:
 
 def _combine_field_value(
     fmt: str,
-    field_zones: list[OcrTemplateZone],
-    zone_values: dict[int, str | None],
+    field_zones: list[tuple[OcrTemplateZone, str | None]],
 ) -> str:
     """Combine the OCR values of all zones targeting one field.
 
@@ -248,11 +228,10 @@ def _combine_field_value(
     and literal text is kept; separators left dangling by an empty token are
     cleaned up. Without a format, the zone values are joined in order by a space.
     """
-    values = {z.name: (zone_values.get(id(z)) or "") for z in field_zones}
+    values = {zone.name: value or "" for zone, value in field_zones}
 
     if not fmt:
-        parts = [zone_values.get(id(z)) or "" for z in field_zones]
-        return " ".join(p for p in parts if p).strip()
+        return " ".join(value for _, value in field_zones if value).strip()
 
     def _replace(match: re.Match) -> str:
         return values.get(match.group(1).strip(), "")
@@ -531,14 +510,6 @@ def _apply_transform(text: str, transform: str, date_format: str = "") -> str:
     return text
 
 
-def _zone_target_label(zone: OcrTemplateZone) -> str:
-    """Human label of a zone's write target (for logging)."""
-    target = getattr(zone, "target", None) or "custom_field"
-    if target == "custom_field":
-        return zone.custom_field.name if zone.custom_field_id else "(no field)"
-    return {"title": "Title", "asn": "ASN", "created": "Created"}.get(target, target)
-
-
 def _parse_created_datetime(value: str):
     """Parse an extracted value into a tz-aware datetime for document.created.
 
@@ -568,70 +539,68 @@ def _parse_created_datetime(value: str):
     return None
 
 
-def _write_zone_value(
-    document: Document,
+def _prepare_field_update(
     zone: OcrTemplateZone,
     value: str,
-) -> None:
-    """Write an extracted value to the zone's target — a custom field, or a
-    built-in document field (title / archive_serial_number / created)."""
-    target = getattr(zone, "target", None) or "custom_field"
-
-    if target == "custom_field":
-        if zone.custom_field_id:
-            _write_custom_field(document, zone.custom_field, value)
-        else:
-            logger.debug("Zone OCR: zone '%s' has no custom field set", zone.name)
-        return
-
-    if target == "title":
-        document.title = value[:128]
-        document.save(update_fields=["title"])
-    elif target == "asn":
+) -> FieldUpdate | None:
+    """Convert extracted text and enforce the destination field's constraints."""
+    custom_field = None
+    if zone.target == OcrTemplateZone.TargetType.CUSTOM_FIELD:
+        if not zone.custom_field_id:
+            return None
+        custom_field = zone.custom_field
+        field_name = CustomFieldInstance.get_value_field_name(custom_field.data_type)
+        typed_value = _convert_value(value, custom_field.data_type)
+        model = CustomFieldInstance
+    elif zone.target == OcrTemplateZone.TargetType.TITLE:
+        field_name, typed_value = "title", value[:128]
+        model = Document
+    elif zone.target == OcrTemplateZone.TargetType.ASN:
         digits = re.sub(r"[^\d]", "", value)
-        if not digits:
-            logger.debug(
-                "Zone OCR: ASN zone '%s' produced no digits (%r)",
-                zone.name,
-                value[:50],
-            )
-            return
-        document.archive_serial_number = int(digits)
-        document.save(update_fields=["archive_serial_number"])
-    elif target == "created":
+        field_name = "archive_serial_number"
+        typed_value = int(digits) if digits else None
+        model = Document
+    elif zone.target == OcrTemplateZone.TargetType.CREATED:
         parsed = _parse_created_datetime(value)
-        if parsed is None:
-            logger.debug(
-                "Zone OCR: created zone '%s' could not parse a date (%r)",
-                zone.name,
-                value[:50],
-            )
-            return
-        document.created = parsed
-        document.save(update_fields=["created"])
-
-
-def _write_custom_field(
-    document: Document,
-    custom_field: CustomField,
-    value: str,
-) -> None:
-    """Write an extracted value to a document's custom field."""
-    typed_value = _convert_value(value, custom_field.data_type)
+        field_name, typed_value = "created", parsed.date() if parsed else None
+        model = Document
+    else:
+        return None
     if typed_value is None:
-        logger.debug(
-            "Zone OCR: skipping custom field '%s' — value conversion returned None",
-            custom_field.name,
-        )
-        return
+        return None
+    try:
+        typed_value = model._meta.get_field(field_name).clean(typed_value, None)
+    except ValidationError:
+        logger.info("Zone OCR: invalid value for zone '%s'", zone.name)
+        return None
+    return FieldUpdate(field_name, typed_value, custom_field)
 
-    value_field_name = CustomFieldInstance.get_value_field_name(custom_field.data_type)
 
-    CustomFieldInstance.objects.update_or_create(
-        document=document,
-        field=custom_field,
-        defaults={value_field_name: typed_value},
-    )
+def _apply_field_updates(document: Document, updates: list[FieldUpdate]) -> None:
+    builtin_values = {
+        update.field_name: update.value
+        for update in updates
+        if update.custom_field is None
+    }
+    original_values = {name: getattr(document, name) for name in builtin_values}
+    try:
+        with transaction.atomic():
+            for update in updates:
+                if update.custom_field is not None:
+                    CustomFieldInstance.objects.update_or_create(
+                        document=document,
+                        field=update.custom_field,
+                        defaults={update.field_name: update.value},
+                    )
+            if builtin_values:
+                for name, value in builtin_values.items():
+                    setattr(document, name, value)
+                document.save(update_fields=list(builtin_values))
+    except Exception:
+        # Database rollback does not restore attributes on the model instance.
+        for name, value in original_values.items():
+            setattr(document, name, value)
+        raise
 
 
 def _convert_value(value: str, data_type: str) -> object | None: