294718: Enhance Batch model and admin functionality

src/country_workspace/admin/batch.py

@@ -14,7 +14,7 @@
 
 @admin.register(Batch)
 class BatchAdmin(BaseModelAdmin):
-    list_display = ("name", "import_date", "imported_by", "program", "source")
+    list_display = ("name", "import_date", "imported_by", "program", "source", "status")
     list_filter = (
         # "country_office",
         # "program",
@@ -23,7 +23,7 @@ class BatchAdmin(BaseModelAdmin):
         ("imported_by", AutoCompleteFilter),
         "source",
     )
-    readonly_fields = ("country_office", "program", "imported_by")
+    readonly_fields = ("country_office", "program", "imported_by", "status")
     search_fields = ("name",)
     ordering = ("-import_date",)
 
@@ -44,8 +44,12 @@ def _get_beneficiary_labels(self, batch: Batch) -> tuple[str, str]:
     @button(change_list=False, label="All Beneficiaries")
     def beneficiaries(self, request: HttpRequest, pk: str) -> HttpResponse:
         batch: Batch = self.get_object(request, pk)
-        households = batch.household_set.all()
-        individuals = batch.individual_set.all()
+        if batch.status == Batch.BatchStatus.COMPLETE:
+            households = batch.household_set.all()
+            individuals = batch.individual_set.all()
+        else:
+            households = batch.household_set.none()
+            individuals = batch.individual_set.none()
         context = {
             "batch": batch,
             "households": households,

src/country_workspace/config/fragments/constance.py

@@ -53,6 +53,7 @@
     "KOBO_PROJECT_VIEW_ID": (KOBO_PROJECT_VIEW_ID, "Kobo Project View ID", str),
     "KOBO_KF_URL": (KOBO_KF_URL, "Kobo Server address", str),
     "KOBO_CACHE_TTL": (86400, "Kobo data cache TTL", int),
+    "KOBO_IMPORT_TIMEBOX_MINUTES": (5, "Max minutes before Kobo import reschedules itself", int),
     "CACHE_TIMEOUT": (86400, "Cache Redis TTL", int),
     "CACHE_BY_VERSION": (False, "Invalidate Cache on CW version change", bool),
     "CONCURRENCY_GUARD": (
@@ -114,6 +115,7 @@
         "KOBO_FIELDS_TO_IGNORE",
         "KOBO_HH_FIELDS_TO_IGNORE",
         "KOBO_IND_FIELDS_TO_IGNORE",
+        "KOBO_IMPORT_TIMEBOX_MINUTES",
     ),
     "Remote System Settings - Aurora": (
         "AURORA_API_TOKEN",

src/country_workspace/contrib/aurora/import_processing.py

@@ -44,7 +44,10 @@ def import_data(job: AsyncJob) -> ImportResult:
         country_office=job.program.country_office,
         imported_by=job.owner,
         source=Batch.BatchSource.AURORA,
+        status=Batch.BatchStatus.LOADING,
     )
+    job.batch = batch
+    job.save(update_fields=["batch"])
 
     total_people = 0
     total_households = 0
@@ -53,6 +56,9 @@ def import_data(job: AsyncJob) -> ImportResult:
         imported = import_result(batch, result, config)
         total_people += imported.people
         total_households += imported.households
+
+    batch.status = Batch.BatchStatus.COMPLETE
+    batch.save(update_fields=["status"])
     return ImportResult(people=total_people, households=total_households)
 
 

src/country_workspace/contrib/kobo/api/client/helpers.py

@@ -60,12 +60,13 @@ def handle_paginated_response[T, U](
     collection_mapper: Callable[[raw_common.ListResponse], list[T]],
     item_mapper: Callable[[T], U],
 ) -> Generator[U, None, None]:
-    while url:
-        response = data_getter(url)
+    next_url: str | None = url
+    while next_url:
+        response = data_getter(next_url)
         response.raise_for_status()
-        data: raw_common.ListResponse = response.json()
+        data = cast("raw_common.ListResponse", response.json())
         yield from map(item_mapper, collection_mapper(data))
-        url = data["next"]
+        next_url = cast("str | None", data["next"])
 
 
 def get_raw_asset_list(data: raw_common.ListResponse) -> list[raw_asset_list.Asset]:
@@ -83,7 +84,7 @@ def get_asset_list(data_getter: DataGetter, url: str) -> Generator[Asset, None,
 def get_submission_list(data_getter: DataGetter, url: str, min_id: int | None = None) -> Iterable[Submission]:
     import json
 
-    query_params = {
+    query_params: dict[str, Any] = {
         START_PARAMETER_NAME: START_PARAMETER_VALUE,
         LIMIT_PARAMETER_NAME: LIMIT_PARAMETER_VALUE,
     }
@@ -92,14 +93,12 @@ def get_submission_list(data_getter: DataGetter, url: str, min_id: int | None =
         query_params["query"] = json.dumps({"_id": {"$gt": min_id}})
 
     url_with_params = change_url(url, query=query_params)
-    return map(
-        partial(download_attachments, data_getter),
-        handle_paginated_response(data_getter, url_with_params, get_raw_submission_list, Submission),
-    )
+    downloader = partial(download_attachments, cast("Callable[[str], Any]", data_getter))
+    return map(downloader, handle_paginated_response(data_getter, url_with_params, get_raw_submission_list, Submission))
 
 
 def get_asset(data_getter: DataGetter, raw: raw_asset_list.Asset) -> Asset:
     response = data_getter(raw["url"])
     response.raise_for_status()
-    raw_asset_data: raw_asset.Asset = response.json()
+    raw_asset_data = cast("raw_asset.Asset", response.json())
     return Asset(raw_asset_data, partial(get_submission_list, data_getter, raw_asset_data["data"]))

src/country_workspace/contrib/kobo/api/data/asset.py

@@ -1,14 +1,12 @@
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Iterable
 
 from country_workspace.contrib.kobo.api.data.common import Raw
 from country_workspace.contrib.kobo.api.data.submission import Submission
 from country_workspace.contrib.kobo.api.raw import asset as raw_asset
 
 
 class Asset(Raw[raw_asset.Asset]):
-    def __init__(
-        self, raw: raw_asset.Asset, submissions: Callable[[int | None], Generator[Submission, None, None]]
-    ) -> None:
+    def __init__(self, raw: raw_asset.Asset, submissions: Callable[[int | None], Iterable[Submission]]) -> None:
         super().__init__(raw)
         self._submissions = submissions
 
@@ -20,7 +18,7 @@ def uid(self) -> str:
     def name(self) -> str:
         return self._raw["name"]
 
-    def submissions(self, min_id: int | None = None) -> Generator[Submission, None, None]:
+    def submissions(self, min_id: int | None = None) -> Iterable[Submission]:
         yield from self._submissions(min_id)
 
     def __str__(self) -> str: