270 lines
10 KiB
Python
270 lines
10 KiB
Python
"""
|
|
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 |