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

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

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

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import random
+import time
+from functools import wraps
+from typing import Any, Callable, Optional, Type
+
+from loguru import logger
+
+from fides.config import CONFIG
+
+
+class RetryableError(Exception):
+    """Base exception for errors that should trigger retries."""
+
+
+class TransientError(RetryableError):
+    """Exception for transient errors that should be retried."""
+
+
+class PermanentError(RetryableError):
+    """Exception for permanent errors that should not be retried."""
+
+
+class RetryConfig:
+    """Configuration for retry behavior."""
+
+    def __init__(
+        self,
+        max_retries: int = 3,
+        base_delay: float = 1.0,
+        max_delay: float = 60.0,
+        backoff_factor: float = 2.0,
+        jitter: bool = True,
+        retry_on_exceptions: Optional[tuple[Type[Exception], ...]] = None,
+    ):
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+        self.backoff_factor = backoff_factor
+        self.jitter = jitter
+        self.retry_on_exceptions = retry_on_exceptions or (Exception,)
+
+
+def is_transient_error(error: Exception) -> bool:
+    """
+    Determine if an error is transient and should be retried.
+
+    This is a cloud-agnostic implementation that can be extended
+    with provider-specific logic.
+    """
+    # Check if this is our custom TransientError
+    if isinstance(error, TransientError):
+        return True
+
+    error_str = str(error).lower()
+
+    # Common transient error patterns across cloud providers
+    transient_patterns = [
+        "timeout",
+        "timed out",
+        "connection",
+        "network",
+        "temporary",
+        "throttling",
+        "rate limit",
+        "too many requests",
+        "service unavailable",
+        "internal server error",
+        "bad gateway",
+        "gateway timeout",
+        "request timeout",
+        "connection reset",
+        "broken pipe",
+    ]
+
+    return any(pattern in error_str for pattern in transient_patterns)
+
+
+def calculate_backoff_delay(
+    attempt: int,
+    base_delay: float,
+    max_delay: float,
+    backoff_factor: float,
+    jitter: bool = True,
+) -> float:
+    """
+    Calculate exponential backoff delay with optional jitter.
+
+    Args:
+        attempt: Current attempt number (0-based)
+        base_delay: Base delay in seconds
+        max_delay: Maximum delay in seconds
+        backoff_factor: Multiplier for each attempt
+        jitter: Whether to add random jitter to prevent thundering herd
+
+    Returns:
+        Delay in seconds before next retry
+    """
+    if attempt == 0:
+        return 0
+
+    # Exponential backoff: base_delay * (backoff_factor ^ attempt)
+    delay = base_delay * (backoff_factor ** (attempt - 1))
+
+    # Cap at maximum delay
+    delay = min(delay, max_delay)
+
+    # Add jitter to prevent multiple retries from synchronizing
+    if jitter:
+        jitter_amount = delay * 0.1  # 10% jitter
+        delay += random.uniform(-jitter_amount, jitter_amount)
+        delay = max(0, delay)  # Ensure non-negative
+
+    return delay
+
+
+def retry_with_backoff(
+    retry_config: Optional[RetryConfig] = None,
+    operation_name: str = "operation",
+) -> Callable:
+    """
+    Decorator for retrying operations with exponential backoff.
+
+    This is a cloud-agnostic retry mechanism that can be extended
+    with provider-specific retry logic.
+
+    Args:
+        retry_config: Configuration for retry behavior
+        operation_name: Name of the operation for logging
+
+    Returns:
+        Decorated function with retry logic
+    """
+    if retry_config is None:
+        # Use default configuration from settings
+        settings = CONFIG.execution
+        retry_config = RetryConfig(
+            max_retries=settings.task_retry_count,
+            base_delay=settings.task_retry_delay,
+            backoff_factor=settings.task_retry_backoff,
+        )
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            last_exception: Optional[Exception] = None
+
+            for attempt in range(retry_config.max_retries + 1):
+                try:
+                    # Track retry attempts
+                    if attempt > 0:
+                        logger.debug(
+                            "Retry attempt {}/{} for {}",
+                            attempt,
+                            retry_config.max_retries,
+                            operation_name,
+                        )
+
+                    return func(*args, **kwargs)
+
+                except Exception as e:
+                    last_exception = e
+
+                    # Check if this is a permanent error that shouldn't be retried
+                    if isinstance(e, PermanentError):
+                        logger.error(
+                            "Permanent error in {}: {}. Not retrying.",
+                            operation_name,
+                            e,
+                        )
+                        raise
+
+                    # Check if we should retry this error
+                    should_retry = (
+                        attempt < retry_config.max_retries
+                        and isinstance(e, retry_config.retry_on_exceptions)
+                        and is_transient_error(e)
+                    )
+
+                    if not should_retry:
+                        logger.error(
+                            "Non-retryable error in {}: {}. Not retrying.",
+                            operation_name,
+                            e,
+                        )
+                        raise
+
+                    # Calculate delay for next retry
+                    delay = calculate_backoff_delay(
+                        attempt + 1,
+                        retry_config.base_delay,
+                        retry_config.max_delay,
+                        retry_config.backoff_factor,
+                        retry_config.jitter,
+                    )
+
+                    logger.warning(
+                        "Transient error in {} (attempt {}/{}): {}. "
+                        "Retrying in {:.1f} seconds...",
+                        operation_name,
+                        attempt + 1,
+                        retry_config.max_retries + 1,
+                        e,
+                        delay,
+                    )
+
+                    # Sleep before retry
+                    time.sleep(delay)
+
+            # If we get here, all retries were exhausted
+            logger.error(
+                "Operation {} failed after {} retry attempts. Last error: {}",
+                operation_name,
+                retry_config.max_retries,
+                last_exception,
+            )
+            if last_exception is not None:
+                raise last_exception
+
+            raise RuntimeError(
+                f"Operation {operation_name} failed after {retry_config.max_retries} retry attempts"
+            )
+
+        return wrapper
+
+    return decorator
+
+
+def retry_cloud_storage_operation(
+    provider: str = "cloud storage",
+    operation_name: str = "storage operation",
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 60.0,
+    backoff_factor: float = 2.0,
+) -> Callable:
+    """
+    Cloud-agnostic retry decorator for storage operations.
+
+    This decorator provides retry logic that works across different
+    cloud storage providers while allowing provider-specific optimizations.
+
+    Args:
+        provider: Name of the cloud storage provider
+        operation_name: Name of the storage operation
+        max_retries: Maximum number of retry attempts
+        base_delay: Base delay in seconds
+        max_delay: Maximum delay in seconds
+        backoff_factor: Multiplier for each attempt
+
+    Returns:
+        Decorated function with cloud-agnostic retry logic
+    """
+    retry_config = RetryConfig(
+        max_retries=max_retries,
+        base_delay=base_delay,
+        max_delay=max_delay,
+        backoff_factor=backoff_factor,
+        jitter=True,
+    )
+
+    return retry_with_backoff(
+        retry_config=retry_config, operation_name=f"{provider} {operation_name}"
+    )
+
+
+def create_retry_config_from_settings() -> RetryConfig:
+    """
+    Create a RetryConfig instance from application settings.
+
+    Returns:
+        RetryConfig configured with application defaults
+    """
+    settings = CONFIG.execution
+    return RetryConfig(
+        max_retries=settings.task_retry_count,
+        base_delay=settings.task_retry_delay,
+        backoff_factor=settings.task_retry_backoff,
+        max_delay=60.0,  # Cap at 1 minute
+        jitter=True,
+    )

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

@@ -0,0 +1,5 @@
+"""S3 streaming storage package."""
+
+from .s3_storage_client import S3StorageClient
+
+__all__ = ["S3StorageClient"]