From b53a636f43ee882150f653e11a0a4e11769fccb5 Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Sat, 9 Nov 2024 13:44:19 -0800 Subject: [PATCH] add some exception handling --- .../websockets/openbb_websockets/client.py | 92 ++++++++++--------- .../websockets/openbb_websockets/models.py | 44 +++++++++ .../openbb_websockets/websockets_router.py | 11 ++- .../models/websocket_connection.py | 8 +- .../openbb_polygon/utils/websocket_client.py | 2 +- 5 files changed, 105 insertions(+), 52 deletions(-) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index 3f042a67b0a..ff909bbb031 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -3,7 +3,7 @@ # pylint: disable=too-many-statements # flake8: noqa: PLR0915 import logging -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, Literal, Optional, Union if TYPE_CHECKING: from openbb_core.provider.abstract.data import Data @@ -23,7 +23,7 @@ class WebSocketClient: The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces. limit : Optional[int] The limit of records to hold in memory. Once the limit is reached, the oldest records are removed. - Default is None. + Default is 300. Set to None to keep all records. results_file : Optional[str] Absolute path to the file for continuous writing. By default, a temporary file is created. table_name : Optional[str] @@ -37,9 +37,11 @@ class WebSocketClient: The authentication token to use for the WebSocket connection. Default is None. Only used for API and Python application endpoints. logger : Optional[logging.Logger] - The logger instance to use this connection. By default, a new logger is created. - kwargs : dict - Additional keyword arguments to pass to the target module. + The pre-configured logger instance to use for this connection. By default, a new logger is created. + kwargs : Optional[dict] + Additional keyword arguments to pass to the target provider module. Keywords and values must not contain spaces. + To pass items to 'websocket.connect()', include them in the 'kwargs' dictionary as, + {'connect_kwargs': {'key': 'value'}}. Properties ---------- @@ -48,9 +50,9 @@ class WebSocketClient: module : str Path to the provider connection script. is_running : bool - Check if the provider connection is running. + Check if the provider connection process is running. is_broadcasting : bool - Check if the broadcast server is running. + Check if the broadcast server process is running. broadcast_address : str URI address for the results broadcast server. results : list @@ -83,7 +85,7 @@ class WebSocketClient: name: str, module: str, symbol: Optional[str] = None, - limit: Optional[int] = None, + limit: Optional[int] = 300, results_file: Optional[str] = None, table_name: Optional[str] = None, save_results: bool = False, @@ -91,7 +93,7 @@ class WebSocketClient: auth_token: Optional[str] = None, logger: Optional[logging.Logger] = None, **kwargs, - ): + ) -> None: """Initialize the WebSocketClient class.""" # pylint: disable=import-outside-toplevel import asyncio # noqa @@ -130,6 +132,7 @@ class WebSocketClient: self._broadcast_thread = None self._broadcast_log_thread = None self._broadcast_message_queue = Queue() + self._exception = None if not results_file: with tempfile.NamedTemporaryFile(delete=False) as temp_file: @@ -144,6 +147,8 @@ class WebSocketClient: atexit.register(self._atexit) + # Set up the SQLite database and table. + # Loop handling is for when the class is used directly instead of from the app or API. try: loop = asyncio.get_event_loop() except (RuntimeError, RuntimeWarning): @@ -156,7 +161,7 @@ class WebSocketClient: except DatabaseError as e: self.logger.error("Error setting up the SQLite database and table: %s", e) - def _atexit(self): + def _atexit(self) -> None: """Clean up the WebSocket client processes at exit.""" # pylint: disable=import-outside-toplevel import os @@ -170,14 +175,14 @@ class WebSocketClient: if os.path.exists(self.results_file): os.remove(self.results_file) - async def _setup_database(self): + async def _setup_database(self) -> None: """Set up the SQLite database and table.""" # pylint: disable=import-outside-toplevel from openbb_websockets.helpers import setup_database return await setup_database(self.results_path, self.table_name) - def _log_provider_output(self, output_queue): + def _log_provider_output(self, output_queue) -> None: """Log output from the provider server queue.""" # pylint: disable=import-outside-toplevel import queue # noqa @@ -191,11 +196,13 @@ class WebSocketClient: if ( "server rejected" in output.lower() or "PROVIDER ERROR" in output + or "Unexpected error" in output ): - self._stop_log_thread_event.set() - self._psutil_process.kill() + err = ChildProcessError(output) + self._exception = err self.logger.error(output) break + output = clean_message(output) output = output + "\n" sys.stdout.write(output + "\n") @@ -203,7 +210,7 @@ class WebSocketClient: except queue.Empty: continue - def _log_broadcast_output(self, output_queue): + def _log_broadcast_output(self, output_queue) -> None: """Log output from the broadcast server queue.""" # pylint: disable=import-outside-toplevel import queue # noqa @@ -245,7 +252,7 @@ class WebSocketClient: except queue.Empty: continue - def connect(self): + def connect(self) -> None: """Connect to the provider WebSocket.""" # pylint: disable=import-outside-toplevel import json # noqa @@ -313,7 +320,7 @@ class WebSocketClient: def send_message( self, message, target: Literal["provider", "broadcast"] = "provider" - ): + ) -> None: """Send a message to the WebSocket process.""" if target == "provider": self._provider_message_queue.put(message) @@ -322,7 +329,7 @@ class WebSocketClient: self._broadcast_message_queue.put(message) read_message_queue(self, self._broadcast_message_queue, target="broadcast") - def disconnect(self): + def disconnect(self) -> None: """Disconnect from the provider WebSocket.""" self._stop_log_thread_event.set() if self._process is None or self.is_running is False: @@ -339,9 +346,11 @@ class WebSocketClient: self._log_thread.join() self._stop_log_thread_event.clear() self.logger.info("Disconnected from the provider WebSocket.") + if hasattr(self, "_exception") and self._exception: + raise self._exception return - def subscribe(self, symbol): + def subscribe(self, symbol) -> None: """Subscribe to a new symbol or list of symbols.""" # pylint: disable=import-outside-toplevel import json @@ -353,7 +362,7 @@ class WebSocketClient: new_symbols = list(set(old_symbols + ticker)) self._symbol = ",".join(new_symbols) - def unsubscribe(self, symbol): + def unsubscribe(self, symbol) -> None: """Unsubscribe from a symbol or list of symbols.""" # pylint: disable=import-outside-toplevel import json @@ -370,22 +379,22 @@ class WebSocketClient: self._symbol = ",".join(new_symbols) @property - def is_running(self): + def is_running(self) -> bool: """Check if the provider connection is running.""" if hasattr(self._psutil_process, "is_running"): return self._psutil_process.is_running() return False @property - def is_broadcasting(self): + def is_broadcasting(self) -> bool: """Check if the broadcast server is running.""" if hasattr(self._psutil_broadcast_process, "is_running"): return self._psutil_broadcast_process.is_running() return False @property - def results(self): - """Retrieve the raw results dumped by the WebSocket stream.""" + def results(self) -> Union[list[dict], None]: + """Retrieve the deserialized results from the results file.""" # pylint: disable=import-outside-toplevel import json # noqa import sqlite3 @@ -397,13 +406,14 @@ class WebSocketClient: cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa for row in cursor: index, message = row - output.append(json.loads(message)) + output.append(json.loads(json.loads(message))) + if output: return output self.logger.info("No results found in %s", self.results_file) - return [] + return @results.deleter def results(self): @@ -448,7 +458,7 @@ class WebSocketClient: self.logger.error("Error clearing results: %s", e) @property - def module(self): + def module(self) -> str: """Path to the provider connection script.""" return self._module @@ -465,12 +475,12 @@ class WebSocketClient: ] @property - def symbol(self): + def symbol(self) -> str: """Symbol(s) requested to subscribe.""" return self._symbol @property - def limit(self): + def limit(self) -> Union[int, None]: """Get the limit of records to hold in memory.""" return self._limit @@ -480,7 +490,7 @@ class WebSocketClient: self._limit = limit @property - def broadcast_address(self): + def broadcast_address(self) -> Union[str, None]: """Get the WebSocket broadcast address.""" return ( self._broadcast_address @@ -493,7 +503,7 @@ class WebSocketClient: host: str = "127.0.0.1", port: int = 6666, **kwargs, - ): + ) -> None: """Broadcast results over a network connection.""" # pylint: disable=import-outside-toplevel import os # noqa @@ -594,15 +604,11 @@ class WebSocketClient: return @property - def transformed_results(self): - """Deserialize the records from the results file.""" - # pylint: disable=import-outside-toplevel - import json - + def transformed_results(self) -> list["Data"]: + """Model validated records from the results file.""" if not self.data_model: raise NotImplementedError("No model provided to transform the results.") - - return [self.data_model.model_validate(json.loads(d)) for d in self.results] + return [self.data_model.model_validate(d) for d in self.results] def __repr__(self): """Return the WebSocketClient representation.""" @@ -617,7 +623,7 @@ class WebSocketClient: ) -def non_blocking_websocket(client, output_queue, provider_message_queue): +def non_blocking_websocket(client, output_queue, provider_message_queue) -> None: """Communicate with the threaded process.""" try: while not client._stop_log_thread_event.is_set(): @@ -639,8 +645,8 @@ def non_blocking_websocket(client, output_queue, provider_message_queue): def send_message( client, message, target: Literal["provider", "broadcast"] = "provider" -): - """Send a message to the WebSocket process.""" +) -> None: + """Send a message to the WebSocketConnection process.""" try: if target == "provider": if client._process and client._process.stdin: @@ -661,7 +667,7 @@ def send_message( def read_message_queue( client, message_queue, target: Literal["provider", "broadcast"] = "provider" ): - """Read messages from the queue and send them to the WebSocket process.""" + """Read messages from the queue and send them to the WebSocketConnection process.""" while not message_queue.empty(): try: if target == "provider": @@ -681,7 +687,7 @@ def read_message_queue( break -def non_blocking_broadcast(client, output_queue, broadcast_message_queue): +def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> None: """Continuously read the output from the broadcast process and log it to the main thread.""" try: while not client._stop_broadcasting_event.is_set(): diff --git a/openbb_platform/extensions/websockets/openbb_websockets/models.py b/openbb_platform/extensions/websockets/openbb_websockets/models.py index e50833f0961..19d59769d5c 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/models.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/models.py @@ -123,3 +123,47 @@ class WebSocketConnection(Data): if not isinstance(v, WebSocketClient): raise ValueError("Client must be an instance of WebSocketClient.") return v + + +class WebSocketConnectionStatus(Data): + """Data model for WebSocketConnection status information.""" + + name: str = Field( + description="Name assigned to the client connection.", + ) + auth_required: bool = Field( + description="True when 'auth_token' is supplied at initialization." + " When True, interactions with the client from the Python or API" + + " endpoints requires it to be supplied as a query parameter.", + ) + subscribed_symbols: str = Field( + description="Symbols subscribed to by the client connection.", + ) + is_running: bool = Field( + description="Whether the client connection is running.", + ) + provider_pid: Optional[int] = Field( + default=None, + description="Process ID of the provider connection.", + ) + is_broadcasting: bool = Field( + description="Whether the client connection is broadcasting.", + ) + broadcast_address: Optional[str] = Field( + default=None, + description="URI to the broadcast server.", + ) + broadcast_pid: Optional[int] = Field( + default=None, + description="Process ID of the broadcast server.", + ) + results_file: Optional[str] = Field( + description="Absolute path to the file for continuous writing.", + ) + table_name: Optional[str] = Field( + default=None, + description="Name of the SQL table to write the results to.", + ) + save_results: bool = Field( + description="Whether to save the results after the session ends.", + ) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py index 18bbbb85bc1..29d5c0446e8 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py @@ -25,6 +25,7 @@ from openbb_websockets.helpers import ( connected_clients, get_status, ) +from openbb_websockets.models import WebSocketConnectionStatus router = Router("", description="WebSockets Router") sys.stdout = StdOutSink() @@ -38,7 +39,7 @@ async def create_connection( provider_choices: ProviderChoices, standard_params: StandardParams, extra_params: ExtraParams, -) -> OBBject: +) -> OBBject[WebSocketConnectionStatus]: """Create a new provider websocket connection.""" name = extra_params.name if name in connected_clients: @@ -67,7 +68,7 @@ async def create_connection( connected_clients[client_name] = client results = await get_status(client_name) - obbject.results = results + obbject.results = WebSocketConnectionStatus(**results) return obbject @@ -205,7 +206,9 @@ async def unsubscribe( @router.command( methods=["GET"], ) -async def get_client_status(name: str = "all") -> OBBject[list[dict]]: +async def get_client_status( + name: str = "all", +) -> OBBject[list[WebSocketConnectionStatus]]: """Get the status of a client, or all client connections. Parameters @@ -226,7 +229,7 @@ async def get_client_status(name: str = "all") -> OBBject[list[dict]]: ] else: connections = [await get_status(name)] - return OBBject(results=connections) + return OBBject(results=[WebSocketConnectionStatus(**d) for d in connections]) @router.command( diff --git a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py index 5972f42cea6..4753f560062 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -193,8 +193,6 @@ class PolygonCryptoAggsWebSocketData(WebSocketData): def _validate_model(cls, values): """Validate the model.""" _ = values.pop("s", None) - if values.get("z") and values["z"] == 0 or not values.get("z"): - _ = values.pop("z", None) return values @@ -839,13 +837,15 @@ class PolygonWebSocketFetcher( try: client.connect() - - except Exception as e: # pylint: disable=broad-except + except Exception as e: client.disconnect() raise OpenBBError(e) from e time.sleep(1) + if client._exception: + raise client._exception from client._exception + if client.is_running: return client diff --git a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py index 993b86eb6e6..a8b8bbce56b 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -195,7 +195,7 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim sys.exit(1) except Exception as e: - msg = f"PROVIDER ERROR: Unexpected error -> {e}" + msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e.__str__()}" logger.error(msg) sys.exit(1)