"""Utilities for S3 key template substitution and helpers.""" import base64 import math import os import re from datetime import datetime, timezone from typing import Optional from abacus_file_upload.constants import ( BYTES_IN_GB, BYTES_IN_KB, BYTES_IN_MB, BYTES_IN_TB, S3_MAX_CHUNK_SIZE_BYTES, S3_MIN_CHUNK_SIZE_BYTES, ) def calculate_optimal_chunk_size( file_size_bytes: int, min_chunk_size_bytes: int = 1 ) -> int: """Calculate optimal chunk size for multipart upload. Ensures: - Chunk size >= min_chunk_size_bytes - Chunk size >= S3_MIN_CHUNK_SIZE_BYTES - Chunk size <= S3_MAX_CHUNK_SIZE_BYTES Args: file_size_bytes: Total file size in bytes min_chunk_size_bytes: Minimum chunk size Returns: Optimal chunk size in bytes """ # Check if file size is at or below minimum chunk size if file_size_bytes <= S3_MIN_CHUNK_SIZE_BYTES: return S3_MIN_CHUNK_SIZE_BYTES # Calculate square root to balance chunk size and part count chunk_size_bytes = math.ceil(math.sqrt(file_size_bytes)) # Enforce given limits chunk_size_bytes = max(chunk_size_bytes, min_chunk_size_bytes) # Enforce hard limits chunk_size_bytes = max(chunk_size_bytes, S3_MIN_CHUNK_SIZE_BYTES) chunk_size_bytes = min(chunk_size_bytes, S3_MAX_CHUNK_SIZE_BYTES) return chunk_size_bytes def convert_md5_hex_to_base64(md5_hex: str) -> str: """Convert hex MD5 to base64 format for S3 ContentMD5. S3 requires ContentMD5 to be base64-encoded binary MD5 digest. Args: md5_hex: MD5 hash as 32 hex characters (e.g., '5d41402abc4b2a76b9719d911017c592') Returns: Base64-encoded MD5 suitable for S3 ContentMD5 parameter """ # Convert hex to binary md5_binary = bytes.fromhex(md5_hex) # Encode to base64 md5_base64 = base64.b64encode(md5_binary).decode('ascii') return md5_base64 def format_bytes(size_bytes: Optional[int]) -> str: """Format bytes into human-readable string with appropriate unit. Args: size_bytes: Size in bytes (None returns 'unlimited') Returns: Formatted string with appropriate unit (B, KB, MB, GB, TB) rounded to 2 decimal places Examples: >>> format_bytes(0) '0 B' >>> format_bytes(1024) '1.00 KB' >>> format_bytes(1536) '1.50 KB' >>> format_bytes(1048576) '1.00 MB' >>> format_bytes(5368709120) '5.00 GB' >>> format_bytes(None) 'unlimited' """ if size_bytes is None: return 'unlimited' if size_bytes == 0: return '0 B' sign = '' if size_bytes >= 0 else '-' size_bytes = abs(size_bytes) # Define units and their thresholds units = ['B', 'KB', 'MB', 'GB', 'TB'] unit_thresholds = [1, BYTES_IN_KB, BYTES_IN_MB, BYTES_IN_GB, BYTES_IN_TB] # Find the appropriate unit unit_index = 1 while ( unit_index < len(unit_thresholds) and size_bytes >= unit_thresholds[unit_index] ): unit_index += 1 # Format with 2 decimal places size = float(size_bytes) / unit_thresholds[unit_index - 1] return f'{sign}{size:.2f} {units[unit_index - 1]}' def generate_key( template: str, file_key: str, upload_type: str, file_name: str, metadata: Optional[dict] = None, upload_time: Optional[datetime] = None, ) -> str: """Generate a key from a template with variable substitution. Args: template: Key template string with variables like {upload_type}, {year}, etc. file_key: Unique upload UUID upload_type: Type of upload file_name: Original filename metadata: Optional metadata dictionary (may contain entity_id) upload_time: Upload timestamp (defaults to current time) Returns: Generated key with all variables substituted Supported template variables: - {year}, {month}, {day}: Date components from upload time (zero-padded) - {file_key}: The unique upload UUID - {entity_id}: Optional entity ID from metadata (removed if not present) - {filename}: Original filename without extension - {ext}: File extension (without dot) - {upload_type}: The upload type (usually not needed since bucket separates types) Examples: >>> generate_key( ... "{year}/{month}/{file_key}.{ext}", ... "abc-123", ... "flowthrough", ... "report.csv" ... ) '2024/12/abc-123.csv' >>> generate_key( ... "{entity_id}/{file_key}.{ext}", ... "abc-123", ... "adjustment", ... "file.csv", ... metadata={"entity_id": 456} ... ) '456/abc-123.csv' """ if upload_time is None: upload_time = datetime.now(timezone.utc) # Extract filename and extension filename_without_ext, ext = get_file_parts(file_name) # Build substitution dictionary variables = { 'upload_type': upload_type, 'year': upload_time.strftime('%Y'), 'month': upload_time.strftime('%m'), 'day': upload_time.strftime('%d'), 'file_key': file_key, 'filename': filename_without_ext, 'ext': ext, } # Add entity_id if present in metadata if metadata and 'entity_id' in metadata: variables['entity_id'] = str(metadata['entity_id']) # Custom dict that returns placeholder for missing keys class FormatDict(dict): def __missing__(self, key): # Return the placeholder unchanged for optional variables if key in ['entity_id']: # List of optional variables return '{' + key + '}' # Raise error for truly invalid variables raise KeyError(key) # Perform substitution try: s3_key = template.format_map(FormatDict(**variables)) except KeyError as e: raise ValueError(f'Invalid template variable: {e.args[0]}') # Clean up any remaining unreplaced variables (e.g., {entity_id} when not provided) # This allows templates to have optional segments # Remove segments that contain unreplaced variables # Pattern: matches path segments containing {variable} s3_key = re.sub(r'/{[^}]+}', '', s3_key) # Remove /{unreplaced} s3_key = re.sub(r'{[^}]+}/', '', s3_key) # Remove {unreplaced}/ s3_key = re.sub(r'{[^}]+}', '', s3_key) # Remove standalone {unreplaced} # Clean up any double slashes that may have resulted s3_key = re.sub(r'/+', '/', s3_key) # Remove leading/trailing slashes s3_key = s3_key.strip('/') return s3_key def get_file_parts(filename: str) -> tuple[str, Optional[str]]: """Extract filename and extension from filename. Args: filename: Original filename Returns: Tuple of (filename_without_ext, extension) Extension is without leading dot and lowercase, or None if no extension Examples: >>> get_file_parts('report.csv') ('report', 'csv') >>> get_file_parts('file.tar.gz') ('file.tar', 'gz') >>> get_file_parts('.gitignore') ('.gitignore', None) >>> get_file_parts('no_extension') ('no_extension', None) """ filename_without_ext, ext = os.path.splitext(filename) ext = ext.lstrip('.').lower() if ext else None return filename_without_ext, ext def get_file_extension(filename: str) -> Optional[str]: """Extract file extension from filename. Args: filename: Original filename Returns: File extension without leading dot and lowercase, or None if no extension """ _, ext = get_file_parts(filename) return ext def validate_file_type( file_type: str | None, allowed_types: Optional[list[str]] ) -> bool: """Validate file type against allowed types. Args: file_type: File extension (lowercase without dot) allowed_types: List of allowed file extensions (lowercase without dot) Returns: True if file type is allowed, False otherwise """ if not allowed_types: return True # No restrictions if not file_type: return False return file_type in [ext for ext in allowed_types]