Spaces:
Runtime error
Runtime error
| """ | |
| Consolidated S3 client for all storage operations. | |
| Provides: | |
| - Environment-based key prefixing for staging/production isolation | |
| - Unified upload interface for transcripts, logs, and state | |
| - Validation of required environment variables | |
| """ | |
| import os | |
| from typing import Optional | |
| import boto3 | |
| from loguru import logger | |
| def get_s3_key_with_prefix(key: str) -> str: | |
| """ | |
| Add S3_PREFIX to key for data isolation between environments. | |
| If S3_PREFIX is empty or not set, returns key unchanged for backward | |
| compatibility with existing production data. | |
| Args: | |
| key: Original S3 key (e.g., 'transcripts/session.vtt') | |
| Returns: | |
| Prefixed key if S3_PREFIX is set (e.g., 'staging/transcripts/session.vtt') | |
| Original key if S3_PREFIX is empty (backward compatible for production) | |
| Examples: | |
| # Production (S3_PREFIX not set) | |
| >>> os.environ.pop('S3_PREFIX', None) | |
| >>> get_s3_key_with_prefix('transcripts/session.vtt') | |
| 'transcripts/session.vtt' | |
| # Staging (S3_PREFIX='staging') | |
| >>> os.environ['S3_PREFIX'] = 'staging' | |
| >>> get_s3_key_with_prefix('transcripts/session.vtt') | |
| 'staging/transcripts/session.vtt' | |
| """ | |
| prefix = os.environ.get("S3_PREFIX", "").strip() | |
| if not prefix: | |
| return key | |
| # Ensure prefix doesn't start/end with slashes | |
| prefix = prefix.strip("/") | |
| # Ensure key doesn't start with slash | |
| key = key.lstrip("/") | |
| return f"{prefix}/{key}" | |
| class S3Client: | |
| """ | |
| Unified S3 client for all storage operations. | |
| Centralizes S3 operations to avoid code duplication across modules. | |
| Validates environment variables and provides consistent error handling. | |
| """ | |
| REQUIRED_ENV_VARS = [ | |
| "AWS_ACCESS_KEY_ID", | |
| "AWS_SECRET_ACCESS_KEY", | |
| "BUCKET_REGION", | |
| "BUCKET_NAME", | |
| ] | |
| def __init__(self): | |
| self._client: Optional[boto3.client] = None | |
| self._bucket: Optional[str] = None | |
| def _ensure_initialized(self) -> bool: | |
| """Lazily initialize S3 client when first needed.""" | |
| if self._client is not None: | |
| return True | |
| missing_vars = [var for var in self.REQUIRED_ENV_VARS if not os.environ.get(var)] | |
| if missing_vars: | |
| logger.warning(f"S3Client: Missing environment variables: {missing_vars}") | |
| return False | |
| try: | |
| self._client = boto3.client( | |
| "s3", | |
| aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"], | |
| aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"], | |
| region_name=os.environ["BUCKET_REGION"], | |
| ) | |
| self._bucket = os.environ["BUCKET_NAME"] | |
| return True | |
| except Exception as e: | |
| logger.error(f"S3Client: Failed to initialize: {e}") | |
| return False | |
| def upload( | |
| self, | |
| key: str, | |
| body: str, | |
| content_type: str = "application/octet-stream", | |
| use_prefix: bool = True, | |
| ) -> bool: | |
| """ | |
| Upload content to S3. | |
| Args: | |
| key: S3 object key | |
| body: Content to upload | |
| content_type: MIME type of the content | |
| use_prefix: Whether to apply S3_PREFIX for environment isolation | |
| Returns: | |
| True if upload succeeded, False otherwise | |
| """ | |
| if not self._ensure_initialized(): | |
| return False | |
| try: | |
| final_key = get_s3_key_with_prefix(key) if use_prefix else key | |
| self._client.put_object( | |
| Bucket=self._bucket, | |
| Key=final_key, | |
| Body=body, | |
| ContentType=content_type, | |
| ) | |
| logger.info(f"S3Client: Uploaded to s3://{self._bucket}/{final_key}") | |
| return True | |
| except Exception as e: | |
| logger.error(f"S3Client: Upload failed for {key}: {e}") | |
| return False | |
| def upload_transcript(self, key: str, vtt_content: str) -> bool: | |
| """Upload VTT transcript to S3.""" | |
| return self.upload(key, vtt_content, content_type="text/vtt") | |
| def upload_json(self, key: str, json_content: str) -> bool: | |
| """Upload JSON content to S3.""" | |
| return self.upload(key, json_content, content_type="application/json") | |
| def upload_log(self, key: str, log_content: str) -> bool: | |
| """Upload log file to S3.""" | |
| if not key.endswith(".log"): | |
| key = f"{key}.log" | |
| return self.upload(key, log_content, content_type="text/plain") | |
| def upload_state(self, base_key: str, state_json: str) -> bool: | |
| """ | |
| Upload flow state to S3. | |
| Args: | |
| base_key: Base key (typically transcript key), will add _state.json suffix | |
| state_json: JSON string of flow state | |
| """ | |
| if base_key.endswith(".vtt"): | |
| state_key = base_key.replace(".vtt", "_state.json") | |
| else: | |
| state_key = f"{base_key}_state.json" | |
| return self.upload_json(state_key, state_json) | |
| # Global singleton for convenience | |
| _s3_client: Optional[S3Client] = None | |
| def get_s3_client() -> S3Client: | |
| """Get the global S3 client instance.""" | |
| global _s3_client | |
| if _s3_client is None: | |
| _s3_client = S3Client() | |
| return _s3_client | |