vasily/TCPDashboard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance logging capabilities across data collection components

Browse Source 
- Added optional logger parameter to various classes including `BaseDataCollector`, `CollectorManager`, `RealTimeCandleProcessor`, and `BatchCandleProcessor` to support conditional logging.
- Implemented error-only logging mode, allowing components to log only error and critical messages when specified.
- Updated logging calls to utilize new helper methods for improved readability and maintainability.
- Enhanced documentation to include details on the new logging system and its usage across components.
- Ensured that child components inherit the logger from their parent components for consistent logging behavior.
This commit is contained in:
Vasily.onl
2025-06-01 14:42:29 +08:00
parent 0697be75da
commit bc13cfcbe0
11 changed files with 1179 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
data/exchanges/okx/collector.py
View File

@@ -23,7 +23,6 @@ from .websocket import (
from .data_processor import OKXDataProcessor
from database.connection import get_db_manager, get_raw_data_manager
from database.models import MarketData, RawTrade
from utils.logger import get_logger
@dataclass
@@ -52,7 +51,9 @@ class OKXCollector(BaseDataCollector):
component_name: Optional[str] = None,
auto_restart: bool = True,
health_check_interval: float = 30.0,
store_raw_data: bool = True):
store_raw_data: bool = True,
logger = None,
log_errors_only: bool = False):
"""
Initialize OKX collector for a single trading pair.
@@ -63,6 +64,8 @@ class OKXCollector(BaseDataCollector):
auto_restart: Enable automatic restart on failures
health_check_interval: Seconds between health checks
store_raw_data: Whether to store raw data for debugging
logger: Logger instance for conditional logging (None for no logging)
log_errors_only: If True and logger provided, only log error-level messages
"""
# Default data types if not specified
if data_types is None:
@@ -79,7 +82,9 @@ class OKXCollector(BaseDataCollector):
data_types=data_types,
component_name=component_name,
auto_restart=auto_restart,
health_check_interval=health_check_interval
health_check_interval=health_check_interval,
logger=logger,
log_errors_only=log_errors_only
)
# OKX-specific settings
@@ -90,7 +95,7 @@ class OKXCollector(BaseDataCollector):
self._ws_client: Optional[OKXWebSocketClient] = None
# Data processor using new common framework
self._data_processor = OKXDataProcessor(symbol, component_name=f"{component_name}_processor")
self._data_processor = OKXDataProcessor(symbol, component_name=f"{component_name}_processor", logger=logger)
# Add callbacks for processed data
self._data_processor.add_trade_callback(self._on_trade_processed)
@@ -113,8 +118,9 @@ class OKXCollector(BaseDataCollector):
DataType.TICKER: OKXChannelType.TICKERS.value
}
self.logger.info(f"Initialized OKX collector for {symbol} with data types: {[dt.value for dt in data_types]}")
self.logger.info(f"Using common data processing framework")
if logger:
logger.info(f"{component_name}: Initialized OKX collector for {symbol} with data types: {[dt.value for dt in data_types]}")
logger.info(f"{component_name}: Using common data processing framework")
async def connect(self) -> bool:
"""
@@ -124,7 +130,8 @@ class OKXCollector(BaseDataCollector):
True if connection successful, False otherwise
"""
try:
self.logger.info(f"Connecting OKX collector for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Connecting OKX collector for {self.symbol}")
# Initialize database managers
self._db_manager = get_db_manager()
@@ -146,29 +153,35 @@ class OKXCollector(BaseDataCollector):
# Connect to WebSocket
if not await self._ws_client.connect(use_public=True):
self.logger.error("Failed to connect to OKX WebSocket")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to connect to OKX WebSocket")
return False
self.logger.info(f"Successfully connected OKX collector for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Successfully connected OKX collector for {self.symbol}")
return True
except Exception as e:
self.logger.error(f"Error connecting OKX collector for {self.symbol}: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error connecting OKX collector for {self.symbol}: {e}")
return False
async def disconnect(self) -> None:
"""Disconnect from OKX WebSocket API."""
try:
self.logger.info(f"Disconnecting OKX collector for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Disconnecting OKX collector for {self.symbol}")
if self._ws_client:
await self._ws_client.disconnect()
self._ws_client = None
self.logger.info(f"Disconnected OKX collector for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Disconnected OKX collector for {self.symbol}")
except Exception as e:
self.logger.error(f"Error disconnecting OKX collector for {self.symbol}: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error disconnecting OKX collector for {self.symbol}: {e}")
async def subscribe_to_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
"""
@@ -182,12 +195,14 @@ class OKXCollector(BaseDataCollector):
True if subscription successful, False otherwise
"""
if not self._ws_client or not self._ws_client.is_connected:
self.logger.error("WebSocket client not connected")
if self.logger:
self.logger.error(f"{self.component_name}: WebSocket client not connected")
return False
# Validate symbol
if self.symbol not in symbols:
self.logger.warning(f"Symbol {self.symbol} not in subscription list: {symbols}")
if self.logger:
self.logger.warning(f"{self.component_name}: Symbol {self.symbol} not in subscription list: {symbols}")
return False
try:
@@ -202,25 +217,31 @@ class OKXCollector(BaseDataCollector):
enabled=True
)
subscriptions.append(subscription)
self.logger.debug(f"Added subscription: {channel} for {self.symbol}")
if self.logger:
self.logger.debug(f"{self.component_name}: Added subscription: {channel} for {self.symbol}")
else:
self.logger.warning(f"Unsupported data type: {data_type}")
if self.logger:
self.logger.warning(f"{self.component_name}: Unsupported data type: {data_type}")
if not subscriptions:
self.logger.warning("No valid subscriptions to create")
if self.logger:
self.logger.warning(f"{self.component_name}: No valid subscriptions to create")
return False
# Subscribe to channels
success = await self._ws_client.subscribe(subscriptions)
if success:
self.logger.info(f"Successfully subscribed to {len(subscriptions)} channels for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Successfully subscribed to {len(subscriptions)} channels for {self.symbol}")
return True
else:
self.logger.error(f"Failed to subscribe to channels for {self.symbol}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to subscribe to channels for {self.symbol}")
return False
except Exception as e:
self.logger.error(f"Error subscribing to data for {self.symbol}: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error subscribing to data for {self.symbol}: {e}")
return False
async def unsubscribe_from_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
@@ -235,7 +256,8 @@ class OKXCollector(BaseDataCollector):
True if unsubscription successful, False otherwise
"""
if not self._ws_client or not self._ws_client.is_connected:
self.logger.warning("WebSocket client not connected")
if self.logger:
self.logger.warning(f"{self.component_name}: WebSocket client not connected")
return True # Consider it successful if not connected
try:
@@ -257,14 +279,17 @@ class OKXCollector(BaseDataCollector):
# Unsubscribe from channels
success = await self._ws_client.unsubscribe(subscriptions)
if success:
self.logger.info(f"Successfully unsubscribed from {len(subscriptions)} channels for {self.symbol}")
if self.logger:
self.logger.info(f"{self.component_name}: Successfully unsubscribed from {len(subscriptions)} channels for {self.symbol}")
return True
else:
self.logger.error(f"Failed to unsubscribe from channels for {self.symbol}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to unsubscribe from channels for {self.symbol}")
return False
except Exception as e:
self.logger.error(f"Error unsubscribing from data for {self.symbol}: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error unsubscribing from data for {self.symbol}: {e}")
return False
async def _process_message(self, message: Any) -> Optional[MarketDataPoint]:
@@ -278,7 +303,8 @@ class OKXCollector(BaseDataCollector):
MarketDataPoint if processing successful, None otherwise
"""
if not isinstance(message, dict):
self.logger.warning(f"Received non-dict message: {type(message)}")
if self.logger:
self.logger.warning(f"{self.component_name}: Received non-dict message: {type(message)}")
return None
try:
@@ -291,11 +317,13 @@ class OKXCollector(BaseDataCollector):
if not success:
self._error_count += 1
self.logger.error(f"Message processing failed: {errors}")
if self.logger:
self.logger.error(f"{self.component_name}: Message processing failed: {errors}")
return None
if errors:
self.logger.warning(f"Message processing warnings: {errors}")
if self.logger:
self.logger.warning(f"{self.component_name}: Message processing warnings: {errors}")
# Store raw data if enabled (for debugging/compliance)
if self.store_raw_data and 'data' in message and 'arg' in message:
@@ -310,7 +338,8 @@ class OKXCollector(BaseDataCollector):
except Exception as e:
self._error_count += 1
self.logger.error(f"Error processing message: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error processing message: {e}")
return None
async def _handle_messages(self) -> None:
@@ -340,10 +369,12 @@ class OKXCollector(BaseDataCollector):
raw_data=data_point.data
)
session.add(raw_trade)
self.logger.debug(f"Stored raw data: {data_point.data_type.value} for {data_point.symbol}")
if self.logger:
self.logger.debug(f"{self.component_name}: Stored raw data: {data_point.data_type.value} for {data_point.symbol}")
except Exception as e:
self.logger.error(f"Error storing raw market data: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error storing raw market data: {e}")
async def _store_completed_candle(self, candle: OHLCVCandle) -> None:
"""
@@ -371,10 +402,12 @@ class OKXCollector(BaseDataCollector):
trades_count=candle.trade_count
)
session.add(market_data)
self.logger.info(f"Stored completed candle: {candle.symbol} {candle.timeframe} at {candle.start_time}")
if self.logger:
self.logger.info(f"{self.component_name}: Stored completed candle: {candle.symbol} {candle.timeframe} at {candle.start_time}")
except Exception as e:
self.logger.error(f"Error storing completed candle: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error storing completed candle: {e}")
async def _store_raw_data(self, channel: str, raw_message: Dict[str, Any]) -> None:
"""
@@ -399,7 +432,8 @@ class OKXCollector(BaseDataCollector):
)
except Exception as e:
self.logger.error(f"Error storing raw WebSocket data: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error storing raw WebSocket data: {e}")
def _on_message(self, message: Dict[str, Any]) -> None:
"""
@@ -412,7 +446,8 @@ class OKXCollector(BaseDataCollector):
# Process message asynchronously
asyncio.create_task(self._process_message(message))
except Exception as e:
self.logger.error(f"Error handling WebSocket message: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error handling WebSocket message: {e}")
def _on_trade_processed(self, trade: StandardizedTrade) -> None:
"""
@@ -422,7 +457,8 @@ class OKXCollector(BaseDataCollector):
trade: Processed standardized trade
"""
self._processed_trades += 1
self.logger.debug(f"Processed trade: {trade.symbol} {trade.side} {trade.size}@{trade.price}")
if self.logger:
self.logger.debug(f"{self.component_name}: Processed trade: {trade.symbol} {trade.side} {trade.size}@{trade.price}")
def _on_candle_processed(self, candle: OHLCVCandle) -> None:
"""
@@ -432,7 +468,8 @@ class OKXCollector(BaseDataCollector):
candle: Completed OHLCV candle
"""
self._processed_candles += 1
self.logger.info(f"Completed candle: {candle.symbol} {candle.timeframe} O:{candle.open} H:{candle.high} L:{candle.low} C:{candle.close} V:{candle.volume}")
if self.logger:
self.logger.info(f"{self.component_name}: Completed candle: {candle.symbol} {candle.timeframe} O:{candle.open} H:{candle.high} L:{candle.low} C:{candle.close} V:{candle.volume}")
# Store completed candle in market_data table
if candle.is_complete:

67
data/exchanges/okx/data_processor.py
View File

@@ -24,7 +24,6 @@ from ...common import (
UnifiedDataTransformer,
create_standardized_trade
)
from utils.logger import get_logger
class OKXMessageType(Enum):
@@ -81,9 +80,9 @@ class OKXDataValidator(BaseDataValidator):
symbol patterns, and data structures.
"""
def __init__(self, component_name: str = "okx_data_validator"):
def __init__(self, component_name: str = "okx_data_validator", logger = None):
"""Initialize OKX data validator."""
super().__init__("okx", component_name)
super().__init__("okx", component_name, logger)
# OKX-specific patterns
self._symbol_pattern = re.compile(r'^[A-Z0-9]+-[A-Z0-9]+$') # BTC-USDT, ETH-USDC
@@ -95,7 +94,8 @@ class OKXDataValidator(BaseDataValidator):
'candle1m', 'candle5m', 'candle15m', 'candle1H', 'candle4H', 'candle1D'
}
self.logger.debug("Initialized OKX data validator")
if self.logger:
self.logger.debug("Initialized OKX data validator")
def validate_symbol_format(self, symbol: str) -> ValidationResult:
"""Validate OKX symbol format (e.g., BTC-USDT)."""
@@ -423,9 +423,9 @@ class OKXDataTransformer(BaseDataTransformer):
This class handles transformation of OKX data formats to standardized formats.
"""
def __init__(self, component_name: str = "okx_data_transformer"):
def __init__(self, component_name: str = "okx_data_transformer", logger = None):
"""Initialize OKX data transformer."""
super().__init__("okx", component_name)
super().__init__("okx", component_name, logger)
def transform_trade_data(self, raw_data: Dict[str, Any], symbol: str) -> Optional[StandardizedTrade]:
"""Transform OKX trade data to standardized format."""
@@ -442,7 +442,8 @@ class OKXDataTransformer(BaseDataTransformer):
is_milliseconds=True
)
except Exception as e:
self.logger.error(f"Error transforming OKX trade data: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error transforming OKX trade data: {e}")
return None
def transform_orderbook_data(self, raw_data: Dict[str, Any], symbol: str) -> Optional[Dict[str, Any]]:
@@ -458,7 +459,8 @@ class OKXDataTransformer(BaseDataTransformer):
'raw_data': raw_data
}
except Exception as e:
self.logger.error(f"Error transforming OKX orderbook data: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error transforming OKX orderbook data: {e}")
return None
def transform_ticker_data(self, raw_data: Dict[str, Any], symbol: str) -> Optional[Dict[str, Any]]:
@@ -497,7 +499,8 @@ class OKXDataTransformer(BaseDataTransformer):
return ticker_data
except Exception as e:
self.logger.error(f"Error transforming OKX ticker data: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error transforming OKX ticker data: {e}")
return None
@@ -512,7 +515,8 @@ class OKXDataProcessor:
def __init__(self,
symbol: str,
config: Optional[CandleProcessingConfig] = None,
component_name: str = "okx_data_processor"):
component_name: str = "okx_data_processor",
logger = None):
"""
Initialize OKX data processor.
@@ -523,17 +527,17 @@ class OKXDataProcessor:
"""
self.symbol = symbol
self.component_name = component_name
self.logger = get_logger(self.component_name)
self.logger = logger
# Core components using common utilities
self.validator = OKXDataValidator(f"{component_name}_validator")
self.transformer = OKXDataTransformer(f"{component_name}_transformer")
self.unified_transformer = UnifiedDataTransformer(self.transformer, f"{component_name}_unified")
self.validator = OKXDataValidator(f"{component_name}_validator", self.logger)
self.transformer = OKXDataTransformer(f"{component_name}_transformer", self.logger)
self.unified_transformer = UnifiedDataTransformer(self.transformer, f"{component_name}_unified", self.logger)
# Real-time candle processing using common utilities
self.config = config or CandleProcessingConfig()
self.candle_processor = RealTimeCandleProcessor(
symbol, "okx", self.config, f"{component_name}_candles"
symbol, "okx", self.config, f"{component_name}_candles", self.logger
)
# Callbacks
@@ -543,7 +547,8 @@ class OKXDataProcessor:
# Connect candle processor callbacks
self.candle_processor.add_candle_callback(self._emit_candle_to_callbacks)
self.logger.info(f"Initialized OKX data processor for {symbol} with real-time candle processing")
if self.logger:
self.logger.info(f"{self.component_name}: Initialized OKX data processor for {symbol} with real-time candle processing")
def add_trade_callback(self, callback: callable) -> None:
"""Add callback for processed trades."""
@@ -571,12 +576,14 @@ class OKXDataProcessor:
validation_result = self.validator.validate_websocket_message(message)
if not validation_result.is_valid:
self.logger.error(f"Message validation failed: {validation_result.errors}")
if self.logger:
self.logger.error(f"{self.component_name}: Message validation failed: {validation_result.errors}")
return False, [], validation_result.errors
# Log warnings if any
if validation_result.warnings:
self.logger.warning(f"Message validation warnings: {validation_result.warnings}")
if self.logger:
self.logger.warning(f"{self.component_name}: Message validation warnings: {validation_result.warnings}")
# Process data if it's a data message
if 'data' in message and 'arg' in message:
@@ -586,8 +593,9 @@ class OKXDataProcessor:
return True, [], []
except Exception as e:
error_msg = f"Exception during message validation and processing: {str(e)}"
self.logger.error(error_msg)
error_msg = f"{self.component_name}: Exception during message validation and processing: {str(e)}"
if self.logger:
self.logger.error(error_msg)
return False, [], [error_msg]
def _process_data_message(self, message: Dict[str, Any], expected_symbol: Optional[str] = None) -> Tuple[bool, List[MarketDataPoint], List[str]]:
@@ -626,7 +634,8 @@ class OKXDataProcessor:
continue
if validation_result.warnings:
self.logger.warning(f"Data validation warnings: {validation_result.warnings}")
if self.logger:
self.logger.warning(f"{self.component_name}: Data validation warnings: {validation_result.warnings}")
# Create MarketDataPoint using sanitized data
sanitized_data = validation_result.sanitized_data or data_item
@@ -650,13 +659,14 @@ class OKXDataProcessor:
self._process_real_time_trade(sanitized_data)
except Exception as e:
self.logger.error(f"Error processing data item: {e}")
errors.append(f"Error processing data item: {str(e)}")
if self.logger:
self.logger.error(f"{self.component_name}: Error processing data item: {e}")
errors.append(f"{self.component_name}: Error processing data item: {str(e)}")
return len(errors) == 0, market_data_points, errors
except Exception as e:
error_msg = f"Exception during data message processing: {str(e)}"
error_msg = f"{self.component_name}: Exception during data message processing: {str(e)}"
errors.append(error_msg)
return False, [], errors
@@ -675,12 +685,14 @@ class OKXDataProcessor:
try:
callback(standardized_trade)
except Exception as e:
self.logger.error(f"Error in trade callback: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error in trade callback: {e}")
# Note: Candle callbacks are handled by _emit_candle_to_callbacks
except Exception as e:
self.logger.error(f"Error processing real-time trade: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error processing real-time trade: {e}")
def _emit_candle_to_callbacks(self, candle: OHLCVCandle) -> None:
"""Emit candle to all registered callbacks."""
@@ -688,7 +700,8 @@ class OKXDataProcessor:
try:
callback(candle)
except Exception as e:
self.logger.error(f"Error in candle callback: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error in candle callback: {e}")
def _channel_to_data_type(self, channel: str) -> Optional[DataType]:
"""Convert OKX channel name to DataType enum."""

145
data/exchanges/okx/websocket.py
View File

@@ -17,8 +17,6 @@ from dataclasses import dataclass
import websockets
from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidURI
from utils.logger import get_logger
class OKXChannelType(Enum):
"""OKX WebSocket channel types."""
@@ -91,7 +89,8 @@ class OKXWebSocketClient:
ping_interval: float = 25.0,
pong_timeout: float = 10.0,
max_reconnect_attempts: int = 5,
reconnect_delay: float = 5.0):
reconnect_delay: float = 5.0,
logger = None):
"""
Initialize OKX WebSocket client.
@@ -109,7 +108,7 @@ class OKXWebSocketClient:
self.reconnect_delay = reconnect_delay
# Initialize logger
self.logger = get_logger(self.component_name, verbose=True)
self.logger = logger
# Connection management
self._websocket: Optional[Any] = None # Changed to Any to handle different websocket types
@@ -138,7 +137,8 @@ class OKXWebSocketClient:
'last_message_time': None
}
self.logger.info(f"Initialized OKX WebSocket client: {component_name}")
if self.logger:
self.logger.info(f"{self.component_name}: Initialized OKX WebSocket client")
@property
def is_connected(self) -> bool:
@@ -184,7 +184,8 @@ class OKXWebSocketClient:
True if connection successful, False otherwise
"""
if self.is_connected:
self.logger.warning("Already connected to OKX WebSocket")
if self.logger:
self.logger.warning("Already connected to OKX WebSocket")
return True
url = self.PUBLIC_WS_URL if use_public else self.PRIVATE_WS_URL
@@ -194,7 +195,8 @@ class OKXWebSocketClient:
self._connection_state = ConnectionState.CONNECTING
try:
self.logger.info(f"Connecting to OKX WebSocket (attempt {attempt + 1}/{self.max_reconnect_attempts}): {url}")
if self.logger:
self.logger.info(f"{self.component_name}: Connecting to OKX WebSocket (attempt {attempt + 1}/{self.max_reconnect_attempts}): {url}")
# Create SSL context for secure connection
ssl_context = ssl.create_default_context()
@@ -219,25 +221,30 @@ class OKXWebSocketClient:
# Start background tasks
await self._start_background_tasks()
self.logger.info("Successfully connected to OKX WebSocket")
if self.logger:
self.logger.info(f"{self.component_name}: Successfully connected to OKX WebSocket")
return True
except (InvalidURI, InvalidHandshake) as e:
self.logger.error(f"Invalid WebSocket configuration: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Invalid WebSocket configuration: {e}")
self._connection_state = ConnectionState.ERROR
return False
except Exception as e:
attempt_num = attempt + 1
self.logger.error(f"Connection attempt {attempt_num} failed: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Connection attempt {attempt_num} failed: {e}")
if attempt_num < self.max_reconnect_attempts:
# Exponential backoff with jitter
delay = self.reconnect_delay * (2 ** attempt) + (0.1 * attempt)
self.logger.info(f"Retrying connection in {delay:.1f} seconds...")
if self.logger:
self.logger.info(f"{self.component_name}: Retrying connection in {delay:.1f} seconds...")
await asyncio.sleep(delay)
else:
self.logger.error(f"All {self.max_reconnect_attempts} connection attempts failed")
if self.logger:
self.logger.error(f"{self.component_name}: All {self.max_reconnect_attempts} connection attempts failed")
self._connection_state = ConnectionState.ERROR
return False
@@ -248,7 +255,8 @@ class OKXWebSocketClient:
if not self._websocket:
return
self.logger.info("Disconnecting from OKX WebSocket")
if self.logger:
self.logger.info(f"{self.component_name}: Disconnecting from OKX WebSocket")
self._connection_state = ConnectionState.DISCONNECTED
# Cancel background tasks
@@ -258,12 +266,14 @@ class OKXWebSocketClient:
try:
await self._websocket.close()
except Exception as e:
self.logger.warning(f"Error closing WebSocket: {e}")
if self.logger:
self.logger.warning(f"{self.component_name}: Error closing WebSocket: {e}")
self._websocket = None
self._is_authenticated = False
self.logger.info("Disconnected from OKX WebSocket")
if self.logger:
self.logger.info(f"{self.component_name}: Disconnected from OKX WebSocket")
async def subscribe(self, subscriptions: List[OKXSubscription]) -> bool:
"""
@@ -276,7 +286,8 @@ class OKXWebSocketClient:
True if subscription successful, False otherwise
"""
if not self.is_connected:
self.logger.error("Cannot subscribe: WebSocket not connected")
if self.logger:
self.logger.error("Cannot subscribe: WebSocket not connected")
return False
try:
@@ -295,11 +306,13 @@ class OKXWebSocketClient:
key = f"{sub.channel}:{sub.inst_id}"
self._subscriptions[key] = sub
self.logger.info(f"Subscribed to {len(subscriptions)} channels")
if self.logger:
self.logger.info(f"{self.component_name}: Subscribed to {len(subscriptions)} channels")
return True
except Exception as e:
self.logger.error(f"Failed to subscribe to channels: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to subscribe to channels: {e}")
return False
async def unsubscribe(self, subscriptions: List[OKXSubscription]) -> bool:
@@ -313,7 +326,8 @@ class OKXWebSocketClient:
True if unsubscription successful, False otherwise
"""
if not self.is_connected:
self.logger.error("Cannot unsubscribe: WebSocket not connected")
if self.logger:
self.logger.error("Cannot unsubscribe: WebSocket not connected")
return False
try:
@@ -332,11 +346,13 @@ class OKXWebSocketClient:
key = f"{sub.channel}:{sub.inst_id}"
self._subscriptions.pop(key, None)
self.logger.info(f"Unsubscribed from {len(subscriptions)} channels")
if self.logger:
self.logger.info(f"{self.component_name}: Unsubscribed from {len(subscriptions)} channels")
return True
except Exception as e:
self.logger.error(f"Failed to unsubscribe from channels: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to unsubscribe from channels: {e}")
return False
def add_message_callback(self, callback: Callable[[Dict[str, Any]], None]) -> None:
@@ -347,7 +363,8 @@ class OKXWebSocketClient:
callback: Function to call when message received
"""
self._message_callbacks.append(callback)
self.logger.debug(f"Added message callback: {callback.__name__}")
if self.logger:
self.logger.debug(f"{self.component_name}: Added message callback: {callback.__name__}")
def remove_message_callback(self, callback: Callable[[Dict[str, Any]], None]) -> None:
"""
@@ -358,7 +375,8 @@ class OKXWebSocketClient:
"""
if callback in self._message_callbacks:
self._message_callbacks.remove(callback)
self.logger.debug(f"Removed message callback: {callback.__name__}")
if self.logger:
self.logger.debug(f"{self.component_name}: Removed message callback: {callback.__name__}")
async def _start_background_tasks(self) -> None:
"""Start background tasks for ping and message handling."""
@@ -368,7 +386,8 @@ class OKXWebSocketClient:
# Start message handler task
self._message_handler_task = asyncio.create_task(self._message_handler())
self.logger.debug("Started background tasks")
if self.logger:
self.logger.debug(f"{self.component_name}: Started background tasks")
async def _stop_background_tasks(self) -> None:
"""Stop background tasks."""
@@ -385,7 +404,8 @@ class OKXWebSocketClient:
self._ping_task = None
self._message_handler_task = None
self.logger.debug("Stopped background tasks")
if self.logger:
self.logger.debug(f"{self.component_name}: Stopped background tasks")
async def _ping_loop(self) -> None:
"""Background task for sending ping messages."""
@@ -401,7 +421,8 @@ class OKXWebSocketClient:
# Check for pong timeout
if (self._last_ping_time > self._last_pong_time and
current_time - self._last_ping_time > self.pong_timeout):
self.logger.warning("Pong timeout - connection may be stale")
if self.logger:
self.logger.warning(f"{self.component_name}: Pong timeout - connection may be stale")
# Don't immediately disconnect, let connection error handling deal with it
await asyncio.sleep(1) # Check every second
@@ -409,7 +430,8 @@ class OKXWebSocketClient:
except asyncio.CancelledError:
break
except Exception as e:
self.logger.error(f"Error in ping loop: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error in ping loop: {e}")
await asyncio.sleep(5)
async def _message_handler(self) -> None:
@@ -432,32 +454,38 @@ class OKXWebSocketClient:
await self._process_message(message)
except ConnectionClosed as e:
self.logger.warning(f"WebSocket connection closed: {e}")
if self.logger:
self.logger.warning(f"{self.component_name}: WebSocket connection closed: {e}")
self._connection_state = ConnectionState.DISCONNECTED
# Attempt automatic reconnection if enabled
if self._reconnect_attempts < self.max_reconnect_attempts:
self._reconnect_attempts += 1
self.logger.info(f"Attempting automatic reconnection ({self._reconnect_attempts}/{self.max_reconnect_attempts})")
if self.logger:
self.logger.info(f"{self.component_name}: Attempting automatic reconnection ({self._reconnect_attempts}/{self.max_reconnect_attempts})")
# Stop current tasks
await self._stop_background_tasks()
# Attempt reconnection
if await self.reconnect():
self.logger.info("Automatic reconnection successful")
if self.logger:
self.logger.info(f"{self.component_name}: Automatic reconnection successful")
continue
else:
self.logger.error("Automatic reconnection failed")
if self.logger:
self.logger.error(f"{self.component_name}: Automatic reconnection failed")
break
else:
self.logger.error("Max reconnection attempts exceeded")
if self.logger:
self.logger.error(f"{self.component_name}: Max reconnection attempts exceeded")
break
except asyncio.CancelledError:
break
except Exception as e:
self.logger.error(f"Error in message handler: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error in message handler: {e}")
await asyncio.sleep(1)
async def _send_message(self, message: Dict[str, Any]) -> None:
@@ -474,14 +502,17 @@ class OKXWebSocketClient:
message_str = json.dumps(message)
await self._websocket.send(message_str)
self._stats['messages_sent'] += 1
self.logger.debug(f"Sent message: {message}")
if self.logger:
self.logger.debug(f"{self.component_name}: Sent message: {message}")
except ConnectionClosed as e:
self.logger.error(f"Connection closed while sending message: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Connection closed while sending message: {e}")
self._connection_state = ConnectionState.DISCONNECTED
raise OKXConnectionError(f"Connection closed: {e}")
except Exception as e:
self.logger.error(f"Failed to send message: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to send message: {e}")
raise OKXConnectionError(f"Failed to send message: {e}")
async def _send_ping(self) -> None:
@@ -493,14 +524,17 @@ class OKXWebSocketClient:
# OKX expects a simple "ping" string, not JSON
await self._websocket.send("ping")
self._stats['pings_sent'] += 1
self.logger.debug("Sent ping to OKX")
if self.logger:
self.logger.debug(f"{self.component_name}: Sent ping to OKX")
except ConnectionClosed as e:
self.logger.error(f"Connection closed while sending ping: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Connection closed while sending ping: {e}")
self._connection_state = ConnectionState.DISCONNECTED
raise OKXConnectionError(f"Connection closed: {e}")
except Exception as e:
self.logger.error(f"Failed to send ping: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to send ping: {e}")
raise OKXConnectionError(f"Failed to send ping: {e}")
async def _process_message(self, message: str) -> None:
@@ -519,7 +553,8 @@ class OKXWebSocketClient:
if message.strip() == "pong":
self._last_pong_time = time.time()
self._stats['pongs_received'] += 1
self.logger.debug("Received pong from OKX")
if self.logger:
self.logger.debug(f"{self.component_name}: Received pong from OKX")
return
# Parse JSON message for all other responses
@@ -529,21 +564,25 @@ class OKXWebSocketClient:
if data.get('event') == 'pong':
self._last_pong_time = time.time()
self._stats['pongs_received'] += 1
self.logger.debug("Received pong from OKX (JSON format)")
if self.logger:
self.logger.debug(f"{self.component_name}: Received pong from OKX (JSON format)")
return
# Handle subscription confirmations
if data.get('event') == 'subscribe':
self.logger.info(f"Subscription confirmed: {data}")
if self.logger:
self.logger.info(f"{self.component_name}: Subscription confirmed: {data}")
return
if data.get('event') == 'unsubscribe':
self.logger.info(f"Unsubscription confirmed: {data}")
if self.logger:
self.logger.info(f"{self.component_name}: Unsubscription confirmed: {data}")
return
# Handle error messages
if data.get('event') == 'error':
self.logger.error(f"OKX error: {data}")
if self.logger:
self.logger.error(f"{self.component_name}: OKX error: {data}")
return
# Process data messages
@@ -553,19 +592,23 @@ class OKXWebSocketClient:
try:
callback(data)
except Exception as e:
self.logger.error(f"Error in message callback {callback.__name__}: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error in message callback {callback.__name__}: {e}")
except json.JSONDecodeError as e:
# Check if it's a simple string response we haven't handled
if message.strip() in ["ping", "pong"]:
self.logger.debug(f"Received simple message: {message.strip()}")
if self.logger:
self.logger.debug(f"{self.component_name}: Received simple message: {message.strip()}")
if message.strip() == "pong":
self._last_pong_time = time.time()
self._stats['pongs_received'] += 1
else:
self.logger.error(f"Failed to parse JSON message: {e}, message: {message}")
if self.logger:
self.logger.error(f"{self.component_name}: Failed to parse JSON message: {e}, message: {message}")
except Exception as e:
self.logger.error(f"Error processing message: {e}")
if self.logger:
self.logger.error(f"{self.component_name}: Error processing message: {e}")
def get_stats(self) -> Dict[str, Any]:
"""Get connection statistics."""
@@ -588,7 +631,8 @@ class OKXWebSocketClient:
Returns:
True if reconnection successful, False otherwise
"""
self.logger.info("Attempting to reconnect to OKX WebSocket")
if self.logger:
self.logger.info(f"{self.component_name}: Attempting to reconnect to OKX WebSocket")
self._connection_state = ConnectionState.RECONNECTING
self._stats['reconnections'] += 1
@@ -605,7 +649,8 @@ class OKXWebSocketClient:
# Re-subscribe to previous subscriptions
if self._subscriptions:
subscriptions = list(self._subscriptions.values())
self.logger.info(f"Re-subscribing to {len(subscriptions)} channels")
if self.logger:
self.logger.info(f"{self.component_name}: Re-subscribing to {len(subscriptions)} channels")
await self.subscribe(subscriptions)
return success