ethyca-fides 2.69.1rc0__py2.py3-none-any.whl → 2.76.0rc2__py2.py3-none-any.whl

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

@@ -3,11 +3,13 @@ from __future__ import annotations
 
 import csv
 import json
+import time
+from copy import deepcopy
 from datetime import datetime
 from io import BytesIO, StringIO
 from itertools import chain
 from typing import Any, Generator, Iterable, Optional, Tuple
-from urllib.parse import urlparse
+from urllib.parse import unquote, urlparse
 
 from fideslang.validation import AnyHttpUrlString
 from loguru import logger
@@ -25,22 +27,21 @@ from fides.api.service.storage.streaming.dsr_storage import (
 )
 from fides.api.service.storage.streaming.retry import retry_cloud_storage_operation
 from fides.api.service.storage.streaming.schemas import (
-    CHUNK_SIZE_THRESHOLD,
+    DEFAULT_CHUNK_SIZE,
+    MAX_FILE_SIZE,
     AttachmentInfo,
     AttachmentProcessingInfo,
-    PackageSplitConfig,
+    SmartOpenStreamingStorageConfig,
     StorageUploadConfig,
     StreamingBufferConfig,
 )
 from fides.api.service.storage.streaming.smart_open_client import SmartOpenStorageClient
 from fides.api.service.storage.util import (
-    convert_processed_attachments_to_attachment_processing_info,
     determine_dataset_name_from_path,
-    extract_storage_key_from_attachment,
     get_unique_filename,
     process_attachments_contextually,
     resolve_attachment_storage_path,
-    resolve_base_path_from_context,
+    resolve_path_from_context,
 )
 
 DEFAULT_ATTACHMENT_NAME = "attachment"
@@ -66,16 +67,22 @@ class SmartOpenStreamingStorage:
     def __init__(
         self,
         storage_client: SmartOpenStorageClient,
-        chunk_size: int = CHUNK_SIZE_THRESHOLD,
+        chunk_size: int = DEFAULT_CHUNK_SIZE,
     ):
         """Initialize with a smart-open storage client.
 
         Args:
             storage_client: Smart-open based storage client
-            chunk_size: Size of chunks for streaming attachments (default: 8KB)
+            chunk_size: Size of chunks for streaming attachments (default: 5MB)
+
+        Raises:
+            ValidationError: If chunk_size is outside valid range (1KB - 2GB)
         """
+        # Validate parameters using Pydantic schema
+        config = SmartOpenStreamingStorageConfig(chunk_size=chunk_size)
+
         self.storage_client = storage_client
-        self.chunk_size = chunk_size
+        self.chunk_size = config.chunk_size
         # Track used filenames per dataset to match DSR report builder behavior
         # Maps dataset_name -> set of used filenames
         self.used_filenames_per_dataset: dict[str, set[str]] = {}
@@ -98,10 +105,15 @@ class SmartOpenStreamingStorage:
         Raises:
             ValueError: If URL cannot be parsed
         """
+        if storage_key is None or storage_key == "":
+            logger.error(f"Storage key cannot be empty: {storage_key}")
+            raise ValueError("Storage key cannot be empty")
+
         if storage_key.startswith("s3://"):
             # Extract bucket from S3 URL: s3://bucket/path
             parts = storage_key.split("/")
             if len(parts) < 4:
+                logger.error(f"Invalid S3 URL format: {storage_key}")
                 raise ValueError(f"Invalid S3 URL format: {storage_key}")
             return parts[2], "/".join(parts[3:])
 
@@ -111,16 +123,16 @@ class SmartOpenStreamingStorage:
             parts = clean_url.split(S3_AMAZONAWS_COM_DOMAIN)
             if len(parts) == 2:
                 bucket = parts[0].replace("https://", "").replace("http://", "")
-                key = parts[1].lstrip(
-                    "/"
-                )  # Strip leading forward slash for S3 compatibility
+                key = unquote(
+                    parts[1].lstrip("/")
+                )  # URL-decode and strip leading slash for S3 compatibility
                 return bucket, key
 
         # Handle generic HTTP(S) URLs
         if storage_key.startswith(("http://", "https://")):
             parsed = urlparse(storage_key)
             bucket = parsed.netloc
-            key = parsed.path.lstrip("/")
+            key = unquote(parsed.path.lstrip("/"))  # URL-decode the path
             return bucket, key
 
         raise ValueError(f"Could not parse storage URL: {storage_key}")
@@ -143,103 +155,6 @@ class SmartOpenStreamingStorage:
                 [content_bytes]
             )
 
-    def build_attachments_list(
-        self, data: dict, config: PackageSplitConfig
-    ) -> list[tuple[str, dict, int]]:
-        """
-        Build a list of attachments from the data.
-
-        Args:
-            data: The data to build the attachments list from
-            config: The configuration for package splitting
-
-        Returns:
-            A list of AttachmentInfo objects
-        """
-        attachments_list = []
-        for key, value in data.items():
-            if not isinstance(value, list):
-                continue
-
-            for item in value:
-                attachments = item.get("attachments", [])
-                if not isinstance(attachments, list):
-                    attachments = []
-
-                attachment_count = len(attachments)
-
-                # Only include items that have attachments
-                if attachment_count > 0:
-                    # If a single item has more attachments than the limit, we need to split it
-                    if attachment_count > config.max_attachments:
-                        # Split the item into multiple sub-items
-                        for i in range(0, attachment_count, config.max_attachments):
-                            sub_attachments = attachments[
-                                i : i + config.max_attachments
-                            ]
-                            sub_item = item.copy()
-                            sub_item["attachments"] = sub_attachments
-                            attachments_list.append(
-                                (key, sub_item, len(sub_attachments))
-                            )
-                    else:
-                        attachments_list.append((key, item, attachment_count))
-
-        return attachments_list
-
-    def split_data_into_packages(
-        self, data: dict, config: Optional[PackageSplitConfig] = None
-    ) -> list[dict]:
-        """Split large datasets into multiple smaller packages.
-
-        Uses a best-fit decreasing algorithm to optimize package distribution:
-        1. Sort items by attachment count (largest first)
-        2. Try to fit each item in the package with the most remaining space
-        3. Create new packages only when necessary
-        4. Handle items that exceed the max_attachments limit by splitting them
-
-        Args:
-            data: The data to split
-            config: Configuration for package splitting (defaults to PackageSplitConfig())
-
-        Returns:
-            List of data packages
-        """
-        # Use default config if none provided
-        if config is None:
-            config = PackageSplitConfig()
-
-        # Collect all items with their attachment counts
-        all_items = self.build_attachments_list(data, config)
-
-        # Sort by attachment count (largest first) for better space utilization
-        all_items.sort(key=lambda x: x[2], reverse=True)
-
-        packages: list[dict[str, Any]] = []
-        package_attachment_counts: list[int] = []
-
-        for key, item, attachment_count in all_items:
-            # Try to find a package with enough space
-            package_found = False
-
-            for i, current_count in enumerate(package_attachment_counts):
-                if current_count + attachment_count <= config.max_attachments:
-                    # Add to existing package
-                    if key not in packages[i]:
-                        packages[i][key] = []
-                    packages[i][key].append(item)
-                    package_attachment_counts[i] += attachment_count
-                    package_found = True
-                    break
-
-            if not package_found:
-                # Create new package - this item cannot fit in any existing package
-                new_package = {key: [item]}
-                packages.append(new_package)
-                package_attachment_counts.append(attachment_count)
-
-        return packages
-
     def _validate_attachment(
         self, attachment: dict
     ) -> Optional[AttachmentProcessingInfo]:
@@ -253,7 +168,9 @@ class SmartOpenStreamingStorage:
         """
         try:
             # Extract storage key using shared utility
-            storage_key = extract_storage_key_from_attachment(attachment)
+            storage_key = (
+                attachment.get("download_url") or attachment.get("file_name") or ""
+            )
             if not storage_key:
                 return None
 
@@ -266,7 +183,7 @@ class SmartOpenStreamingStorage:
             )
 
             # Resolve base path using shared utility
-            base_path = resolve_base_path_from_context(attachment)
+            base_path = resolve_path_from_context(attachment)
 
             # Create AttachmentProcessingInfo
             processing_info = AttachmentProcessingInfo(
@@ -294,26 +211,82 @@ class SmartOpenStreamingStorage:
         Returns:
             Iterator that yields chunks of the attachment content
         """
+        for arg in [bucket, key, storage_key]:
+            if arg is None or arg == "":
+                logger.error(f"{arg} cannot be empty: {arg}")
+                raise ValueError(f"{arg} cannot be empty")
+
         try:
             with self.storage_client.stream_read(bucket, key) as content_stream:
                 # Stream in chunks instead of reading entire file
-                chunk_count = 0
-                total_bytes = 0
-                while True:
-                    chunk = content_stream.read(self.chunk_size)
+                chunk_count = total_bytes = 0
+                max_chunks = (
+                    MAX_FILE_SIZE // self.chunk_size + 1
+                )  # Safety limit to prevent infinite loops
+
+                size_based_timeout = MAX_FILE_SIZE // (10 * 1024 * 1024)  # 1s per 10MB
+                timeout = 300 + size_based_timeout  # 5 minutes base + 1s per 10MB
+                start_time = time.time()
+
+                # Log the calculated timeout for debugging
+                logger.debug(
+                    f"Starting stream for {storage_key} with timeout: {timeout}s "
+                    f"(base: 300s + size-based: {size_based_timeout}s)"
+                )
+
+                while chunk_count < max_chunks and total_bytes < MAX_FILE_SIZE:
+                    elapsed_time = time.time() - start_time
+                    if elapsed_time >= timeout:
+                        raise TimeoutError(
+                            f"Timeout reached ({timeout}s) while streaming attachment {storage_key}."
+                        )
+
+                    try:
+                        chunk = content_stream.read(self.chunk_size)
+                    except Exception as read_error:
+                        logger.error(
+                            f"Error reading chunk from stream for {storage_key}: {read_error}"
+                        )
+                        raise StorageUploadError(
+                            f"Stream read error for {storage_key}: {read_error}"
+                        ) from read_error
+
                     if not chunk:
+                        # End of stream reached normally
+                        logger.debug(
+                            f"Successfully streamed attachment {storage_key}: "
+                            f"{total_bytes} bytes in {chunk_count} chunks"
+                        )
                         break
+
                     chunk_count += 1
                     total_bytes += len(chunk)
                     yield chunk
 
+                # Log if we hit limits
+                if chunk_count >= max_chunks:
+                    logger.warning(
+                        f"Maximum chunk count ({max_chunks}) reached for attachment {storage_key}. "
+                        f"Streamed {total_bytes} bytes. Stream may be incomplete."
+                    )
+                elif total_bytes >= MAX_FILE_SIZE:
+                    logger.warning(
+                        f"Maximum file size ({MAX_FILE_SIZE} bytes) reached for attachment {storage_key}. "
+                        f"Streamed {total_bytes} bytes in {chunk_count} chunks. Stream may be incomplete."
+                    )
+
         except Exception as e:
-            logger.warning(f"Failed to stream attachment {storage_key}: {e}")
-            # Yield empty content on failure
-            yield b""
+            logger.error(f"Failed to stream attachment {storage_key}: {e}")
+            raise StorageUploadError(
+                f"Failed to stream attachment {storage_key}: {e}"
+            ) from e
 
     def _collect_and_validate_attachments(
-        self, data: dict
+        self,
+        data: dict,
+        used_filenames_data: set[str],
+        used_filenames_attachments: set[str],
+        processed_attachments: dict[tuple[str, str], str],
     ) -> list[AttachmentProcessingInfo]:
         """Collect and validate attachments using the same contextual approach as DSR report builder.
 
@@ -326,24 +299,30 @@ class SmartOpenStreamingStorage:
         Returns:
             List of validated AttachmentProcessingInfo objects
         """
-        # Initialize tracking structures (similar to DSR report builder)
-        used_filenames_data: set[str] = set()
-        used_filenames_attachments: set[str] = set()
-        processed_attachments: dict[tuple[str, str], str] = {}
+        validated_attachments = []
 
         # Use the shared contextual processing function
+        # Note: This method should only be used when DSR report builder is not available
+        # For HTML format, use _collect_and_validate_attachments_from_dsr_builder instead
         processed_attachments_list = process_attachments_contextually(
-            data,
-            used_filenames_data,
-            used_filenames_attachments,
-            processed_attachments,
-            enable_streaming=True,  # Always use streaming mode for storage
+            data=data,
+            used_filenames_data=used_filenames_data,
+            used_filenames_attachments=used_filenames_attachments,
+            processed_attachments=processed_attachments,
+            enable_streaming=True,
         )
 
-        # Convert to AttachmentProcessingInfo objects using shared utility
-        return convert_processed_attachments_to_attachment_processing_info(
-            processed_attachments_list, self._validate_attachment
-        )
+        for processed_attachment in processed_attachments_list:
+            # Add context information to the attachment data
+            attachment_with_context = deepcopy(processed_attachment["attachment"])
+            attachment_with_context["_context"] = processed_attachment["context"]
+
+            # Validate and convert to AttachmentProcessingInfo
+            validated = self._validate_attachment(attachment_with_context)
+            if validated:
+                validated_attachments.append(validated)
+
+        return validated_attachments
 
     def _collect_and_validate_attachments_from_dsr_builder(
         self, data: dict, dsr_builder: "DSRReportBuilder"
@@ -372,17 +351,11 @@ class SmartOpenStreamingStorage:
             else:
                 used_filenames_data.update(filenames)
 
-        processed_attachments_list = process_attachments_contextually(
+        return self._collect_and_validate_attachments(
             data,
             used_filenames_data,
             used_filenames_attachments,
             dsr_builder.processed_attachments,
-            enable_streaming=True,  # Always use streaming mode for storage
-        )
-
-        # Convert to AttachmentProcessingInfo objects using shared utility
-        return convert_processed_attachments_to_attachment_processing_info(
-            processed_attachments_list, self._validate_attachment
         )
 
     @retry_cloud_storage_operation(
@@ -404,7 +377,7 @@ class SmartOpenStreamingStorage:
         """Upload data to cloud storage using smart-open streaming for memory efficiency.
 
         This function leverages smart-open's streaming capabilities while maintaining
-        our DSR-specific business logic for package splitting and attachment processing.
+        our DSR-specific business logic for package and attachment processing.
         All data is streamed directly from source to destination without local storage.
 
         Args:
@@ -557,8 +530,9 @@ class SmartOpenStreamingStorage:
             raise StorageUploadError(f"Failed to generate DSR report: {e}") from e
 
         # Use the DSR report builder's processed attachments to avoid duplicates
+        # Use the redacted data from the DSR report builder instead of the original data
         all_attachments = self._collect_and_validate_attachments_from_dsr_builder(
-            data, dsr_builder
+            dsr_builder.dsr_data, dsr_builder
         )
 
         if not all_attachments:
@@ -595,7 +569,9 @@ class SmartOpenStreamingStorage:
         )
 
         # Create ZIP generator with attachment files
-        attachment_files_generator = self._create_attachment_files(all_attachments)
+        attachment_files_generator = self._create_attachment_files(
+            all_attachments, buffer_config
+        )
 
         # Combine both generators and stream the complete ZIP to storage
         combined_entries = chain(attachment_files_generator, dsr_files_generator)
@@ -654,7 +630,12 @@ class SmartOpenStreamingStorage:
             resp_format: Response format (csv, json)
         """
         # Collect and validate all attachments using shared contextual processing
-        all_attachments = self._collect_and_validate_attachments(data)
+        all_attachments = self._collect_and_validate_attachments(
+            data=data,
+            used_filenames_data=set(),
+            used_filenames_attachments=set(),
+            processed_attachments={},
+        )
 
         if not all_attachments:
             # No attachments, just upload the data
@@ -673,6 +654,7 @@ class SmartOpenStreamingStorage:
             max_workers,
             batch_size,
             resp_format,
+            buffer_config,
         )
 
         # Use smart-open's streaming upload capability
@@ -718,6 +700,7 @@ class SmartOpenStreamingStorage:
         max_workers: int,
         batch_size: int,
         resp_format: str,
+        buffer_config: Optional[StreamingBufferConfig] = None,
     ) -> Generator[Tuple[str, datetime, int, Any, Iterable[bytes]], None, None]:
         """Create a generator for ZIP file contents including data and attachments.
 
@@ -743,7 +726,9 @@ class SmartOpenStreamingStorage:
             yield from self._convert_to_stream_zip_format(data_files_generator)
 
         # Then, yield attachment files (already in stream_zip format, stream directly)
-        attachment_files_generator = self._create_attachment_files(all_attachments)
+        attachment_files_generator = self._create_attachment_files(
+            all_attachments, buffer_config
+        )
         yield from attachment_files_generator
 
     def _create_data_files(
@@ -785,9 +770,62 @@ class SmartOpenStreamingStorage:
                     data_content = json.dumps(value, default=str).encode("utf-8")
                     yield f"{key}.json", BytesIO(data_content), {}
 
+    def _handle_attachment_error(
+        self,
+        all_attachments: list[AttachmentProcessingInfo],
+        failed_attachments: list[dict[str, Optional[str]]],
+    ) -> Generator[Tuple[str, datetime, int, Any, Iterable[bytes]], None, None]:
+        """Handle attachment errors and create a summary file."""
+
+        try:
+            # Calculate success rate with division by zero protection
+            total_attempted = len(all_attachments)
+            total_failed = len(failed_attachments)
+            success_rate = "N/A"
+
+            if total_attempted > 0:
+                success_rate = (
+                    f"{((total_attempted - total_failed) / total_attempted * 100):.1f}%"
+                )
+
+            error_summary = {
+                "failed_attachments": failed_attachments,
+                "total_failed": total_failed,
+                "total_attempted": total_attempted,
+                "success_rate": success_rate,
+                "timestamp": datetime.now().isoformat(),
+            }
+
+            error_summary_content = json.dumps(error_summary, indent=2).encode("utf-8")
+            yield (
+                "errors/attachment_failures_summary.json",
+                datetime.now(),
+                DEFAULT_FILE_MODE,
+                _ZIP_32_TYPE(),
+                iter([error_summary_content]),
+            )
+        except Exception as summary_error:
+            logger.error(f"Failed to create error summary: {summary_error}")
+            # Create a minimal error summary as fallback
+            fallback_summary = {
+                "error": "Failed to generate detailed error summary",
+                "total_failed": len(failed_attachments),
+                "total_attempted": len(all_attachments),
+                "timestamp": datetime.now().isoformat(),
+            }
+            fallback_content = json.dumps(fallback_summary, indent=2).encode("utf-8")
+            yield (
+                "errors/attachment_failures_summary.json",
+                datetime.now(),
+                DEFAULT_FILE_MODE,
+                _ZIP_32_TYPE(),
+                iter([fallback_content]),
+            )
+
     def _create_attachment_files(
         self,
         all_attachments: list[AttachmentProcessingInfo],
+        buffer_config: Optional[StreamingBufferConfig] = None,
     ) -> Generator[Tuple[str, datetime, int, Any, Iterable[bytes]], None, None]:
         """Create attachment files for the ZIP using true cloud-to-cloud streaming.
 
@@ -797,13 +835,66 @@ class SmartOpenStreamingStorage:
 
         Args:
             all_attachments: List of validated attachments
+            buffer_config: Configuration for error handling behavior
 
         Returns:
             Generator yielding attachment file entries in stream_zip format
         """
+        if buffer_config is None:
+            buffer_config = StreamingBufferConfig()
+
+        failed_attachments = []
+
         for attachment_info in all_attachments:
-            result = self._process_attachment_safely(attachment_info)
-            yield result
+            try:
+                result = self._process_attachment_safely(attachment_info)
+                yield result
+            except StorageUploadError as e:
+                # Log the failure
+                failed_attachments.append(
+                    {
+                        "attachment": attachment_info.attachment.file_name,
+                        "storage_key": attachment_info.attachment.storage_key,
+                        "error": str(e),
+                    }
+                )
+                logger.error(
+                    f"Failed to process attachment {attachment_info.attachment.file_name} "
+                    f"({attachment_info.attachment.storage_key}): {e}"
+                )
+
+                # If fail_fast is enabled, re-raise the exception
+                if buffer_config.fail_fast_on_attachment_errors:
+                    raise
+
+                # Create a placeholder file with error information if error details are enabled
+                if buffer_config.include_error_details:
+                    error_content = (
+                        f"Error: Failed to retrieve attachment - {e}".encode("utf-8")
+                    )
+                    error_filename = (
+                        f"ERROR_{attachment_info.attachment.file_name or 'unknown'}"
+                    )
+                    yield (
+                        f"errors/{error_filename}",
+                        datetime.now(),
+                        DEFAULT_FILE_MODE,
+                        _ZIP_32_TYPE(),
+                        iter([error_content]),
+                    )
+
+        # Log summary of failed attachments
+        if failed_attachments:
+            logger.warning(
+                f"Failed to process {len(failed_attachments)} attachments: "
+                f"{[att['attachment'] for att in failed_attachments]}"
+            )
+
+            # Create a summary error file with details about all failures if error details are enabled
+            if buffer_config.include_error_details:
+                yield from self._handle_attachment_error(
+                    all_attachments, failed_attachments
+                )
 
     def _transform_data_for_access_package(
         self, data: dict[str, Any], all_attachments: list[AttachmentProcessingInfo]