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

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

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

@@ -0,0 +1,113 @@
+"""S3-specific storage client implementation."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from fideslang.validation import AnyHttpUrlString
+from loguru import logger
+
+from fides.api.schemas.storage.storage import AWSAuthMethod, StorageSecrets
+from fides.api.service.storage.s3 import create_presigned_url_for_s3
+from fides.api.service.storage.streaming.base_storage_client import BaseStorageClient
+from fides.api.util.aws_util import get_s3_client
+
+
+class S3StorageClient(BaseStorageClient):
+    """S3-specific storage client implementation.
+
+    Handles S3-specific URI construction, transport parameters, and presigned URL
+    generation for the smart-open storage client.
+    """
+
+    def __init__(self, storage_secrets: dict[StorageSecrets, Any]):
+        """Initialize the storage client with secrets.
+
+        Args:
+            storage_secrets: Provider-specific storage credentials and configuration using StorageSecrets enum keys
+        """
+        super().__init__(storage_secrets)
+        self.storage_secrets: dict[StorageSecrets, Any] = storage_secrets
+
+    def build_uri(self, bucket: str, key: str) -> str:
+        """Build the S3 URI for the storage location.
+
+        Args:
+            bucket: S3 bucket name
+            key: Object key/path
+
+        Returns:
+            S3 URI string for smart-open
+        """
+        # Handle custom endpoint (e.g., MinIO) - endpoint_url is not in StorageSecrets enum
+        # For now, we'll use standard S3 URI since we only have enum keys
+        return f"s3://{bucket}/{key}"
+
+    def get_transport_params(self) -> dict[str, Any]:
+        """Get S3-specific transport parameters for smart-open.
+
+        Returns:
+            Dictionary of S3 transport parameters for smart-open
+        """
+        params = {}
+
+        if StorageSecrets.AWS_ACCESS_KEY_ID in self.storage_secrets:
+            params["access_key"] = self.storage_secrets[
+                StorageSecrets.AWS_ACCESS_KEY_ID
+            ]
+        if StorageSecrets.AWS_SECRET_ACCESS_KEY in self.storage_secrets:
+            params["secret_key"] = self.storage_secrets[
+                StorageSecrets.AWS_SECRET_ACCESS_KEY
+            ]
+        if StorageSecrets.REGION_NAME in self.storage_secrets:
+            params["region"] = self.storage_secrets[StorageSecrets.REGION_NAME]
+        if StorageSecrets.AWS_ASSUME_ROLE in self.storage_secrets:
+            params["assume_role_arn"] = self.storage_secrets[
+                StorageSecrets.AWS_ASSUME_ROLE
+            ]
+
+        return params
+
+    def generate_presigned_url(
+        self, bucket: str, key: str, ttl_seconds: Optional[int] = None
+    ) -> AnyHttpUrlString:
+        """Generate an S3 presigned URL for the object.
+
+        Args:
+            bucket: S3 bucket name
+            key: Object key/path
+            ttl_seconds: Time to live in seconds
+
+        Returns:
+            S3 presigned URL for the object
+
+        Raises:
+            Exception: If presigned URL generation fails
+        """
+        try:
+            # Storage secrets are already in the right format for get_s3_client
+            # get_s3_client expects dict[StorageSecrets, Any] with enum keys
+            s3_secrets = self.storage_secrets
+
+            # Determine auth method based on available credentials
+            # If AWS credentials are present, use SECRET_KEYS, otherwise use AUTOMATIC
+            if (
+                StorageSecrets.AWS_ACCESS_KEY_ID in self.storage_secrets
+                and StorageSecrets.AWS_SECRET_ACCESS_KEY in self.storage_secrets
+                and self.storage_secrets[StorageSecrets.AWS_ACCESS_KEY_ID]
+                and self.storage_secrets[StorageSecrets.AWS_SECRET_ACCESS_KEY]
+            ):
+                auth_method = AWSAuthMethod.SECRET_KEYS.value
+            else:
+                auth_method = AWSAuthMethod.AUTOMATIC.value
+
+            # Extract assume_role_arn if present for role assumption
+            assume_role_arn = None
+            if StorageSecrets.AWS_ASSUME_ROLE in self.storage_secrets:
+                assume_role_arn = self.storage_secrets[StorageSecrets.AWS_ASSUME_ROLE]
+
+            s3_client = get_s3_client(auth_method, s3_secrets, assume_role_arn)
+            return create_presigned_url_for_s3(s3_client, bucket, key, ttl_seconds)
+        except Exception as e:
+            logger.error(f"Failed to generate S3 presigned URL: {e}")
+            raise

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

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import Any, Optional, Union
+
+from fideslang.validation import AnyHttpUrlString
+from loguru import logger
+
+from fides.api.common_exceptions import StorageUploadError
+from fides.api.models.privacy_request import PrivacyRequest
+from fides.api.schemas.storage.storage import StorageSecrets, StorageSecretsS3
+from fides.api.service.storage.s3 import generic_upload_to_s3
+from fides.api.service.storage.streaming.retry import retry_cloud_storage_operation
+from fides.api.service.storage.streaming.schemas import StorageUploadConfig
+from fides.api.service.storage.streaming.smart_open_client import SmartOpenStorageClient
+from fides.api.service.storage.streaming.smart_open_streaming_storage import (
+    SmartOpenStreamingStorage,
+)
+
+
+@retry_cloud_storage_operation(
+    provider="s3_streaming",
+    operation_name="upload_to_s3_streaming",
+    max_retries=2,
+    base_delay=2.0,
+    max_delay=30.0,
+)
+def upload_to_s3_streaming(
+    storage_secrets: Union[StorageSecretsS3, dict[StorageSecrets, Any]],
+    data: dict,
+    bucket_name: str,
+    file_key: str,
+    resp_format: str,
+    privacy_request: PrivacyRequest,
+    document: Optional[BytesIO],
+    auth_method: str,
+    max_workers: int = 5,
+) -> Optional[AnyHttpUrlString]:
+    """Uploads arbitrary data to S3 using smart-open streaming for memory efficiency.
+
+    This function now uses smart-open for efficient cloud storage operations while maintaining
+    our DSR-specific business logic for package splitting and attachment processing.
+    """
+    formatted_secrets = format_secrets(storage_secrets)
+
+    if document is not None:
+        _, response = generic_upload_to_s3(
+            formatted_secrets,  # type: ignore[arg-type]
+            bucket_name,
+            file_key,
+            auth_method,
+            document,
+        )
+        return response
+
+    try:
+        storage_client = SmartOpenStorageClient("s3", formatted_secrets)
+
+        # Create upload config for the streaming interface
+        upload_config = StorageUploadConfig(
+            bucket_name=bucket_name,
+            file_key=file_key,
+            resp_format=resp_format,
+            max_workers=max_workers,
+        )
+
+        # Use the smart-open streaming implementation
+        streaming_storage = SmartOpenStreamingStorage(storage_client)
+        result = streaming_storage.upload_to_storage_streaming(
+            data,
+            upload_config,
+            privacy_request,
+            document,
+        )
+
+        logger.debug(
+            "Successfully uploaded streaming archive to S3 using smart-open: {}",
+            file_key,
+        )
+        return result
+
+    except Exception as e:
+        logger.error("Unexpected error during smart-open streaming upload: {}", e)
+        raise StorageUploadError(
+            f"Unexpected error during smart-open streaming upload: {e}"
+        )
+
+
+def _process_storage_secrets_input(
+    storage_secrets: Union[StorageSecretsS3, dict[StorageSecrets, Any]]  # type: ignore[misc]
+) -> dict[str, Any]:
+    """Process input and convert to string-keyed dictionary."""
+    final_secrets: dict[str, Any] = {}
+
+    if isinstance(storage_secrets, StorageSecretsS3):
+        # Convert StorageSecretsS3 model directly to string keys
+        for key, value in storage_secrets.model_dump().items():
+            if value is not None:
+                final_secrets[key] = value
+    else:
+        # Process dict input, converting enum keys to strings if needed
+        for key, value in (storage_secrets or {}).items():  # type: ignore[assignment]
+            if isinstance(key, str):
+                final_secrets[key] = value
+            elif isinstance(key, StorageSecrets):
+                final_secrets[key.value] = value
+
+    return final_secrets
+
+
+def _validate_aws_credentials(final_secrets: dict[str, Any]) -> None:
+    """Validate that required AWS credentials are present."""
+    has_access_key = (
+        "aws_access_key_id" in final_secrets and final_secrets["aws_access_key_id"]
+    )
+    has_secret_key = (
+        "aws_secret_access_key" in final_secrets
+        and final_secrets["aws_secret_access_key"]
+    )
+
+    # If we have any AWS credentials, we need both for SECRET_KEYS auth
+    if has_access_key or has_secret_key:
+        if not has_access_key:
+            raise ValueError(
+                "Missing required AWS credentials for SECRET_KEYS auth: aws_access_key_id. "
+                "Storage configuration must include valid AWS access key and secret key."
+            )
+        if not has_secret_key:
+            raise ValueError(
+                "Missing required AWS credentials for SECRET_KEYS auth: aws_secret_access_key. "
+                "Storage configuration must include valid AWS access key and secret key."
+            )
+    else:
+        # AUTOMATIC authentication - check if region is provided
+        has_region = "region_name" in final_secrets and final_secrets["region_name"]
+        if not has_region:
+            raise ValueError(
+                "Missing required region_name for AUTOMATIC authentication. "
+                "Storage configuration must include a valid AWS region."
+            )
+
+
+def _set_default_region(final_secrets: dict[str, Any]) -> None:
+    """Set default region if missing."""
+    if "region_name" not in final_secrets or not final_secrets["region_name"]:
+        logger.debug("Setting default region to 'us-east-1'")
+        final_secrets["region_name"] = "us-east-1"
+
+
+def format_secrets(
+    storage_secrets: Union[StorageSecretsS3, dict[StorageSecrets, Any]]  # type: ignore[misc]
+) -> dict[str, Any]:
+    """
+    Returns the correct format for the S3StorageClient.
+
+    This function handles multiple credential formats and processes them through several stages:
+    1. Input processing: Accepts either StorageSecretsS3 models (from API) or raw dicts (from database)
+    2. Key normalization: Converts all keys to string format for consistency
+    3. Validation: Ensures required AWS credentials are present based on auth method
+    4. Default setting: Sets default region if missing (after validation)
+    5. Return: Returns string-keyed dict ready for S3StorageClient
+
+    Input formats:
+    - StorageSecretsS3: Used by storage API endpoints (e.g., put_config_secrets)
+    - dict[StorageSecrets, Any]: Used by storage models (e.g., StorageConfig.secrets)
+    - dict[str, Any]: Used by database queries (e.g., StorageConfig.get_by)
+
+    Authentication methods:
+    - SECRET_KEYS: Requires aws_access_key_id, aws_secret_access_key, and region_name
+    - AUTOMATIC: Requires only region_name (relies on AWS SDK credential chain)
+    - Role assumption: Can be used with either auth method via assume_role_arn
+
+    Output format:
+    - dict[str, Any]: Required by S3StorageClient.generate_presigned_url() and other AWS operations
+
+    Returns:
+        dict[str, Any]: Credentials with string keys like 'aws_access_key_id'
+
+    Raises:
+        ValueError: If required AWS credentials are missing for the chosen auth method
+    """
+    logger.debug("format_secrets called with type: {}", type(storage_secrets).__name__)
+
+    # Stage 1: Process input and create final format directly
+    final_secrets = _process_storage_secrets_input(storage_secrets)
+
+    # Stage 2: Validate required credentials BEFORE setting defaults
+    _validate_aws_credentials(final_secrets)
+
+    # Stage 3: Set default region if missing (after validation)
+    _set_default_region(final_secrets)
+
+    logger.debug(
+        "format_secrets completed successfully with {} keys", len(final_secrets)
+    )
+    return final_secrets