ethyca-fides 2.71.1b1__py2.py3-none-any.whl → 2.74.0rc1__py2.py3-none-any.whl

fides/api/service/privacy_request/request_service.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-from asyncio import sleep
 from datetime import datetime, timedelta
 from typing import Any, Dict, List, Optional, Set
 
@@ -11,7 +10,7 @@ from sqlalchemy import text
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.elements import TextClause
 
-from fides.api.common_exceptions import PrivacyRequestError, PrivacyRequestNotFound
+from fides.api.common_exceptions import PrivacyRequestError
 from fides.api.models.privacy_request import (
     EXITED_EXECUTION_LOG_STATUSES,
     PrivacyRequest,
@@ -21,10 +20,7 @@ from fides.api.models.privacy_request.request_task import AsyncTaskType
 from fides.api.models.worker_task import ExecutionLogStatus
 from fides.api.schemas.drp_privacy_request import DrpPrivacyRequestCreate
 from fides.api.schemas.policy import ActionType
-from fides.api.schemas.privacy_request import (
-    PrivacyRequestResponse,
-    PrivacyRequestStatus,
-)
+from fides.api.schemas.privacy_request import PrivacyRequestStatus
 from fides.api.schemas.redis_cache import Identity
 from fides.api.tasks import DSR_QUEUE_NAME, DatabaseTask, celery_app
 from fides.api.tasks.scheduled.scheduler import scheduler
@@ -38,7 +34,6 @@ from fides.api.util.cache import (
     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
 from fides.config import CONFIG
 
 PRIVACY_REQUEST_STATUS_CHANGE_POLL = "privacy_request_status_change_poll"
@@ -98,62 +93,6 @@ def get_async_client() -> AsyncClient:
     return AsyncClient()
 
 
-async def poll_server_for_completion(
-    privacy_request_id: str,
-    server_url: str,
-    token: str,
-    *,
-    poll_interval_seconds: int = 30,
-    timeout_seconds: int = 1800,  # 30 minutes
-    client: AsyncClient | None = None,
-) -> PrivacyRequestResponse:
-    """Poll a server for privacy request completion.
-
-    Requests will report complete with if they have a status of canceled, complete,
-    denied, or error. By default the polling will time out if not completed in 30
-    minutes, time can be overridden by setting the timeout_seconds.
-    """
-    url = (
-        f"{server_url}{V1_URL_PREFIX}{PRIVACY_REQUESTS}?request_id={privacy_request_id}"
-    )
-    start_time = datetime.now()
-    elapsed_time = 0.0
-    while elapsed_time < timeout_seconds:
-        if client:
-            response = await client.get(
-                url, headers={"Authorization": f"Bearer {token}"}
-            )
-        else:
-            async_client = get_async_client()
-            response = await async_client.get(
-                url, headers={"Authorization": f"Bearer {token}"}
-            )
-        response.raise_for_status()
-
-        # Privacy requests are returned paginated. Since this is searching for a specific
-        # privacy request there should only be one value present in items.
-        items = response.json()["items"]
-        if not items:
-            raise PrivacyRequestNotFound(
-                f"No privacy request found with id '{privacy_request_id}'"
-            )
-        status = PrivacyRequestResponse(**items[0])
-        if status.status and status.status in (
-            PrivacyRequestStatus.complete,
-            PrivacyRequestStatus.canceled,
-            PrivacyRequestStatus.error,
-            PrivacyRequestStatus.denied,
-        ):
-            return status
-
-        await sleep(poll_interval_seconds)
-        time_delta = datetime.now() - start_time
-        elapsed_time = time_delta.seconds
-    raise TimeoutError(
-        f"Timeout of {timeout_seconds} seconds has been exceeded while waiting for privacy request {privacy_request_id}"
-    )
-
-
 def initiate_poll_for_exited_privacy_request_tasks() -> None:
     """Initiates scheduler to check if a Privacy Request's status needs to be flipped when all
     Request Tasks have had a chance to run"""
@@ -364,7 +303,7 @@ def initiate_interrupted_task_requeue_poll() -> None:
     )
 
 
-def initiate_async_tasks_status_polling() -> None:
+def initiate_polling_task_requeue() -> None:
     """Initiates scheduler to check for and requeue pending polling async tasks"""
     if CONFIG.test_mode:
         return
@@ -375,13 +314,13 @@ def initiate_async_tasks_status_polling() -> None:
 
     logger.info("Initiating scheduler for async tasks status polling")
     scheduler.add_job(
-        func=poll_async_tasks_status,
+        func=requeue_polling_tasks,
         trigger="interval",
         kwargs={},
         id=ASYNC_TASKS_STATUS_POLLING,
         coalesce=True,
         replace_existing=True,
-        seconds=CONFIG.execution.async_tasks_status_polling_interval_seconds,
+        seconds=CONFIG.execution.async_polling_interval_hours * 3600,
     )
 
 
@@ -541,6 +480,27 @@ def _get_request_task_ids_in_progress(
     return [task[0] for task in request_tasks_in_progress]
 
 
+def _has_polling_tasks(db: Session, privacy_request_id: str) -> bool:
+    """
+    Check if a privacy request has any polling async tasks.
+
+    Args:
+        db: Database session
+        privacy_request_id: The ID of the privacy request to check
+
+    Returns:
+        bool: True if the privacy request has polling tasks, False otherwise
+    """
+    return db.query(
+        db.query(RequestTask)
+        .filter(
+            RequestTask.privacy_request_id == privacy_request_id,
+            RequestTask.async_type == AsyncTaskType.polling,
+        )
+        .exists()
+    ).scalar()
+
+
 # pylint: disable=too-many-branches
 # pylint: disable=too-many-statements
 @celery_app.task(base=DatabaseTask, bind=True)
@@ -680,6 +640,17 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
                                 should_requeue = False
                                 break
 
+                            # Check if the Privacy request has polling tasks
+                            if _has_polling_tasks(db, privacy_request.id):
+                                # If the polling request task has no cached task ID, it's stuck
+                                logger.warning(
+                                    f"No task ID found for request task {request_task_id} "
+                                    f"(privacy request {privacy_request.id}) Contains polling tasks - "
+                                    f"keeping request in current status as it may be waiting for polling task to complete"
+                                )
+                                should_requeue = False
+                                break
+
                             # For other statuses, cancel the entire privacy request
                             _cancel_interrupted_tasks_and_error_privacy_request(
                                 db,
@@ -706,7 +677,7 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
 
 
 @celery_app.task(base=DatabaseTask, bind=True)
-def poll_async_tasks_status(self: DatabaseTask) -> None:
+def requeue_polling_tasks(self: DatabaseTask) -> None:
     """
     Poll the status of async tasks that are awaiting processing.
     """
@@ -720,19 +691,21 @@ def poll_async_tasks_status(self: DatabaseTask) -> None:
         with self.get_new_session() as db:
             logger.debug("Polling for async tasks status")
 
-            # Get all tasks that are awaiting processing and are from polling async tasks
+            # Get all tasks that are polling and are from polling async tasks
             async_tasks = (
                 db.query(RequestTask)
-                .filter(RequestTask.status == ExecutionLogStatus.awaiting_processing)
+                .filter(RequestTask.status == ExecutionLogStatus.polling)
                 .filter(RequestTask.async_type == AsyncTaskType.polling)
                 .all()
             )
-            logger.info(f"Found {len(async_tasks)} async tasks awaiting processing")
+            logger.info(f"Found {len(async_tasks)} async polling tasks")
 
             if async_tasks:
-                from fides.api.service.async_dsr.async_dsr_service import (  # pylint: disable=cyclic-import
-                    requeue_polling_request,
-                )
+                # Avoiding cyclic imports
+                from fides.api.task.execute_request_tasks import queue_request_task
 
                 for async_task in async_tasks:
-                    requeue_polling_request(db, async_task)
+                    logger.info(
+                        f"Requeuing polling task {async_task.id} for processing"
+                    )
+                    queue_request_task(async_task, privacy_request_proceed=True)

fides/api/service/saas_request/saas_request_override_factory.py

@@ -1,6 +1,6 @@
 from enum import Enum
 from inspect import Signature, signature
-from typing import Callable, Dict, List, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 
 from loguru import logger
 
@@ -12,6 +12,9 @@ from fides.api.schemas.consentable_item import ConsentableItem, ConsentWebhookRe
 from fides.api.schemas.saas.shared_schemas import ConsentPropagationStatus
 from fides.api.util.collection_util import Row
 
+if TYPE_CHECKING:
+    from fides.api.schemas.saas.async_polling_configuration import PollingResult
+
 
 # at some point this should likely be formalized more centrally...
 class SaaSRequestType(Enum):
@@ -30,6 +33,10 @@ class SaaSRequestType(Enum):
     UPDATE_CONSENT = "update_consent"
     PROCESS_CONSENT_WEBHOOK = "process_consent_webhook"
 
+    # Async polling request types
+    POLLING_STATUS = "polling_status"
+    POLLING_RESULT = "polling_result"
+
 
 RequestOverrideFunction = Callable[
     ...,
@@ -39,6 +46,8 @@ RequestOverrideFunction = Callable[
         List[Row],
         ConsentPropagationStatus,
         int,
+        bool,  # For polling status overrides
+        "PollingResult",  # For polling result overrides - string literal to avoid circular import
         None,
     ],
 ]
@@ -105,6 +114,10 @@ class SaaSRequestOverrideFactory:
                     validate_update_consent_function(override_function)
                 elif request_type == SaaSRequestType.PROCESS_CONSENT_WEBHOOK:
                     validate_process_consent_webhook_function(override_function)
+                elif request_type == SaaSRequestType.POLLING_STATUS:
+                    validate_polling_status_override_function(override_function)
+                elif request_type == SaaSRequestType.POLLING_RESULT:
+                    validate_polling_result_override_function(override_function)
                 else:
                     raise ValueError(
                         f"Invalid SaaSRequestType '{request_type}' provided for SaaS request override function"
@@ -264,5 +277,62 @@ def validate_process_consent_webhook_function(f: Callable) -> None:
         )
 
 
+def validate_polling_status_override_function(f: Callable) -> None:
+    """
+    Perform some basic checks on the user-provided SaaS request override function
+    that will be used for polling status checks.
+
+    The validation is not overly strict to allow for some flexibility in
+    the functions that are used for overrides, but we check to ensure that
+    the function meets the framework's basic expectations.
+
+    Specifically, the validation checks that function's return type is `bool`
+    and that it declares at least 4 parameters.
+    """
+    sig: Signature = signature(f)
+    if sig.return_annotation is not bool:
+        raise InvalidSaaSRequestOverrideException(
+            "Polling status override function must return a bool"
+        )
+    if len(sig.parameters) < 4:
+        raise InvalidSaaSRequestOverrideException(
+            "Polling status override function must declare at least 4 parameters"
+        )
+
+
+def validate_polling_result_override_function(f: Callable) -> None:
+    """
+    Perform some basic checks on the user-provided SaaS request override function
+    that will be used for polling result retrieval.
+
+    The validation is not overly strict to allow for some flexibility in
+    the functions that are used for overrides, but we check to ensure that
+    the function meets the framework's basic expectations.
+
+    Specifically, the validation checks that function's return type is `PollingResult`
+    and that it declares at least 4 parameters.
+    """
+    sig: Signature = signature(f)
+
+    # Import PollingResult here to avoid circular imports at module level
+    from fides.api.schemas.saas.async_polling_configuration import PollingResult
+
+    # Check return type annotation - handle both direct class and string literals
+    return_annotation = sig.return_annotation
+    if return_annotation not in (
+        PollingResult,
+        "PollingResult",
+        Optional[PollingResult],
+        "Optional[PollingResult]",
+    ):
+        raise InvalidSaaSRequestOverrideException(
+            "Polling result override function must return a PollingResult"
+        )
+    if len(sig.parameters) < 4:
+        raise InvalidSaaSRequestOverrideException(
+            "Polling result override function must declare at least 4 parameters"
+        )
+
+
 # TODO: Avoid running this on import?
 register = SaaSRequestOverrideFactory.register

fides/api/task/execute_request_tasks.py

@@ -22,7 +22,10 @@ from fides.api.common_exceptions import (
 from fides.api.graph.config import TERMINATOR_ADDRESS, CollectionAddress
 from fides.api.models.connectionconfig import ConnectionConfig
 from fides.api.models.privacy_request import ExecutionLog, PrivacyRequest, RequestTask
-from fides.api.models.worker_task import ExecutionLogStatus
+from fides.api.models.worker_task import (
+    RESUMABLE_EXECUTION_LOG_STATUSES,
+    ExecutionLogStatus,
+)
 from fides.api.schemas.policy import ActionType, CurrentStep
 from fides.api.schemas.privacy_request import PrivacyRequestStatus
 from fides.api.task.graph_task import (
@@ -162,7 +165,7 @@ def can_run_task_body(
     if request_task.is_root_task:
         # Shouldn't be possible but adding as a catch-all
         return False
-    if request_task.status != ExecutionLogStatus.pending:
+    if request_task.status not in RESUMABLE_EXECUTION_LOG_STATUSES:
         logger_method(request_task)(
             "Skipping {} task {} with status {}.",
             request_task.action_type,
@@ -274,15 +277,18 @@ def run_access_node(
 
     try:
         with self.get_new_session() as session:
-            privacy_request, request_task, upstream_results = (
-                run_prerequisite_task_checks(
-                    session, privacy_request_id, privacy_request_task_id
-                )
+            (
+                privacy_request,
+                request_task,
+                upstream_results,
+            ) = run_prerequisite_task_checks(
+                session, privacy_request_id, privacy_request_task_id
             )
             with logger.contextualize(
                 privacy_request_source=(
                     privacy_request.source.value if privacy_request.source else None
-                )
+                ),
+                collection=request_task.collection_address,
             ):
                 log_task_starting(request_task)
 
@@ -343,7 +349,8 @@ def run_erasure_node(
         with logger.contextualize(
             privacy_request_source=(
                 privacy_request.source.value if privacy_request.source else None
-            )
+            ),
+            collection=request_task.collection_address,
         ):
             log_task_starting(request_task)
 
@@ -416,7 +423,8 @@ def run_consent_node(
         with logger.contextualize(
             privacy_request_source=(
                 privacy_request.source.value if privacy_request.source else None
-            )
+            ),
+            collection=request_task.collection_address,
         ):
             log_task_starting(request_task)
 

fides/api/task/filter_results.py

@@ -6,10 +6,43 @@ from loguru import logger
 
 from fides.api.graph.config import CollectionAddress, FieldPath
 from fides.api.graph.graph import DatasetGraph
+from fides.api.service.storage.util import is_attachment_field
 from fides.api.task.manual.manual_task_address import ManualTaskAddress
 from fides.api.util.collection_util import Row
 
 
+def should_skip_data_category_filtering(
+    node_address: str, results: List[Dict[str, Optional[Any]]]
+) -> bool:
+    """
+    Determine if results should skip data category filtering.
+
+    This includes:
+    1. Manual task data - controlled by field definitions, not data categories
+    2. Attachment data - should be preserved regardless of data categories (from polling or manual tasks)
+
+    :param node_address: The collection address being processed
+    :param results: The results data to check
+    :return: True if filtering should be skipped, False otherwise
+    """
+    # Check if this is a manual task address
+    if ManualTaskAddress.is_manual_task_address(node_address):
+        return True
+
+    # Check if results contain attachment data
+    if not results:
+        return False
+
+    # Verify it looks like actual attachment data
+    for result in results:
+        if isinstance(result, dict) and any(
+            is_attachment_field(value) for value in result.values()
+        ):
+            return True
+
+    return False
+
+
 def filter_data_categories(
     access_request_results: Dict[str, List[Dict[str, Optional[Any]]]],
     target_categories: Set[str],
@@ -38,8 +71,8 @@ def filter_data_categories(
         if not results:
             continue
 
-        # Skip manual task data - it doesn't need filtering since it's controlled by field definitions
-        if ManualTaskAddress.is_manual_task_address(node_address):
+        # Skip data that doesn't need filtering (manual tasks and attachment data)
+        if should_skip_data_category_filtering(node_address, results):
             filtered_access_results[node_address].extend(results)
             continue
 

fides/api/task/graph_task.py

@@ -13,6 +13,7 @@ from sqlalchemy.orm import Session
 from fides.api.api.deps import get_autoclose_db_session as get_db
 from fides.api.common_exceptions import (
     ActionDisabled,
+    AwaitingAsyncProcessing,
     AwaitingAsyncTask,
     CollectionDisabled,
     NotSupportedForCollection,
@@ -61,7 +62,6 @@ from fides.api.util.collection_util import (
 from fides.api.util.consent_util import (
     add_errored_system_status_for_consent_reporting_on_preferences,
 )
-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
@@ -96,6 +96,7 @@ def retry(
     If we exceed the number of TASK_RETRY_COUNT retries, we re-raise the exception to stop execution of the privacy request.
     """
 
+    # pylint: disable=too-many-return-statements
     def decorator(func: Callable) -> Callable:
         @wraps(func)
         def result(*args: Any, **kwargs: Any) -> Any:
@@ -115,6 +116,17 @@ def retry(
                         self.log_start(action_type)
                     # Run access or erasure request
                     return func(*args, **kwargs)
+                except AwaitingAsyncProcessing as ex:
+                    logger.warning(
+                        "Request Task {} {} {} async processing in progress",
+                        self.request_task.id if self.request_task.id else None,
+                        method_name,
+                        self.execution_node.address,
+                    )
+                    # Log the async processing status and exit without retrying.
+                    self.log_async_processing(action_type, ex)
+                    # Request Task put in "polling" status and exited, external system is processing
+                    return None
                 except AwaitingAsyncTask as ex:
                     logger.warning(
                         "Request Task {} {} {} awaiting async continuation",
@@ -416,9 +428,13 @@ class GraphTask(ABC):  # pylint: disable=too-many-instance-attributes
             # created to keep these in sync.
             # TODO remove conditional above alongside deprecating DSR 2.0
             if self.request_task.id:
-                # Merge the request task to ensure we keep the changes made
-                # to self.request_task that haven't yet been committed to the database
+                # Merge the request_task into the current session to make it persistent,
+                # then refresh its `async_type` to load the latest state from the
+                # database. This is crucial for async tasks where `async_type` might be
+                # updated by another process, and avoids overwriting local data like
+                # `access_data`.
                 request_task = db.merge(self.request_task)
+                db.refresh(request_task, attribute_names=["async_type"])
                 request_task.update_status(db, status)
                 self.request_task = request_task
 
@@ -453,6 +469,21 @@ class GraphTask(ABC):  # pylint: disable=too-many-instance-attributes
             message, [], action_type, ExecutionLogStatus.awaiting_processing
         )
 
+    def log_async_processing(
+        self,
+        action_type: ActionType,
+        ex: Optional[BaseException],
+        extra_message: Optional[str] = None,
+    ) -> None:
+        """On async processing activities - external system is actively processing, Fides is polling"""
+        logger.info("Polling for async results on node {}", self.key)
+
+        message = str(ex)
+        if extra_message:
+            message = f"{message}. {extra_message}"
+
+        self.update_status(message, [], action_type, ExecutionLogStatus.polling)
+
     def log_skipped(self, action_type: ActionType, ex: str) -> None:
         """Log that a collection was skipped.  For now, this is because a collection has been disabled."""
         logger.info("Skipping node {}", self.key)
@@ -469,11 +500,11 @@ class GraphTask(ABC):  # pylint: disable=too-many-instance-attributes
     ) -> None:
         """On completion activities"""
         if ex:
-            logger.warning(
-                "Ending {}, {} with failure {}",
+            logger.error(
+                "Ending {}, {} with failure: {}",
                 self.resources.request.id,
                 self.key,
-                Pii(ex),
+                ex,
             )
             self.update_status(str(ex), [], action_type, ExecutionLogStatus.error)
             # For DSR 3.0, Hooking into the GraphTask.log_end method to also mark the current
@@ -854,7 +885,6 @@ class GraphTask(ABC):  # pylint: disable=too-many-instance-attributes
         privacy_request_id = self.resources.request.id
 
         with get_db() as db:
-
             privacy_preferences = db.query(PrivacyPreferenceHistory).filter(
                 PrivacyPreferenceHistory.privacy_request_id == privacy_request_id
             )
@@ -975,7 +1005,6 @@ def build_affected_field_logs(
     policy_id = policy.id
 
     with get_db() as db:
-
         rules = db.query(Rule).filter(Rule.policy_id == policy_id)
 
         targeted_field_paths: Dict[FieldAddress, str] = {}

fides/api/tasks/__init__.py

@@ -86,7 +86,6 @@ class DatabaseTask(Task):  # pylint: disable=W0223
         # to prevent session overlap when requests are executing concurrently
         # when in task_always_eager mode (i.e. without proper workers)
         new_session = self._sessionmaker()
-        logger.debug(f"DatabaseTaskSession ID: {id(new_session)}. Self ID: {id(self)}")
         return new_session