TCPDashboard/database/repositories/raw_trade_repository.py

"""Repository for raw_trades table operations."""

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

from ..models import RawTrade
from data.base_collector import MarketDataPoint
from .base_repository import BaseRepository, DatabaseOperationError


class RawTradeRepository(BaseRepository):
    """Repository for raw_trades table operations."""
    
    def insert_market_data_point(self, data_point: MarketDataPoint) -> bool:
        """
        Insert a market data point into raw_trades table.
        """
        try:
            with self.get_session() as session:
                query = text("""
                    INSERT INTO raw_trades (
                        exchange, symbol, timestamp, data_type, raw_data, created_at
                    ) VALUES (
                        :exchange, :symbol, :timestamp, :data_type, :raw_data, NOW()
                    )
                """)
                
                session.execute(query, {
                    'exchange': data_point.exchange,
                    'symbol': data_point.symbol,
                    'timestamp': data_point.timestamp,
                    'data_type': data_point.data_type.value,
                    'raw_data': json.dumps(data_point.data)
                })
                
                session.commit()
                
                self.log_debug(f"Stored raw {data_point.data_type.value} data for {data_point.symbol}")
                return True
                
        except Exception as e:
            self.log_error(f"Error storing raw data for {data_point.symbol}: {e}")
            raise DatabaseOperationError(f"Failed to store raw data: {e}")
    
    def insert_raw_websocket_data(self, 
                                 exchange: str,
                                 symbol: str, 
                                 data_type: str,
                                 raw_data: Dict[str, Any],
                                 timestamp: Optional[datetime] = None) -> bool:
        """
        Insert raw WebSocket data for debugging purposes.
        """
        try:
            if timestamp is None:
                timestamp = datetime.now()
            
            with self.get_session() as session:
                query = text("""
                    INSERT INTO raw_trades (
                        exchange, symbol, timestamp, data_type, raw_data, created_at
                    ) VALUES (
                        :exchange, :symbol, :timestamp, :data_type, :raw_data, NOW()
                    )
                """)
                
                session.execute(query, {
                    'exchange': exchange,
                    'symbol': symbol,
                    'timestamp': timestamp,
                    'data_type': data_type,
                    'raw_data': json.dumps(raw_data)
                })
                
                session.commit()
                
                self.log_debug(f"Stored raw WebSocket data: {data_type} for {symbol}")
                return True
                
        except Exception as e:
            self.log_error(f"Error storing raw WebSocket data for {symbol}: {e}")
            raise DatabaseOperationError(f"Failed to store raw WebSocket data: {e}")
    
    def get_raw_trades(self,
                      symbol: str,
                      data_type: str,
                      start_time: datetime,
                      end_time: datetime,
                      exchange: str = "okx",
                      limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        Retrieve raw trades from the database.
        """
        try:
            with self.get_session() as session:
                query = text("""
                    SELECT id, exchange, symbol, timestamp, data_type, raw_data, created_at
                    FROM raw_trades
                    WHERE exchange = :exchange 
                      AND symbol = :symbol 
                      AND data_type = :data_type
                      AND timestamp >= :start_time 
                      AND timestamp <= :end_time
                    ORDER BY timestamp ASC
                """)
                
                if limit:
                    query_str = str(query.compile(compile_kwargs={"literal_binds": True})) + f" LIMIT {limit}"
                    query = text(query_str)

                result = session.execute(query, {
                    'exchange': exchange,
                    'symbol': symbol,
                    'data_type': data_type,
                    'start_time': start_time,
                    'end_time': end_time
                })
                
                trades = [dict(row._mapping) for row in result]
                
                self.log_info(f"Retrieved {len(trades)} raw trades for {symbol} {data_type}")
                return trades
                
        except Exception as e:
            self.log_error(f"Error retrieving raw trades for {symbol}: {e}")
            raise DatabaseOperationError(f"Failed to retrieve raw trades: {e}")