207255_Celery_tasks_refactor

src/hope/admin/household.py

@@ -242,15 +242,15 @@ def has_add_permission(self, request: HttpRequest, obj: Household | None = None)
 
 class HouseholdWithdrawFromListMixin:
     @staticmethod
-    def get_household_queryset_from_list(household_id_list: list, program: Program) -> QuerySet:
+    def get_household_queryset_from_list(household_id_list: list[str], program: Program) -> QuerySet:
         return Household.objects.filter(
             unicef_id__in=household_id_list,
             withdrawn=False,
             program=program,
         )
 
     @transaction.atomic
-    def mass_withdraw_households_from_list_bulk(self, household_id_list: list, tag: str, program: Program) -> None:
+    def mass_withdraw_households_from_list_bulk(self, household_id_list: list[str], tag: str, program: Program) -> None:
         households = self.get_household_queryset_from_list(household_id_list, program)
         individuals = Individual.objects.filter(household__in=households, withdrawn=False, duplicate=False)
 
@@ -631,7 +631,7 @@ def mass_enroll_to_another_program(self, request: HttpRequest, qs: QuerySet) ->
                 program_for_enroll = form.cleaned_data["program_for_enroll"]
                 households_ids = list(qs.distinct("unicef_id").values_list("id", flat=True))
                 enroll_households_to_program_task.delay(
-                    households_ids=households_ids,
+                    households_ids=[str(_id) for (_id) in households_ids],
                     program_for_enroll_id=str(program_for_enroll.id),
                     user_id=str(request.user.id),
                 )

src/hope/admin/individual.py

@@ -226,16 +226,15 @@ def sanity_check(self, request: HttpRequest, pk: UUID) -> TemplateResponse:
     def revalidate_phone_number_sync(self, request: HttpRequest, queryset: QuerySet) -> None:
         try:
             ids = queryset.values_list("id", flat=True)
-            revalidate_phone_number_task(ids)
+            revalidate_phone_number_task([str(_id) for _id in ids])
             self.message_user(request, f"Updated {len(ids)} records", messages.SUCCESS)
         except Error as e:
             self.message_user(request, str(e), messages.ERROR)
 
     revalidate_phone_number_sync.short_description = "Re-validate phone number (sync)"
 
     def revalidate_phone_number_async(self, request: HttpRequest, queryset: QuerySet) -> None:
-        ids = list(queryset.values_list("id", flat=True))
-        revalidate_phone_number_task.delay(ids)
+        revalidate_phone_number_task.delay([str(_id) for _id in queryset.values_list("id", flat=True)])
         self.message_user(request, "Updating in progress", messages.SUCCESS)
 
     revalidate_phone_number_async.short_description = "Re-validate phone number (async)"

src/hope/admin/registration_data.py

@@ -347,7 +347,7 @@ def enroll_to_program(self, request: HttpRequest, pk: UUID) -> HttpResponse | No
                 program_for_enroll = form.cleaned_data["program_for_enroll"]
                 households_ids = list(qs.distinct("unicef_id").values_list("id", flat=True))
                 enroll_households_to_program_task.delay(
-                    households_ids=households_ids,
+                    households_ids=[str(_id) for (_id) in households_ids],
                     program_for_enroll_id=str(program_for_enroll.id),
                     user_id=str(request.user.id),
                 )

src/hope/apps/account/celery_tasks.py

@@ -4,25 +4,44 @@
 
 from django.db.models import Q
 from django.utils import timezone
+from django_celery_boost.models import AsyncJobModel
 
 from hope.apps.account.signals import _invalidate_user_permissions_cache
 from hope.apps.core.celery import app
 from hope.apps.utils.logs import log_start_and_end
 from hope.apps.utils.sentry import sentry_tags
+from hope.models import AsyncRetryJob
 
 logger = logging.getLogger(__name__)
 
 
+def invalidate_permissions_cache_for_user_if_expired_role_action(job: AsyncRetryJob) -> bool:
+    # Invalidate permissions cache for users with roles that expired a day before
+    from hope.models import User
+
+    try:
+        day_ago = timezone.now() - datetime.timedelta(days=1)
+        users = User.objects.filter(
+            Q(role_assignments__expiry_date=day_ago.date()) | Q(partner__role_assignments__expiry_date=day_ago.date())
+        ).distinct()
+        _invalidate_user_permissions_cache(users)
+        return True
+    except Exception:
+        logger.exception("Failed to invalidate permissions cache for users with expired roles")
+        raise
+
+
 @app.task(bind=True, default_retry_delay=60, max_retries=3)
 @log_start_and_end
 @sentry_tags
 def invalidate_permissions_cache_for_user_if_expired_role(self: Any) -> bool:
-    # Invalidate permissions cache for users with roles that expired a day before
-    from hope.models import User
-
-    day_ago = timezone.now() - datetime.timedelta(days=1)
-    users = User.objects.filter(
-        Q(role_assignments__expiry_date=day_ago.date()) | Q(partner__role_assignments__expiry_date=day_ago.date())
-    ).distinct()
-    _invalidate_user_permissions_cache(users)
+    job = AsyncRetryJob.objects.create(
+        owner=None,
+        type=AsyncJobModel.JobType.JOB_TASK,
+        action="hope.apps.account.celery_tasks.invalidate_permissions_cache_for_user_if_expired_role_action",
+        config={},
+        group_key="invalidate_permissions_cache_for_user_if_expired_role",
+        description="Invalidate permissions cache for users with expired roles",
+    )
+    job.queue()
     return True

src/hope/apps/accountability/celery_tasks.py

@@ -1,5 +1,7 @@
 import logging
 
+from django_celery_boost.models import AsyncJobModel
+
 from hope.apps.accountability.services.export_survey_sample_service import (
     ExportSurveySampleService,
 )
@@ -8,17 +10,17 @@
 from hope.apps.core.utils import send_email_notification
 from hope.apps.utils.logs import log_start_and_end
 from hope.apps.utils.sentry import sentry_tags, set_sentry_business_area_tag
-from hope.models import BusinessArea, Survey
+from hope.models import AsyncJob, BusinessArea, Survey
 
 logger = logging.getLogger(__name__)
 
 
-@app.task
-@log_start_and_end
-@sentry_tags
-def export_survey_sample_task(survey_id: str, user_id: str) -> None:
+def export_survey_sample_task_action(job: AsyncJob) -> None:
     from hope.models import User
 
+    survey_id = job.config["survey_id"]
+    user_id = job.config["user_id"]
+
     try:
         survey = Survey.objects.get(id=survey_id)
         user = User.objects.get(pk=user_id)
@@ -30,43 +32,92 @@ def export_survey_sample_task(survey_id: str, user_id: str) -> None:
         if survey.business_area.enable_email_notification:
             send_email_notification(service, user)
 
-    except Exception as e:
-        logger.warning(e)
+    except Exception as exc:
+        job.errors = {
+            "error": str(exc),
+        }
+        job.save(update_fields=["errors"])
+        logger.exception("Failed to export survey sample")
+        raise
+
+
+def send_survey_to_users_action(job: AsyncJob) -> None:
+    survey_id = job.config["survey_id"]
+    try:
+        survey = Survey.objects.get(id=survey_id)
+        set_sentry_business_area_tag(survey.business_area.name)
+        if survey.category == Survey.CATEGORY_MANUAL:
+            return
+        phone_numbers = survey.recipients.filter(head_of_household__phone_no_valid=True).values_list(
+            "head_of_household__phone_no", flat=True
+        )
+        if survey.category == Survey.CATEGORY_SMS:
+            api = RapidProAPI(survey.business_area.slug, RapidProAPI.MODE_MESSAGE)
+            api.broadcast_message(phone_numbers, survey.body)
+            return
+        business_area = BusinessArea.objects.get(id=survey.business_area_id)
+        api = RapidProAPI(business_area.slug, RapidProAPI.MODE_VERIFICATION)
+
+        already_received = [
+            phone_number
+            for successful_call in survey.successful_rapid_pro_calls
+            for phone_number in successful_call["urns"]
+        ]
+        phone_numbers = [phone_number for phone_number in phone_numbers if phone_number not in already_received]
+
+        successful_flows, error = api.start_flow(survey.flow_id, phone_numbers)
+        if error:
+            job.errors = {
+                "start_flow_error": str(error),
+            }
+            job.save(update_fields=["errors"])
+
+        for successful_flow in successful_flows:
+            survey.successful_rapid_pro_calls.append(
+                {
+                    "flow_uuid": successful_flow.response["uuid"],
+                    "urns": list(map(str, successful_flow.urns)),
+                }
+            )
+        survey.save()
+    except Exception as exc:
+        job.errors = {
+            "error": str(exc),
+        }
+        job.save(update_fields=["errors"])
+        logger.exception("Failed to send survey to users")
         raise
 
 
+@app.task
+@log_start_and_end
+@sentry_tags
+def export_survey_sample_task(survey_id: str, user_id: str) -> None:
+    survey = Survey.objects.get(id=survey_id)
+    job = AsyncJob.objects.create(
+        owner_id=user_id,
+        program=survey.program,
+        type=AsyncJobModel.JobType.JOB_TASK,
+        action="hope.apps.accountability.celery_tasks.export_survey_sample_task_action",
+        config={"survey_id": str(survey_id), "user_id": str(user_id)},
+        group_key=f"export_survey_sample_task:{survey_id}",
+        description=f"Export survey sample for survey {survey_id}",
+    )
+    job.queue()
+
+
 @app.task
 @log_start_and_end
 @sentry_tags
 def send_survey_to_users(survey_id: str) -> None:
     survey = Survey.objects.get(id=survey_id)
-    set_sentry_business_area_tag(survey.business_area.name)
-    if survey.category == Survey.CATEGORY_MANUAL:
-        return
-    phone_numbers = survey.recipients.filter(head_of_household__phone_no_valid=True).values_list(
-        "head_of_household__phone_no", flat=True
+    job = AsyncJob.objects.create(
+        owner=survey.created_by,
+        program=survey.program,
+        type=AsyncJobModel.JobType.JOB_TASK,
+        action="hope.apps.accountability.celery_tasks.send_survey_to_users_action",
+        config={"survey_id": str(survey_id)},
+        group_key=f"send_survey_to_users:{survey_id}",
+        description=f"Send survey to users for survey {survey_id}",
     )
-    if survey.category == Survey.CATEGORY_SMS:
-        api = RapidProAPI(survey.business_area.slug, RapidProAPI.MODE_MESSAGE)
-        api.broadcast_message(phone_numbers, survey.body)
-        return
-    business_area = BusinessArea.objects.get(id=survey.business_area_id)
-    api = RapidProAPI(business_area.slug, RapidProAPI.MODE_VERIFICATION)
-
-    already_received = [
-        phone_number
-        for successful_call in survey.successful_rapid_pro_calls
-        for phone_number in successful_call["urns"]
-    ]
-    phone_numbers = [phone_number for phone_number in phone_numbers if phone_number not in already_received]
-
-    successful_flows, error = api.start_flow(survey.flow_id, phone_numbers)
-
-    for successful_flow in successful_flows:
-        survey.successful_rapid_pro_calls.append(
-            {
-                "flow_uuid": successful_flow.response["uuid"],
-                "urns": list(map(str, successful_flow.urns)),
-            }
-        )
-    survey.save()
+    job.queue()