Re-queue interrupted Celery tasks (#5800)

Co-authored-by: Robert Keyser <[email protected]>

Author: galvana

CHANGELOG.md

@@ -26,6 +26,7 @@ Changes can also be flagged with a GitHub label for tracking purposes. The URL o
 - Added support for queue-specific Celery workers [#5761](https://github.com/ethyca/fides/pull/5761)
 - Added support for AWS SES as an email provider [#5804](https://github.com/ethyca/fides/pull/5804)
 - Nested identity query support for BigQuery [#5814](https://github.com/ethyca/fides/pull/5814)
+- Added job that automatically requeues interrupted tasks for in progress privacy requests [#5800](https://github.com/ethyca/fides/pull/5800)
 
 ### Changed
 - Improved dataset validation for namespace metadata and dataset reachability [#5744](https://github.com/ethyca/fides/pull/5744)

docker-compose.yml

@@ -191,6 +191,18 @@ services:
         target: /usr/local/etc/redis
         read_only: False
 
+  flower:
+    image: mher/flower:2.0
+    command: celery flower --broker=redis://:redispassword@redis:6379/0
+    ports:
+      - "5555:5555"
+    depends_on:
+      - redis
+      - worker
+    environment:
+      - CELERY_BROKER_URL=redis://:redispassword@redis:6379/0
+      - CELERY_RESULT_BACKEND=redis://:redispassword@redis:6379/0
+
 volumes:
   postgres: null
 

noxfiles/dev_nox.py

@@ -36,6 +36,7 @@ def shell(session: Session) -> None:
         )
 
 
+# pylint: disable=too-many-branches
 @nox_session()
 def dev(session: Session) -> None:
     """
@@ -51,6 +52,7 @@ def dev(session: Session) -> None:
         - pc = Build and run the Privacy Center
         - remote_debug = Run with remote debugging enabled (see docker-compose.remote-debug.yml)
         - worker = Run a Fides worker
+        - flower = Run Flower monitoring dashboard for Celery
         - child = Run a Fides child node
         - <datastore(s)> = Run a test datastore (e.g. 'mssql', 'mongodb')
 
@@ -61,21 +63,20 @@ def dev(session: Session) -> None:
     build(session, "dev")
     session.notify("teardown")
 
-    if "worker" in session.posargs:
-        session.run("docker", "compose", "up", "--wait", "worker", external=True)
+    workers = ["worker", "worker-privacy-preferences", "worker-dsr"]
 
-    if "worker-privacy-preferences" in session.posargs:
-        session.run(
-            "docker",
-            "compose",
-            "up",
-            "--wait",
-            "worker-privacy-preferences",
-            external=True,
-        )
+    for worker in workers:
+        if worker in session.posargs:
+            session.run("docker", "compose", "up", "--wait", worker, external=True)
 
-    if "worker-dsr" in session.posargs:
-        session.run("docker", "compose", "up", "--wait", "worker-dsr", external=True)
+    if "flower" in session.posargs:
+        # Only start Flower if worker is also enabled
+        if any(worker in session.posargs for worker in workers):
+            session.run("docker", "compose", "up", "-d", "flower", external=True)
+        else:
+            session.error(
+                "Flower requires the worker service. Please add 'worker' to your arguments."
+            )
 
     datastores = [
         datastore for datastore in session.posargs if datastore in ALL_DATASTORES

requirements.txt

@@ -19,6 +19,7 @@ fastapi[all]==0.111.0
 fastapi-pagination[sqlalchemy]==0.12.25
 fideslog==1.2.10
 firebase-admin==5.3.0
+flower==2.0.1
 GitPython==3.1.41
 httpx==0.23.1
 hvac==0.11.2

src/fides/api/api/v1/endpoints/privacy_request_endpoints.py

@@ -84,7 +84,7 @@
 )
 from fides.api.schemas.dataset import CollectionAddressResponse, DryRunDatasetResponse
 from fides.api.schemas.external_https import PrivacyRequestResumeFormat
-from fides.api.schemas.policy import ActionType, CurrentStep
+from fides.api.schemas.policy import ActionType
 from fides.api.schemas.privacy_request import (
     BulkPostPrivacyRequests,
     BulkReviewResponse,
@@ -170,7 +170,10 @@
 )
 from fides.service.messaging.messaging_service import MessagingService
 from fides.service.privacy_request.privacy_request_service import (
+    PrivacyRequestError,
     PrivacyRequestService,
+    _process_privacy_request_restart,
+    _requeue_privacy_request,
     _trigger_pre_approval_webhooks,
     queue_privacy_request,
 )
@@ -1867,35 +1870,6 @@ def resume_privacy_request_from_requires_input(
     return privacy_request  # type: ignore[return-value]
 
 
-def _process_privacy_request_restart(
-    privacy_request: PrivacyRequest,
-    failed_step: Optional[CurrentStep],
-    db: Session,
-) -> PrivacyRequestResponse:
-    """If failed_step is provided, restart the DSR within that step. Otherwise,
-    restart the privacy request from the beginning."""
-    if failed_step:
-        logger.info(
-            "Restarting failed privacy request '{}' from '{}'",
-            privacy_request.id,
-            failed_step,
-        )
-    else:
-        logger.info(
-            "Restarting failed privacy request '{}' from the beginning",
-            privacy_request.id,
-        )
-
-    privacy_request.status = PrivacyRequestStatus.in_processing
-    privacy_request.save(db=db)
-    queue_privacy_request(
-        privacy_request_id=privacy_request.id,
-        from_step=failed_step.value if failed_step else None,
-    )
-
-    return privacy_request  # type: ignore[return-value]
-
-
 @router.get(
     REQUEST_TASKS,
     dependencies=[Security(verify_oauth_client, scopes=[PRIVACY_REQUEST_READ])],
@@ -1936,52 +1910,18 @@ def requeue_privacy_request(
 
     Don't use this unless the Privacy Request is stuck.
     """
-    pr: PrivacyRequest = get_privacy_request_or_error(db, privacy_request_id)
+    privacy_request: PrivacyRequest = get_privacy_request_or_error(
+        db, privacy_request_id
+    )
 
-    if pr.status not in [
-        PrivacyRequestStatus.approved,
-        PrivacyRequestStatus.in_processing,
-    ]:
+    try:
+        return _requeue_privacy_request(db, privacy_request)
+    except PrivacyRequestError as exc:
         raise HTTPException(
             status_code=HTTP_400_BAD_REQUEST,
-            detail=f"Request failed. Cannot re-queue privacy request {pr.id} with status {pr.status.value}",
+            detail=exc.message,
         )
 
-    # Both DSR 2.0 and 3.0 cache checkpoint details
-    checkpoint_details: Optional[CheckpointActionRequired] = (
-        pr.get_failed_checkpoint_details()
-    )
-    resume_step = checkpoint_details.step if checkpoint_details else None
-
-    # DSR 3.0 additionally stores Request Tasks in the application db that can be used to infer
-    # a resume checkpoint in the event the cache has expired.
-    if not resume_step and pr.request_tasks.count():
-        if pr.consent_tasks.count():
-            resume_step = CurrentStep.consent
-        elif pr.erasure_tasks.count():
-            # Checking if access terminator task was completed, because erasure tasks are created
-            # at the same time as the access tasks
-            terminator_access_task = pr.get_terminate_task_by_action(ActionType.access)
-            resume_step = (
-                CurrentStep.erasure
-                if terminator_access_task.status == ExecutionLogStatus.complete
-                else CurrentStep.access
-            )
-        elif pr.access_tasks.count():
-            resume_step = CurrentStep.access
-
-    logger.info(
-        "Manually re-queuing Privacy Request {} from step {}",
-        pr,
-        resume_step.value if resume_step else None,
-    )
-
-    return _process_privacy_request_restart(
-        pr,
-        resume_step,
-        db,
-    )
-
 
 @router.post(
     REQUEST_TASK_CALLBACK,

src/fides/api/main.py

@@ -40,9 +40,13 @@
     initiate_scheduled_batch_email_send,
 )
 from fides.api.service.privacy_request.request_service import (
+    initiate_interrupted_task_requeue_poll,
     initiate_poll_for_exited_privacy_request_tasks,
     initiate_scheduled_dsr_data_removal,
 )
+
+# pylint: disable=wildcard-import, unused-wildcard-import
+from fides.api.service.saas_request.override_implementations import *
 from fides.api.tasks.scheduled.scheduler import async_scheduler, scheduler
 from fides.api.ui import (
     get_admin_index_as_response,
@@ -91,6 +95,7 @@ async def lifespan(wrapped_app: FastAPI) -> AsyncGenerator[None, None]:
     initiate_scheduled_batch_email_send()
     initiate_poll_for_exited_privacy_request_tasks()
     initiate_scheduled_dsr_data_removal()
+    initiate_interrupted_task_requeue_poll()
     initiate_bcrypt_migration_task()
 
     logger.debug("Sending startup analytics events...")