Validate records in chunks (#271)

Author: datamik

src/country_workspace/config/fragments/constance.py

@@ -64,6 +64,7 @@
     ),
     "MAILJET_API_KEY": (MAILJET_API_KEY, "Mailjet API key", str),
     "MAILJET_SECRET_KEY": (MAILJET_SECRET_KEY, "Mailjet secret key", "write_only_input"),
+    "CHUNK_SIZE_FOR_VALIDATION_TASK": (500, "Number of records to process per chunk in validation tasks", int),
 }
 
 CONSTANCE_CONFIG_FIELDSETS = {
@@ -83,6 +84,7 @@
         "MAILJET_SECRET_KEY",
     ),
     "Data consistency": ("CONCURRENCY_GUARD",),
+    "System": ("CHUNK_SIZE_FOR_VALIDATION_TASK",),
 }
 
 # Mapping of config keys to masked default display values in the Constance admin UI.

src/country_workspace/workspaces/admin/cleaners/actions.py

@@ -17,7 +17,7 @@
 from .calculate_checksum import calculate_checksum_impl
 from .mass_update import MassUpdateForm, mass_update_impl
 from .regex import RegexUpdateForm, regex_update_impl
-from .validate import validate_queryset
+from .validate import create_validation_jobs
 
 if TYPE_CHECKING:
     from django.db.models import QuerySet
@@ -32,19 +32,15 @@ def validate_records(
     queryset: "QuerySet[Beneficiary]",
 ) -> None:
     if model_admin._check_empty_queryset(request, queryset):
-        return None
-    opts = queryset.model._meta
-    job = AsyncJob.objects.create(
+        return
+    create_validation_jobs(
         description=validate_records.short_description,
-        type=AsyncJob.JobType.ACTION,
         owner=request.user,
-        action=fqn(validate_queryset),
         program=state.program,
-        config={"pks": list(queryset.values_list("pk", flat=True)), "model_name": opts.label},
+        queryset=queryset,
     )
-    job.queue()
     model_admin.message_user(request, "Task scheduled", messages.SUCCESS)
-    return job
+    return
 
 
 @admin.action(description="Mass update record fields", permissions=["mass_update"])

src/country_workspace/workspaces/admin/cleaners/validate.py

@@ -2,11 +2,14 @@
 from typing import Any
 from itertools import batched
 from collections.abc import Iterable
+
+from concurrency.utils import fqn
+from constance import config
 from django.db.models import Model, QuerySet, Prefetch
 from django.db.models.query import prefetch_related_objects
 
 from country_workspace.context import batch_ctx
-from country_workspace.models import AsyncJob, Household, Individual
+from country_workspace.models import AsyncJob, Household, Individual, Program
 from country_workspace.state import state
 
 
@@ -50,17 +53,6 @@ def validate_queryset(queryset: QuerySet[Model], chunk_size: int = 2000, **kwarg
     return {"valid": valid, "invalid": invalid}
 
 
-def validate_program(job: AsyncJob) -> dict[str, int]:
-    try:
-        program = job.program
-        qs = program.households.all() if program.beneficiary_group.master_detail else program.individuals.all()
-        return validate_queryset(qs)
-
-    except Exception as e:  # pragma: no cover
-        logger.error("Error during program validation: %s", e)
-        raise
-
-
 def _validate_and_count(objs: Iterable[Model]) -> tuple[int, int]:
     valid = invalid = 0
     for obj in objs:
@@ -70,3 +62,18 @@ def _validate_and_count(objs: Iterable[Model]) -> tuple[int, int]:
             else:
                 invalid += 1
     return valid, invalid
+
+
+def create_validation_jobs(description: str, owner: str, program: Program, queryset: QuerySet) -> AsyncJob:
+    opts = queryset.model._meta
+    queryset = queryset.values_list("pk", flat=True).order_by("pk")
+    for chunk in batched(queryset, config.CHUNK_SIZE_FOR_VALIDATION_TASK):
+        job = AsyncJob.objects.create(
+            description=f"{description} (PKs {chunk[0]} - {chunk[-1]})",
+            type=AsyncJob.JobType.ACTION,
+            owner=owner,
+            action=fqn(validate_queryset),
+            program=program,
+            config={"pks": chunk, "model_name": opts.label},
+        )
+        job.queue()

src/country_workspace/workspaces/admin/hh_ind.py

@@ -2,7 +2,6 @@
 
 from admin_extra_buttons.decorators import button
 from adminfilters.mixin import AdminAutoCompleteSearchMixin
-from concurrency.utils import fqn
 from django.contrib import messages
 from django.contrib.admin.utils import unquote
 from django.core.exceptions import PermissionDenied
@@ -16,12 +15,11 @@
 from django.utils.translation import gettext_lazy as _
 
 from ...cache.manager import cache_manager
-from ...models import AsyncJob
 from ...state import state
 from ..options import WorkspaceModelAdmin
 from ..models import CountryHousehold, CountryIndividual
 from .cleaners import actions
-from .cleaners.validate import validate_program
+from .cleaners.validate import create_validation_jobs
 from ...utils.flex_fields import Base64ImageField, get_checker_fields
 
 if TYPE_CHECKING:
@@ -168,15 +166,13 @@ def validate_single(self, request: HttpRequest, pk: str) -> "HttpResponse":
     )
     def validate_program(self, request: HttpRequest) -> "HttpResponse":
         program = state.program
-        job = AsyncJob.objects.create(
+        queryset = program.households.all() if program.beneficiary_group.master_detail else program.individuals.all()
+        create_validation_jobs(
             description="Validate Entire Programme",
-            type=AsyncJob.JobType.TASK,
             owner=request.user,
-            action=fqn(validate_program),
             program=program,
-            config={},
+            queryset=queryset,
         )
-        job.queue()
         self.message_user(request, _("Task scheduled"), messages.SUCCESS)
 
     @button(html_attrs={"title": "Shows raw data as stored, ready to be sent to HOPE"})

tests/admin/test_admin_cleaners.py

@@ -1,8 +1,4 @@
 import pytest
-from pytest_mock import MockerFixture
-from unittest.mock import Mock
-
-from country_workspace.workspaces.admin.cleaners.validate import validate_program
 from country_workspace.workspaces.models import CountryHousehold, CountryIndividual, CountryProgram
 
 type Beneficiary = CountryHousehold | CountryIndividual
@@ -31,27 +27,3 @@ def beneficiary(program: CountryProgram, master_detail: bool) -> Beneficiary:
         batch__country_office=program.country_office,
         household=None,
     )
-
-
-@pytest.mark.parametrize(
-    ("return_value", "expected"),
-    [(True, {"valid": 1, "invalid": 0}), (False, {"valid": 0, "invalid": 1})],
-    ids=["valid", "invalid"],
-)
-def test_validate_program_success(
-    mocker: MockerFixture, beneficiary: Beneficiary, master_detail: bool, return_value: bool, expected: dict[str, int]
-) -> None:
-    model_path = "country_workspace.models.Household" if master_detail else "country_workspace.models.Individual"
-    mocker.patch(f"{model_path}.validate_with_checker", return_value=return_value)
-    assert validate_program(Mock(program=beneficiary.program)) == expected
-
-
-def test_validate_program_exception(mocker: MockerFixture, beneficiary: Beneficiary, master_detail: bool) -> None:
-    model_path = "country_workspace.models.Household" if master_detail else "country_workspace.models.Individual"
-    mocker.patch(f"{model_path}.validate_with_checker", side_effect=Exception("Test error"))
-    with pytest.raises(Exception, match="Test error"):
-        validate_program(Mock(program=beneficiary.program))
-
-
-def test_validate_program_empty_queryset(program: CountryProgram) -> None:
-    assert validate_program(Mock(program=program)) == {"valid": 0, "invalid": 0}