gogo2/COBY/interfaces/storage_manager.py

"""
Interface for data storage and retrieval.
"""

from abc import ABC, abstractmethod
from datetime import datetime
from typing import List, Dict, Optional, Any
from ..models.core import OrderBookSnapshot, TradeEvent, HeatmapData, SystemMetrics


class StorageManager(ABC):
    """Manages data persistence and retrieval"""
    
    @abstractmethod
    async def store_orderbook(self, data: OrderBookSnapshot) -> bool:
        """
        Store order book snapshot to database.
        
        Args:
            data: Order book snapshot to store
            
        Returns:
            bool: True if stored successfully, False otherwise
        """
        pass
    
    @abstractmethod
    async def store_trade(self, data: TradeEvent) -> bool:
        """
        Store trade event to database.
        
        Args:
            data: Trade event to store
            
        Returns:
            bool: True if stored successfully, False otherwise
        """
        pass
    
    @abstractmethod
    async def store_heatmap(self, data: HeatmapData) -> bool:
        """
        Store heatmap data to database.
        
        Args:
            data: Heatmap data to store
            
        Returns:
            bool: True if stored successfully, False otherwise
        """
        pass
    
    @abstractmethod
    async def store_metrics(self, data: SystemMetrics) -> bool:
        """
        Store system metrics to database.
        
        Args:
            data: System metrics to store
            
        Returns:
            bool: True if stored successfully, False otherwise
        """
        pass
    
    @abstractmethod
    async def get_historical_orderbooks(self, symbol: str, exchange: str,
                                      start: datetime, end: datetime,
                                      limit: Optional[int] = None) -> List[OrderBookSnapshot]:
        """
        Retrieve historical order book data.
        
        Args:
            symbol: Trading symbol
            exchange: Exchange name
            start: Start timestamp
            end: End timestamp
            limit: Maximum number of records to return
            
        Returns:
            List[OrderBookSnapshot]: Historical order book data
        """
        pass
    
    @abstractmethod
    async def get_historical_trades(self, symbol: str, exchange: str,
                                  start: datetime, end: datetime,
                                  limit: Optional[int] = None) -> List[TradeEvent]:
        """
        Retrieve historical trade data.
        
        Args:
            symbol: Trading symbol
            exchange: Exchange name
            start: Start timestamp
            end: End timestamp
            limit: Maximum number of records to return
            
        Returns:
            List[TradeEvent]: Historical trade data
        """
        pass
    
    @abstractmethod
    async def get_latest_orderbook(self, symbol: str, exchange: str) -> Optional[OrderBookSnapshot]:
        """
        Get latest order book snapshot.
        
        Args:
            symbol: Trading symbol
            exchange: Exchange name
            
        Returns:
            OrderBookSnapshot: Latest order book or None if not found
        """
        pass
    
    @abstractmethod
    async def get_latest_heatmap(self, symbol: str, bucket_size: float) -> Optional[HeatmapData]:
        """
        Get latest heatmap data.
        
        Args:
            symbol: Trading symbol
            bucket_size: Price bucket size
            
        Returns:
            HeatmapData: Latest heatmap or None if not found
        """
        pass
    
    @abstractmethod
    async def get_ohlcv_data(self, symbol: str, exchange: str, timeframe: str,
                           start: datetime, end: datetime) -> List[Dict[str, Any]]:
        """
        Get OHLCV candlestick data.
        
        Args:
            symbol: Trading symbol
            exchange: Exchange name
            timeframe: Timeframe (e.g., '1m', '5m', '1h')
            start: Start timestamp
            end: End timestamp
            
        Returns:
            List[Dict]: OHLCV data
        """
        pass
    
    @abstractmethod
    async def batch_store_orderbooks(self, data: List[OrderBookSnapshot]) -> int:
        """
        Store multiple order book snapshots in batch.
        
        Args:
            data: List of order book snapshots
            
        Returns:
            int: Number of records stored successfully
        """
        pass
    
    @abstractmethod
    async def batch_store_trades(self, data: List[TradeEvent]) -> int:
        """
        Store multiple trade events in batch.
        
        Args:
            data: List of trade events
            
        Returns:
            int: Number of records stored successfully
        """
        pass
    
    @abstractmethod
    def setup_database_schema(self) -> None:
        """
        Set up database schema and tables.
        Should be idempotent - safe to call multiple times.
        """
        pass
    
    @abstractmethod
    async def cleanup_old_data(self, retention_days: int) -> int:
        """
        Clean up old data based on retention policy.
        
        Args:
            retention_days: Number of days to retain data
            
        Returns:
            int: Number of records deleted
        """
        pass
    
    @abstractmethod
    async def get_storage_stats(self) -> Dict[str, Any]:
        """
        Get storage statistics.
        
        Returns:
            Dict: Storage statistics (table sizes, record counts, etc.)
        """
        pass
    
    @abstractmethod
    async def health_check(self) -> bool:
        """
        Check storage system health.
        
        Returns:
            bool: True if healthy, False otherwise
        """
        pass