From d579caf1ea20a53b7f5b7acaf5f4660bccd8eef7 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Mon, 23 Jun 2025 09:19:28 -0700
Subject: [PATCH] Enhancement: better try to catch db errors before unlink

---
 src/documents/consumer.py            | 68 +++++++++++++++++-----------
 src/documents/tests/test_consumer.py | 19 ++++++++
 2 files changed, 61 insertions(+), 26 deletions(-)

diff --git a/src/documents/consumer.py b/src/documents/consumer.py
index 31db723d9..e25f438db 100644
--- a/src/documents/consumer.py
+++ b/src/documents/consumer.py
@@ -467,6 +467,8 @@ class ConsumerPlugin(
         )
         # now that everything is done, we can start to store the document
         # in the system. This will be a transaction and reasonably fast.
+        success = False
+        result = None
         try:
             with transaction.atomic():
                 # store the document.
@@ -531,7 +533,22 @@ class ConsumerPlugin(
                 # renaming logic to acquire the lock as well.
                 # This triggers things like file renaming
                 document.save()
+                success = True
 
+        except Exception as e:
+            # save the exception for later
+            try:
+                self._fail(
+                    str(e),
+                    f"The following error occurred while storing document "
+                    f"{self.filename} after parsing: {e}",
+                    exc_info=True,
+                    exception=e,
+                )
+            except Exception as fail_exc:
+                stored_exception = fail_exc
+        finally:
+            if success:
                 # Delete the file only if it was successfully consumed
                 self.log.debug(f"Deleting file {self.working_copy}")
                 self.input_doc.original_file.unlink()
@@ -549,34 +566,33 @@ class ConsumerPlugin(
                     self.log.debug(f"Deleting file {shadow_file}")
                     Path(shadow_file).unlink()
 
-        except Exception as e:
-            self._fail(
-                str(e),
-                f"The following error occurred while storing document "
-                f"{self.filename} after parsing: {e}",
-                exc_info=True,
-                exception=e,
-            )
-        finally:
+                self.run_post_consume_script(document)
+
+                self.log.info(f"Document {document} consumption finished")
+
+                self._send_progress(
+                    100,
+                    100,
+                    ProgressStatusOptions.SUCCESS,
+                    ConsumerStatusShortMessage.FINISHED,
+                    document.id,
+                )
+
+                # Return the most up to date fields
+                document.refresh_from_db()
+
+                result = f"Success. New document id {document.pk} created"
+            elif stored_exception:
+                raise stored_exception
+            else:
+                self._fail(
+                    ConsumerStatusShortMessage.FAILED,
+                    f"Error occurred while saving {self.filename}.",
+                )
+
             document_parser.cleanup()
             tempdir.cleanup()
-
-        self.run_post_consume_script(document)
-
-        self.log.info(f"Document {document} consumption finished")
-
-        self._send_progress(
-            100,
-            100,
-            ProgressStatusOptions.SUCCESS,
-            ConsumerStatusShortMessage.FINISHED,
-            document.id,
-        )
-
-        # Return the most up to date fields
-        document.refresh_from_db()
-
-        return f"Success. New document id {document.pk} created"
+            return result
 
     def _parse_title_placeholders(self, title: str) -> str:
         local_added = timezone.localtime(timezone.now())
diff --git a/src/documents/tests/test_consumer.py b/src/documents/tests/test_consumer.py
index f0fdc02c7..bed07fc8a 100644
--- a/src/documents/tests/test_consumer.py
+++ b/src/documents/tests/test_consumer.py
@@ -633,6 +633,25 @@ class TestConsumer(
         # Database empty
         self.assertEqual(Document.objects.all().count(), 0)
 
+    @mock.patch("documents.consumer.ConsumerPlugin._store")
+    @mock.patch("documents.consumer.document_consumption_finished.send")
+    @mock.patch("documents.consumer.generate_unique_filename")
+    def testSaveFailsStillCaught(self, m_filename, m_signal, m_store):
+        filename = self.get_test_file()
+        m_store.return_value = None
+        m_filename.side_effect = AttributeError("BOOM")
+
+        with self.get_consumer(filename) as consumer:
+            with self.assertRaisesMessage(
+                ConsumerError,
+                "sample.pdf: The following error occurred while storing document sample.pdf after parsing: BOOM",
+            ):
+                consumer.run()
+
+        self._assert_first_last_send_progress(last_status="FAILED")
+        self.assertIsFile(filename)
+        self.assertEqual(Document.objects.count(), 0)
+
     @override_settings(FILENAME_FORMAT="{correspondent}/{title}")
     def testFilenameHandling(self):
         with self.get_consumer(