From 9ab9c6e769ea862fa45f6d05e10eae2db302a16c Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Mon, 9 Dec 2024 11:45:23 -0800 Subject: [PATCH] lots of linting --- .../websockets/openbb_websockets/broadcast.py | 58 ++++--- .../websockets/openbb_websockets/client.py | 161 +++++++----------- .../websockets/openbb_websockets/helpers.py | 3 +- .../websockets/openbb_websockets/listen.py | 40 ++--- .../openbb_websockets/websockets_router.py | 13 +- .../openbb_fmp/models/websocket_connection.py | 5 + .../fmp/openbb_fmp/utils/websocket_client.py | 41 +++-- .../models/websocket_connection.py | 13 +- .../openbb_intrinio/utils/websocket_client.py | 85 ++++----- .../models/websocket_connection.py | 7 +- .../models/websocket_connection.py | 12 +- .../openbb_tiingo/utils/websocket_client.py | 53 +++--- 12 files changed, 240 insertions(+), 251 deletions(-) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py index 576b602a559..dee00bd9f4f 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py @@ -10,7 +10,7 @@ import uvicorn from fastapi import FastAPI, WebSocket, WebSocketDisconnect from starlette.websockets import WebSocketState -from openbb_websockets.helpers import get_logger, parse_kwargs +from openbb_websockets.helpers import get_logger, parse_kwargs, setup_database connected_clients: set = set() @@ -65,7 +65,10 @@ async def websocket_endpoint( # noqa: PLR0915 if ( broadcast_server.auth_token is not None - and auth_token != broadcast_server._decrypt_value(broadcast_server.auth_token) + and auth_token + != broadcast_server._decrypt_value( # pylint: disable=protected-access + broadcast_server.auth_token + ) ): await websocket.accept() await websocket.send_text( @@ -92,10 +95,9 @@ async def websocket_endpoint( # noqa: PLR0915 except WebSocketDisconnect: pass - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"Unexpected error: {e.__class__.__name__} -> {e}" broadcast_server.logger.error(msg) - pass finally: if broadcast_server in connected_clients: connected_clients.remove(broadcast_server) @@ -106,7 +108,7 @@ async def websocket_endpoint( # noqa: PLR0915 await stdin_task except asyncio.CancelledError: broadcast_server.logger.info("INFO: A listener task was cancelled.") - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"Unexpected error while cancelling stream task: {e.__class__.__name__} -> {e}" broadcast_server.logger.error(msg) if websocket.client_state != WebSocketState.DISCONNECTED: @@ -117,7 +119,7 @@ async def websocket_endpoint( # noqa: PLR0915 broadcast_server.logger.error(msg) -class BroadcastServer: +class BroadcastServer: # pylint: disable=too-many-instance-attributes """Stream new results from a continuously written SQLite database. Not intended to be used directly, it is initialized by the server app when it accepts a new connection. @@ -144,6 +146,7 @@ class BroadcastServer: self._iv = os.urandom(16) self.auth_token = self._encrypt_value(auth_token) if auth_token else None self.websocket = None + self.kwargs = kwargs def _encrypt_value(self, value: str) -> str: """Encrypt the value for storage.""" @@ -159,8 +162,11 @@ class BroadcastServer: return decrypt_value(self._key, self._iv, value) - async def stream_results(self): # noqa: PLR0915 + async def stream_results( # noqa: PLR0915 # pylint: disable=too-many-branches + self, + ): """Continuously read the database and send new messages as JSON via WebSocket.""" + # pylint: disable=import-outside-toplevel import sqlite3 # noqa from openbb_core.app.model.abstract.error import OpenBBError @@ -168,16 +174,15 @@ class BroadcastServer: last_id = 0 if not file_path.exists(): - self.logger.error(f"Results file not found: {file_path}") + self.logger.error("Results file not found: %s", str(file_path)) return - else: - conn = sqlite3.connect(self.results_file) - cursor = conn.cursor() - cursor.execute(f"SELECT MAX(id) FROM {self.table_name}") # noqa:S608 - last_id = cursor.fetchone()[0] or 0 - conn.close() + conn = sqlite3.connect(self.results_file) + cursor = conn.cursor() + cursor.execute(f"SELECT MAX(id) FROM {self.table_name}") # noqa:S608 + last_id = cursor.fetchone()[0] or 0 + conn.close() - try: + try: # pylint: disable=too-many-nested-blocks while True: try: if file_path.exists(): @@ -192,7 +197,7 @@ class BroadcastServer: if rows: for row in rows: - index, message = row + _, message = row await self.broadcast(json.dumps(json.loads(message))) last_id = max(row[0] for row in rows) else: @@ -209,17 +214,15 @@ class BroadcastServer: "Results file was removed by the parent process." ) break - else: - raise OpenBBError(e) from e + raise OpenBBError(e) from e except asyncio.CancelledError: break except WebSocketDisconnect: pass - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"Unexpected error: {e.__class__.__name__} -> {e}" self.logger.error(msg) - finally: - return + return async def broadcast(self, message: str): """Broadcast a message to all connected connected clients.""" @@ -229,14 +232,15 @@ class BroadcastServer: await client.websocket.send_text(message) except WebSocketDisconnect: disconnected_clients.add(client) - except Exception as e: - self.logger.error(f"Unexpected error: {e}") + except Exception as e: # pylint: disable=broad-except + msg = f"Unexpected error: {e.__class__.__name__} -> {e}" + self.logger.error(msg) disconnected_clients.add(client) # Remove disconnected connected clients for client in disconnected_clients: connected_clients.remove(client) - def start_app(self, host: str = "127.0.0.1", port: int = 6666, **kwargs): + def start_app(self, host: str = "127.0.0.1", port: int = 6666): """Start the FastAPI app with Uvicorn.""" uvicorn.run( self._app, @@ -251,7 +255,6 @@ def create_broadcast_server( table_name: str, sleep_time: float = 0.25, auth_token: Optional[str] = None, - **kwargs, ): """Create a new BroadcastServer instance.""" return BroadcastServer(results_file, table_name, sleep_time, auth_token) @@ -288,5 +291,8 @@ if __name__ == "__main__": raise ValueError("Results file path is required for Broadcast server.") if not Path(RESULTS_FILE).absolute().exists(): - raise FileNotFoundError(f"Results file not found: {RESULTS_FILE}") + # pylint: disable=import-outside-toplevel + from openbb_core.provider.utils.helpers import run_async + + run_async(setup_database, RESULTS_FILE, TABLE_NAME) main() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index 183b6fc957b..ac9b258e34f 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -1,6 +1,6 @@ """WebSocket Client module for interacting with a provider websocket in a non-blocking pattern.""" -# pylint: disable=too-many-statements +# pylint: disable=too-many-statements,protected-access # flake8: noqa: PLR0915 import logging @@ -10,7 +10,7 @@ if TYPE_CHECKING: from openbb_core.provider.abstract.data import Data -class WebSocketClient: +class WebSocketClient: # pylint: disable=too-many-instance-attributes """Client for interacting with a websocket server in a non-blocking pattern. Parameters @@ -79,7 +79,7 @@ class WebSocketClient: Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets. """ - def __init__( # noqa: PLR0913 + def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals self, name: str, module: str, @@ -95,8 +95,7 @@ class WebSocketClient: ) -> None: """Initialize the WebSocketClient class.""" # pylint: disable=import-outside-toplevel - import asyncio # noqa - import atexit + import atexit # noqa import os import tempfile import threading @@ -105,6 +104,7 @@ class WebSocketClient: from pathlib import Path from openbb_core.app.model.abstract.error import OpenBBError from openbb_websockets.helpers import get_logger + from openbb_websockets.helpers import encrypt_value self.name = name self.module = module.replace(".py", "") # type: ignore @@ -115,13 +115,14 @@ class WebSocketClient: self._symbol = symbol self._key = os.urandom(32) self._iv = os.urandom(16) - self._auth_token = self._encrypt_value(auth_token) if auth_token else None + self._auth_token = ( + encrypt_value(self._key, self._iv, auth_token) if auth_token else None + ) # strings in kwargs are encrypted before storing in the class but unencrypted when passed to the provider module. if kwargs: for k, v in kwargs.items(): if isinstance(v, str): - encrypted_value = self._encrypt_value(v) - kwargs[k] = encrypted_value + kwargs[k] = encrypt_value(self._key, self._iv, v) else: kwargs[k] = v @@ -144,7 +145,6 @@ class WebSocketClient: if not results_file: with tempfile.NamedTemporaryFile(delete=False) as temp_file: - pass temp_file_path = temp_file.name self.results_path = Path(temp_file_path).absolute() self.results_file = temp_file_path @@ -162,20 +162,6 @@ class WebSocketClient: self.logger.error(msg) self._exception = OpenBBError(msg) - def _encrypt_value(self, value): - """Encrypt a value before storing.""" - # pylint: disable=import-outside-toplevel - from openbb_websockets.helpers import encrypt_value - - return encrypt_value(self._key, self._iv, value) - - def _decrypt_value(self, encrypted_value): - """Decrypt the value for use.""" - # pylint: disable=import-outside-toplevel - from openbb_websockets.helpers import decrypt_value - - return decrypt_value(self._key, self._iv, encrypted_value) - def _atexit(self) -> None: """Clean up the WebSocket client processes at exit.""" # pylint: disable=import-outside-toplevel @@ -195,34 +181,10 @@ class WebSocketClient: def _setup_database(self) -> None: """Set up the SQLite database and table.""" # pylint: disable=import-outside-toplevel - import asyncio # noqa - import threading - import time + from openbb_core.provider.utils.helpers import run_async # noqa from openbb_websockets.helpers import setup_database - def run_in_new_loop(): - """Run setup in new event loop.""" - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - loop.run_until_complete( - setup_database(self.results_path, self.table_name) - ) - finally: - loop.close() - - def run_in_thread(): - """Run setup in separate thread.""" - thread = threading.Thread(target=run_in_new_loop) - thread.start() - time.sleep(0.01) - thread.join() - - try: - loop = asyncio.get_running_loop() # noqa - run_in_thread() - except RuntimeError: - run_in_new_loop() + run_async(setup_database, self.results_path, self.table_name) def _log_provider_output(self, output_queue) -> None: """Log output from the provider logger, handling exceptions, errors, and messages that are not data.""" @@ -346,17 +308,17 @@ class WebSocketClient: except queue.Empty: continue - def connect(self) -> None: + def connect(self) -> None: # pylint: disable=too-many-locals """Connect to the provider WebSocket.""" # pylint: disable=import-outside-toplevel - import json # noqa - import os + import os # noqa import psutil import queue import subprocess import threading import time from openbb_core.app.model.abstract.error import OpenBBError + from openbb_websockets.helpers import decrypt_value if self.is_running: self.logger.info("Provider connection already running.") @@ -381,7 +343,9 @@ class WebSocketClient: if kwargs: for k, v in kwargs.items(): if isinstance(v, str): - unencrypted_value = self._decrypt_value(v) + unencrypted_value = decrypt_value( + self._key, self._iv, v # pylint: disable=protected-access + ) kwargs[k] = unencrypted_value else: kwargs[k] = v @@ -396,7 +360,7 @@ class WebSocketClient: if kwarg not in command: command.extend([kwarg]) - self._process = subprocess.Popen( # noqa + self._process = subprocess.Popen( # noqa # pylint: disable=consider-using-with command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, @@ -430,9 +394,9 @@ class WebSocketClient: time.sleep(2) if self._exception is not None: - exc = getattr(self, "_exception", None) - self._exception = None - raise OpenBBError(exc) + with self._exception as exc: + self._exception = None + raise OpenBBError(exc) if not self.is_running: self.logger.error("The provider server failed to start.") @@ -542,7 +506,7 @@ class WebSocketClient: try: cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa for row in cursor: - index, message = row + _, message = row if self.data_model: message = json.loads(message) if isinstance(message, (str, bytes)): @@ -554,7 +518,7 @@ class WebSocketClient: else: output.append(json.loads(json.loads(message))) except Exception as e: - raise OpenBBError(f"Error retrieving results: {e}") + raise OpenBBError(f"Error retrieving results: {e}") from e if output: return output @@ -580,8 +544,9 @@ class WebSocketClient: self.table_name, self.results_file, ) - except Exception as e: - self.logger.error("Error clearing results: %s", e) + except Exception as e: # pylint: disable=broad-except + msg = f"Error clearing results: {e.__class__.__name__}: {e}" + self.logger.error(msg) @property def module(self) -> list: @@ -624,7 +589,7 @@ class WebSocketClient: else None ) - def start_broadcasting( + def start_broadcasting( # pylint: disable=too-many-locals self, host: str = "127.0.0.1", port: int = 6666, @@ -639,14 +604,14 @@ class WebSocketClient: import psutil import queue from openbb_platform_api.utils.api import check_port + from openbb_websockets.helpers import decrypt_value if ( self._broadcast_process is not None and self._broadcast_process.poll() is None ): - self.logger.info( - f"WebSocket broadcast already running on: {self._broadcast_address}" - ) + msg = f"WebSocket broadcast already running on: {self._broadcast_address}" + self.logger.info(msg) return open_port = check_port(host, port) @@ -663,20 +628,22 @@ class WebSocketClient: f"port={open_port}", f"results_file={self.results_file}", f"table_name={self.table_name}", - f"auth_token={self._decrypt_value(self._auth_token) if self._auth_token else None}", + f"auth_token={decrypt_value(self._key, self._iv, self._auth_token) if self._auth_token else None}", ] if kwargs: - for kwarg in kwargs: - command.extend([f"{kwarg}={kwargs[kwarg]}"]) + for k, v in kwargs.items(): + command.extend([f"{k}={v}"]) - self._broadcast_process = subprocess.Popen( # noqa - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - stdin=subprocess.PIPE, - env=os.environ, - text=True, - bufsize=1, + self._broadcast_process = ( + subprocess.Popen( # noqa # pylint: disable=consider-using-with + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + env=os.environ, + text=True, + bufsize=1, + ) ) self._psutil_broadcast_process = psutil.Process(self._broadcast_process.pid) output_queue: queue.Queue = queue.Queue() @@ -786,7 +753,7 @@ def send_message( client._broadcast_process.stdin.flush() else: client.logger.error("Broadcast process is not running.") - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"Error sending message to the {target} process: {e.__class__.__name__} -> {e}" client.logger.error(msg) @@ -796,28 +763,22 @@ def read_message_queue( ): """Read messages from the queue and send them to the WebSocketConnection process.""" while not message_queue.empty(): - try: - message = message_queue.get(timeout=1) - if message: - try: - if ( - target == "provider" - and not client._stop_log_thread_event.is_set() - ): - send_message(client, message, target="provider") - elif ( - target == "broadcast" - and not client._stop_broadcasting_event.is_set() - ): - send_message(client, message, target="broadcast") - except Exception as e: - err = ( - f"Error while attempting to transmit from the outgoing message queue: {e.__class__.__name__} " - f"-> {e} -> {message}" - ) - client.logger.error(err) - finally: - break + message = message_queue.get(timeout=1) + if message: + try: + if target == "provider" and not client._stop_log_thread_event.is_set(): + send_message(client, message, target="provider") + elif ( + target == "broadcast" + and not client._stop_broadcasting_event.is_set() + ): + send_message(client, message, target="broadcast") + except Exception as e: # pylint: disable=broad-except + err = ( + f"Error while attempting to transmit from the outgoing message queue: {e.__class__.__name__} " + f"-> {e} -> {message}" + ) + client.logger.error(err) def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> None: @@ -832,7 +793,7 @@ def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> Non break if output: output_queue.put(output.strip()) - except Exception as e: + except Exception as e: # pylint: disable=broad-except client.logger.error( f"Unexpected error in non_blocking_broadcast: {e.__class__.__name__} -> {e}" ) diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py index 3f8168acba5..c2888b7585f 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py @@ -1,5 +1,7 @@ """WebSockets helpers.""" +# pylint: disable=protected-access + import logging import re from typing import Any, Optional @@ -24,7 +26,6 @@ def clean_message(message: str) -> str: def get_logger(name, level=logging.INFO): """Get a logger instance.""" # pylint: disable=import-outside-toplevel - import logging import uuid logger = logging.getLogger(f"{name}-{uuid.uuid4()}") diff --git a/openbb_platform/extensions/websockets/openbb_websockets/listen.py b/openbb_platform/extensions/websockets/openbb_websockets/listen.py index 61473df05ec..db28a80766b 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/listen.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/listen.py @@ -10,11 +10,14 @@ class Listener: self.loop = None self.websocket = None self.current_task = None + self.logger = None self.kwargs: dict = {} if kwargs: self.kwargs = kwargs - async def listen(self, url, **kwargs): # noqa: PLR0915 + async def listen( # noqa: PLR0915 # pylint: disable=too-many-branches,too-many-statements,too-many-locals + self, url, **kwargs + ): """Listen for WebSocket messages.""" # pylint: disable=import-outside-toplevel import asyncio # noqa @@ -47,9 +50,8 @@ class Listener: async with websockets.connect(url, **kwargs) as websocket: self.websocket = websocket url = clean_message(url) - self.logger.info( - f"\nListening for messages from {clean_message(url)}" - ) + msg = f"\nListening for messages from {clean_message(url)}" + self.logger.info(msg) for handler in self.logger.handlers: handler.flush() async for message in websocket: @@ -71,30 +73,30 @@ class Listener: websockets.ConnectionClosedError, asyncio.IncompleteReadError, ): - self.logger.error( - f"The process hosting {clean_message(url)} was terminated." - ) + msg = f"The process hosting {clean_message(url)} was terminated." + self.logger.error(msg) break except websockets.exceptions.InvalidURI as error: - self.logger.error(f"Invalid URI -> {error}") + msg = f"Invalid URI -> {error}" + self.logger.error(msg) break except InvalidStatusCode as error: - self.logger.error(f"Invalid status code -> {error}") + msg = f"Invalid status code -> {error}" + self.logger.error(msg) break except OSError as error: if "Multiple exceptions" in str(error): err = str(error).split("Multiple exceptions:")[1].strip() err = err.split("[")[-1].strip().replace("]", ":") - self.logger.error( - f"An error occurred while attempting to connect to: {clean_message(url)} -> {err}" - ) + msg = f"An error occurred while attempting to connect to: {clean_message(url)} -> {err}" + self.logger.error(msg) else: - self.logger.error( - f"An error occurred while attempting to connect to: {clean_message(url)} -> {error}" - ) + msg = f"An error occurred while attempting to connect to: {clean_message(url)} -> {error}" + self.logger.error(msg) break - except Exception as error: - self.logger.error(f"An unexpected error occurred: {error}") + except Exception as error: # pylint: disable=broad-except + msg = f"Unexpected error -> {error.__class__.__name__}: {error}" + self.logger.error(msg) raise OpenBBError(error) from error finally: if self.websocket: @@ -154,7 +156,5 @@ def listen(url, **kwargs): try: listener = Listener(**kwargs) listener.run(url, **kwargs) - except Exception as e: + except Exception as e: # pylint: disable=broad-except raise OpenBBError(e) from e - finally: - return diff --git a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py index f75af9b37f3..720061a702f 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py @@ -1,6 +1,6 @@ """Websockets Router.""" -# pylint: disable=unused-argument +# pylint: disable=unused-argument,protected-access,unused-import import asyncio import sys @@ -34,6 +34,17 @@ sys.stdout = StdOutSink() @router.command( model="WebSocketConnection", + examples=[ + APIEx( + parameters={ + "name": "client1", + "provider": "fmp", + "asset_type": "crypto", + "symbol": "btcusd,ethusd,solusd", + "start_broadcast": True, + } + ) + ], ) async def create_connection( cc: CommandContext, diff --git a/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py b/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py index c9fd7a5d765..b9779e5f0d2 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py +++ b/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py @@ -1,5 +1,7 @@ """FMP WebSocket model.""" +# pylint: disable=unused-argument,protected-access + from datetime import datetime from typing import Any, Literal, Optional @@ -92,11 +94,13 @@ class FmpWebSocketData(WebSocketData): ) @field_validator("symbol", mode="before") + @classmethod def _validate_symbol(cls, v): """Validate the symbol.""" return v.upper() @field_validator("type", mode="before", check_fields=False) + @classmethod def _valiidate_data_type(cls, v): """Validate the data type.""" return ( @@ -104,6 +108,7 @@ class FmpWebSocketData(WebSocketData): ) @field_validator("date", mode="before", check_fields=False) + @classmethod def _validate_date(cls, v): """Validate the date.""" # pylint: disable=import-outside-toplevel diff --git a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py index 352c7b3f5e4..5090088e0df 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -1,4 +1,4 @@ -"""FMP WebSocket server.""" +"""FMP WebSocket client.""" import asyncio import json @@ -24,14 +24,15 @@ kwargs = parse_kwargs() queue = MessageQueue() command_queue = MessageQueue() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) +kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) -async def login(websocket, api_key): +async def login(websocket): """Login to the WebSocket.""" login_event = { "event": "login", "data": { - "apiKey": api_key, + "apiKey": kwargs["api_key"], }, } try: @@ -97,7 +98,7 @@ async def process_stdin_queue(websocket): await subscribe(websocket, symbol, event) -async def process_message(message, results_path, table_name, limit): +async def process_message(message): """Process the message and write to the database.""" result: dict = {} message = json.loads(message) if isinstance(message, str) else message @@ -124,24 +125,27 @@ async def process_message(message, results_path, table_name, limit): except ValidationError: raise e from e if result: - await write_to_db(result, results_path, table_name, limit) + await write_to_db( + result, + kwargs["results_file"], + kwargs["table_name"], + kwargs.get("limit"), + ) -async def connect_and_stream(url, symbol, api_key, results_path, table_name, limit): +async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" handler_task = asyncio.create_task( - queue.process_queue( - lambda message: process_message(message, results_path, table_name, limit) - ) + queue.process_queue(lambda message: process_message(message)) ) stdin_task = asyncio.create_task(read_stdin_and_queue_commands()) try: - websocket = await websockets.connect(url, **CONNECT_KWARGS) - await login(websocket, api_key) - await subscribe(websocket, symbol, "subscribe") + websocket = await websockets.connect(kwargs["url"], **CONNECT_KWARGS) + await login(websocket) + await subscribe(websocket, kwargs["symbol"], "subscribe") while True: ws_task = asyncio.create_task(websocket.recv()) @@ -166,13 +170,13 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds.") await asyncio.sleep(5) - await connect_and_stream(url, symbol, api_key, results_path, table_name, limit) + await connect_and_stream() except websockets.WebSocketException as e: logger.error(e) sys.exit(1) - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -194,14 +198,7 @@ if __name__ == "__main__": loop.add_signal_handler(sig, handle_termination_signal, logger) asyncio.run_coroutine_threadsafe( - connect_and_stream( - kwargs["url"], - kwargs["symbol"], - kwargs["api_key"], - os.path.abspath(kwargs["results_file"]), - kwargs["table_name"], - kwargs.get("limit", None), - ), + connect_and_stream(), loop, ) loop.run_forever() diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py index cc926dd0b6e..4f5f89b98e1 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py @@ -1,5 +1,7 @@ """Intrinio WebSocket model.""" +# pylint: disable=unused-argument + from datetime import datetime from typing import Any, Literal, Optional @@ -82,6 +84,7 @@ class IntrinioWebSocketData(WebSocketData): ) @field_validator("date", mode="before", check_fields=False) + @classmethod def _validate_date(cls, v): """Validate the date.""" # pylint: disable=import-outside-toplevel @@ -100,6 +103,12 @@ class IntrinioWebSocketData(WebSocketData): return dt.astimezone(timezone("America/New_York")) + @field_validator("condition", mode="before", check_fields=False) + @classmethod + def _validate_condition(cls, v): + """Strip the empty spaces from the condition.""" + return v.strip().replace(" ", "") if v and isinstance(v, str) else None + class IntrinioWebSocketConnection(WebSocketConnection): """Intrinio WebSocket connection model.""" @@ -152,8 +161,8 @@ class IntrinioWebSocketFetcher( try: client.connect() await asyncio.sleep(2) - if client._exception: - raise client._exception + if client._exception: # pylint: disable=protected-access + raise client._exception # pylint: disable=protected-access except OpenBBError as e: if client.is_running: client.disconnect() diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py index 48575f9d6e1..49764c6dcd5 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -1,10 +1,13 @@ """Intrinio WebSocket server.""" +# pylint: disable=unused-argument + import asyncio import json import signal import sys +from openbb_core.provider.utils.helpers import run_async from openbb_intrinio.models.websocket_connection import IntrinioWebSocketData from openbb_intrinio.utils.stocks_client import IntrinioRealtimeClient, Quote, Trade from openbb_websockets.helpers import ( @@ -23,45 +26,6 @@ command_queue = MessageQueue() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) -async def subscribe(client, symbol, event): - """Subscribe or unsubscribe to a symbol.""" - ticker = symbol.split(",") if isinstance(symbol, str) else symbol - try: - if event == "subscribe": - client.join(ticker) - elif event == "unsubscribe": - client.leave(ticker) - except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" - logger.error(msg) - - -async def read_stdin_and_queue_commands(): - """Read from stdin and queue commands.""" - while True: - line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) - sys.stdin.flush() - - if not line: - break - - try: - command = json.loads(line.strip()) - await command_queue.enqueue(command) - except json.JSONDecodeError: - logger.error("Invalid JSON received from stdin -> %s", line.strip()) - - -async def process_stdin_queue(client): - """Process the command queue.""" - while True: - command = await command_queue.dequeue() - symbol = ["lobby" if d == "*" else d.upper() for d in command.get("symbol", [])] - event = command.get("event") - if symbol and event: - await subscribe(client, symbol, event) - - async def process_message(message): """Process the message and write to the database.""" result: dict = {} @@ -91,7 +55,7 @@ async def process_message(message): def on_message(message, backlog): """Process the message and write to the database.""" - asyncio.run(process_message(message)) + run_async(process_message, message) options = { @@ -108,6 +72,45 @@ client = IntrinioRealtimeClient( ) +async def subscribe(symbol, event): + """Subscribe or unsubscribe to a symbol.""" + ticker = symbol.split(",") if isinstance(symbol, str) else symbol + try: + if event == "subscribe": + client.join(ticker) + elif event == "unsubscribe": + client.leave(ticker) + except Exception as e: # pylint: disable=broad-except + msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" + logger.error(msg) + + +async def read_stdin_and_queue_commands(): + """Read from stdin and queue commands.""" + while True: + line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) + sys.stdin.flush() + + if not line: + break + + try: + command = json.loads(line.strip()) + await command_queue.enqueue(command) + except json.JSONDecodeError: + logger.error("Invalid JSON received from stdin -> %s", line.strip()) + + +async def process_stdin_queue(): + """Process the command queue.""" + while True: + command = await command_queue.dequeue() + symbol = ["lobby" if d == "*" else d.upper() for d in command.get("symbol", [])] + event = command.get("event") + if symbol and event: + await subscribe(symbol, event) + + async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" symbol = kwargs.pop("symbol", "lobby") @@ -115,7 +118,7 @@ async def connect_and_stream(): asyncio.create_task(read_stdin_and_queue_commands()) client.connect() client.join(symbol) - asyncio.create_task(process_stdin_queue(client)) + asyncio.create_task(process_stdin_queue()) if __name__ == "__main__": 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 250f8a21331..99db223ac12 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -371,6 +371,7 @@ class PolygonCryptoTradeWebSocketData(WebSocketData): return CRYPTO_EXCHANGE_MAP.get(v, str(v)) @model_validator(mode="before") + @classmethod def _validate_model(cls, values): """Validate the model.""" _ = values.pop("i", None) @@ -1184,7 +1185,7 @@ class PolygonWebSocketData(Data): options_symbol = data.get("sym", "").startswith("O:") or data.get( "symbol", "" ).startswith("O:") - + model = "" if options_symbol: model = OPTIONS_MODEL_MAP.get(data.get("ev", "")) or OPTIONS_MODEL_MAP.get( data.get("type", "") @@ -1259,8 +1260,8 @@ class PolygonWebSocketFetcher( client.disconnect() raise e from e - if client._exception: - raise client._exception from client._exception + if client._exception: # pylint: disable=protected-access + raise client._exception # pylint: disable=protected-access if client.is_running: return {"client": client} diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py index ffeb3a3781d..7afe4319918 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py @@ -1,5 +1,7 @@ """Tiingo WebSocket model.""" +# pylint: disable=unused-argument + from datetime import datetime from typing import Any, Literal, Optional @@ -166,11 +168,13 @@ class TiingoWebSocketData(WebSocketData): ) @field_validator("symbol", mode="before", check_fields=False) + @classmethod def _validate_symbol(cls, v): """Validate the symbol.""" return v.upper() @field_validator("type", mode="before", check_fields=False) + @classmethod def _valiidate_data_type(cls, v): """Validate the data type.""" return ( @@ -178,6 +182,7 @@ class TiingoWebSocketData(WebSocketData): ) @field_validator("date", "timestamp", mode="before", check_fields=False) + @classmethod def _validate_date(cls, v): """Validate the date.""" # pylint: disable=import-outside-toplevel @@ -284,10 +289,9 @@ class TiingoWebSocketFetcher( await sleep(1) - if client._exception: - exc = getattr(client, "_exception", None) - client._exception = None - client._atexit() + if getattr(client, "_exception", None): + exc = client._exception # pylint: disable=protected-access + client._exception = None # pylint: disable=protected-access raise OpenBBError(exc) if client.is_running: diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py index 25243c9dc55..a384962dfe9 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -73,6 +73,7 @@ queue = MessageQueue() logger = get_logger("openbb.websocket.tiingo") kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) +kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) # Subscribe and unsubscribe events are handled in a separate connection using the subscription_id set by the login event. @@ -123,7 +124,7 @@ async def read_stdin_and_update_symbols(): await update_symbols(symbol, event) -async def process_message(message, results_path, table_name, limit): +async def process_message(message): # pylint: disable=too-many-branches """Process the message and write to the database.""" result: dict = {} data_message: dict = {} @@ -148,7 +149,7 @@ async def process_message(message, results_path, table_name, limit): msg = f"PROVIDER INFO: Authorization: {response.get('message')}" logger.info(msg) if message.get("data", {}).get("subscriptionId"): - global SUBSCRIPTION_ID # noqa: PLW0603 + global SUBSCRIPTION_ID # noqa: PLW0603 # pylint: disable=global-statement SUBSCRIPTION_ID = message["data"]["subscriptionId"] if "tickers" in response.get("data", {}): @@ -188,31 +189,32 @@ async def process_message(message, results_path, table_name, limit): raise e from e if result: - await write_to_db(result, results_path, table_name, limit) + await write_to_db( + result, + kwargs["results_file"], + kwargs["table_name"], + kwargs.get("limit"), + ) return -async def connect_and_stream( - url, symbol, threshold_level, api_key, results_path, table_name, limit -): +async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" handler_task = asyncio.create_task( - queue.process_queue( - lambda message: process_message(message, results_path, table_name, limit) - ) + queue.process_queue(lambda message: process_message(message)) ) - stdin_task = asyncio.create_task(read_stdin_and_update_symbols()) + ticker: list = [] - if isinstance(symbol, str): - ticker = symbol.lower().split(",") + if isinstance(kwargs["symbol"], str): + ticker = kwargs["symbol"].lower().split(",") subscribe_event = { "eventName": "subscribe", - "authorization": api_key, + "authorization": kwargs["api_key"], "eventData": { - "thresholdLevel": threshold_level, + "thresholdLevel": kwargs["threshold_level"], "tickers": ticker, }, } @@ -224,7 +226,7 @@ async def connect_and_stream( try: try: - async with websockets.connect(url, **connect_kwargs) as websocket: + async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket: logger.info("PROVIDER INFO: WebSocket connection established.") await websocket.send(json.dumps(subscribe_event)) while True: @@ -241,15 +243,13 @@ async def connect_and_stream( # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds...") await asyncio.sleep(5) - await connect_and_stream( - url, symbol, threshold_level, api_key, results_path, table_name, limit - ) + await connect_and_stream() except websockets.WebSocketException as e: logger.info(str(e)) sys.exit(0) - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -270,24 +270,15 @@ if __name__ == "__main__": loop.add_signal_handler(signal.SIGTERM, handle_termination_signal, logger) asyncio.run_coroutine_threadsafe( - connect_and_stream( - kwargs["url"], - kwargs["symbol"], - kwargs["threshold_level"], - kwargs["api_key"], - os.path.abspath(kwargs["results_file"]), - kwargs["table_name"], - kwargs.get("limit", None), - ), + connect_and_stream(), loop, ) loop.run_forever() except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error -> {e.__class__.__name__}: {e}" - logger.error(msg) + ERR = f"Unexpected error -> {e.__class__.__name__}: {e}" + logger.error(ERR) finally: loop.stop() - loop.close sys.exit(0)