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

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

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

@@ -0,0 +1,173 @@
+from typing import Annotated, Any, Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+# multipart upload requirements
+MIN_PART_SIZE = 5 * 1024 * 1024  # 5MB minimum part size for AWS S3 uploads
+ZIP_BUFFER_THRESHOLD = 5 * 1024 * 1024  # 5MB threshold for zip buffer uploads
+CHUNK_SIZE_THRESHOLD = 5 * 1024 * 1024  # 5MB threshold for chunk-based uploads
+LARGE_FILE_THRESHOLD = 25 * 1024 * 1024  # 25MB threshold for large file handling
+
+
+class AttachmentInfo(BaseModel):
+    """Schema for attachment information in streaming storage"""
+
+    storage_key: str = Field(
+        ..., description="Storage key for the attachment (cloud-agnostic)"
+    )
+    file_name: Optional[str] = Field(None, description="Human-readable filename")
+    size: Optional[int] = Field(None, ge=0, description="File size in bytes")
+    content_type: Optional[str] = Field(None, description="MIME type of the file")
+
+    @field_validator("storage_key")
+    @classmethod
+    def validate_storage_key(cls, v: str) -> str:
+        if not v or not v.strip():
+            raise ValueError("storage_key cannot be empty or whitespace")
+        return v.strip()
+
+    @field_validator("file_name")
+    @classmethod
+    def validate_file_name(cls, v: Optional[str]) -> Optional[str]:
+        if v is not None and (not v or not v.strip()):
+            raise ValueError("file_name cannot be empty or whitespace if provided")
+        return v.strip() if v else None
+
+
+class StorageUploadConfig(BaseModel):
+    """Configuration for storage upload operations"""
+
+    bucket_name: str = Field(..., description="Storage bucket name")
+    file_key: str = Field(..., description="File key in storage")
+    max_workers: int = Field(
+        default=5, ge=1, le=20, description="Maximum parallel workers"
+    )
+    resp_format: str = Field(..., description="Response format (csv, json, html)")
+
+    @field_validator("bucket_name", "file_key")
+    @classmethod
+    def validate_storage_identifiers(cls, v: str) -> str:
+        if not v or not v.strip():
+            raise ValueError("Storage identifier cannot be empty or whitespace")
+        return v.strip()
+
+
+class PackageSplitConfig(BaseModel):
+    """Configuration for package splitting operations"""
+
+    max_attachments: int = Field(
+        default=100, ge=1, le=1000, description="Max attachments per package"
+    )
+
+    @field_validator("max_attachments")
+    @classmethod
+    def validate_max_attachments(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError("max_attachments must be at least 1")
+        return v
+
+
+class StreamingBufferConfig(BaseModel):
+    """Configuration for streaming buffer operations"""
+
+    zip_buffer_threshold: int = Field(
+        default=5 * 1024 * 1024,
+        ge=1024 * 1024,
+        description="Zip buffer threshold in bytes (default: 5MB)",
+    )
+    stream_buffer_threshold: int = Field(
+        default=1024 * 1024,
+        ge=1024 * 1024,
+        description="Stream buffer threshold in bytes (default: 1MB)",
+    )
+    chunk_size_threshold: int = Field(
+        default=1024 * 1024,
+        ge=1024 * 1024,
+        description="Chunk size threshold in bytes (default: 1MB)",
+    )
+
+
+class AttachmentProcessingInfo(BaseModel):
+    """Information for processing a single attachment"""
+
+    attachment: AttachmentInfo = Field(..., description="Attachment to process")
+    base_path: str = Field(..., description="Base path for the attachment in the zip")
+    item: dict[str, Any] = Field(..., description="Associated item data")
+
+    @field_validator("base_path")
+    @classmethod
+    def validate_base_path(cls, v: str) -> str:
+        if not v or not v.strip():
+            raise ValueError("base_path cannot be empty or whitespace")
+        return v.strip()
+
+
+class MultipartUploadResponse(BaseModel):
+    """Response from initiating a multipart upload"""
+
+    upload_id: Annotated[
+        str, Field(..., description="Unique identifier for the multipart upload")
+    ]
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None, description="Additional metadata from the storage provider"
+    )
+
+    model_config = {
+        "extra": "allow",  # Allow extra fields from storage providers
+    }
+
+    @field_validator("upload_id")
+    @classmethod
+    def validate_upload_id(cls, v: Any) -> str:
+        """Validate upload ID is not empty or whitespace."""
+        if not isinstance(v, str):
+            raise ValueError("Upload ID must be a string")
+        if not v or not v.strip():
+            raise ValueError("Upload ID cannot be empty or whitespace")
+        return v.strip()
+
+    def is_valid_upload_id(self) -> bool:
+        """Check if the upload ID is valid (not empty or whitespace)"""
+        return bool(self.upload_id and self.upload_id.strip())
+
+    def get_provider_metadata(self, key: str, default: Any = None) -> Any:
+        """Get metadata value by key with fallback to default"""
+        if self.metadata is not None:
+            return self.metadata.get(key, default)  # pylint: disable=no-member
+        return default
+
+
+class UploadPartResponse(BaseModel):
+    """Response from uploading a part"""
+
+    etag: Annotated[str, Field(..., description="ETag/checksum of the uploaded part")]
+    part_number: Annotated[
+        int, Field(..., description="Sequential number of the part (1-10,000)")
+    ]
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None, description="Additional metadata from the storage provider"
+    )
+
+    model_config = {
+        "extra": "allow",  # Allow extra fields from storage providers
+    }
+
+    @field_validator("part_number")
+    @classmethod
+    def validate_part_number(cls, v: Any) -> int:
+        """Validate part number is within AWS S3 limits (1-10,000)."""
+        if not isinstance(v, int):
+            raise ValueError("Part number must be an integer")
+        if v < 1 or v > 10000:
+            raise ValueError("Part number must be between 1 and 10,000")
+        return v
+
+    def is_valid_etag(self) -> bool:
+        """Check if the ETag is valid (not empty and properly formatted)"""
+        return bool(self.etag and self.etag.strip())
+
+    def get_provider_metadata(self, key: str, default: Any = None) -> Any:
+        """Get metadata value by key with fallback to default"""
+        if self.metadata is not None:
+            return self.metadata.get(key, default)  # pylint: disable=no-member
+        return default