Add complete time series aggregation example and refactor OKXCollector for repository pattern
- Introduced `example_complete_series_aggregation.py` to demonstrate time series aggregation, emitting candles even when no trades occur. - Implemented `CompleteSeriesProcessor` extending `RealTimeCandleProcessor` to handle time-based candle emission and empty candle creation. - Refactored `OKXCollector` to utilize the new repository pattern for database operations, enhancing modularity and maintainability. - Updated database operations to centralize data handling through `DatabaseOperations`, improving error handling and logging. - Enhanced documentation to include details on the new aggregation example and repository pattern implementation, ensuring clarity for users.
This commit is contained in:
513
database/operations.py
Normal file
513
database/operations.py
Normal file
@@ -0,0 +1,513 @@
|
||||
"""
|
||||
Database Operations Module
|
||||
|
||||
This module provides centralized database operations for all tables,
|
||||
following the Repository pattern to abstract SQL complexity from business logic.
|
||||
|
||||
Benefits:
|
||||
- Centralized database operations
|
||||
- Clean API for different tables
|
||||
- Easy to test and maintain
|
||||
- Database implementation can change without affecting business logic
|
||||
- Consistent error handling and logging
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from typing import List, Optional, Dict, Any, Union
|
||||
from contextlib import contextmanager
|
||||
import logging
|
||||
import json
|
||||
from sqlalchemy import text
|
||||
|
||||
from .connection import get_db_manager
|
||||
from .models import MarketData, RawTrade
|
||||
from data.common.data_types import OHLCVCandle, StandardizedTrade
|
||||
from data.base_collector import MarketDataPoint, DataType
|
||||
|
||||
|
||||
class DatabaseOperationError(Exception):
|
||||
"""Custom exception for database operation errors."""
|
||||
pass
|
||||
|
||||
|
||||
class BaseRepository:
|
||||
"""Base class for all repository classes."""
|
||||
|
||||
def __init__(self, logger: Optional[logging.Logger] = None):
|
||||
"""Initialize repository with optional logger."""
|
||||
self.logger = logger
|
||||
self._db_manager = get_db_manager()
|
||||
self._db_manager.initialize()
|
||||
|
||||
def log_info(self, message: str) -> None:
|
||||
"""Log info message if logger is available."""
|
||||
if self.logger:
|
||||
self.logger.info(message)
|
||||
|
||||
def log_error(self, message: str) -> None:
|
||||
"""Log error message if logger is available."""
|
||||
if self.logger:
|
||||
self.logger.error(message)
|
||||
|
||||
@contextmanager
|
||||
def get_session(self):
|
||||
"""Get database session with automatic cleanup."""
|
||||
if not self._db_manager:
|
||||
raise DatabaseOperationError("Database manager not initialized")
|
||||
|
||||
with self._db_manager.get_session() as session:
|
||||
yield session
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
candle: OHLCV candle to store
|
||||
force_update: If True, update existing records; if False, ignore duplicates
|
||||
|
||||
Returns:
|
||||
True if operation successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Use right-aligned timestamp (end_time) following industry standard
|
||||
candle_timestamp = candle.end_time
|
||||
|
||||
with self.get_session() as session:
|
||||
if force_update:
|
||||
# Update existing records with new data
|
||||
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:
|
||||
# Ignore duplicates, keep existing records
|
||||
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_info(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.
|
||||
|
||||
Args:
|
||||
symbol: Trading symbol
|
||||
timeframe: Candle timeframe
|
||||
start_time: Start timestamp
|
||||
end_time: End timestamp
|
||||
exchange: Exchange name
|
||||
|
||||
Returns:
|
||||
List of candle dictionaries
|
||||
"""
|
||||
try:
|
||||
with self.get_session() as session:
|
||||
query = text("""
|
||||
SELECT exchange, symbol, timeframe, timestamp,
|
||||
open, high, low, close, volume, trades_count,
|
||||
created_at, updated_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 = []
|
||||
for row in result:
|
||||
candles.append({
|
||||
'exchange': row.exchange,
|
||||
'symbol': row.symbol,
|
||||
'timeframe': row.timeframe,
|
||||
'timestamp': row.timestamp,
|
||||
'open': row.open,
|
||||
'high': row.high,
|
||||
'low': row.low,
|
||||
'close': row.close,
|
||||
'volume': row.volume,
|
||||
'trades_count': row.trades_count,
|
||||
'created_at': row.created_at,
|
||||
'updated_at': row.updated_at
|
||||
})
|
||||
|
||||
self.log_info(f"Retrieved {len(candles)} candles for {symbol} {timeframe}")
|
||||
return candles
|
||||
|
||||
except Exception as e:
|
||||
self.log_error(f"Error retrieving candles for {symbol} {timeframe}: {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.
|
||||
|
||||
Args:
|
||||
symbol: Trading symbol
|
||||
timeframe: Candle timeframe
|
||||
exchange: Exchange name
|
||||
|
||||
Returns:
|
||||
Latest candle dictionary or None
|
||||
"""
|
||||
try:
|
||||
with self.get_session() as session:
|
||||
query = text("""
|
||||
SELECT exchange, symbol, timeframe, timestamp,
|
||||
open, high, low, close, volume, trades_count,
|
||||
created_at, updated_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 {
|
||||
'exchange': row.exchange,
|
||||
'symbol': row.symbol,
|
||||
'timeframe': row.timeframe,
|
||||
'timestamp': row.timestamp,
|
||||
'open': row.open,
|
||||
'high': row.high,
|
||||
'low': row.low,
|
||||
'close': row.close,
|
||||
'volume': row.volume,
|
||||
'trades_count': row.trades_count,
|
||||
'created_at': row.created_at,
|
||||
'updated_at': row.updated_at
|
||||
}
|
||||
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}")
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
data_point: Market data point to store
|
||||
|
||||
Returns:
|
||||
True if operation successful, False otherwise
|
||||
"""
|
||||
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_info(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.
|
||||
|
||||
Args:
|
||||
exchange: Exchange name
|
||||
symbol: Trading symbol
|
||||
data_type: Type of data (e.g., 'raw_trades', 'raw_orderbook')
|
||||
raw_data: Raw data dictionary
|
||||
timestamp: Optional timestamp (defaults to now)
|
||||
|
||||
Returns:
|
||||
True if operation successful, False otherwise
|
||||
"""
|
||||
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_info(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.
|
||||
|
||||
Args:
|
||||
symbol: Trading symbol
|
||||
data_type: Data type filter
|
||||
start_time: Start timestamp
|
||||
end_time: End timestamp
|
||||
exchange: Exchange name
|
||||
limit: Maximum number of records to return
|
||||
|
||||
Returns:
|
||||
List of raw trade dictionaries
|
||||
"""
|
||||
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 += f" LIMIT {limit}"
|
||||
|
||||
result = session.execute(query, {
|
||||
'exchange': exchange,
|
||||
'symbol': symbol,
|
||||
'data_type': data_type,
|
||||
'start_time': start_time,
|
||||
'end_time': end_time
|
||||
})
|
||||
|
||||
trades = []
|
||||
for row in result:
|
||||
trades.append({
|
||||
'id': row.id,
|
||||
'exchange': row.exchange,
|
||||
'symbol': row.symbol,
|
||||
'timestamp': row.timestamp,
|
||||
'data_type': row.data_type,
|
||||
'raw_data': row.raw_data,
|
||||
'created_at': row.created_at
|
||||
})
|
||||
|
||||
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}")
|
||||
|
||||
|
||||
class DatabaseOperations:
|
||||
"""
|
||||
Main database operations manager that provides access to all repositories.
|
||||
|
||||
This is the main entry point for database operations, providing a
|
||||
centralized interface to all table-specific repositories.
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Optional[logging.Logger] = None):
|
||||
"""Initialize database operations with optional logger."""
|
||||
self.logger = logger
|
||||
|
||||
# Initialize repositories
|
||||
self.market_data = MarketDataRepository(logger)
|
||||
self.raw_trades = RawTradeRepository(logger)
|
||||
|
||||
def health_check(self) -> bool:
|
||||
"""
|
||||
Perform a health check on database connections.
|
||||
|
||||
Returns:
|
||||
True if database is healthy, False otherwise
|
||||
"""
|
||||
try:
|
||||
with self.market_data.get_session() as session:
|
||||
# Simple query to test connection
|
||||
result = session.execute(text("SELECT 1"))
|
||||
return result.fetchone() is not None
|
||||
except Exception as e:
|
||||
if self.logger:
|
||||
self.logger.error(f"Database health check failed: {e}")
|
||||
return False
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get database statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary containing database statistics
|
||||
"""
|
||||
try:
|
||||
stats = {
|
||||
'healthy': self.health_check(),
|
||||
'repositories': {
|
||||
'market_data': 'MarketDataRepository',
|
||||
'raw_trades': 'RawTradeRepository'
|
||||
}
|
||||
}
|
||||
|
||||
# Get table counts
|
||||
with self.market_data.get_session() as session:
|
||||
# Market data count
|
||||
result = session.execute(text("SELECT COUNT(*) FROM market_data"))
|
||||
stats['candle_count'] = result.fetchone()[0]
|
||||
|
||||
# Raw trades count
|
||||
result = session.execute(text("SELECT COUNT(*) FROM raw_trades"))
|
||||
stats['raw_trade_count'] = result.fetchone()[0]
|
||||
|
||||
return stats
|
||||
|
||||
except Exception as e:
|
||||
if self.logger:
|
||||
self.logger.error(f"Error getting database stats: {e}")
|
||||
return {'healthy': False, 'error': str(e)}
|
||||
|
||||
|
||||
# Singleton instance for global access
|
||||
_db_operations_instance: Optional[DatabaseOperations] = None
|
||||
|
||||
|
||||
def get_database_operations(logger: Optional[logging.Logger] = None) -> DatabaseOperations:
|
||||
"""
|
||||
Get the global database operations instance.
|
||||
|
||||
Args:
|
||||
logger: Optional logger for database operations
|
||||
|
||||
Returns:
|
||||
DatabaseOperations instance
|
||||
"""
|
||||
global _db_operations_instance
|
||||
|
||||
if _db_operations_instance is None:
|
||||
_db_operations_instance = DatabaseOperations(logger)
|
||||
|
||||
return _db_operations_instance
|
||||
|
||||
|
||||
__all__ = [
|
||||
'DatabaseOperationError',
|
||||
'MarketDataRepository',
|
||||
'RawTradeRepository',
|
||||
'DatabaseOperations',
|
||||
'get_database_operations'
|
||||
]
|
||||
@@ -27,7 +27,7 @@ CREATE TABLE market_data (
|
||||
volume DECIMAL(18,8) NOT NULL,
|
||||
trades_count INTEGER, -- number of trades in this candle
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
-- updated_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
CONSTRAINT unique_market_data UNIQUE(exchange, symbol, timeframe, timestamp)
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user