ethyca-fides 2.71.0rc4__py2.py3-none-any.whl → 2.71.1__py2.py3-none-any.whl

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

@@ -0,0 +1,88 @@
+"""
+Pure HTTP utility for executing polling requests.
+
+This module contains low-level HTTP request execution for async DSR polling,
+with no business logic or orchestration dependencies.
+"""
+
+from typing import Any, Dict, Optional
+
+from requests import Response
+
+from fides.api.common_exceptions import PrivacyRequestError
+from fides.api.schemas.saas.async_polling_configuration import (
+    PollingResultRequest,
+    PollingStatusRequest,
+)
+from fides.api.service.connectors.saas.authenticated_client import AuthenticatedClient
+from fides.api.util.saas_util import map_param_values
+
+
+class PollingRequestHandler:
+    """
+    Pure HTTP utility for executing polling requests.
+
+    Handles status checks and result retrieval with no business logic dependencies.
+    """
+
+    def __init__(
+        self,
+        status_request: PollingStatusRequest,
+        result_request: Optional[PollingResultRequest] = None,
+    ):
+        self.status_request = status_request
+        self.result_request = result_request
+
+    def get_status_response(
+        self,
+        client: AuthenticatedClient,
+        param_values: Dict[str, Any],
+    ) -> Response:
+        """Execute HTTP status request and return raw response."""
+        if not self.status_request:
+            raise PrivacyRequestError(
+                "status_request is not configured in the async polling configuration"
+            )
+
+        prepared_status_request = map_param_values(
+            action="status",
+            context="polling request",
+            current_request=self.status_request,
+            param_values=param_values,
+        )
+
+        response: Response = client.send(prepared_status_request)
+
+        if not response.ok:
+            raise PrivacyRequestError(
+                f"Status request failed with status code {response.status_code}: {response.text}"
+            )
+
+        return response
+
+    def get_result_response(
+        self,
+        client: AuthenticatedClient,
+        param_values: Dict[str, Any],
+    ) -> Response:
+        """Execute HTTP result request and return raw response."""
+        if not self.result_request:
+            raise PrivacyRequestError(
+                "result_request is not configured in the async polling configuration"
+            )
+
+        prepared_result_request = map_param_values(
+            action="result",
+            context="polling request",
+            current_request=self.result_request,
+            param_values=param_values,
+        )
+
+        response: Response = client.send(prepared_result_request)
+
+        if not response.ok:
+            raise PrivacyRequestError(
+                f"Result request failed with status code {response.status_code}: {response.text}"
+            )
+
+        return response

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