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

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

@@ -0,0 +1,261 @@
+"""
+Pure response processing utility for async DSR polling.
+
+This module handles data type inference, attachment classification,
+and response parsing with no business logic dependencies.
+"""
+
+import mimetypes
+import os
+from email.message import Message
+from typing import Any, List, Optional
+from urllib.parse import urlparse
+
+import pydash
+from loguru import logger
+from requests import Response
+
+from fides.api.common_exceptions import PrivacyRequestError
+from fides.api.schemas.saas.async_polling_configuration import (
+    PollingResult,
+    PollingResultType,
+    SupportedDataType,
+)
+from fides.api.util.collection_util import Row
+
+CONTENT_TYPE = "content-type"
+CONTENT_DISPOSITION = "content-disposition"
+
+# Content type mappings for data type detection
+CONTENT_TYPE_MAP = {
+    "application/json": SupportedDataType.json,
+    "text/csv": SupportedDataType.csv,
+    "application/csv": SupportedDataType.csv,
+}
+
+# File extensions that indicate attachments
+ATTACHMENT_EXTENSIONS = {
+    ".zip",
+    ".pdf",
+    ".tar",
+    ".gz",
+    ".xlsx",
+    ".xls",
+    ".xml",
+    ".csv",
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".gif",
+    ".bmp",
+    ".tiff",
+    ".webp",
+    ".svg",
+}
+
+# Content types that should always be attachments
+ATTACHMENT_CONTENT_TYPES = {
+    "application/octet-stream",
+    "application/zip",
+    "application/pdf",
+    "application/xml",
+    "text/xml",
+    "image/",
+    "video/",
+    "text/csv",
+    "application/csv",
+}
+
+# Initialize mimetypes module for content type inference
+mimetypes.init()
+
+
+class PollingResponseProcessor:
+    """Pure utility for processing async polling responses with smart data type inference."""
+
+    @classmethod
+    def process_result_response(
+        cls, request_path: str, response: Response, result_path: Optional[str] = None
+    ) -> PollingResult:
+        """Process response with smart data type inference."""
+        inferred_type = cls._infer_data_type(request_path, response)
+
+        if cls._should_store_as_attachment(response, request_path, inferred_type):
+            return cls._build_attachment_result(response, request_path, inferred_type)
+
+        rows = cls._parse_to_rows(response, inferred_type, result_path)
+        return PollingResult(
+            data=rows,
+            result_type=PollingResultType.rows,
+            metadata={
+                "inferred_type": inferred_type.value,
+                "content_type": response.headers.get(CONTENT_TYPE, ""),
+                "row_count": len(rows),
+                "parsed_to_rows": True,
+            },
+        )
+
+    @staticmethod
+    def evaluate_status_response(
+        response: Response,
+        status_path: str,
+        status_completed_value: Optional[Any] = None,
+    ) -> bool:
+        """Process status response and extract completion status."""
+        # Check if response is JSON
+        content_type = response.headers.get(CONTENT_TYPE, "").lower()
+        if "application/json" not in content_type:
+            return False
+
+        try:
+            status_value = pydash.get(response.json(), status_path)
+        except ValueError as e:
+            logger.error(f"Invalid JSON response in status check: {e}")
+            return False
+
+        # Handle list values - check first element
+        if isinstance(status_value, list) and status_value:
+            status_value = status_value[0]
+
+        # Direct comparison if completed value specified
+        if status_completed_value is not None:
+            return status_value == status_completed_value
+
+        # Otherwise, check truthiness
+        return bool(status_value)
+
+    @staticmethod
+    def _infer_data_type(request_path: str, response: Response) -> SupportedDataType:
+        """Infer data type from response characteristics."""
+        content_type = response.headers.get(CONTENT_TYPE, "").lower()
+
+        # Check for attachment content types first
+        if any(ct in content_type for ct in ATTACHMENT_CONTENT_TYPES):
+            return SupportedDataType.attachment
+
+        # Check content-type header (most reliable)
+        for content_type_pattern, data_type in CONTENT_TYPE_MAP.items():
+            if content_type_pattern in content_type:
+                return data_type
+
+        # Check URL extension
+        path = urlparse(request_path.lower()).path
+        if path.endswith(".csv"):
+            return SupportedDataType.csv
+
+        # Try parsing as JSON
+        try:
+            response.json()
+            return SupportedDataType.json
+        except (ValueError, TypeError):
+            pass
+
+        return SupportedDataType.attachment  # Preserve unknown types as raw data
+
+    @staticmethod
+    def _should_store_as_attachment(
+        response: Response, request_path: str, inferred_type: SupportedDataType
+    ) -> bool:
+        """Determine if response should be treated as an attachment."""
+        content_type = response.headers.get(CONTENT_TYPE, "").lower()
+
+        # Check attachment-specific indicators
+        if inferred_type == SupportedDataType.attachment:
+            return True
+
+        if "attachment" in response.headers.get(CONTENT_DISPOSITION, "").lower():
+            return True
+
+        if any(ct in content_type for ct in ATTACHMENT_CONTENT_TYPES):
+            return True
+
+        # Check file extension in URL
+        path = urlparse(request_path.lower()).path
+        if any(path.endswith(ext) for ext in ATTACHMENT_EXTENSIONS):
+            return True
+
+        return False
+
+    @staticmethod
+    def _extract_filename(response: Response, request_url: str) -> str:
+        """Extract filename from response or URL."""
+        # Try Content-Disposition header using email.message for proper parsing
+        disposition = response.headers.get(CONTENT_DISPOSITION)
+        if disposition:
+            msg = Message()
+            msg[CONTENT_DISPOSITION] = disposition
+            filename = msg.get_filename()
+            if filename:
+                return filename
+
+        # Extract from URL path using os.path.basename
+        path = urlparse(request_url).path
+        if path:
+            filename = os.path.basename(path)
+            if filename:
+                return filename
+
+        return "polling_result"
+
+    @staticmethod
+    def _build_attachment_result(
+        response: Response, request_url: str, inferred_type: SupportedDataType
+    ) -> PollingResult:
+        """Build a PollingResult for attachment data."""
+        filename = PollingResponseProcessor._extract_filename(response, request_url)
+
+        # Get content type from header, or infer from filename extension
+        content_type = response.headers.get(CONTENT_TYPE)
+
+        # If no content type or it's generic, try to infer from filename extension using mimetypes
+        if not content_type or content_type in (
+            "application/octet-stream",
+            "text/plain",
+        ):
+            guessed_type, _ = mimetypes.guess_type(filename)
+            if guessed_type:
+                content_type = guessed_type
+            else:
+                content_type = content_type or "application/octet-stream"
+
+        return PollingResult(
+            data=response.content,
+            result_type=PollingResultType.attachment,
+            metadata={
+                "inferred_type": inferred_type.value,
+                "content_type": content_type,
+                "filename": filename,
+                "size": len(response.content),
+                "preserved_as_attachment": True,
+            },
+        )
+
+    @staticmethod
+    def _parse_to_rows(
+        response: Response,
+        data_type: SupportedDataType,
+        result_path: Optional[str] = None,
+    ) -> List[Row]:
+        """Parse response to List[Row] based on data type."""
+        if data_type == SupportedDataType.json:
+            try:
+                data = response.json()
+                if result_path:
+                    data = pydash.get(data, result_path)
+                    if data is None:
+                        raise PrivacyRequestError(
+                            f"Could not extract data from response using path: {result_path}"
+                        )
+
+                if isinstance(data, list):
+                    return data
+                if isinstance(data, dict):
+                    return [data]
+                raise PrivacyRequestError(
+                    f"Expected list or dict from result request, got: {type(data)}"
+                )
+
+            except ValueError as e:
+                raise PrivacyRequestError(f"Invalid JSON response: {e}")
+
+        raise PrivacyRequestError(f"Cannot parse {data_type} to rows")

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