aiecs 1.0.0__py3-none-any.whl

aiecs/infrastructure/persistence/file_storage.py

@@ -0,0 +1,671 @@
+"""
+File Storage Implementation with Google Cloud Storage
+
+Provides file storage capabilities using Google Cloud Storage as the backend,
+with support for local fallback and caching.
+"""
+
+import os
+import json
+import logging
+import asyncio
+import aiofiles
+from typing import Dict, List, Any, Optional, Union, BinaryIO
+from datetime import datetime, timedelta
+from pathlib import Path
+import hashlib
+import gzip
+import pickle
+
+try:
+    from google.cloud import storage
+    from google.cloud.exceptions import NotFound, GoogleCloudError
+    from google.auth.exceptions import DefaultCredentialsError
+    GCS_AVAILABLE = True
+except ImportError:
+    GCS_AVAILABLE = False
+    storage = None
+    NotFound = Exception
+    GoogleCloudError = Exception
+    DefaultCredentialsError = Exception
+
+from ..monitoring.executor_metrics import ExecutorMetrics
+
+logger = logging.getLogger(__name__)
+
+
+class FileStorageError(Exception):
+    """Base exception for file storage operations."""
+    pass
+
+
+class FileStorageConfig:
+    """Configuration for file storage."""
+
+    def __init__(self, config: Dict[str, Any]):
+        # Google Cloud Storage settings
+        self.gcs_bucket_name = config.get('gcs_bucket_name', 'multi-task-storage')
+        self.gcs_project_id = config.get('gcs_project_id')
+        self.gcs_credentials_path = config.get('gcs_credentials_path')
+        self.gcs_location = config.get('gcs_location', 'US')
+
+        # Local storage fallback
+        self.local_storage_path = config.get('local_storage_path', './storage')
+        self.enable_local_fallback = config.get('enable_local_fallback', True)
+
+        # Cache settings
+        self.enable_cache = config.get('enable_cache', True)
+        self.cache_ttl_seconds = config.get('cache_ttl_seconds', 3600)
+        self.max_cache_size_mb = config.get('max_cache_size_mb', 100)
+
+        # Performance settings
+        self.chunk_size = config.get('chunk_size', 8192)
+        self.max_retries = config.get('max_retries', 3)
+        self.timeout_seconds = config.get('timeout_seconds', 30)
+
+        # Compression settings
+        self.enable_compression = config.get('enable_compression', True)
+        self.compression_threshold_bytes = config.get('compression_threshold_bytes', 1024)
+
+        # Security settings
+        self.enable_encryption = config.get('enable_encryption', False)
+        self.encryption_key = config.get('encryption_key')
+
+
+class FileStorage:
+    """
+    File storage implementation with Google Cloud Storage backend.
+
+    Features:
+    - Google Cloud Storage as primary backend
+    - Local filesystem fallback
+    - In-memory caching with TTL
+    - Automatic compression for large files
+    - Retry logic with exponential backoff
+    - Metrics collection
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = FileStorageConfig(config)
+        self._gcs_client = None
+        self._gcs_bucket = None
+        self._cache = {}
+        self._cache_timestamps = {}
+        self._initialized = False
+
+        # Metrics
+        self.metrics = ExecutorMetrics(enable_metrics=True)
+
+        # Ensure local storage directory exists
+        if self.config.enable_local_fallback:
+            Path(self.config.local_storage_path).mkdir(parents=True, exist_ok=True)
+
+    async def initialize(self) -> bool:
+        """
+        Initialize the file storage system.
+
+        Returns:
+            True if initialization was successful
+        """
+        try:
+            if GCS_AVAILABLE:
+                await self._init_gcs()
+            else:
+                logger.warning("Google Cloud Storage not available, using local storage only")
+
+            self._initialized = True
+            logger.info("File storage initialized successfully")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to initialize file storage: {e}")
+            if not self.config.enable_local_fallback:
+                raise FileStorageError(f"Storage initialization failed: {e}")
+
+            logger.info("Falling back to local storage only")
+            self._initialized = True
+            return True
+
+    async def _init_gcs(self):
+        """Initialize Google Cloud Storage client."""
+        try:
+            # Set credentials if provided
+            if self.config.gcs_credentials_path:
+                os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = self.config.gcs_credentials_path
+
+            # Create client
+            self._gcs_client = storage.Client(project=self.config.gcs_project_id)
+
+            # Get or create bucket
+            try:
+                self._gcs_bucket = self._gcs_client.bucket(self.config.gcs_bucket_name)
+                # Test bucket access
+                self._gcs_bucket.reload()
+                logger.info(f"Connected to GCS bucket: {self.config.gcs_bucket_name}")
+
+            except NotFound:
+                # Create bucket if it doesn't exist
+                self._gcs_bucket = self._gcs_client.create_bucket(
+                    self.config.gcs_bucket_name,
+                    location=self.config.gcs_location
+                )
+                logger.info(f"Created GCS bucket: {self.config.gcs_bucket_name}")
+
+        except DefaultCredentialsError:
+            logger.warning("GCS credentials not found, using local storage only")
+            self._gcs_client = None
+            self._gcs_bucket = None
+
+        except Exception as e:
+            logger.error(f"Failed to initialize GCS: {e}")
+            self._gcs_client = None
+            self._gcs_bucket = None
+
+    async def store(self, key: str, data: Union[str, bytes, Dict[str, Any]],
+                   metadata: Optional[Dict[str, Any]] = None) -> bool:
+        """
+        Store data with the given key.
+
+        Args:
+            key: Storage key
+            data: Data to store
+            metadata: Optional metadata
+
+        Returns:
+            True if storage was successful
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        start_time = datetime.utcnow()
+
+        try:
+            # Serialize data
+            serialized_data = await self._serialize_data(data)
+
+            # Compress if enabled and data is large enough
+            if (self.config.enable_compression and
+                len(serialized_data) > self.config.compression_threshold_bytes):
+                serialized_data = gzip.compress(serialized_data)
+                compressed = True
+            else:
+                compressed = False
+
+            # Store in cache
+            if self.config.enable_cache:
+                self._cache[key] = {
+                    'data': data,
+                    'metadata': metadata,
+                    'compressed': compressed
+                }
+                self._cache_timestamps[key] = datetime.utcnow()
+                await self._cleanup_cache()
+
+            # Store in GCS if available
+            if self._gcs_bucket:
+                success = await self._store_gcs(key, serialized_data, metadata, compressed)
+                if success:
+                    self.metrics.record_operation('gcs_store_success', 1)
+                    duration = (datetime.utcnow() - start_time).total_seconds()
+                    self.metrics.record_duration('gcs_store_duration', duration)
+                    return True
+
+            # Fallback to local storage
+            if self.config.enable_local_fallback:
+                success = await self._store_local(key, serialized_data, metadata, compressed)
+                if success:
+                    self.metrics.record_operation('local_store_success', 1)
+                    duration = (datetime.utcnow() - start_time).total_seconds()
+                    self.metrics.record_duration('local_store_duration', duration)
+                    return True
+
+            self.metrics.record_operation('store_failure', 1)
+            return False
+
+        except Exception as e:
+            logger.error(f"Failed to store data for key {key}: {e}")
+            self.metrics.record_operation('store_error', 1)
+            raise FileStorageError(f"Storage failed: {e}")
+
+    async def retrieve(self, key: str) -> Optional[Union[str, bytes, Dict[str, Any]]]:
+        """
+        Retrieve data by key.
+
+        Args:
+            key: Storage key
+
+        Returns:
+            The stored data if found, None otherwise
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        start_time = datetime.utcnow()
+
+        try:
+            # Check cache first
+            if self.config.enable_cache and key in self._cache:
+                cache_time = self._cache_timestamps.get(key)
+                if cache_time and (datetime.utcnow() - cache_time).total_seconds() < self.config.cache_ttl_seconds:
+                    self.metrics.record_operation('cache_hit', 1)
+                    return self._cache[key]['data']
+                else:
+                    # Remove expired cache entry
+                    self._cache.pop(key, None)
+                    self._cache_timestamps.pop(key, None)
+
+            # Try GCS first
+            if self._gcs_bucket:
+                data = await self._retrieve_gcs(key)
+                if data is not None:
+                    self.metrics.record_operation('gcs_retrieve_success', 1)
+                    duration = (datetime.utcnow() - start_time).total_seconds()
+                    self.metrics.record_duration('gcs_retrieve_duration', duration)
+
+                    # Update cache
+                    if self.config.enable_cache:
+                        self._cache[key] = {'data': data, 'metadata': {}}
+                        self._cache_timestamps[key] = datetime.utcnow()
+
+                    return data
+
+            # Fallback to local storage
+            if self.config.enable_local_fallback:
+                data = await self._retrieve_local(key)
+                if data is not None:
+                    self.metrics.record_operation('local_retrieve_success', 1)
+                    duration = (datetime.utcnow() - start_time).total_seconds()
+                    self.metrics.record_duration('local_retrieve_duration', duration)
+
+                    # Update cache
+                    if self.config.enable_cache:
+                        self._cache[key] = {'data': data, 'metadata': {}}
+                        self._cache_timestamps[key] = datetime.utcnow()
+
+                    return data
+
+            self.metrics.record_operation('retrieve_not_found', 1)
+            return None
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve data for key {key}: {e}")
+            self.metrics.record_operation('retrieve_error', 1)
+            raise FileStorageError(f"Retrieval failed: {e}")
+
+    async def delete(self, key: str) -> bool:
+        """
+        Delete data by key.
+
+        Args:
+            key: Storage key
+
+        Returns:
+            True if deletion was successful
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            success = True
+
+            # Remove from cache
+            if self.config.enable_cache:
+                self._cache.pop(key, None)
+                self._cache_timestamps.pop(key, None)
+
+            # Delete from GCS
+            if self._gcs_bucket:
+                gcs_success = await self._delete_gcs(key)
+                if gcs_success:
+                    self.metrics.record_operation('gcs_delete_success', 1)
+                else:
+                    success = False
+
+            # Delete from local storage
+            if self.config.enable_local_fallback:
+                local_success = await self._delete_local(key)
+                if local_success:
+                    self.metrics.record_operation('local_delete_success', 1)
+                else:
+                    success = False
+
+            if success:
+                self.metrics.record_operation('delete_success', 1)
+            else:
+                self.metrics.record_operation('delete_failure', 1)
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to delete data for key {key}: {e}")
+            self.metrics.record_operation('delete_error', 1)
+            raise FileStorageError(f"Deletion failed: {e}")
+
+    async def exists(self, key: str) -> bool:
+        """
+        Check if data exists for the given key.
+
+        Args:
+            key: Storage key
+
+        Returns:
+            True if data exists
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            # Check cache first
+            if self.config.enable_cache and key in self._cache:
+                cache_time = self._cache_timestamps.get(key)
+                if cache_time and (datetime.utcnow() - cache_time).total_seconds() < self.config.cache_ttl_seconds:
+                    return True
+
+            # Check GCS
+            if self._gcs_bucket:
+                if await self._exists_gcs(key):
+                    return True
+
+            # Check local storage
+            if self.config.enable_local_fallback:
+                return await self._exists_local(key)
+
+            return False
+
+        except Exception as e:
+            logger.error(f"Failed to check existence for key {key}: {e}")
+            raise FileStorageError(f"Existence check failed: {e}")
+
+    async def list_keys(self, prefix: Optional[str] = None, limit: Optional[int] = None) -> List[str]:
+        """
+        List storage keys with optional prefix filtering.
+
+        Args:
+            prefix: Optional key prefix filter
+            limit: Maximum number of keys to return
+
+        Returns:
+            List of storage keys
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            keys = set()
+
+            # Get keys from GCS
+            if self._gcs_bucket:
+                gcs_keys = await self._list_keys_gcs(prefix, limit)
+                keys.update(gcs_keys)
+
+            # Get keys from local storage
+            if self.config.enable_local_fallback:
+                local_keys = await self._list_keys_local(prefix, limit)
+                keys.update(local_keys)
+
+            # Apply limit if specified
+            keys_list = list(keys)
+            if limit:
+                keys_list = keys_list[:limit]
+
+            return keys_list
+
+        except Exception as e:
+            logger.error(f"Failed to list keys: {e}")
+            raise FileStorageError(f"Key listing failed: {e}")
+
+    # GCS implementation methods
+
+    async def _store_gcs(self, key: str, data: bytes, metadata: Optional[Dict[str, Any]], compressed: bool) -> bool:
+        """Store data in Google Cloud Storage."""
+        try:
+            blob = self._gcs_bucket.blob(key)
+
+            # Set metadata
+            if metadata:
+                blob.metadata = metadata
+            if compressed:
+                blob.content_encoding = 'gzip'
+
+            # Upload data
+            blob.upload_from_string(data)
+            return True
+
+        except Exception as e:
+            logger.error(f"GCS store failed for key {key}: {e}")
+            return False
+
+    async def _retrieve_gcs(self, key: str) -> Optional[Any]:
+        """Retrieve data from Google Cloud Storage."""
+        try:
+            blob = self._gcs_bucket.blob(key)
+
+            if not blob.exists():
+                return None
+
+            # Download data
+            data = blob.download_as_bytes()
+
+            # Decompress if needed
+            if blob.content_encoding == 'gzip':
+                data = gzip.decompress(data)
+
+            # Deserialize data
+            return await self._deserialize_data(data)
+
+        except NotFound:
+            return None
+        except Exception as e:
+            logger.error(f"GCS retrieve failed for key {key}: {e}")
+            return None
+
+    async def _delete_gcs(self, key: str) -> bool:
+        """Delete data from Google Cloud Storage."""
+        try:
+            blob = self._gcs_bucket.blob(key)
+            blob.delete()
+            return True
+
+        except NotFound:
+            return True  # Already deleted
+        except Exception as e:
+            logger.error(f"GCS delete failed for key {key}: {e}")
+            return False
+
+    async def _exists_gcs(self, key: str) -> bool:
+        """Check if data exists in Google Cloud Storage."""
+        try:
+            blob = self._gcs_bucket.blob(key)
+            return blob.exists()
+
+        except Exception as e:
+            logger.error(f"GCS exists check failed for key {key}: {e}")
+            return False
+
+    async def _list_keys_gcs(self, prefix: Optional[str], limit: Optional[int]) -> List[str]:
+        """List keys from Google Cloud Storage."""
+        try:
+            blobs = self._gcs_bucket.list_blobs(prefix=prefix, max_results=limit)
+            return [blob.name for blob in blobs]
+
+        except Exception as e:
+            logger.error(f"GCS list keys failed: {e}")
+            return []
+
+    # Local storage implementation methods
+
+    async def _store_local(self, key: str, data: bytes, metadata: Optional[Dict[str, Any]], compressed: bool) -> bool:
+        """Store data in local filesystem."""
+        try:
+            file_path = Path(self.config.local_storage_path) / key
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+
+            async with aiofiles.open(file_path, 'wb') as f:
+                await f.write(data)
+
+            # Store metadata separately
+            if metadata:
+                metadata_path = file_path.with_suffix('.metadata')
+                metadata_with_compression = {**metadata, 'compressed': compressed}
+                async with aiofiles.open(metadata_path, 'w') as f:
+                    await f.write(json.dumps(metadata_with_compression))
+
+            return True
+
+        except Exception as e:
+            logger.error(f"Local store failed for key {key}: {e}")
+            return False
+
+    async def _retrieve_local(self, key: str) -> Optional[Any]:
+        """Retrieve data from local filesystem."""
+        try:
+            file_path = Path(self.config.local_storage_path) / key
+
+            if not file_path.exists():
+                return None
+
+            async with aiofiles.open(file_path, 'rb') as f:
+                data = await f.read()
+
+            # Check for compression metadata
+            metadata_path = file_path.with_suffix('.metadata')
+            compressed = False
+            if metadata_path.exists():
+                async with aiofiles.open(metadata_path, 'r') as f:
+                    metadata = json.loads(await f.read())
+                    compressed = metadata.get('compressed', False)
+
+            # Decompress if needed
+            if compressed:
+                data = gzip.decompress(data)
+
+            # Deserialize data
+            return await self._deserialize_data(data)
+
+        except Exception as e:
+            logger.error(f"Local retrieve failed for key {key}: {e}")
+            return None
+
+    async def _delete_local(self, key: str) -> bool:
+        """Delete data from local filesystem."""
+        try:
+            file_path = Path(self.config.local_storage_path) / key
+            metadata_path = file_path.with_suffix('.metadata')
+
+            success = True
+            if file_path.exists():
+                file_path.unlink()
+
+            if metadata_path.exists():
+                metadata_path.unlink()
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Local delete failed for key {key}: {e}")
+            return False
+
+    async def _exists_local(self, key: str) -> bool:
+        """Check if data exists in local filesystem."""
+        try:
+            file_path = Path(self.config.local_storage_path) / key
+            return file_path.exists()
+
+        except Exception as e:
+            logger.error(f"Local exists check failed for key {key}: {e}")
+            return False
+
+    async def _list_keys_local(self, prefix: Optional[str], limit: Optional[int]) -> List[str]:
+        """List keys from local filesystem."""
+        try:
+            storage_path = Path(self.config.local_storage_path)
+            if not storage_path.exists():
+                return []
+
+            keys = []
+            for file_path in storage_path.rglob('*'):
+                if file_path.is_file() and not file_path.name.endswith('.metadata'):
+                    key = str(file_path.relative_to(storage_path))
+                    if not prefix or key.startswith(prefix):
+                        keys.append(key)
+                        if limit and len(keys) >= limit:
+                            break
+
+            return keys
+
+        except Exception as e:
+            logger.error(f"Local list keys failed: {e}")
+            return []
+
+    # Utility methods
+
+    async def _serialize_data(self, data: Union[str, bytes, Dict[str, Any]]) -> bytes:
+        """Serialize data for storage."""
+        if isinstance(data, bytes):
+            return data
+        elif isinstance(data, str):
+            return data.encode('utf-8')
+        else:
+            # Use pickle for complex objects
+            return pickle.dumps(data)
+
+    async def _deserialize_data(self, data: bytes) -> Any:
+        """Deserialize data from storage."""
+        try:
+            # Try to deserialize as pickle first
+            return pickle.loads(data)
+        except:
+            try:
+                # Try as JSON
+                return json.loads(data.decode('utf-8'))
+            except:
+                # Return as string
+                return data.decode('utf-8')
+
+    async def _cleanup_cache(self):
+        """Clean up expired cache entries."""
+        if not self.config.enable_cache:
+            return
+
+        current_time = datetime.utcnow()
+        expired_keys = []
+
+        for key, timestamp in self._cache_timestamps.items():
+            if (current_time - timestamp).total_seconds() > self.config.cache_ttl_seconds:
+                expired_keys.append(key)
+
+        for key in expired_keys:
+            self._cache.pop(key, None)
+            self._cache_timestamps.pop(key, None)
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get storage statistics."""
+        return {
+            'initialized': self._initialized,
+            'gcs_available': self._gcs_bucket is not None,
+            'local_fallback_enabled': self.config.enable_local_fallback,
+            'cache_enabled': self.config.enable_cache,
+            'cache_size': len(self._cache),
+            'metrics': self.metrics.get_metrics_summary() if hasattr(self.metrics, 'get_metrics_summary') else {}
+        }
+
+
+# Global instance
+_file_storage_instance = None
+
+def get_file_storage(config: Optional[Dict[str, Any]] = None) -> FileStorage:
+    """Get the global file storage instance."""
+    global _file_storage_instance
+    if _file_storage_instance is None:
+        if config is None:
+            from aiecs.config.config import get_settings
+            settings = get_settings()
+            config = settings.file_storage_config
+        _file_storage_instance = FileStorage(config)
+    return _file_storage_instance
+
+async def initialize_file_storage(config: Optional[Dict[str, Any]] = None) -> FileStorage:
+    """Initialize and return the file storage instance."""
+    storage = get_file_storage(config)
+    await storage.initialize()
+    return storage