Always set page text positions and prevent data clobber

documentcloud/documents/models/document.py

@@ -486,42 +486,34 @@ def set_page_text(self, page_text_infos):
 
     def _set_page_positions(self, pages, file_names, file_contents):
         """Handle grafting page positions back into the document"""
-
         current_pdf = pymupdf.open(stream=storage.open(self.doc_path, "rb").read())
         start_page = pages[0]["page_number"]
         stop_page = pages[-1]["page_number"]
 
+        # always write the position JSON - this is cheap and does not cause the
+        # memory issues that gating below guards against
+        self._write_position_json(pages, file_names, file_contents)
+
         visible_text = self._check_visible_text(current_pdf, start_page, stop_page)
         logger.info(
             "[SET PAGE TEXT] %d - visible text detected: %s", self.pk, visible_text
         )
         if visible_text:
             # merging when we need to flatten visible text causes excessive memory usage
+            current_pdf.close()
             return None
-
         grafted_pdf, base_pdf_stream = self._init_graft_pdf(
             current_pdf,
             start_page,
             stop_page,
             visible_text,
         )
-
         for page in pages:
             page_number = page["page_number"]
             if page.get("positions"):
-                logger.info(
-                    "[SET PAGE TEXT] %d - positions page %d", self.pk, page_number
-                )
-                file_names.append(
-                    path.page_text_position_path(self.pk, self.slug, page_number)
-                )
-                positions = [{**p.pop("metadata", {}), **p} for p in page["positions"]]
-                file_contents.append(json.dumps(positions).encode("utf-8"))
-
                 logger.info("[SET PAGE TEXT] %d - graft page %d", self.pk, page_number)
                 # create the overlay file
                 graft_page(page["positions"], grafted_pdf[page_number - start_page])
-
         # merge the overlay pages back onto the original document
         if visible_text:
             contents = self._merge_overlay_visible(
@@ -539,9 +531,26 @@ def _set_page_positions(self, pages, file_names, file_contents):
             )
         current_pdf.close()
         grafted_pdf.close()
-
         return contents
 
+    def _write_position_json(self, pages, file_names, file_contents):
+        """Stage the per-page position JSON files for upload.
+
+        Always safe to run - writing the position files is cheap and does not
+        trigger the memory issues associated with grafting into the PDF.
+        """
+        for page in pages:
+            page_number = page["page_number"]
+            if page.get("positions"):
+                logger.info(
+                    "[SET PAGE TEXT] %d - positions page %d", self.pk, page_number
+                )
+                file_names.append(
+                    path.page_text_position_path(self.pk, self.slug, page_number)
+                )
+                positions = [{**p.pop("metadata", {}), **p} for p in page["positions"]]
+                file_contents.append(json.dumps(positions).encode("utf-8"))
+
     def _merge_overlay(self, base_pdf_stream, grafted_pdf, start_page, stop_page):
         """Merge the text only overlay pages back in to the base PDF"""
         base_pdf = Pdf.open(base_pdf_stream)

documentcloud/documents/tasks.py

@@ -282,7 +282,7 @@ def set_page_text(document_pk, page_text_infos):
     logger.info("[SET PAGE TEXT] %d - setting status to readable", document_pk)
     with transaction.atomic():
         document.status = Status.readable
-        document.save()
+        document.save(update_fields=["status", "solr_dirty"])
         document.index_on_commit(field_updates={"status": "set"})
     kwargs = {"field_updates": {}}
     try:
@@ -298,7 +298,7 @@ def set_page_text(document_pk, page_text_infos):
         with transaction.atomic():
             logger.info("[SET PAGE TEXT] %d - setting status to success", document_pk)
             document.status = Status.success
-            document.save()
+            document.save(update_fields=["status", "solr_dirty"])
             kwargs["field_updates"]["status"] = "set"
             document.index_on_commit(**kwargs)
 

documentcloud/documents/tests/test_tasks.py

@@ -0,0 +1,29 @@
+# Standard Library
+from unittest.mock import patch
+
+# Third Party
+import pytest
+
+# DocumentCloud
+from documentcloud.documents.choices import Status
+from documentcloud.documents.models import Document
+from documentcloud.documents.tasks import set_page_text
+from documentcloud.documents.tests.factories import DocumentFactory
+
+
+@pytest.mark.django_db
+def test_set_page_text_does_not_clobber_data():
+    doc = DocumentFactory(status=Status.pending, slug="test-doc")
+
+    def fake_set_page_text(self, *args, **kwargs):
+        # simulate the Add-On writing a key/value pair into data via a
+        # separate DB row update while this task holds a stale instance
+        Document.objects.filter(pk=self.pk).update(data={"_tag": "applied"})
+        return {"pages": [], "updated": 123}
+
+    with patch.object(Document, "set_page_text", fake_set_page_text):
+        set_page_text(doc.pk, [{"page_number": 0, "text": "hi"}])
+
+    doc.refresh_from_db()
+    assert doc.data == {"_tag": "applied"}  # survives with the fix
+    assert doc.status == Status.success