diff --git a/documentcloud/documents/models/document.py b/documentcloud/documents/models/document.py index 5e4b7a41..4b427e05 100644 --- a/documentcloud/documents/models/document.py +++ b/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) diff --git a/documentcloud/documents/tasks.py b/documentcloud/documents/tasks.py index d3ac75de..4b52f51d 100644 --- a/documentcloud/documents/tasks.py +++ b/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) diff --git a/documentcloud/documents/tests/test_tasks.py b/documentcloud/documents/tests/test_tasks.py new file mode 100644 index 00000000..8f6ec770 --- /dev/null +++ b/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