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

fides/api/schemas/saas/async_polling_configuration.py

@@ -0,0 +1,81 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from fides.api.schemas.saas.saas_config import SaaSRequest
+from fides.api.schemas.saas.strategy_configuration import StrategyConfiguration
+from fides.api.util.collection_util import Row
+
+
+class SupportedDataType(Enum):
+    """Supported data types for polling async DSR result requests."""
+
+    # Structured data types that can be parsed into rows
+    json = "json"  # Parsed into List[Row] from JSON response
+    csv = "csv"  # Parsed into List[Row] from CSV response
+    # Binary/non-parseable data stored as raw bytes
+    attachment = "attachment"  # Binary files (.zip, .pdf, .xml, etc.) stored as bytes
+
+
+class PollingResultType(Enum):
+    """Types of results from async polling operations."""
+
+    rows = "rows"  # Structured data parsed into List[Row]
+    attachment = "attachment"  # Binary file data stored as bytes
+
+
+class PollingResult(BaseModel):
+    """
+    Flexible result container for async polling operations.
+    Handles both structured data and file attachments.
+    """
+
+    data: Union[List[Row], bytes]
+    result_type: PollingResultType
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class PollingStatusRequest(SaaSRequest):
+    """
+    Extended SaaSRequest for checking async job status.
+    Uses request_override for custom status checking logic or standard fields for simple cases.
+    """
+
+    status_path: Optional[str] = None
+    status_completed_value: Optional[Union[str, bool, int]] = None
+
+    @model_validator(mode="after")
+    def validate_status_fields(self) -> "PollingStatusRequest":
+        """Ensure required fields are present unless using an override."""
+        if self.request_override:
+            return self
+
+        if not self.status_path:
+            raise ValueError("status_path is required when request_override is not set")
+        if self.status_completed_value is None:
+            raise ValueError(
+                "status_completed_value is required when request_override is not set"
+            )
+        return self
+
+
+class PollingResultRequest(SaaSRequest):
+    """
+    Extended SaaSRequest for retrieving async job results.
+    Uses request_override for custom result retrieval or standard HTTP request for simple cases.
+    Data type is automatically inferred from response.
+    """
+
+    result_path: Optional[str] = None
+
+
+class AsyncPollingConfiguration(StrategyConfiguration):
+    """
+    Simplified configuration for polling async DSR requests.
+    The main read request serves as the initial request.
+    """
+
+    status_request: PollingStatusRequest
+    # result_request is optional for delete/update operations
+    result_request: Optional[PollingResultRequest] = None

fides/api/schemas/saas/saas_config.py

@@ -2,7 +2,9 @@ from typing import Any, Dict, List, Optional, Set, Union
 
 from fideslang.models import FidesCollectionKey, FidesDatasetReference
 from fideslang.validation import FidesKey
-from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as PydanticField
+from pydantic import field_validator, model_validator
 
 from fides.api.common_exceptions import ValidationError
 from fides.api.graph.config import (
@@ -115,6 +117,10 @@ class SaaSRequest(BaseModel):
     skip_missing_param_values: Optional[bool] = (
         False  # Skip instead of raising an exception if placeholders can't be populated in body
     )
+    correlation_id_path: Optional[str] = PydanticField(
+        default=None,
+        description="The path to the correlation ID in the response. For use with async polling.",
+    )
     model_config = ConfigDict(
         from_attributes=True, use_enum_values=True, extra="forbid"
     )
@@ -213,7 +219,7 @@ class SaaSRequest(BaseModel):
 class ReadSaaSRequest(SaaSRequest):
     """
     An extension of the base SaaSRequest that allows the inclusion of an output template
-    that is used to format each collection result.
+    that is used to format each collection result, and correlation_id_path for async polling.
     """
 
     output: Optional[str] = None
@@ -230,7 +236,8 @@ class ReadSaaSRequest(SaaSRequest):
                 raise ValueError(
                     "A read request must specify a method if a path is provided and no request_override is specified"
                 )
-        else:
+
+        if self.request_override:
             allowed_fields = {
                 "request_override",
                 "param_values",

fides/api/schemas/saas/strategy_configuration.py

@@ -178,15 +178,3 @@ class OAuth2ClientCredentialsConfiguration(OAuth2BaseConfiguration):
     """
 
     refresh_request: Optional[SaaSRequest] = Field(exclude=True)
-
-
-class PollingAsyncDSRConfiguration(StrategyConfiguration):
-    """
-    Configuration for polling async DSR requests.
-    """
-
-    status_request: SaaSRequest
-    status_path: str
-    status_completed_value: Optional[str] = None
-    result_request: SaaSRequest
-    result_path: str

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

@@ -0,0 +1,155 @@
+"""Handler for storing and managing attachments from async polling operations."""
+
+from io import BytesIO
+from typing import List, Union
+
+from loguru import logger
+from sqlalchemy.orm import Session
+
+from fides.api.common_exceptions import PrivacyRequestError
+from fides.api.models.attachment import (
+    Attachment,
+    AttachmentReference,
+    AttachmentReferenceType,
+    AttachmentType,
+)
+from fides.api.models.privacy_request.request_task import RequestTask
+from fides.api.models.storage import get_active_default_storage_config
+from fides.api.util.collection_util import Row
+
+
+class PollingAttachmentHandler:
+    """Utility class for handling attachment storage and metadata in polling operations."""
+
+    @staticmethod
+    def store_attachment(
+        session: Session,
+        request_task: RequestTask,
+        attachment_data: bytes,
+        filename: str,
+    ) -> str:
+        """
+        Store polling attachment data and return attachment ID.
+
+        This utility function handles the storage of attachment data
+        from polling results and creates the necessary database records.
+
+        Args:
+            session: Database session
+            request_task: The request task associated with this attachment
+            attachment_data: The binary attachment data to store
+            filename: The filename for the attachment
+
+        Returns:
+            str: The ID of the created attachment
+
+        Raises:
+            PrivacyRequestError: If storage configuration is not found or storage fails
+        """
+        try:
+            # Get active storage config
+            storage_config = get_active_default_storage_config(session)
+            if not storage_config:
+                raise PrivacyRequestError("No active storage configuration found")
+
+            # Create attachment record and upload to storage
+            attachment = Attachment.create_and_upload(
+                db=session,
+                data={
+                    "file_name": filename,
+                    "attachment_type": AttachmentType.include_with_access_package,
+                    "storage_key": storage_config.key,
+                },
+                attachment_file=BytesIO(attachment_data),
+            )
+
+            # Create attachment references
+            AttachmentReference.create(
+                db=session,
+                data={
+                    "attachment_id": attachment.id,
+                    "reference_id": request_task.id,
+                    "reference_type": AttachmentReferenceType.request_task,
+                },
+            )
+
+            AttachmentReference.create(
+                db=session,
+                data={
+                    "attachment_id": attachment.id,
+                    "reference_id": request_task.privacy_request.id,
+                    "reference_type": AttachmentReferenceType.privacy_request,
+                },
+            )
+
+            logger.info(
+                f"Successfully stored polling attachment {attachment.id} for request_task {request_task.id}"
+            )
+            return attachment.id
+
+        except Exception as e:
+            logger.error(f"Failed to store polling attachment: {e}")
+            raise PrivacyRequestError(f"Failed to store polling attachment: {e}")
+
+    @staticmethod
+    def ensure_attachment_bytes(data: Union[List[Row], bytes]) -> bytes:
+        """
+        Ensure attachment polling results provide bytes content.
+
+        Args:
+            data: The data that should be bytes
+
+        Returns:
+            bytes: The validated bytes data
+
+        Raises:
+            PrivacyRequestError: If data is not bytes
+        """
+        if isinstance(data, bytes):
+            return data
+        raise PrivacyRequestError("Expected bytes data for attachment polling result")
+
+    @staticmethod
+    def add_metadata_to_rows(db: Session, attachment_id: str, rows: List[Row]) -> None:
+        """
+        Add attachment metadata to rows collection (like manual tasks do).
+
+        Args:
+            db: Database session
+            attachment_id: The ID of the attachment to add metadata for
+            rows: The list of rows to add the attachment metadata to
+        """
+        attachment_record = (
+            db.query(Attachment).filter(Attachment.id == attachment_id).first()
+        )
+
+        if attachment_record:
+            try:
+                size, url = attachment_record.retrieve_attachment()
+                attachment_info = {
+                    "file_name": attachment_record.file_name,
+                    "download_url": str(url) if url else None,
+                    "file_size": size,
+                }
+            except Exception as exc:
+                logger.warning(
+                    f"Could not retrieve attachment content for {attachment_record.file_name}: {exc}"
+                )
+                attachment_info = {
+                    "file_name": attachment_record.file_name,
+                    "download_url": None,
+                    "file_size": None,
+                }
+
+            # Add attachment to the polling results
+            attachments_item = None
+            for item in rows:
+                if isinstance(item, dict) and "retrieved_attachments" in item:
+                    attachments_item = item
+                    break
+
+            if attachments_item is None:
+                attachments_item = {"retrieved_attachments": []}
+                rows.append(attachments_item)
+
+            attachments_item["retrieved_attachments"].append(attachment_info)

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")