Reconnect on default DB at end of export

arkindex/documents/export/__init__.py

@@ -11,6 +11,7 @@ from pathlib import Path
 from django.conf import settings
 from django.core.mail import send_mail
 from django.db import connections
+from django.db.utils import InterfaceError, OperationalError
 from django.template.loader import render_to_string
 from django.utils.text import slugify
 from django_rq import job
@@ -118,6 +119,22 @@ def send_email(subject, template_name, corpus_export, **context):
         logger.error(f"Failed to send email to {corpus_export.user.email}")
 
 
+def update_state(corpus_export: CorpusExport, state: CorpusExportState):
+    """
+    Make sure the corpus export instance from default DB is still available for updates
+    Sometimes the DB connection may drop after a really long export, especially with remote DB exports
+    """
+    try:
+        corpus_export.state = state
+        corpus_export.save()
+    except (InterfaceError, OperationalError) as e:
+        logger.warning(f"Database connection has been lost, retrying: {e}")
+        connections["default"].connect()
+        corpus_export.refresh_from_db(using="default")
+        corpus_export.state = state
+        corpus_export.save()
+
+
 @job("export", timeout=settings.RQ_TIMEOUTS["export_corpus"])
 def export_corpus(corpus_export: CorpusExport) -> None:
     _, db_path = tempfile.mkstemp(suffix=".db")
@@ -181,8 +198,8 @@ def export_corpus(corpus_export: CorpusExport) -> None:
             },
         )
 
-        corpus_export.state = CorpusExportState.Done
-        corpus_export.save()
+        # Safely update state with auto-reconnect to db
+        update_state(corpus_export, CorpusExportState.Done)
 
         send_email(
             "Arkindex project export completed",
@@ -190,8 +207,9 @@ def export_corpus(corpus_export: CorpusExport) -> None:
             corpus_export,
         )
     except Exception as e:
-        corpus_export.state = CorpusExportState.Failed
-        corpus_export.save()
+        # Safely update state with auto-reconnect to db
+        update_state(corpus_export, CorpusExportState.Failed)
+
         send_email(
             "Arkindex project export failed",
             "export_error.html",