ethyca-fides 2.67.0rc2__py2.py3-none-any.whl → 2.67.0rc4__py2.py3-none-any.whl

fides/api/service/privacy_request/request_service.py

@@ -3,11 +3,12 @@ from __future__ import annotations
 import json
 from asyncio import sleep
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
 from httpx import AsyncClient
 from loguru import logger
 from sqlalchemy import text
+from sqlalchemy.orm import Session
 from sqlalchemy.sql.elements import TextClause
 
 from fides.api.common_exceptions import PrivacyRequestNotFound
@@ -31,6 +32,9 @@ from fides.api.util.cache import (
     celery_tasks_in_flight,
     get_async_task_tracking_cache_key,
     get_cache,
+    get_privacy_request_retry_count,
+    increment_privacy_request_retry_count,
+    reset_privacy_request_retry_count,
 )
 from fides.api.util.lock import redis_lock
 from fides.common.api.v1.urn_registry import PRIVACY_REQUESTS, V1_URL_PREFIX
@@ -350,10 +354,17 @@ def initiate_interrupted_task_requeue_poll() -> None:
 
 
 def get_cached_task_id(entity_id: str) -> Optional[str]:
-    """Gets the cached task ID for a privacy request or request task by ID."""
+    """Gets the cached task ID for a privacy request or request task by ID.
+
+    Raises Exception if cache operations fail, allowing callers to handle cache failures appropriately.
+    """
     cache: FidesopsRedis = get_cache()
-    task_id = cache.get(get_async_task_tracking_cache_key(entity_id))
-    return task_id
+    try:
+        task_id = cache.get(get_async_task_tracking_cache_key(entity_id))
+        return task_id
+    except Exception as exc:
+        logger.error(f"Failed to get cached task ID for entity {entity_id}: {exc}")
+        raise
 
 
 REQUEUE_INTERRUPTED_TASKS_LOCK = "requeue_interrupted_tasks_lock"
@@ -393,6 +404,115 @@ def _get_task_ids_from_dsr_queue(
     return queued_tasks_ids
 
 
+def _cancel_interrupted_tasks_and_error_privacy_request(
+    db: Session, privacy_request: PrivacyRequest, error_message: Optional[str] = None
+) -> None:
+    """
+    Cancel all tasks associated with an interrupted privacy request and set the privacy request to error state.
+
+    This function:
+    1. Logs the error message (either provided or default)
+    2. Revokes the main privacy request task and all associated request tasks
+    3. Sets the privacy request status to error
+    4. Creates an error log entry
+
+    Args:
+        db: Database session
+        privacy_request: The privacy request to cancel and error
+        error_message: Optional error message to log. If not provided, uses default message.
+    """
+    if error_message:
+        logger.error(error_message)
+    else:
+        logger.error(
+            f"Canceling interrupted tasks and marking privacy request {privacy_request.id} as error"
+        )
+
+    # Cancel all associated Celery tasks
+    privacy_request.cancel_celery_tasks()
+
+    # Set privacy request to error state using the existing method
+    try:
+        privacy_request.error_processing(db)
+        logger.info(
+            f"Privacy request {privacy_request.id} marked as error due to task interruption"
+        )
+    except Exception as exc:
+        logger.error(
+            f"Failed to mark privacy request {privacy_request.id} as error: {exc}"
+        )
+
+
+def _handle_privacy_request_requeue(
+    db: Session, privacy_request: PrivacyRequest
+) -> None:
+    """Handle retry logic for a privacy request - either requeue or cancel based on retry count."""
+    try:
+        # Check retry count and either requeue or cancel based on limit
+        current_retry_count = get_privacy_request_retry_count(privacy_request.id)
+        max_retries = CONFIG.execution.privacy_request_requeue_retry_count
+
+        if current_retry_count < max_retries:
+            # Increment retry count and attempt requeue
+            new_retry_count = increment_privacy_request_retry_count(privacy_request.id)
+            logger.info(
+                f"Requeuing privacy request {privacy_request.id} "
+                f"(attempt {new_retry_count}/{max_retries})"
+            )
+
+            from fides.service.privacy_request.privacy_request_service import (  # pylint: disable=cyclic-import
+                PrivacyRequestError,
+                _requeue_privacy_request,
+            )
+
+            try:
+                _requeue_privacy_request(db, privacy_request)
+            except PrivacyRequestError as exc:
+                # If requeue fails, cancel tasks and set to error state
+                _cancel_interrupted_tasks_and_error_privacy_request(
+                    db, privacy_request, exc.message
+                )
+        else:
+            # Exceeded retry limit, cancel tasks and set to error state
+            _cancel_interrupted_tasks_and_error_privacy_request(
+                db,
+                privacy_request,
+                f"Privacy request {privacy_request.id} exceeded max retry attempts "
+                f"({max_retries}), canceling tasks and setting to error state",
+            )
+            # Reset retry count since we're giving up
+            reset_privacy_request_retry_count(privacy_request.id)
+
+    except Exception as cache_exc:
+        # If cache operations fail (Redis down, network issues, etc.), fail safe by canceling
+        _cancel_interrupted_tasks_and_error_privacy_request(
+            db,
+            privacy_request,
+            f"Cache operation failed for privacy request {privacy_request.id}, "
+            f"failing safe by canceling tasks: {cache_exc}",
+        )
+
+
+def _get_request_task_ids_in_progress(
+    db: Session, privacy_request_id: str
+) -> List[str]:
+    """Get the IDs of request tasks that are currently in progress for a privacy request."""
+    request_tasks_in_progress = (
+        db.query(RequestTask.id)
+        .filter(RequestTask.privacy_request_id == privacy_request_id)
+        .filter(
+            RequestTask.status.in_(
+                [
+                    ExecutionLogStatus.in_processing,
+                    ExecutionLogStatus.pending,
+                ]
+            )
+        )
+        .all()
+    )
+    return [task[0] for task in request_tasks_in_progress]
+
+
 # pylint: disable=too-many-branches
 @celery_app.task(base=DatabaseTask, bind=True)
 def requeue_interrupted_tasks(self: DatabaseTask) -> None:
@@ -442,17 +562,40 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
             )
 
             # Get task IDs from the queue in a memory-efficient way
-            queued_tasks_ids = _get_task_ids_from_dsr_queue(redis_conn)
+            try:
+                queued_tasks_ids = _get_task_ids_from_dsr_queue(redis_conn)
+            except Exception as queue_exc:
+                logger.warning(
+                    f"Failed to get task IDs from queue, skipping queue state checks: {queue_exc}"
+                )
+                return
 
             # Check each privacy request
             for privacy_request in in_progress_requests:
                 should_requeue = False
                 logger.debug(f"Checking tasks for privacy request {privacy_request.id}")
 
-                task_id = get_cached_task_id(privacy_request.id)
+                try:
+                    task_id = get_cached_task_id(privacy_request.id)
+                except Exception as cache_exc:
+                    # If we can't get the task ID due to cache failure, fail safe by canceling
+                    _cancel_interrupted_tasks_and_error_privacy_request(
+                        db,
+                        privacy_request,
+                        f"Cache failure when getting task ID for privacy request {privacy_request.id}, "
+                        f"failing safe by canceling tasks: {cache_exc}",
+                    )
+                    continue
 
                 # If the task ID is not cached, we can't check if it's running
+                # This means the request is stuck - cancel it
                 if not task_id:
+                    _cancel_interrupted_tasks_and_error_privacy_request(
+                        db,
+                        privacy_request,
+                        f"No task ID found for privacy request {privacy_request.id}, "
+                        f"request is stuck without a running task - canceling",
+                    )
                     continue
 
                 # Check if the main privacy request task is active
@@ -470,30 +613,36 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
                         )
                         should_requeue = True
 
-                    request_tasks_in_progress = (
-                        db.query(RequestTask.id)
-                        .filter(RequestTask.privacy_request_id == privacy_request.id)
-                        .filter(
-                            RequestTask.status.in_(
-                                [
-                                    ExecutionLogStatus.in_processing,
-                                    ExecutionLogStatus.pending,
-                                ]
-                            )
-                        )
-                        .all()
+                    request_task_ids_in_progress = _get_request_task_ids_in_progress(
+                        db, privacy_request.id
                     )
-                    request_task_ids_in_progress = [
-                        task[0] for task in request_tasks_in_progress
-                    ]
 
                     # Check each individual request task
                     for request_task_id in request_task_ids_in_progress:
-                        subtask_id = get_cached_task_id(request_task_id)
+                        try:
+                            subtask_id = get_cached_task_id(request_task_id)
+                        except Exception as cache_exc:
+                            # If we can't get the subtask ID due to cache failure, fail safe by canceling
+                            _cancel_interrupted_tasks_and_error_privacy_request(
+                                db,
+                                privacy_request,
+                                f"Cache failure when getting subtask ID for request task {request_task_id} "
+                                f"(privacy request {privacy_request.id}), failing safe by canceling tasks: {cache_exc}",
+                            )
+                            should_requeue = False
+                            break
 
                         # If the task ID is not cached, we can't check if it's running
+                        # This means the subtask is stuck - cancel the entire privacy request
                         if not subtask_id:
-                            continue
+                            _cancel_interrupted_tasks_and_error_privacy_request(
+                                db,
+                                privacy_request,
+                                f"No task ID found for request task {request_task_id} "
+                                f"(privacy request {privacy_request.id}), subtask is stuck - canceling privacy request",
+                            )
+                            should_requeue = False
+                            break
 
                         if (
                             subtask_id not in queued_tasks_ids
@@ -507,12 +656,4 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
 
                 # Requeue the privacy request if needed
                 if should_requeue:
-                    from fides.service.privacy_request.privacy_request_service import (  # pylint: disable=cyclic-import
-                        PrivacyRequestError,
-                        _requeue_privacy_request,
-                    )
-
-                    try:
-                        _requeue_privacy_request(db, privacy_request)
-                    except PrivacyRequestError as exc:
-                        logger.error(exc.message)
+                    _handle_privacy_request_requeue(db, privacy_request)

fides/api/task/execute_request_tasks.py

@@ -36,6 +36,7 @@ from fides.api.tasks import DSR_QUEUE_NAME, DatabaseTask, celery_app
 from fides.api.util.cache import cache_task_tracking_key
 from fides.api.util.collection_util import Row
 from fides.api.util.logger_context_utils import LoggerContextKeys, log_context
+from fides.api.util.memory_watchdog import memory_limiter
 
 # DSR 3.0 task functions
 
@@ -255,6 +256,7 @@ def queue_downstream_tasks(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,
@@ -319,6 +321,7 @@ def run_access_node(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,
@@ -391,6 +394,7 @@ def run_erasure_node(
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
+@memory_limiter
 @log_context(
     capture_args={
         "privacy_request_id": LoggerContextKeys.privacy_request_id,

fides/api/task/graph_task.py

@@ -62,6 +62,7 @@ from fides.api.util.consent_util import (
 )
 from fides.api.util.logger import Pii
 from fides.api.util.logger_context_utils import LoggerContextKeys
+from fides.api.util.memory_watchdog import MemoryLimitExceeded
 from fides.api.util.saas_util import FIDESOPS_GROUPED_INPUTS
 from fides.config import CONFIG
 
@@ -71,6 +72,16 @@ EMPTY_REQUEST = PrivacyRequest()
 EMPTY_REQUEST_TASK = RequestTask()
 
 
+def _is_memory_limit_exceeded(exception: BaseException) -> bool:
+    """Check if the exception or any exception in its chain is a MemoryLimitExceeded."""
+    current_exception: Optional[BaseException] = exception
+    while current_exception:
+        if isinstance(current_exception, MemoryLimitExceeded):
+            return True
+        current_exception = current_exception.__cause__ or current_exception.__context__
+    return False
+
+
 def retry(
     action_type: ActionType,
     default_return: Any,
@@ -146,7 +157,31 @@ def retry(
                     self.log_skipped(action_type, exc)
                     self.cache_system_status_for_preferences()
                     return default_return
+                except MemoryLimitExceeded as ex:
+                    # Hard failure – mark task & downstream as errored and abort.
+                    logger.error(
+                        "Memory watchdog exceeded ({}%). Aborting {} {} without retry.",
+                        ex.memory_percent,
+                        method_name,
+                        self.execution_node.address,
+                    )
+                    # Persist error status and create execution logs before raising
+                    self.log_end(action_type, ex)
+                    self.add_error_status_for_consent_reporting()
+                    raise
                 except BaseException as ex:  # pylint: disable=W0703
+                    # Check if this exception was caused by memory limit exceeded
+                    if _is_memory_limit_exceeded(ex):
+                        logger.error(
+                            "Memory watchdog exceeded (wrapped exception). Aborting {} {} without retry.",
+                            method_name,
+                            self.execution_node.address,
+                        )
+                        # Persist error status and create execution logs before raising
+                        self.log_end(action_type, ex)
+                        self.add_error_status_for_consent_reporting()
+                        raise
+
                     traceback.print_exc()
                     func_delay *= CONFIG.execution.task_retry_backoff
                     logger.warning(
@@ -555,12 +590,21 @@ class GraphTask(ABC):  # pylint: disable=too-many-instance-attributes
         # For access request results, mutate rows in-place to remove non-matching
         # array elements.  We already iterated over `output` above, so reuse the same
         # loop structure to keep cache locality.
+        logger.info(
+            "Filtering {} rows in {} for matching array elements.",
+            len(output),
+            self.execution_node.address,
+        )
         for row in output:
+            filter_element_match(row, post_processed_node_input_data)
+
+        if len(output) > 0:
             logger.info(
-                "Filtering row in {} for matching array elements.",
+                "Filtering completed for {} rows in {}. Post-processed node size: {}",
+                len(output),
                 self.execution_node.address,
+                len(post_processed_node_input_data),
             )
-            filter_element_match(row, post_processed_node_input_data)
 
         if self.request_task.id:
             # Saves intermediate access results for DSR 3.0 directly on the Request Task

fides/api/util/cache.py

@@ -334,6 +334,62 @@ def cache_task_tracking_key(request_id: str, celery_task_id: str) -> None:
         )
 
 
+def get_privacy_request_retry_cache_key(privacy_request_id: str) -> str:
+    """Get cache key for tracking privacy request requeue retry attempts."""
+    return f"id-{privacy_request_id}-privacy-request-retry-count"
+
+
+def get_privacy_request_retry_count(privacy_request_id: str) -> int:
+    """Get the current retry count for a privacy request requeue attempts.
+
+    Raises Exception if cache operations fail, allowing callers to handle cache failures appropriately.
+    """
+    cache: FidesopsRedis = get_cache()
+    try:
+        retry_count = cache.get(get_privacy_request_retry_cache_key(privacy_request_id))
+        return int(retry_count) if retry_count else 0
+    except Exception as exc:
+        logger.error(
+            f"Failed to get retry count for privacy request {privacy_request_id}: {exc}"
+        )
+        raise
+
+
+def increment_privacy_request_retry_count(privacy_request_id: str) -> int:
+    """Increment and return the retry count for a privacy request requeue attempts.
+
+    Raises Exception if cache operations fail, allowing callers to handle cache failures appropriately.
+    """
+    cache: FidesopsRedis = get_cache()
+    cache_key = get_privacy_request_retry_cache_key(privacy_request_id)
+
+    try:
+        # Increment the counter, will be 1 if key doesn't exist
+        new_count = cache.incr(cache_key)
+        # Set expiry to prevent cache buildup (24 hours)
+        cache.expire(cache_key, 86400)
+        return new_count
+    except Exception as exc:
+        logger.error(
+            f"Failed to increment retry count for privacy request {privacy_request_id}: {exc}"
+        )
+        raise
+
+
+def reset_privacy_request_retry_count(privacy_request_id: str) -> None:
+    """Reset the retry count for a privacy request requeue attempts.
+
+    Silently fails if cache operations fail since this is cleanup.
+    """
+    cache: FidesopsRedis = get_cache()
+    try:
+        cache.delete(get_privacy_request_retry_cache_key(privacy_request_id))
+    except Exception as exc:
+        logger.warning(
+            f"Failed to reset retry count for privacy request {privacy_request_id}: {exc}"
+        )
+
+
 def celery_tasks_in_flight(celery_task_ids: List[str]) -> bool:
     """Returns True if supplied Celery Tasks appear to be in-flight"""
     if not celery_task_ids: