""" Comprehensive storage manager that integrates TimescaleDB, connection pooling, and schema management. """ import asyncio import logging from typing import Dict, List, Optional, Any from datetime import datetime, timedelta from .timescale_manager import TimescaleManager from .connection_pool import ConnectionPoolManager from .schema import SchemaManager from ..models.core import OrderBookSnapshot, TradeEvent, HeatmapData from ..config import Config from ..utils.exceptions import DatabaseError, ConnectionError logger = logging.getLogger(__name__) class StorageManager: """Unified storage manager for all database operations.""" def __init__(self, config: Config): self.config = config self.connection_pool = ConnectionPoolManager(config) self.schema_manager = SchemaManager(self.connection_pool) self.timescale_manager = TimescaleManager(config) self._initialized = False async def initialize(self) -> None: """Initialize all storage components.""" try: logger.info("Initializing storage manager...") # Initialize connection pool await self.connection_pool.initialize() # Set up database schema await self.schema_manager.setup_complete_schema() # Initialize TimescaleDB manager with existing pool self.timescale_manager.pool = self.connection_pool.pool self._initialized = True logger.info("Storage manager initialized successfully") except Exception as e: logger.error(f"Failed to initialize storage manager: {e}") raise ConnectionError(f"Storage initialization failed: {e}") async def close(self) -> None: """Close all storage connections.""" if self.timescale_manager: await self.timescale_manager.close() if self.connection_pool: await self.connection_pool.close() logger.info("Storage manager closed") def is_healthy(self) -> bool: """Check if storage system is healthy.""" return self._initialized and self.connection_pool.is_healthy() # Order book operations async def store_orderbook(self, data: OrderBookSnapshot) -> bool: """Store order book snapshot.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.store_orderbook(data) async def batch_store_orderbooks(self, data_list: List[OrderBookSnapshot]) -> bool: """Store multiple order book snapshots.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.batch_store_orderbooks(data_list) async def get_latest_orderbook(self, symbol: str, exchange: Optional[str] = None) -> Optional[Dict]: """Get latest order book snapshot.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.get_latest_orderbook(symbol, exchange) # Trade operations async def store_trade(self, data: TradeEvent) -> bool: """Store trade event.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.store_trade(data) async def batch_store_trades(self, data_list: List[TradeEvent]) -> bool: """Store multiple trade events.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.batch_store_trades(data_list) # Heatmap operations async def store_heatmap_data(self, data: HeatmapData) -> bool: """Store heatmap data.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.store_heatmap_data(data) async def get_heatmap_data(self, symbol: str, bucket_size: float, start: Optional[datetime] = None) -> List[Dict]: """Get heatmap data for visualization.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.get_heatmap_data(symbol, bucket_size, start) # Historical data operations async def get_historical_data(self, symbol: str, start: datetime, end: datetime, data_type: str = 'orderbook') -> List[Dict]: """Get historical data for a symbol within time range.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") return await self.timescale_manager.get_historical_data(symbol, start, end, data_type) # System operations async def get_system_stats(self) -> Dict[str, Any]: """Get comprehensive system statistics.""" if not self._initialized: return {"status": "not_initialized"} try: # Get database stats db_stats = await self.timescale_manager.get_database_stats() # Get connection pool stats pool_stats = self.connection_pool.get_pool_stats() # Get schema info schema_info = await self.schema_manager.get_schema_info() return { "status": "healthy" if self.is_healthy() else "unhealthy", "database": db_stats, "connection_pool": pool_stats, "schema": schema_info, "initialized": self._initialized } except Exception as e: logger.error(f"Failed to get system stats: {e}") return {"status": "error", "error": str(e)} async def health_check(self) -> Dict[str, Any]: """Perform comprehensive health check.""" health_status = { "healthy": True, "components": {}, "timestamp": datetime.utcnow().isoformat() } try: # Check connection pool pool_healthy = self.connection_pool.is_healthy() health_status["components"]["connection_pool"] = { "healthy": pool_healthy, "stats": self.connection_pool.get_pool_stats() } # Test database connection try: async with self.connection_pool.acquire() as conn: await conn.execute('SELECT 1') health_status["components"]["database"] = {"healthy": True} except Exception as e: health_status["components"]["database"] = { "healthy": False, "error": str(e) } health_status["healthy"] = False # Check schema version try: current_version = await self.schema_manager.get_current_schema_version() health_status["components"]["schema"] = { "healthy": True, "version": current_version } except Exception as e: health_status["components"]["schema"] = { "healthy": False, "error": str(e) } health_status["healthy"] = False # Overall health health_status["healthy"] = all( comp.get("healthy", False) for comp in health_status["components"].values() ) except Exception as e: logger.error(f"Health check failed: {e}") health_status["healthy"] = False health_status["error"] = str(e) return health_status # Maintenance operations async def cleanup_old_data(self, days: int = 30) -> Dict[str, int]: """Clean up old data beyond retention period.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") try: cutoff_date = datetime.utcnow() - timedelta(days=days) cleanup_stats = {} async with self.connection_pool.acquire() as conn: # Clean up old order book snapshots result = await conn.execute(""" DELETE FROM order_book_snapshots WHERE timestamp < $1 """, cutoff_date) cleanup_stats["order_book_snapshots"] = int(result.split()[-1]) # Clean up old trade events result = await conn.execute(""" DELETE FROM trade_events WHERE timestamp < $1 """, cutoff_date) cleanup_stats["trade_events"] = int(result.split()[-1]) # Clean up old heatmap data result = await conn.execute(""" DELETE FROM heatmap_data WHERE timestamp < $1 """, cutoff_date) cleanup_stats["heatmap_data"] = int(result.split()[-1]) logger.info(f"Cleaned up old data: {cleanup_stats}") return cleanup_stats except Exception as e: logger.error(f"Failed to cleanup old data: {e}") raise DatabaseError(f"Data cleanup failed: {e}") async def optimize_database(self) -> bool: """Run database optimization tasks.""" if not self._initialized: raise DatabaseError("Storage manager not initialized") try: async with self.connection_pool.acquire() as conn: # Analyze tables for better query planning tables = ['order_book_snapshots', 'trade_events', 'heatmap_data', 'ohlcv_data'] for table in tables: await conn.execute(f"ANALYZE {table}") # Vacuum tables to reclaim space for table in tables: await conn.execute(f"VACUUM {table}") logger.info("Database optimization completed") return True except Exception as e: logger.error(f"Database optimization failed: {e}") return False