ethyca-fides 2.71.1b1__py2.py3-none-any.whl → 2.71.2b0__py2.py3-none-any.whl

fides/api/service/async_dsr/handlers/polling_sub_request_handler.py

@@ -0,0 +1,123 @@
+"""Manager for polling sub-request operations and lifecycle management."""
+
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict
+
+from loguru import logger
+from sqlalchemy.orm import Session
+
+from fides.api.common_exceptions import PrivacyRequestError
+from fides.api.models.privacy_request.request_task import (
+    RequestTask,
+    RequestTaskSubRequest,
+)
+from fides.api.models.worker_task import ExecutionLogStatus
+
+
+class PollingSubRequestHandler:
+    """Utility class for managing polling sub-request lifecycle and status checking."""
+
+    @staticmethod
+    def create_sub_request(
+        session: Session,
+        request_task: RequestTask,
+        param_values_map: Dict[str, Any],
+    ) -> RequestTaskSubRequest:
+        """
+        Create a new sub-request for tracking async polling operations.
+
+        Args:
+            session: Database session
+            request_task: The parent request task
+            param_values_map: Parameter values including correlation_id for polling
+
+        Returns:
+            RequestTaskSubRequest: The created sub-request
+        """
+        sub_request = RequestTaskSubRequest.create(
+            session,
+            data={
+                "request_task_id": request_task.id,
+                "param_values": param_values_map,
+                "status": ExecutionLogStatus.pending.value,
+            },
+        )
+        logger.info(
+            f"Created sub-request {sub_request.id} for task '{request_task.id}'"
+        )
+        return sub_request
+
+    @staticmethod
+    def check_completion(polling_task: RequestTask) -> bool:
+        """
+        Check if all sub-requests for a polling task are complete.
+
+        Args:
+            polling_task: The polling task to check
+
+        Returns:
+            bool: True if all sub-requests are complete, False if still in progress
+        """
+        # Get all sub-requests and categorize by status
+        all_sub_requests = polling_task.sub_requests
+        completed_sub_requests = [
+            sub_request
+            for sub_request in all_sub_requests
+            if sub_request.status == ExecutionLogStatus.complete.value
+        ]
+        failed_sub_requests = [
+            sub_request
+            for sub_request in all_sub_requests
+            if sub_request.status == ExecutionLogStatus.error.value
+        ]
+
+        if (
+            len(completed_sub_requests) == len(all_sub_requests)
+            and len(all_sub_requests) > 0
+        ):
+            # All sub-requests completed successfully - aggregate results
+            logger.info(
+                f"All sub-requests completed successfully for task {polling_task.id}"
+            )
+            return True
+
+        # Still polling - some sub-requests are pending
+        logger.info(
+            f"Polling task {polling_task.id}: {len(completed_sub_requests)}/{len(all_sub_requests)} sub-requests complete, {len(failed_sub_requests)} failed"
+        )
+        return False
+
+    @staticmethod
+    def check_timeout(polling_task: RequestTask, timeout_days: int) -> None:
+        """
+        Check if any sub-requests have exceeded the polling timeout.
+
+        Args:
+            polling_task: The polling task to check
+            timeout_days: Timeout threshold in days
+
+        Raises:
+            PrivacyRequestError: If any sub-request has timed out
+        """
+        timeout_seconds = timeout_days * 24 * 60 * 60  # Convert days to seconds
+
+        # Check timeout for incomplete sub-requests only
+        timed_out_sub_requests = []
+
+        for sub_request in polling_task.sub_requests:
+            if sub_request.status != ExecutionLogStatus.complete.value:
+                # Check if this sub-request has timed out
+                if sub_request.created_at:
+                    timeout_threshold = sub_request.created_at + timedelta(
+                        seconds=timeout_seconds
+                    )
+                    current_time = datetime.now(timezone.utc)
+                    if current_time > timeout_threshold:
+                        timed_out_sub_requests.append(sub_request)
+
+        if timed_out_sub_requests:
+            sub_request_ids = [sr.id for sr in timed_out_sub_requests]
+            raise PrivacyRequestError(
+                f"Polling timeout exceeded for sub-requests {sub_request_ids} "
+                f"in task {polling_task.id}. Timeout interval: {timeout_days} days"
+            )

fides/api/service/async_dsr/strategies/async_dsr_strategy.py

@@ -0,0 +1,52 @@
+from abc import abstractmethod
+from typing import Any, Dict, List
+
+from sqlalchemy.orm import Session
+
+from fides.api.models.privacy_request.request_task import AsyncTaskType, RequestTask
+from fides.api.service.connectors.query_configs.saas_query_config import SaaSQueryConfig
+from fides.api.service.connectors.saas.authenticated_client import AuthenticatedClient
+from fides.api.service.strategy import Strategy
+from fides.api.util.collection_util import Row
+
+
+class AsyncDSRStrategy(Strategy):
+    """Abstract base class for async DSR strategies"""
+
+    type: AsyncTaskType
+    session: Session
+
+    @abstractmethod
+    def async_retrieve_data(
+        self,
+        client: AuthenticatedClient,
+        request_task_id: str,
+        query_config: SaaSQueryConfig,
+        input_data: Dict[str, List[Any]],
+    ) -> List[Row]:
+        """
+        Execute async retrieve data.
+        """
+
+    @abstractmethod
+    def async_mask_data(
+        self,
+        client: AuthenticatedClient,
+        request_task_id: str,
+        query_config: SaaSQueryConfig,
+        rows: List[Row],
+    ) -> int:
+        """
+        Execute async mask data.
+        """
+
+    def _get_request_task(self, request_task_id: str) -> RequestTask:
+        """Get request task by ID or raise ValueError if not found."""
+        request_task = (
+            self.session.query(RequestTask)
+            .filter(RequestTask.id == request_task_id)
+            .first()
+        )
+        if not request_task:
+            raise ValueError(f"RequestTask with ID {request_task_id} not found")
+        return request_task

fides/api/service/async_dsr/strategies/async_dsr_strategy_callback.py

@@ -0,0 +1,199 @@
+from typing import Any, Dict, List
+
+from loguru import logger
+from sqlalchemy.orm import Session
+
+from fides.api.common_exceptions import AwaitingAsyncTask
+from fides.api.models.privacy_request.request_task import AsyncTaskType, RequestTask
+from fides.api.service.async_dsr.strategies.async_dsr_strategy import AsyncDSRStrategy
+from fides.api.service.async_dsr.utils import AsyncPhase, get_async_phase
+from fides.api.service.connectors.query_configs.saas_query_config import SaaSQueryConfig
+from fides.api.service.connectors.saas.authenticated_client import AuthenticatedClient
+from fides.api.util.collection_util import Row
+
+
+class AsyncCallbackStrategy(AsyncDSRStrategy):
+    """
+    Enhanced strategy for callback async DSR requests.
+    Works for both access and erasure operations with internal phase-based organization.
+    """
+
+    type = AsyncTaskType.callback
+
+    def __init__(self, session: Session):
+        self.session = session
+
+    def async_retrieve_data(
+        self,
+        client: AuthenticatedClient,
+        request_task_id: str,
+        query_config: SaaSQueryConfig,
+        input_data: Dict[str, List[Any]],
+    ) -> List[Row]:
+        """Execute async retrieve data with internal phase routing."""
+        request_task = self._get_request_task(request_task_id)
+        async_phase = get_async_phase(request_task, query_config)
+
+        if async_phase == AsyncPhase.initial_async:
+            return self._initial_request_access(
+                client, request_task, query_config, input_data
+            )
+        if async_phase == AsyncPhase.callback_completion:
+            return self._callback_completion_access(request_task)
+
+        logger.warning(
+            f"Unexpected async phase '{async_phase}' for callback access task {request_task.id}"
+        )
+        return []
+
+    def async_mask_data(
+        self,
+        client: AuthenticatedClient,
+        request_task_id: str,
+        query_config: SaaSQueryConfig,
+        rows: List[Row],
+    ) -> int:
+        """Execute async mask data with internal phase routing."""
+        request_task = self._get_request_task(request_task_id)
+        async_phase = get_async_phase(request_task, query_config)
+
+        if async_phase == AsyncPhase.initial_async:
+            return self._initial_request_erasure(
+                client, request_task, query_config, rows
+            )
+        if async_phase == AsyncPhase.callback_completion:
+            return self._callback_completion_erasure(request_task)
+
+        logger.warning(
+            f"Unexpected async phase '{async_phase}' for callback erasure task {request_task.id}"
+        )
+        return 0
+
+    # Private helper methods
+
+    def _initial_request_access(
+        self,
+        client: AuthenticatedClient,
+        request_task: RequestTask,
+        query_config: SaaSQueryConfig,
+        input_data: Dict[str, List[Any]],
+    ) -> List[Row]:
+        """Handle initial setup for access callback requests."""
+        logger.info(f"Initial callback request for access task {request_task.id}")
+
+        read_requests = query_config.get_read_requests_by_identity()
+
+        # Filter to get only the requests that need async processing
+        async_requests_to_process = [
+            req for req in read_requests if req.async_config and request_task.id
+        ]
+
+        # If there are no async requests, we shouldn't be in this handler.
+        if not async_requests_to_process:
+            logger.warning(
+                f"Async callback handler was called, but no async-configured read requests were found for task {request_task.id}."
+            )
+            return []
+
+        # Set async_type once for the request task
+        request_task.async_type = AsyncTaskType.callback
+        self.session.commit()
+
+        # For callback strategy, we execute the initial request to trigger the callback
+        # then wait for the external system to call us back via webhook.
+        # The callback endpoint logic will then queue the request task for it
+        # to be executed again with
+        # AsyncCallbackStrategy.async_retrieve_data -> _callback_completion_access
+
+        # Execute the initial callback request to trigger the async process
+        privacy_request = request_task.privacy_request
+        policy = privacy_request.policy
+
+        for read_request in async_requests_to_process:
+            if read_request.path:  # Only execute if there's an actual request to make
+                try:
+                    # Set the current request context
+                    query_config.current_request = read_request
+                    prepared_request = query_config.generate_query(input_data, policy)
+                    # Execute the initial request to trigger the callback
+                    client.send(prepared_request, read_request.ignore_errors)
+                    logger.info(
+                        f"Executed initial callback request for access task {request_task.id}"
+                    )
+                except ValueError as exc:
+                    if read_request.skip_missing_param_values:
+                        logger.debug(
+                            "Skipping optional access callback request: {}",
+                            exc,
+                        )
+                        continue
+                    raise exc
+
+        # Asynchronous callback access request detected in saas config.
+        # We raise AwaitingAsyncTask to put this task in an awaiting_processing state.
+        raise AwaitingAsyncTask(
+            f"Awaiting callback from {request_task.dataset_name} for access results"
+        )
+
+    def _initial_request_erasure(
+        self,
+        client: AuthenticatedClient,
+        request_task: RequestTask,
+        query_config: SaaSQueryConfig,
+        rows: List[Row],
+    ) -> int:
+        """Handle initial setup for erasure callback requests."""
+        logger.info(f"Initial callback request for erasure task {request_task.id}")
+
+        privacy_request = request_task.privacy_request
+        policy = privacy_request.policy
+
+        # For erasure, we only execute masking requests (delete/update requests)
+        masking_request = query_config.get_masking_request()
+
+        if not masking_request:
+            logger.warning(
+                f"No masking request found for erasure task {request_task.id}"
+            )
+            return 0
+
+        rows_updated = 0
+
+        # Check if the masking request has async_config and set async_type once
+        if masking_request.async_config and request_task.id:
+            request_task.async_type = AsyncTaskType.callback
+            self.session.commit()
+
+            if masking_request.path:
+                for row in rows:
+                    try:
+                        prepared_request = query_config.generate_update_stmt(
+                            row, policy, privacy_request
+                        )
+                        client.send(prepared_request, masking_request.ignore_errors)
+                        rows_updated += 1
+                    except ValueError as exc:
+                        if masking_request.skip_missing_param_values:
+                            logger.debug(
+                                "Skipping optional masking request: {}",
+                                exc,
+                            )
+                            continue
+                        raise exc
+
+            raise AwaitingAsyncTask("Awaiting erasure callback results")
+
+        logger.warning(
+            f"No async configuration found for erasure task {request_task.id}"
+        )
+        return rows_updated
+
+    def _callback_completion_access(self, request_task: RequestTask) -> List[Row]:
+        """Handle completion of access callback requests."""
+        logger.info(f"Access callback succeeded for request task '{request_task.id}'")
+        return request_task.get_access_data()
+
+    def _callback_completion_erasure(self, request_task: RequestTask) -> int:
+        """Handle completion of erasure callback requests."""
+        logger.info(f"Masking callback succeeded for request task '{request_task.id}'")
+        return request_task.rows_masked or 0

fides/api/service/async_dsr/strategies/async_dsr_strategy_factory.py

@@ -0,0 +1,72 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import ValidationError
+from sqlalchemy.orm import Session
+
+from fides.api.common_exceptions import NoSuchStrategyException
+from fides.api.common_exceptions import ValidationError as FidesopsValidationError
+from fides.api.schemas.saas.async_polling_configuration import AsyncPollingConfiguration
+from fides.api.service.async_dsr.strategies.async_dsr_strategy import AsyncDSRStrategy
+from fides.api.service.async_dsr.strategies.async_dsr_strategy_callback import (
+    AsyncCallbackStrategy,
+)
+from fides.api.service.async_dsr.strategies.async_dsr_strategy_polling import (
+    AsyncPollingStrategy,
+)
+
+
+class SupportedAsyncDSRStrategies(Enum):
+    """
+    The supported async DSR strategies for handling asynchronous data subject requests.
+    """
+
+    callback = AsyncCallbackStrategy
+    polling = AsyncPollingStrategy
+
+    @classmethod
+    def __contains__(cls, item: str) -> bool:
+        try:
+            cls[item]
+        except KeyError:
+            return False
+
+        return True
+
+
+def get_strategy(
+    strategy_name: str,
+    session: Session,
+    configuration: Optional[Dict[str, Any]] = None,
+) -> AsyncDSRStrategy:
+    """
+    Returns the enhanced strategy given the name, session, and configuration.
+    Raises NoSuchStrategyException if the strategy does not exist
+    """
+    if not SupportedAsyncDSRStrategies.__contains__(strategy_name):
+        valid_strategies = ", ".join(get_strategy_names())
+        raise NoSuchStrategyException(
+            f"Strategy '{strategy_name}' does not exist. Valid strategies are [{valid_strategies}]"
+        )
+
+    strategy_class = SupportedAsyncDSRStrategies[strategy_name].value
+
+    if strategy_name == "polling":
+        # Polling strategy requires configuration
+        if not configuration:
+            raise FidesopsValidationError("Configuration required for polling strategy")
+        try:
+            strategy_config = AsyncPollingConfiguration(**configuration)
+            return strategy_class(session, strategy_config)
+        except ValidationError as e:
+            raise FidesopsValidationError(message=str(e))
+    elif strategy_name == "callback":
+        # Callback strategy only requires session
+        return strategy_class(session)
+
+    raise NoSuchStrategyException(f"Unsupported strategy: {strategy_name}")
+
+
+def get_strategy_names() -> List[str]:
+    """Returns all supported async DSR strategies"""
+    return [s.name for s in SupportedAsyncDSRStrategies]