Refactor BaseDataCollector to utilize CollectorStateAndTelemetry for improved state management

- Introduced a new `CollectorStateAndTelemetry` class to encapsulate the status, health checks, and statistics of the data collector, promoting modularity and separation of concerns.
- Updated `BaseDataCollector` to replace direct status management with calls to the new telemetry class, enhancing maintainability and readability.
- Refactored logging methods to utilize the telemetry class, ensuring consistent logging practices.
- Modified the `OKXCollector` to integrate with the new telemetry system for improved status reporting and error handling.
- Added comprehensive tests for the `CollectorStateAndTelemetry` class to ensure functionality and reliability.

These changes streamline the data collector's architecture, aligning with project standards for maintainability and performance.
This commit is contained in:
Vasily.onl
2025-06-09 17:27:29 +08:00
parent ec8f5514bb
commit 60434afd5d
6 changed files with 615 additions and 294 deletions

View File

@@ -132,10 +132,9 @@ class OKXCollector(BaseDataCollector):
DataType.TICKER: OKXChannelType.TICKERS.value
}
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 timeframes: {self.timeframes}")
logger.info(f"{component_name}: Using common data processing framework")
self._log_info(f"{component_name}: Initialized OKX collector for {symbol} with data types: {[dt.value for dt in data_types]}")
self._log_info(f"{component_name}: Using timeframes: {self.timeframes}")
self._log_info(f"{component_name}: Using common data processing framework")
async def connect(self) -> bool:
"""
@@ -145,11 +144,10 @@ class OKXCollector(BaseDataCollector):
True if connection successful, False otherwise
"""
try:
if self.logger:
self.logger.info(f"{self.component_name}: Connecting OKX collector for {self.symbol}")
self._log_info(f"Connecting OKX collector for {self.symbol}")
# Initialize database operations using repository pattern
self._db_operations = get_database_operations(self.logger)
self._db_operations = get_database_operations(self.logger) # self.logger needs to be accessible for database operations
# Create WebSocket client
ws_component_name = f"okx_ws_{self.symbol.replace('-', '_').lower()}"
@@ -167,35 +165,31 @@ class OKXCollector(BaseDataCollector):
# Connect to WebSocket
if not await self._ws_client.connect(use_public=True):
if self.logger:
self.logger.error(f"{self.component_name}: Failed to connect to OKX WebSocket")
self._log_error(f"Failed to connect to OKX WebSocket")
return False
if self.logger:
self.logger.info(f"{self.component_name}: Successfully connected OKX collector for {self.symbol}")
self._log_info(f"Successfully connected OKX collector for {self.symbol}")
return True
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error connecting OKX collector for {self.symbol}: {e}")
self._log_error(f"Error connecting OKX collector for {self.symbol}: {e}")
return False
async def disconnect(self) -> None:
"""Disconnect from OKX WebSocket API."""
"""
Disconnect from OKX WebSocket API.
"""
try:
if self.logger:
self.logger.info(f"{self.component_name}: Disconnecting OKX collector for {self.symbol}")
self._log_info(f"Disconnecting OKX collector for {self.symbol}")
if self._ws_client:
await self._ws_client.disconnect()
self._ws_client = None
if self.logger:
self.logger.info(f"{self.component_name}: Disconnected OKX collector for {self.symbol}")
self._log_info(f"Disconnected OKX collector for {self.symbol}")
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error disconnecting OKX collector for {self.symbol}: {e}")
self._log_error(f"Error disconnecting OKX collector for {self.symbol}: {e}")
async def subscribe_to_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
"""
@@ -209,14 +203,12 @@ class OKXCollector(BaseDataCollector):
True if subscription successful, False otherwise
"""
if not self._ws_client or not self._ws_client.is_connected:
if self.logger:
self.logger.error(f"{self.component_name}: WebSocket client not connected")
self._log_error(f"WebSocket client not connected")
return False
# Validate symbol
if self.symbol not in symbols:
if self.logger:
self.logger.warning(f"{self.component_name}: Symbol {self.symbol} not in subscription list: {symbols}")
self._log_warning(f"Symbol {self.symbol} not in subscription list: {symbols}")
return False
try:
@@ -231,31 +223,25 @@ class OKXCollector(BaseDataCollector):
enabled=True
)
subscriptions.append(subscription)
if self.logger:
self.logger.debug(f"{self.component_name}: Added subscription: {channel} for {self.symbol}")
self._log_debug(f"Added subscription: {channel} for {self.symbol}")
else:
if self.logger:
self.logger.warning(f"{self.component_name}: Unsupported data type: {data_type}")
self._log_warning(f"Unsupported data type: {data_type}")
if not subscriptions:
if self.logger:
self.logger.warning(f"{self.component_name}: No valid subscriptions to create")
self._log_warning(f"No valid subscriptions to create")
return False
# Subscribe to channels
success = await self._ws_client.subscribe(subscriptions)
if success:
if self.logger:
self.logger.info(f"{self.component_name}: Successfully subscribed to {len(subscriptions)} channels for {self.symbol}")
self._log_info(f"Successfully subscribed to {len(subscriptions)} channels for {self.symbol}")
return True
else:
if self.logger:
self.logger.error(f"{self.component_name}: Failed to subscribe to channels for {self.symbol}")
self._log_error(f"Failed to subscribe to channels for {self.symbol}")
return False
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error subscribing to data for {self.symbol}: {e}")
self._log_error(f"Error subscribing to data for {self.symbol}: {e}")
return False
async def unsubscribe_from_data(self, symbols: List[str], data_types: List[DataType]) -> bool:
@@ -270,8 +256,7 @@ class OKXCollector(BaseDataCollector):
True if unsubscription successful, False otherwise
"""
if not self._ws_client or not self._ws_client.is_connected:
if self.logger:
self.logger.warning(f"{self.component_name}: WebSocket client not connected")
self._log_warning(f"WebSocket client not connected")
return True # Consider it successful if not connected
try:
@@ -293,17 +278,14 @@ class OKXCollector(BaseDataCollector):
# Unsubscribe from channels
success = await self._ws_client.unsubscribe(subscriptions)
if success:
if self.logger:
self.logger.info(f"{self.component_name}: Successfully unsubscribed from {len(subscriptions)} channels for {self.symbol}")
self._log_info(f"Successfully unsubscribed from {len(subscriptions)} channels for {self.symbol}")
return True
else:
if self.logger:
self.logger.error(f"{self.component_name}: Failed to unsubscribe from channels for {self.symbol}")
self._log_error(f"Failed to unsubscribe from channels for {self.symbol}")
return False
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error unsubscribing from data for {self.symbol}: {e}")
self._log_error(f"Error unsubscribing from data for {self.symbol}: {e}")
return False
async def _process_message(self, message: Any) -> Optional[MarketDataPoint]:
@@ -317,8 +299,7 @@ class OKXCollector(BaseDataCollector):
MarketDataPoint if processing successful, None otherwise
"""
if not isinstance(message, dict):
if self.logger:
self.logger.warning(f"{self.component_name}: Received non-dict message: {type(message)}")
self._log_warning(f"Received non-dict message: {type(message)}")
return None
try:
@@ -331,13 +312,11 @@ class OKXCollector(BaseDataCollector):
if not success:
self._error_count += 1
if self.logger:
self.logger.error(f"{self.component_name}: Message processing failed: {errors}")
self._log_error(f"Message processing failed: {errors}")
return None
if errors:
if self.logger:
self.logger.warning(f"{self.component_name}: Message processing warnings: {errors}")
self._log_warning(f"Message processing warnings: {errors}")
# Store raw data if enabled (for debugging/compliance)
if self.store_raw_data:
@@ -353,22 +332,23 @@ class OKXCollector(BaseDataCollector):
except Exception as e:
self._error_count += 1
if self.logger:
self.logger.error(f"{self.component_name}: Error processing message: {e}")
self._log_error(f"Error processing message: {e}")
return None
async def _handle_messages(self) -> None:
"""Handle message processing in the background."""
"""
Handle message processing in the background.
This method exists for compatibility with BaseDataCollector
"""
# The new data processor handles messages through callbacks
# This method exists for compatibility with BaseDataCollector
# Update heartbeat to indicate the message loop is active
self._last_heartbeat = datetime.now(timezone.utc)
self._state_telemetry.update_heartbeat()
# Check if we're receiving WebSocket messages
if self._ws_client and self._ws_client.is_connected:
# Update last data received timestamp if WebSocket is connected and active
self._last_data_received = datetime.now(timezone.utc)
self._state_telemetry.update_data_received_timestamp()
# Short sleep to prevent busy loop while maintaining heartbeat
await asyncio.sleep(0.1)
@@ -386,15 +366,13 @@ class OKXCollector(BaseDataCollector):
# Store raw market data points in raw_trades table using repository
success = self._db_operations.raw_trades.insert_market_data_point(data_point)
if success and self.logger:
self.logger.debug(f"{self.component_name}: Stored raw data: {data_point.data_type.value} for {data_point.symbol}")
if success:
self._log_debug(f"Stored raw data: {data_point.data_type.value} for {data_point.symbol}")
except DatabaseOperationError as e:
if self.logger:
self.logger.error(f"{self.component_name}: Database error storing raw market data: {e}")
self._log_error(f"Database error storing raw market data: {e}")
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error storing raw market data: {e}")
self._log_error(f"Error storing raw market data: {e}")
async def _store_completed_candle(self, candle: OHLCVCandle) -> None:
"""
@@ -414,21 +392,19 @@ class OKXCollector(BaseDataCollector):
# Store completed candles using repository pattern
success = self._db_operations.market_data.upsert_candle(candle, self.force_update_candles)
if success and self.logger:
if success:
action = "Updated" if self.force_update_candles else "Stored"
self.logger.debug(f"{self.component_name}: {action} candle: {candle.symbol} {candle.timeframe} at {candle.end_time} (force_update={self.force_update_candles}) - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}, Vol: {candle.volume}, Trades: {candle.trade_count}")
self._log_debug(f"{action} candle: {candle.symbol} {candle.timeframe} at {candle.end_time} (force_update={self.force_update_candles}) - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}, Vol: {candle.volume}, Trades: {candle.trade_count}")
except DatabaseOperationError as e:
if self.logger:
self.logger.error(f"{self.component_name}: Database error storing completed candle: {e}")
# Log candle details for debugging
self.logger.error(f"{self.component_name}: Failed candle details: {candle.symbol} {candle.timeframe} {candle.end_time} - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}")
self._log_error(f"Database error storing completed candle: {e}")
# Log candle details for debugging
self._log_error(f"Failed candle details: {candle.symbol} {candle.timeframe} {candle.end_time} - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}")
self._error_count += 1
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error storing completed candle: {e}")
# Log candle details for debugging
self.logger.error(f"{self.component_name}: Failed candle details: {candle.symbol} {candle.timeframe} {candle.end_time} - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}")
self._log_error(f"Error storing completed candle: {e}")
# Log candle details for debugging
self._log_error(f"Failed candle details: {candle.symbol} {candle.timeframe} {candle.end_time} - OHLCV: {candle.open}/{candle.high}/{candle.low}/{candle.close}")
self._error_count += 1
async def _store_raw_data(self, channel: str, raw_message: Dict[str, Any]) -> None:
@@ -452,15 +428,13 @@ class OKXCollector(BaseDataCollector):
raw_data=data_item,
timestamp=datetime.now(timezone.utc)
)
if not success and self.logger:
self.logger.warning(f"{self.component_name}: Failed to store raw WebSocket data for {channel}")
if not success:
self._log_warning(f"Failed to store raw WebSocket data for {channel}")
except DatabaseOperationError as e:
if self.logger:
self.logger.error(f"{self.component_name}: Database error storing raw WebSocket data: {e}")
self._log_error(f"Database error storing raw WebSocket data: {e}")
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error storing raw WebSocket data: {e}")
self._log_error(f"Error storing raw WebSocket data: {e}")
def _on_message(self, message: Dict[str, Any]) -> None:
"""
@@ -471,16 +445,14 @@ class OKXCollector(BaseDataCollector):
"""
try:
# Update heartbeat and data received timestamps
current_time = datetime.now(timezone.utc)
self._last_heartbeat = current_time
self._last_data_received = current_time
self._state_telemetry.update_heartbeat()
self._state_telemetry.update_data_received_timestamp()
self._message_count += 1
# Process message asynchronously
asyncio.create_task(self._process_message(message))
except Exception as e:
if self.logger:
self.logger.error(f"{self.component_name}: Error handling WebSocket message: {e}")
self._log_error(f"Error handling WebSocket message: {e}")
def _on_trade_processed(self, trade: StandardizedTrade) -> None:
"""
@@ -490,8 +462,7 @@ class OKXCollector(BaseDataCollector):
trade: Processed standardized trade
"""
self._processed_trades += 1
if self.logger:
self.logger.debug(f"{self.component_name}: Processed trade: {trade.symbol} {trade.side} {trade.size}@{trade.price}")
self._log_debug(f"Processed trade: {trade.symbol} {trade.side} {trade.size}@{trade.price}")
def _on_candle_processed(self, candle: OHLCVCandle) -> None:
"""
@@ -501,8 +472,7 @@ class OKXCollector(BaseDataCollector):
candle: Completed OHLCV candle
"""
self._processed_candles += 1
if self.logger:
self.logger.debug(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}")
self._log_debug(f"Processed candle: {candle.symbol} {candle.timeframe}")
# Store completed candle in market_data table
if candle.is_complete:
@@ -510,41 +480,22 @@ class OKXCollector(BaseDataCollector):
def get_status(self) -> Dict[str, Any]:
"""
Get current collector status including processing statistics.
Get current collector status and statistics.
Returns:
Dictionary containing collector status information
Dictionary containing status information
"""
base_status = super().get_status()
# Add OKX-specific status
okx_status = {
"symbol": self.symbol,
"websocket_connected": self._ws_client.is_connected if self._ws_client else False,
"websocket_state": self._ws_client.connection_state.value if self._ws_client else "disconnected",
"store_raw_data": self.store_raw_data,
"force_update_candles": self.force_update_candles,
"timeframes": self.timeframes,
"processing_stats": {
"messages_received": self._message_count,
"trades_processed": self._processed_trades,
"candles_processed": self._processed_candles,
"errors": self._error_count
}
}
# Add data processor statistics
if self._data_processor:
okx_status["data_processor_stats"] = self._data_processor.get_processing_stats()
# Add WebSocket statistics
if self._ws_client:
okx_status["websocket_stats"] = self._ws_client.get_stats()
# Merge with base status
base_status.update(okx_status)
return base_status
return self._state_telemetry.get_status()
def get_health_status(self) -> Dict[str, Any]:
"""
Get detailed health status for monitoring.
Returns:
Dictionary containing health information
"""
return self._state_telemetry.get_health_status()
def __repr__(self) -> str:
"""String representation of the collector."""
return f"OKXCollector(symbol='{self.symbol}', status='{self.status.value}', data_types={[dt.value for dt in self.data_types]})"
return f"<{self.__class__.__name__}({self.exchange_name}, {len(self.symbols)} symbols, {self._state_telemetry.status.value})>"