TCPDashboard/database/repositories/market_data_repository.py

"""Repository for market_data table operations."""

from datetime import datetime
from typing import List, Optional, Dict, Any
from sqlalchemy import text

from ..models import MarketData
from data.common.data_types import OHLCVCandle
from .base_repository import BaseRepository, DatabaseOperationError


class MarketDataRepository(BaseRepository):
    """Repository for market_data table operations."""
    
    def upsert_candle(self, candle: OHLCVCandle, force_update: bool = False) -> bool:
        """
        Insert or update a candle in the market_data table.
        """
        try:
            candle_timestamp = candle.end_time
            
            with self.get_session() as session:
                if force_update:
                    query = text("""
                        INSERT INTO market_data (
                            exchange, symbol, timeframe, timestamp,
                            open, high, low, close, volume, trades_count, 
                            created_at
                        ) VALUES (
                            :exchange, :symbol, :timeframe, :timestamp,
                            :open, :high, :low, :close, :volume, :trades_count,
                            NOW()
                        )
                        ON CONFLICT (exchange, symbol, timeframe, timestamp)
                        DO UPDATE SET
                            open = EXCLUDED.open,
                            high = EXCLUDED.high,
                            low = EXCLUDED.low,
                            close = EXCLUDED.close,
                            volume = EXCLUDED.volume,
                            trades_count = EXCLUDED.trades_count
                    """)
                    action = "Updated"
                else:
                    query = text("""
                        INSERT INTO market_data (
                            exchange, symbol, timeframe, timestamp,
                            open, high, low, close, volume, trades_count, 
                            created_at
                        ) VALUES (
                            :exchange, :symbol, :timeframe, :timestamp,
                            :open, :high, :low, :close, :volume, :trades_count,
                            NOW()
                        )
                        ON CONFLICT (exchange, symbol, timeframe, timestamp)
                        DO NOTHING
                    """)
                    action = "Stored"
                
                session.execute(query, {
                    'exchange': candle.exchange,
                    'symbol': candle.symbol,
                    'timeframe': candle.timeframe,
                    'timestamp': candle_timestamp,
                    'open': float(candle.open),
                    'high': float(candle.high),
                    'low': float(candle.low),
                    'close': float(candle.close),
                    'volume': float(candle.volume),
                    'trades_count': candle.trade_count
                })
                
                session.commit()
                
                self.log_debug(f"{action} candle: {candle.symbol} {candle.timeframe} at {candle_timestamp} (force_update={force_update})")
                return True
                
        except Exception as e:
            self.log_error(f"Error storing candle {candle.symbol} {candle.timeframe}: {e}")
            raise DatabaseOperationError(f"Failed to store candle: {e}")
    
    def get_candles(self, 
                   symbol: str, 
                   timeframe: str, 
                   start_time: datetime, 
                   end_time: datetime,
                   exchange: str = "okx") -> List[Dict[str, Any]]:
        """
        Retrieve candles from the database.
        """
        try:
            with self.get_session() as session:
                query = text("""
                    SELECT exchange, symbol, timeframe, timestamp,
                           open, high, low, close, volume, trades_count,
                           created_at
                    FROM market_data
                    WHERE exchange = :exchange 
                      AND symbol = :symbol 
                      AND timeframe = :timeframe
                      AND timestamp >= :start_time 
                      AND timestamp <= :end_time
                    ORDER BY timestamp ASC
                """)
                
                result = session.execute(query, {
                    'exchange': exchange,
                    'symbol': symbol,
                    'timeframe': timeframe,
                    'start_time': start_time,
                    'end_time': end_time
                })
                
                candles = [dict(row._mapping) for row in result]
                
                self.log_debug(f"Retrieved {len(candles)} candles for {symbol} {timeframe}")
                return candles
                
        except Exception as e:
            self.log_error(f"Error retrieving candles: {e}")
            raise DatabaseOperationError(f"Failed to retrieve candles: {e}")
    
    def get_latest_candle(self, symbol: str, timeframe: str, exchange: str = "okx") -> Optional[Dict[str, Any]]:
        """
        Get the latest candle for a symbol and timeframe.
        """
        try:
            with self.get_session() as session:
                query = text("""
                    SELECT exchange, symbol, timeframe, timestamp,
                           open, high, low, close, volume, trades_count,
                           created_at
                    FROM market_data
                    WHERE exchange = :exchange 
                      AND symbol = :symbol 
                      AND timeframe = :timeframe
                    ORDER BY timestamp DESC
                    LIMIT 1
                """)
                
                result = session.execute(query, {
                    'exchange': exchange,
                    'symbol': symbol,
                    'timeframe': timeframe
                })
                
                row = result.fetchone()
                if row:
                    return dict(row._mapping)
                return None
                
        except Exception as e:
            self.log_error(f"Error retrieving latest candle for {symbol} {timeframe}: {e}")
            raise DatabaseOperationError(f"Failed to retrieve latest candle: {e}")