"""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}")