simplify get/catch and add transaction

Author: pld

tally_ho/apps/tally/management/commands/async_pvp_import.py

@@ -8,6 +8,8 @@
 
 import zipfile
 
+from django.db import transaction
+
 from tally_ho.apps.tally.models.pvp_upload_bundle import PvpUploadBundle
 from tally_ho.apps.tally.models.result_form import ResultForm
 from tally_ho.apps.tally.models.user_profile import UserProfile
@@ -28,19 +30,32 @@ def async_pvp_import(bundle_id, user_id):
     The orchestrator flips status to COMPLETED on success or FAILED if
     anything raises.
 
+    Re-invocations on a non-PENDING bundle short-circuit: a celery
+    retry or duplicate enqueue must not overwrite a terminal status
+    (which would also explode on the PvpSubmission unique constraint).
+
     :param bundle_id: PvpUploadBundle.id
     :param user_id:   UserProfile.id (the super admin who confirmed)
     :returns:         the bundle id (round-trip for callers/monitoring)
     """
-    bundle = PvpUploadBundle.objects.get(id=bundle_id)
     user = UserProfile.objects.get(id=user_id)
 
-    bundle.status = PvpBundleStatus.IMPORTING
-    bundle.save(update_fields=["status", "modified_date"])
+    # Atomically claim the bundle: lock the row, verify it's PENDING,
+    # flip to IMPORTING. A duplicate task arriving here finds the row
+    # already IMPORTING/COMPLETED/FAILED and bails before doing work.
+    with transaction.atomic():
+        bundle = PvpUploadBundle.objects.select_for_update().get(
+            id=bundle_id,
+        )
+        if bundle.status != PvpBundleStatus.PENDING:
+            return bundle.id
+        bundle.status = PvpBundleStatus.IMPORTING
+        bundle.save(update_fields=["status", "modified_date"])
 
-    # Catch failures from parse/validation too — otherwise a corrupted
-    # zip leaves the bundle stuck in IMPORTING. import_bundle handles
-    # its own FAILED flip; this except covers everything before it.
+    # All failure handling — parse errors, validation crashes, and
+    # import_bundle rollbacks — is centralized in this except. The
+    # FAILED save happens outside any atomic here (the claim above
+    # already committed), so it sticks even after the re-raise.
     try:
         zip_path = bundle.zip_file.path
         parsed = parse_bundle(zip_path)
@@ -66,9 +81,7 @@ def async_pvp_import(bundle_id, user_id):
                 zip_ref=zip_ref,
             )
     except Exception as exc:
-        bundle.refresh_from_db(fields=["status"])
-        if bundle.status != PvpBundleStatus.FAILED:
-            bundle.status = PvpBundleStatus.FAILED
+        bundle.status = PvpBundleStatus.FAILED
         bundle.error_message = str(exc)
         bundle.save(
             update_fields=["status", "error_message", "modified_date"],

tally_ho/apps/tally/management/commands/create_demo_pvp_bundle.py

@@ -23,9 +23,9 @@
 )
 
 
-CSV_HEADERS = tuple(
-    dict.fromkeys((*REQUIRED_COLUMNS, *RECON_COLUMNS, *IMAGE_COLUMNS))
-)
+# REQUIRED_COLUMNS is a superset — RECON_COLUMNS is already splatted in
+# and IMAGE_COLUMNS members are listed explicitly.
+CSV_HEADERS = REQUIRED_COLUMNS
 STUB_IMAGES = ("demo_sig.jpg", "demo_p1.jpg", "demo_p2.jpg")
 # Minimal valid JPEG (SOI + APP0 JFIF marker + EOI). Parser only checks
 # that the file exists; content does not need to be a real photo.

tally_ho/apps/tally/models/result_form.py

@@ -188,8 +188,12 @@ class Meta:
                               blank=True,
                               related_name='result_forms',
                               on_delete=models.PROTECT)
-    # Set to the PvpSubmission whose round-2 candidate results populated
-    # this form's DATA_ENTRY_1 results.
+    # Set to the PvpSubmission whose candidate results populated this
+    # form. In DE1_ONLY mode the submission's round-2 values become a
+    # single DATA_ENTRY_1 result set (DE2 is left for a human clerk).
+    # In DE1_AND_DE2 mode the submission populates DATA_ENTRY_1,
+    # DATA_ENTRY_2, and FINAL — corrections is skipped because the
+    # device guarantees round1 == round2.
     pvp_submission = models.ForeignKey(
         PvpSubmission,
         null=True,

tally_ho/apps/tally/tests/management/commands/test_async_pvp_import.py

@@ -201,6 +201,30 @@ def test_async_pvp_import_marks_failed_when_parse_raises(self):
         self.assertIsNone(self.bundle.imported_at)
         self.assertEqual(PvpSubmission.objects.count(), 0)
 
+    def test_async_pvp_import_is_noop_when_bundle_not_pending(self):
+        # If the task is re-queued externally (celery retry, manual
+        # re-run, sibling worker race), the second invocation must not
+        # overwrite a terminal status. Without the guard the task would
+        # flip COMPLETED -> IMPORTING and then explode on the duplicate
+        # PvpSubmission unique constraint, ending in spurious FAILED.
+        original_imported_at = self.bundle.imported_at
+        self.bundle.status = PvpBundleStatus.COMPLETED
+        self.bundle.number_of_submissions = 1
+        self.bundle.save(update_fields=[
+            "status", "number_of_submissions", "modified_date",
+        ])
+
+        async_pvp_import.apply(
+            kwargs={"bundle_id": self.bundle.id, "user_id": self.user.id},
+        )
+
+        self.bundle.refresh_from_db()
+        self.assertEqual(self.bundle.status, PvpBundleStatus.COMPLETED)
+        self.assertEqual(self.bundle.number_of_submissions, 1)
+        self.assertEqual(self.bundle.imported_at, original_imported_at)
+        # No new PvpSubmission row should have been created.
+        self.assertEqual(PvpSubmission.objects.count(), 0)
+
     def test_async_pvp_import_stores_error_message_on_failure(self):
         # On failure the bundle must record why so the operator sees
         # the cause on the result page instead of a bare FAILED.

tally_ho/apps/tally/tests/views/test_pvp_views.py

@@ -7,6 +7,7 @@
 import csv
 import io
 import json
+import pathlib
 import shutil
 import tempfile
 import zipfile
@@ -193,6 +194,70 @@ def test_post_invalid_zip_shows_error_no_bundle_created(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(PvpUploadBundle.objects.count(), 0)
 
+    def test_post_invalid_zip_removes_persisted_zip_from_disk(self):
+        # On parse failure the bundle row is deleted, but Django does
+        # not delete the underlying file when the model row goes away.
+        # The view must explicitly clean it up to avoid an orphan in
+        # MEDIA_ROOT.
+        view = PvpUploadView.as_view()
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, mode="w") as zf:
+            zf.writestr("README.txt", b"oops")
+        upload = SimpleUploadedFile(
+            "bundle.zip", buf.getvalue(),
+            content_type="application/zip",
+        )
+        request = self.factory.post("/", data={"zip_file": upload})
+        request.user = User.objects.get(id=self.user.id)
+        request.session = {}
+
+        # Snapshot the bundle directory tree before; anything new here
+        # at the end of the test is a leak.
+        bundles_root = pathlib.Path(self._media_root) / "pvp" / "bundles"
+        before = (
+            set(p.relative_to(bundles_root) for p in bundles_root.rglob("*"))
+            if bundles_root.exists() else set()
+        )
+
+        response = view(request, tally_id=self.tally.id)
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(PvpUploadBundle.objects.count(), 0)
+
+        after_files = (
+            {
+                p for p in bundles_root.rglob("*")
+                if p.is_file()
+            } if bundles_root.exists() else set()
+        )
+        # Compare file paths only; intermediate dirs are not interesting.
+        new_files = {
+            p for p in after_files
+            if p.relative_to(bundles_root) not in before
+        }
+        self.assertEqual(new_files, set())
+
+    def test_post_to_disabled_tally_shows_error_no_bundle_created(self):
+        # Uploading to a DISABLED tally would have every row rejected
+        # downstream as `pvp_disabled` — short-circuit at upload so the
+        # operator sees the actual cause instead of a useless skip list.
+        self.tally.pvp_mode = PvpMode.DISABLED
+        self.tally.save(update_fields=["pvp_mode"])
+
+        view = PvpUploadView.as_view()
+        zip_bytes = _zip_bytes([
+            _csv_row(instance_id="uuid:1", barcode="111",
+                     candidate_id=131301, order=1, r2=12),
+        ])
+        upload = SimpleUploadedFile(
+            "bundle.zip", zip_bytes, content_type="application/zip",
+        )
+        request = self.factory.post("/", data={"zip_file": upload})
+        request.user = User.objects.get(id=self.user.id)
+        request.session = {}
+        response = view(request, tally_id=self.tally.id)
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(PvpUploadBundle.objects.count(), 0)
+
 
 class TestPvpConfirmView(_PvpViewTestBase, TestCase):
     def _make_pending_bundle(self):

tally_ho/apps/tally/views/pvp.py

@@ -12,6 +12,7 @@
 """
 
 from django.contrib import messages
+from django.db import transaction
 from django.http import JsonResponse
 from django.shortcuts import get_object_or_404, redirect
 from django.utils.translation import gettext_lazy as _
@@ -27,6 +28,7 @@
 from tally_ho.apps.tally.models.tally import Tally
 from tally_ho.apps.tally.models.user_profile import UserProfile
 from tally_ho.libs.models.enums.pvp_bundle_status import PvpBundleStatus
+from tally_ho.libs.models.enums.pvp_mode import PvpMode
 from tally_ho.libs.permissions import groups
 from tally_ho.libs.pvp.bundle import (
     DuplicateBarcodeError,
@@ -54,6 +56,17 @@ def get_context_data(self, **kwargs):
     def form_valid(self, form):
         tally_id = self.kwargs["tally_id"]
         tally = get_object_or_404(Tally, id=tally_id)
+
+        # A DISABLED tally would reject every row at confirm-time with
+        # `pvp_disabled`. Short-circuit at upload so the operator sees
+        # the actual cause and isn't asked to confirm an empty import.
+        if tally.pvp_mode == PvpMode.DISABLED:
+            form.add_error(
+                "zip_file",
+                _("PVP is disabled for this tally."),
+            )
+            return self.form_invalid(form)
+
         upload = form.cleaned_data["zip_file"]
 
         # request.user is a Django User; the FK targets UserProfile (the
@@ -84,7 +97,11 @@ def form_valid(self, form):
             RoundIntegrityError,
             UnsafeImageFilenameError,
         ) as exc:
-            bundle.delete()  # zip_file deletion handled by FileField
+            # Django's FileField does not auto-delete the on-disk file
+            # when the model row is deleted; do it explicitly so a
+            # rejected upload does not leak a zip in MEDIA_ROOT.
+            bundle.zip_file.delete(save=False)
+            bundle.delete()
             form.add_error("zip_file", str(exc))
             return self.form_invalid(form)
 
@@ -135,13 +152,20 @@ def get_context_data(self, **kwargs):
     def post(self, request, *args, **kwargs):
         tally_id = self.kwargs["tally_id"]
         bundle_id = self.kwargs["bundle_id"]
-        bundle = get_object_or_404(
+        # Tally membership check up front; the locked row read below is
+        # for the enqueue gate, not authorization.
+        get_object_or_404(
             PvpUploadBundle, id=bundle_id, tally_id=tally_id,
         )
-        # Only PENDING bundles can be enqueued. Re-confirming an already
-        # IMPORTING/COMPLETED/FAILED bundle would otherwise duplicate the
-        # import task and race against the in-flight one.
-        if bundle.status != PvpBundleStatus.PENDING:
+        # Take a row lock so two concurrent confirms can't both observe
+        # PENDING and both enqueue. The lock is released when the
+        # atomic block exits, which happens before we hand off to celery.
+        with transaction.atomic():
+            bundle = PvpUploadBundle.objects.select_for_update().get(
+                id=bundle_id,
+            )
+            already_handled = bundle.status != PvpBundleStatus.PENDING
+        if already_handled:
             messages.info(
                 request,
                 _("PVP import already %(status)s.") % {

tally_ho/libs/pvp/bundle.py

@@ -125,8 +125,13 @@ def parse_bundle(zip_path: str | Path) -> ParsedBundle:
     """Parse a PVP upload bundle into submission-level rows.
 
     Raises:
-      InvalidBundleError: zip can't open, csv missing, required columns gone
-      DuplicateBarcodeError: same barcode appears in more than one submission
+      InvalidBundleError: zip can't open, csv missing, required columns
+        gone, or a value that must parse as an integer doesn't
+      DuplicateBarcodeError: same barcode appears in more than one
+        submission
+      RoundIntegrityError: any candidate row has missing rounds or
+        round1 != round2 (device guarantees they match)
+      UnsafeImageFilenameError: an image filename contains path syntax
     """
     path = Path(zip_path)
     try:

tally_ho/libs/pvp/import_submission.py

@@ -65,7 +65,12 @@
 def import_submission(
     parsed_submission, *, tally, bundle, uploaded_by, zip_ref,
 ):
-    """Import one validated PVP submission as DE1 results + recon."""
+    """Import one validated PVP submission.
+
+    Entry-version writes depend on ``bundle.mode`` — see the module
+    docstring for the DE1_ONLY vs DE1_AND_DE2 breakdown. Wraps the
+    write pass in a reversion revision for audit history.
+    """
     with reversion.create_revision():
         reversion.set_user(uploaded_by)
         reversion.set_comment(
@@ -191,37 +196,32 @@ def import_bundle(*, bundle, submissions, tally, uploaded_by, zip_ref):
     parse-time validation. The bundle row already exists (created at
     upload time); this function walks the validated subset, calls
     `import_submission` for each row inside a single outer atomic
-    transaction, and flips the bundle's terminal status.
+    transaction, and flips the bundle to COMPLETED on success.
 
     All-or-nothing: if any `import_submission` raises, the outer atomic
     rolls back every prior submission's writes (each row's own atomic
     becomes a savepoint of the outer). Image `on_commit` callbacks
     registered by successful submissions are discarded on rollback, so
-    no orphan files land on disk. The bundle status flip to FAILED
-    happens after the rollback so the FAILED row sticks.
+    no orphan files land on disk. Callers are responsible for flipping
+    the bundle to FAILED and persisting the exception message.
     """
     submissions = list(submissions)
-    try:
-        with transaction.atomic():
-            for parsed in submissions:
-                import_submission(
-                    parsed,
-                    tally=tally,
-                    bundle=bundle,
-                    uploaded_by=uploaded_by,
-                    zip_ref=zip_ref,
-                )
-            bundle.number_of_submissions = len(submissions)
-            bundle.imported_at = timezone.now()
-            bundle.status = PvpBundleStatus.COMPLETED
-            bundle.save(update_fields=[
-                "number_of_submissions", "imported_at",
-                "status", "modified_date",
-            ])
-    except Exception:
-        bundle.status = PvpBundleStatus.FAILED
-        bundle.save(update_fields=["status", "modified_date"])
-        raise
+    with transaction.atomic():
+        for parsed in submissions:
+            import_submission(
+                parsed,
+                tally=tally,
+                bundle=bundle,
+                uploaded_by=uploaded_by,
+                zip_ref=zip_ref,
+            )
+        bundle.number_of_submissions = len(submissions)
+        bundle.imported_at = timezone.now()
+        bundle.status = PvpBundleStatus.COMPLETED
+        bundle.save(update_fields=[
+            "number_of_submissions", "imported_at",
+            "status", "modified_date",
+        ])
     return bundle
 
 

tally_ho/libs/tests/pvp/test_import_submission.py

@@ -524,10 +524,12 @@ def test_empty_submission_list_marks_completed_with_zero(self):
         self.assertEqual(self.bundle.number_of_submissions, 0)
         self.assertIsNotNone(self.bundle.imported_at)
 
-    def test_failure_marks_bundle_failed_reraises_and_rolls_back_all(self):
+    def test_failure_reraises_and_rolls_back_all(self):
         # Second submission references a candidate_id that doesn't exist.
         # Under all-or-nothing semantics the first submission must also
         # be rolled back so the bundle leaves no partial trace.
+        # import_bundle re-raises for the caller to handle the FAILED
+        # transition (see async_pvp_import for the terminal flip).
         bad = self._parsed_submission(
             odk_instance_id="uuid:rf-2", barcode="222",
             candidates=(
@@ -551,7 +553,10 @@ def test_failure_marks_bundle_failed_reraises_and_rolls_back_all(self):
         finally:
             zip_ref.close()
         self.bundle.refresh_from_db()
-        self.assertEqual(self.bundle.status, PvpBundleStatus.FAILED)
+        # import_bundle no longer sets FAILED — that's the caller's job.
+        # The bundle stays at whatever status it had when import_bundle
+        # was called (IMPORTING in production, PENDING in this test).
+        self.assertNotEqual(self.bundle.status, PvpBundleStatus.COMPLETED)
         self.assertIsNone(self.bundle.imported_at)
         # All-or-nothing: nothing persisted.
         self.assertEqual(self.bundle.submissions.count(), 0)