diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index ac9b258e34f..d73a4a32bbd 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -237,6 +237,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes if "SymbolError" in output: err = ValueError(output) self._exception = err + sys.stdout.write(output + "\n") + sys.stdout.flush() continue # Other errors are logged to stdout and the process is killed. # If the exception is raised by the parent thread, it will be treated as an unexpected error. @@ -262,7 +264,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes elif output.startswith("INFO:"): output = output.replace("INFO:", "PROVIDER INFO:") - self.logger.info(output) + sys.stdout.write(output + "\n") + sys.stdout.flush() except queue.Empty: continue @@ -394,9 +397,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes time.sleep(2) if self._exception is not None: - with self._exception as exc: - self._exception = None - raise OpenBBError(exc) + exc = getattr(self, "_exception", None) + self._exception = None + raise OpenBBError(exc) if not self.is_running: self.logger.error("The provider server failed to start.") @@ -722,9 +725,10 @@ def non_blocking_websocket(client, output_queue, provider_message_queue) -> None output_queue.put(output.strip()) except Exception as e: - client.logger.error( + msg = ( f"Unexpected error in non_blocking_websocket: {e.__class__.__name__} -> {e}" ) + client.logger.error(msg) raise e from e finally: client._process.stdout.close() @@ -794,9 +798,10 @@ def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> Non if output: output_queue.put(output.strip()) except Exception as e: # pylint: disable=broad-except - client.logger.error( + err = ( f"Unexpected error in non_blocking_broadcast: {e.__class__.__name__} -> {e}" ) + client.logger.error(err) finally: client._broadcast_process.stdout.close() client._broadcast_process.wait() 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 f5a76aa73ea..59e726eb0a3 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -122,8 +122,8 @@ def validate_date(cls, v): v = v / 1e9 # Convert nanoseconds to seconds dt = datetime.fromtimestamp(v, tz=timezone("UTC")) dt = dt.astimezone(timezone("America/New_York")) - - return dt + return dt + return v class PolygonWebSocketQueryParams(WebSocketQueryParams): @@ -353,9 +353,9 @@ class PolygonCryptoTradeWebSocketData(WebSocketData): """Validate the conditions.""" if v is None or isinstance(v, list) and v[0] == 0: return None - if isinstance(v, list) and v[0] == 1: + elif isinstance(v, list) and v[0] == 1: return "sellside" - if isinstance(v, list) and v[0] == 2: + elif isinstance(v, list) and v[0] == 2: return "buyside" return str(v) @@ -1200,7 +1200,7 @@ class PolygonWebSocketData(Data): or MODEL_MAP.get(data.get("type", "")) ) if not model: - return super().__new__(cls) # pylint: disable=E1120 + return super().__new__(cls) return model.model_validate(data) # type: ignore 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 c1749b9bb8a..2a18dc2c517 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -1,4 +1,4 @@ -"""Polygon WebSocket server.""" +"""Polygon WebSocket client.""" import asyncio import json @@ -30,7 +30,7 @@ kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) FEED = kwargs.pop("feed", None) ASSET_TYPE = kwargs.pop("asset_type", None) -kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) +kwargs["results_file"] = os.path.abspath(kwargs.get("results_file")) async def handle_symbol(symbol): @@ -85,9 +85,8 @@ async def handle_symbol(symbol): return ",".join(new_symbols) -async def login(websocket, api_key): - """Login to the WebSocket.""" - login_event = f'{{"action":"auth","params":"{api_key}"}}' +async def login(websocket): + login_event = f'{{"action":"auth","params":"{kwargs["api_key"]}"}}' try: await websocket.send(login_event) res = await websocket.recv() @@ -108,7 +107,7 @@ async def login(websocket, api_key): logger.error(err) sys.exit(1) logger.info("PROVIDER INFO: %s", msg.get("message")) - except Exception as e: # pylint: disable=broad-except + except Exception as e: logger.error("PROVIDER ERROR: %s -> %s", e.__class__.__name__, e.args[0]) sys.exit(1) @@ -123,7 +122,7 @@ async def subscribe(websocket, symbol, event): subscribe_event = f'{{"action":"{event}","params":"{ticker}"}}' try: await websocket.send(subscribe_event) - except Exception as e: # pylint: disable=broad-except + except Exception as e: msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" logger.error(msg) @@ -183,7 +182,7 @@ async def process_message(message): if result: await write_to_db( result, - kwargs["results_path"], + kwargs["results_file"], kwargs["table_name"], kwargs.get("limit"), ) @@ -191,7 +190,7 @@ async def process_message(message): logger.info("PROVIDER INFO: %s", msg) -async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-statements +async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" handler_task = asyncio.create_task( @@ -207,7 +206,8 @@ async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-s try: async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket: - await login(websocket, kwargs["api_key"]) + + await login(websocket) response = await websocket.recv() messages = json.loads(response) await process_message(messages) @@ -262,7 +262,7 @@ async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-s logger.error(msg) sys.exit(1) - except Exception as e: # pylint: disable=broad-except + except Exception as e: msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -293,8 +293,8 @@ if __name__ == "__main__": logger.error("PROVIDER ERROR: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - ERR = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" - logger.error(ERR) + msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" + logger.error(msg) finally: sys.exit(0)