TCPDashboard/database/repositories/raw_trade_repository.py

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

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

from sqlalchemy import desc

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 using the ORM.
        """
        try:
            with self.get_session() as session:
                new_trade = RawTrade(
                    exchange=data_point.exchange,
                    symbol=data_point.symbol,
                    timestamp=data_point.timestamp,
                    data_type=data_point.data_type.value,
                    raw_data=data_point.data
                )
                session.add(new_trade)
                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 using the ORM.
        """
        try:
            with self.get_session() as session:
                new_trade = RawTrade(
                    exchange=exchange,
                    symbol=symbol,
                    timestamp=timestamp or datetime.now(datetime.timezone.utc),
                    data_type=data_type,
                    raw_data=raw_data
                )
                session.add(new_trade)
                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 using the ORM.
        """
        try:
            with self.get_session() as session:
                query = (
                    session.query(RawTrade)
                    .filter(
                        RawTrade.exchange == exchange,
                        RawTrade.symbol == symbol,
                        RawTrade.data_type == data_type,
                        RawTrade.timestamp >= start_time,
                        RawTrade.timestamp <= end_time
                    )
                    .order_by(RawTrade.timestamp.asc())
                )
                
                if limit:
                    query = query.limit(limit)

                results = query.all()

                trades = [
                    {
                        "id": r.id, "exchange": r.exchange, "symbol": r.symbol,
                        "timestamp": r.timestamp, "data_type": r.data_type,
                        "raw_data": r.raw_data, "created_at": r.created_at
                    } for r in results
                ]
                
                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}")