storage manager
This commit is contained in:
Dobromir Popov
270
COBY/storage/storage_manager.py
Normal file
270
COBY/storage/storage_manager.py
Normal file
@ -0,0 +1,270 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user