Enhance transaction management in Kobo asset import

- Updated `import_asset` to ensure each submission is processed in its own transaction, allowing for partial success and rollback of only failed submissions.
- Added detailed comments to clarify transaction handling and watermark persistence for recovery.
- Introduced a new test to verify that earlier submissions remain intact when a later submission fails, ensuring data integrity during import processes.

These changes improve the robustness of the import functionality and enhance error handling in the Kobo sync process.

Author: Arsen P

src/country_workspace/contrib/kobo/sync.py

@@ -293,6 +293,8 @@ def import_asset(
         for submission in submissions_iterator:
             current_submission = submission
             originating_id = get_kobo_originating_id(asset.uid, str(submission.id))
+            # One transaction per submission: if this block fails, only this
+            # submission is rolled back; previously committed data stays.
             with transaction.atomic():
                 household = create_household(batch, submission, config, id_generator, originating_id)
                 individuals = create_individuals(batch, household, submission, config, originating_id)
@@ -319,6 +321,8 @@ def import_asset(
         raise ImportError(error_msg) from e
 
     finally:
+        # Persist watermark so the next run (or retry) resumes after the last
+        # successful submission; runs even when an exception is raised.
         if last_successful_id > last_id:
             SyncLog.objects.update_or_create(
                 name=sync_log_name,
@@ -347,8 +351,22 @@ def id_generator() -> int:
 
 
 def import_data(job: AsyncJob) -> ImportResult:
+    """
+    Import Kobo asset data for a job.
+
+    Transaction management:
+    - Batch setup (create or lock batch, set status to LOADING) runs in a single
+      atomic block so the batch is visible and locked before import starts.
+    - import_asset() runs outside that block; it commits each submission in its
+      own transaction (see import_asset). That way, partial progress is persisted
+      and the watermark (SyncLog) can be updated per submission.
+    - On success, batch status is updated to COMPLETE in a final atomic block.
+    - On incomplete (timebox or reschedule), a new AsyncJob is queued; no batch
+      status change until the run that completes.
+    """
     config: Config = job.config
 
+    # Create or lock the batch and set status to LOADING in one transaction.
     with transaction.atomic():
         batch_id = getattr(job, "batch_id", None)
         if batch_id:
@@ -381,25 +399,28 @@ def import_data(job: AsyncJob) -> ImportResult:
     timebox_minutes = getattr(constance_config, "KOBO_IMPORT_TIMEBOX_MINUTES", 0) or 0
     timebox_seconds = int(timebox_minutes * 60) if timebox_minutes > 0 else None
 
-    import_result = import_asset(
+    # import_asset commits each submission in its own transaction; on error it
+    # raises after updating the watermark so the next run can resume.
+    asset_import_result = import_asset(
         batch,
         asset,
         config,
         id_generator,
         timebox_seconds=timebox_seconds,
     )
-    household_counter += import_result["households"]
-    individual_counter += import_result["individuals"]
+    household_counter += asset_import_result["households"]
+    individual_counter += asset_import_result["individuals"]
 
-    if import_result["completed"] and config.get("validate_after_import"):
+    if asset_import_result["completed"] and config.get("validate_after_import"):
         create_validation_jobs(
             description=f"Validate records for batch {batch.pk}",
             owner=job.owner,
             program=job.program,
             queryset=batch.household_set.all().prefetch_related("members"),  # type: ignore[attr-defined]
         )
 
-    if import_result["completed"]:
+    if asset_import_result["completed"]:
+        # Mark batch complete in a dedicated transaction.
         with transaction.atomic():
             Batch.objects.select_for_update().filter(pk=batch.pk).update(status=Batch.BatchStatus.COMPLETE)
     else:
@@ -417,5 +438,5 @@ def import_data(job: AsyncJob) -> ImportResult:
     return ImportResult(
         households=household_counter,
         individuals=individual_counter,
-        completed=import_result["completed"],
+        completed=asset_import_result["completed"],
     )

tests/contrib/kobo/test_kobo_sync.py

@@ -35,6 +35,8 @@
     DataCheckerFactory,
     FieldsetFactory,
     FlexFieldFactory,
+    HouseholdFactory,
+    IndividualFactory,
     SyncLogFactory,
 )
 from testutils.factories.smart_fields import DataCheckerFieldsetFactory
@@ -392,6 +394,65 @@ def test_import_asset_with_error(mocker: MockerFixture, config: Config) -> None:
     assert sync_log.last_id == "101"
 
 
+@pytest.mark.django_db
+def test_import_asset_on_error_persists_previous_data(mocker: MockerFixture, config: Config) -> None:
+    """
+    When import_asset fails on a later submission, earlier data are persisted:
+    each submission runs in its own transaction, so a failure rolls back only
+    that submission; we assert the first household (and its individuals) remain
+    in the DB and the watermark is updated for recovery.
+    """
+    batch = BatchFactory()
+    program_ct = ContentType.objects.get_for_model(Program)
+    SyncLogFactory(
+        name="kobo_test_asset_uid",
+        content_type=program_ct,
+        object_id=batch.program.id,
+        last_id="100",
+    )
+
+    id_generator_mock = mocker.Mock(name="id_generator")
+
+    def create_household_real(batch, submission, config, id_generator, originating_id):
+        return HouseholdFactory(batch=batch, individuals=[])
+
+    def create_individuals_real(batch, household, submission, config, originating_id):
+        return [IndividualFactory(batch=batch, household=household)]
+
+    mocker.patch(
+        "country_workspace.contrib.kobo.sync.create_household",
+        side_effect=create_household_real,
+    )
+    mocker.patch(
+        "country_workspace.contrib.kobo.sync.create_individuals",
+        side_effect=create_individuals_real,
+    )
+    set_roles_and_relationships_mock = mocker.patch("country_workspace.contrib.kobo.sync.set_roles_and_relationships")
+    set_roles_and_relationships_mock.side_effect = [None, ValueError("fail on second")]
+
+    submission_1 = Mock(spec=dict)
+    submission_1.id = 101
+    submission_2 = Mock(spec=dict)
+    submission_2.id = 102
+    asset_mock = Mock()
+    asset_mock.uid = "test_asset_uid"
+    asset_mock.submissions = Mock(return_value=iter([submission_1, submission_2]))
+
+    with pytest.raises(ImportError, match=r"Successfully imported.*at submission 102"):
+        import_asset(batch, asset_mock, config, id_generator_mock)
+
+    batch.refresh_from_db()
+    assert batch.household_set.count() == 1
+    assert batch.household_set.first().members.count() == 1
+
+    sync_log = SyncLog.objects.get(
+        name="kobo_test_asset_uid",
+        content_type=program_ct,
+        object_id=batch.program.id,
+    )
+    assert sync_log.last_id == "101"
+
+
 @pytest.mark.django_db
 def test_import_asset_no_new_submissions(mocker: MockerFixture, config: Config) -> None:
     batch = BatchFactory()