ethyca-fides 2.68.1b1__py2.py3-none-any.whl → 2.68.1b3__py2.py3-none-any.whl

fides/api/service/privacy_request/request_service.py

@@ -11,12 +11,13 @@ from sqlalchemy import text
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.elements import TextClause
 
-from fides.api.common_exceptions import PrivacyRequestNotFound
+from fides.api.common_exceptions import PrivacyRequestError, PrivacyRequestNotFound
 from fides.api.models.privacy_request import (
     EXITED_EXECUTION_LOG_STATUSES,
     PrivacyRequest,
     RequestTask,
 )
+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
@@ -43,6 +44,7 @@ from fides.config import CONFIG
 PRIVACY_REQUEST_STATUS_CHANGE_POLL = "privacy_request_status_change_poll"
 DSR_DATA_REMOVAL = "dsr_data_removal"
 INTERRUPTED_TASK_REQUEUE_POLL = "interrupted_task_requeue_poll"
+ASYNC_TASKS_STATUS_POLLING = "async_tasks_status_polling"
 
 
 def build_required_privacy_request_kwargs(
@@ -191,7 +193,11 @@ def poll_for_exited_privacy_request_tasks(self: DatabaseTask) -> Set[str]:
             db.query(PrivacyRequest)
             .filter(
                 PrivacyRequest.status.in_(
-                    [PrivacyRequestStatus.in_processing, PrivacyRequestStatus.approved]
+                    [
+                        PrivacyRequestStatus.in_processing,
+                        PrivacyRequestStatus.approved,
+                        PrivacyRequestStatus.requires_input,
+                    ]
                 )
             )
             # Only look at Privacy Requests that haven't been deleted
@@ -353,6 +359,27 @@ def initiate_interrupted_task_requeue_poll() -> None:
     )
 
 
+def initiate_async_tasks_status_polling() -> None:
+    """Initiates scheduler to check for and requeue pending polling async tasks"""
+    if CONFIG.test_mode:
+        return
+
+    assert (
+        scheduler.running
+    ), "Scheduler is not running! Cannot add async tasks status polling job."
+
+    logger.info("Initiating scheduler for async tasks status polling")
+    scheduler.add_job(
+        func=poll_async_tasks_status,
+        trigger="interval",
+        kwargs={},
+        id=ASYNC_TASKS_STATUS_POLLING,
+        coalesce=True,
+        replace_existing=True,
+        seconds=CONFIG.execution.async_tasks_status_polling_interval_seconds,
+    )
+
+
 def get_cached_task_id(entity_id: str) -> Optional[str]:
     """Gets the cached task ID for a privacy request or request task by ID.
 
@@ -461,7 +488,6 @@ def _handle_privacy_request_requeue(
             )
 
             from fides.service.privacy_request.privacy_request_service import (  # pylint: disable=cyclic-import
-                PrivacyRequestError,
                 _requeue_privacy_request,
             )
 
@@ -546,6 +572,7 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
                         [
                             PrivacyRequestStatus.in_processing,
                             PrivacyRequestStatus.approved,
+                            PrivacyRequestStatus.requires_input,
                         ]
                     )
                 )
@@ -657,3 +684,29 @@ def requeue_interrupted_tasks(self: DatabaseTask) -> None:
                 # Requeue the privacy request if needed
                 if should_requeue:
                     _handle_privacy_request_requeue(db, privacy_request)
+
+
+@celery_app.task(base=DatabaseTask, bind=True)
+def poll_async_tasks_status(self: DatabaseTask) -> None:
+    """
+    Poll the status of async tasks that are awaiting processing.
+    """
+
+    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
+        async_tasks = (
+            db.query(RequestTask)
+            .filter(RequestTask.status == ExecutionLogStatus.awaiting_processing)
+            .filter(RequestTask.async_type == AsyncTaskType.polling)
+            .all()
+        )
+        if async_tasks:
+            logger.debug(f"Found {len(async_tasks)} async tasks to poll")
+            from fides.api.service.async_dsr.async_dsr_service import (  # pylint: disable=cyclic-import
+                requeue_polling_request,
+            )
+
+            for async_task in async_tasks:
+                requeue_polling_request(db, async_task)

fides/api/service/storage/storage_uploader_service.py

@@ -14,6 +14,7 @@ from fides.api.schemas.storage.storage import (
     StorageDetails,
     StorageType,
 )
+from fides.api.service.storage.streaming.s3.streaming_s3 import upload_to_s3_streaming
 from fides.api.tasks.storage import upload_to_gcs, upload_to_local, upload_to_s3
 
 
@@ -33,6 +34,8 @@ def upload(
     :param storage_key: Key representing where to upload data
     :return str representing location of upload (url or simply a description of where to find the data)
     """
+    logger.debug("upload called with storage_key: {}", storage_key)
+
     config: Optional[StorageConfig] = StorageConfig.get_by(
         db=db, field="key", value=storage_key
     )
@@ -40,16 +43,41 @@ def upload(
     if config is None:
         logger.warning("Storage type not found: {}", storage_key)
         raise StorageUploadError(f"Storage type not found: {storage_key}")
+
+    logger.debug(
+        "Retrieved storage config: key={}, type={}, has_secrets={}",
+        config.key,
+        config.type,
+        config.secrets is not None,
+    )
+
+    if config.secrets:
+        logger.debug("Storage config secrets type: {}", type(config.secrets))
+        if isinstance(config.secrets, dict):
+            logger.debug("Storage config secrets keys: {}", list(config.secrets.keys()))
+        else:
+            logger.debug("Storage config secrets is not a dict: {}", config.secrets)
+    else:
+        logger.warning("Storage config has no secrets!")
+
     uploader: Any = _get_uploader_from_config_type(config.type)  # type: ignore
+    logger.debug(
+        "Using uploader: {}",
+        uploader.__name__ if hasattr(uploader, "__name__") else type(uploader),
+    )
+
     return uploader(db, config, data, privacy_request)
 
 
-def get_extension(resp_format: ResponseFormat) -> str:
+def get_extension(resp_format: ResponseFormat, enable_streaming: bool = False) -> str:
     """
     Determine file extension for various response formats.
 
     CSV's and HTML reports are zipped together before uploading to s3.
     """
+    if enable_streaming:
+        return "zip"
+
     if resp_format == ResponseFormat.csv:
         return "zip"
 
@@ -62,7 +90,9 @@ def get_extension(resp_format: ResponseFormat) -> str:
     raise NotImplementedError(f"No extension defined for {resp_format}")
 
 
-def _construct_file_key(request_id: str, config: StorageConfig) -> str:
+def _construct_file_key(
+    request_id: str, config: StorageConfig, enable_streaming: bool = False
+) -> str:
     """Constructs file key based on desired naming convention and request id, e.g. 23847234.json"""
     naming = config.details.get(
         StorageDetails.NAMING.value, FileNaming.request_id.value
@@ -70,7 +100,7 @@ def _construct_file_key(request_id: str, config: StorageConfig) -> str:
     if naming != FileNaming.request_id.value:
         raise ValueError(f"File naming of {naming} not supported")
 
-    return f"{request_id}.{get_extension(config.format)}"  # type: ignore
+    return f"{request_id}.{get_extension(config.format, enable_streaming)}"  # type: ignore
 
 
 def _get_uploader_from_config_type(storage_type: StorageType) -> Any:
@@ -88,13 +118,58 @@ def _s3_uploader(
     data: Dict,
     privacy_request: PrivacyRequest,
 ) -> str:
-    """Constructs necessary info needed for s3 before calling upload"""
-    file_key: str = _construct_file_key(privacy_request.id, config)
+    """
+    Constructs necessary info needed for s3 before calling upload.
+    If `enable_streaming` is configured in the storage config, we use a streaming approach for better memory efficiency.
+    Otherwise we fall back to the traditional upload method.
+    """
+    logger.debug(
+        "_s3_uploader called with config: key={}, type={}, has_secrets={}",
+        config.key,
+        config.type,
+        config.secrets is not None,
+    )
+
+    if config.secrets:
+        logger.debug(
+            "Config secrets keys: {}",
+            (
+                list(config.secrets.keys())
+                if isinstance(config.secrets, dict)
+                else "Not a dict"
+            ),
+        )
+        logger.debug("Config secrets type: {}", type(config.secrets))
+    else:
+        logger.warning("Config secrets is None or empty!")
+
+    enable_streaming = config.details.get(StorageDetails.ENABLE_STREAMING.value, False)
+    file_key: str = _construct_file_key(privacy_request.id, config, enable_streaming)
 
     bucket_name = config.details[StorageDetails.BUCKET.value]
     auth_method = config.details[StorageDetails.AUTH_METHOD.value]
     document = None
 
+    if enable_streaming:
+        file_key = f"{privacy_request.id}.zip"
+        # Use streaming upload for better memory efficiency
+        logger.debug("Using streaming S3 upload for {}", file_key)
+        logger.debug("Calling upload_to_s3_streaming with secrets: {}", config.secrets)
+        return upload_to_s3_streaming(
+            config.secrets,  # type: ignore
+            data,
+            bucket_name,
+            file_key,
+            config.format.value,  # type: ignore
+            privacy_request,
+            document,
+            auth_method,
+        )
+
+    file_key = _construct_file_key(privacy_request.id, config)
+
+    # Fall back to traditional upload method
+    logger.debug("Using traditional S3 upload for {}", file_key)
     return upload_to_s3(
         config.secrets,  # type: ignore
         data,

fides/api/service/storage/streaming/__init__.py

@@ -0,0 +1,42 @@
+"""Streaming storage module for efficient cloud-to-cloud data transfer."""
+
+from .base_storage_client import BaseStorageClient
+from .retry import (
+    PermanentError,
+    RetryableError,
+    RetryConfig,
+    TransientError,
+    create_retry_config_from_settings,
+    is_transient_error,
+    retry_cloud_storage_operation,
+    retry_with_backoff,
+)
+from .s3.s3_storage_client import S3StorageClient
+from .schemas import AttachmentInfo, StorageUploadConfig, StreamingBufferConfig
+from .smart_open_client import SmartOpenStorageClient
+from .smart_open_streaming_storage import SmartOpenStreamingStorage
+from .storage_client_factory import StorageClientFactory
+
+__all__ = [
+    # Base classes and interfaces
+    "BaseStorageClient",
+    "StorageClientFactory",
+    # Provider-specific clients
+    "S3StorageClient",
+    # Main clients
+    "SmartOpenStorageClient",
+    "SmartOpenStreamingStorage",
+    # Schemas
+    "StorageUploadConfig",
+    "StreamingBufferConfig",
+    "AttachmentInfo",
+    # Retry utilities
+    "RetryConfig",
+    "RetryableError",
+    "TransientError",
+    "PermanentError",
+    "retry_with_backoff",
+    "retry_cloud_storage_operation",
+    "create_retry_config_from_settings",
+    "is_transient_error",
+]

fides/api/service/storage/streaming/base_storage_client.py

@@ -0,0 +1,61 @@
+"""Base abstract class for cloud storage clients."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from fideslang.validation import AnyHttpUrlString
+
+
+class BaseStorageClient(ABC):
+    """Abstract base class for cloud storage clients.
+
+    This class defines the interface that all provider-specific storage clients
+    must implement. It provides a common contract for URI building, transport
+    parameters, and presigned URL generation.
+    """
+
+    def __init__(self, storage_secrets: Any):
+        """Initialize the storage client with secrets.
+
+        Args:
+            storage_secrets: Provider-specific storage credentials and configuration.
+                           Derived classes will specify the exact type they expect.
+        """
+        self.storage_secrets = storage_secrets
+
+    @abstractmethod
+    def build_uri(self, bucket: str, key: str) -> str:
+        """Build the URI for the storage location.
+
+        Args:
+            bucket: Storage bucket/container name
+            key: Object key/path
+
+        Returns:
+            URI string for smart-open
+        """
+
+    @abstractmethod
+    def get_transport_params(self) -> dict[str, Any]:
+        """Get transport parameters for smart-open.
+
+        Returns:
+            Dictionary of transport parameters for smart-open
+        """
+
+    @abstractmethod
+    def generate_presigned_url(
+        self, bucket: str, key: str, ttl_seconds: Optional[int] = None
+    ) -> AnyHttpUrlString:
+        """Generate a presigned URL for the object.
+
+        Args:
+            bucket: Storage bucket/container name
+            key: Object key/path
+            ttl_seconds: Time to live in seconds
+
+        Returns:
+            Presigned URL for the object
+        """

fides/api/service/storage/streaming/dsr_storage.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import zipfile
+from datetime import datetime
+from io import BytesIO
+from typing import Any, Generator, Tuple
+
+from loguru import logger
+from stream_zip import _ZIP_32_TYPE
+
+from fides.api.service.storage.streaming.schemas import AttachmentProcessingInfo
+from fides.api.service.storage.streaming.smart_open_client import SmartOpenStorageClient
+
+
+def stream_dsr_buffer_to_storage(
+    storage_client: SmartOpenStorageClient,
+    bucket_name: str,
+    file_key: str,
+    dsr_buffer: BytesIO,
+    content_type: str = "application/zip",
+) -> None:
+    """Stream DSR buffer to storage using smart-open streaming.
+
+    This function handles only the storage streaming concern, accepting a pre-generated
+    DSR buffer to maintain clear separation of concerns.
+
+    Args:
+        storage_client: The storage client for streaming uploads
+        bucket_name: Storage bucket name
+        file_key: File key in storage
+        dsr_buffer: Pre-generated DSR report buffer (BytesIO)
+        content_type: MIME type for the uploaded file (defaults to application/zip)
+    """
+    # Get the content from the buffer
+    content = dsr_buffer.getvalue()
+    try:
+        # Use smart-open's streaming upload for efficient memory usage
+        with storage_client.stream_upload(
+            bucket_name, file_key, content_type=content_type
+        ) as upload_stream:
+            upload_stream.write(content)
+
+    except Exception as e:
+        logger.error("Failed to upload DSR report using smart-open streaming: {}", e)
+        raise e
+
+
+def create_dsr_report_files_generator(
+    dsr_buffer: BytesIO,
+    all_attachments: list["AttachmentProcessingInfo"],
+    bucket_name: str,
+    max_workers: int,
+    batch_size: int,
+) -> Generator[
+    Tuple[str, datetime, int, Any, Generator[bytes, None, None]], None, None
+]:
+    """Create a ZIP generator for DSR report HTML files only.
+
+    This method extracts and yields the HTML files from a DSR report buffer.
+    Note: This function only handles the DSR report files (HTML, CSS, etc.).
+    The caller is responsible for combining this with attachment files to create
+    the complete ZIP.
+
+    Args:
+        dsr_buffer: The DSR report buffer (ZIP file as BytesIO)
+        all_attachments: List of validated attachments (used for logging only)
+        bucket_name: Storage bucket name (used for logging only)
+        max_workers: Maximum parallel workers (used for logging only)
+        batch_size: Number of attachments to process in each batch (used for logging only)
+
+    Returns:
+        Generator yielding DSR report files in stream_zip format
+    """
+    logger.debug(
+        f"Creating DSR report files generator with {len(all_attachments)} attachments"
+    )
+
+    # Reset buffer position to ensure we can read from it
+    dsr_buffer.seek(0)
+
+    # Extract and yield the DSR report files from the buffer
+    # The dsr_buffer is already a ZIP file, so we need to extract and re-yield its contents
+    with zipfile.ZipFile(dsr_buffer) as dsr_zip:
+        for file_info in dsr_zip.filelist:
+            if not file_info.is_dir():
+                # Read the file content and yield it in stream_zip format
+                content = dsr_zip.read(file_info.filename)
+
+                def content_generator(
+                    file_content: bytes,
+                ) -> Generator[bytes, None, None]:
+                    yield file_content
+
+                yield file_info.filename, datetime.now(), 0o644, _ZIP_32_TYPE(), content_generator(
+                    content
+                )
+
+    logger.debug("DSR report files extracted and ready for ZIP creation")