From f7faf07ea8032397d588ded1fbc46e42483bc379 Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Tue, 17 Dec 2024 16:36:08 -0800 Subject: [PATCH] fix some imports and make default message queue 10000 --- .../provider/utils/websockets/database.py | 19 +---- .../utils/websockets/message_queue.py | 4 +- .../websockets/openbb_websockets/broadcast.py | 11 +-- .../websockets/openbb_websockets/client.py | 84 ++++++------------- .../websockets/openbb_websockets/listen.py | 5 +- .../fmp/openbb_fmp/utils/websocket_client.py | 18 ++-- .../openbb_intrinio/utils/websocket_client.py | 23 ++--- .../openbb_polygon/utils/websocket_client.py | 16 ++-- .../openbb_tiingo/utils/websocket_client.py | 16 ++-- 9 files changed, 77 insertions(+), 119 deletions(-) diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py index 2a7feff32e2..cefb2545a1e 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -41,8 +41,6 @@ class Database: A Pydantic model to validate the JSON data. Default is None. limit: Optional[int] The maximum number of records to keep in the database. Default is None. - keep_results: Optional[bool] - Whether to persist the file on disk after the Python session is over. Default is False. logger : Optional[logging.Logger] A custom logger to use. If not provided, a new logger will be created. loop: Optional[asyncio.AbstractEventLoop] @@ -75,8 +73,7 @@ class Database: ): """Initialize the ResultsDB class.""" # pylint: disable=import-outside-toplevel - import atexit # noqa - import tempfile + import tempfile # noqa from pathlib import Path from openbb_core.provider.utils.websockets.helpers import get_logger @@ -94,11 +91,9 @@ class Database: self.results_path = Path(results_file).absolute() self.results_file = results_file - self.keep_results = keep_results self.table_name = table_name if table_name else "records" self.limit = limit self.loop = loop - atexit.register(self._atexit) self.kwargs = kwargs if kwargs else {} run_async(self._setup_database) self.data_model = data_model @@ -136,7 +131,6 @@ class Database: f""" CREATE TABLE IF NOT EXISTS {self.table_name} ( id INTEGER PRIMARY KEY AUTOINCREMENT, - symbol TEXT, message TEXT NOT NULL ) """ @@ -150,17 +144,6 @@ class Database: self._exception = e raise OpenBBError(msg) from e - def _atexit(self) -> None: - """Clean up the WebSocket client processes at exit.""" - # pylint: disable=import-outside-toplevel - import os - - self._exception = None - if self.keep_results: - self.logger.info("Results database saved to, %s\n", str(self.results_file)) - if os.path.exists(self.results_file) and not self.keep_results: # type: ignore - os.remove(self.results_file) # type: ignore - async def _write_to_db(self, message) -> None: """Write the WebSocket message to the SQLite database.""" # pylint: disable=import-outside-toplevel diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py index f626af8b51b..d54d590e278 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/message_queue.py @@ -9,7 +9,7 @@ class MessageQueue: def __init__( self, - max_size: int = 5000, + max_size: int = 10000, max_retries=5, backoff_factor=0.5, logger: Optional[logging.Logger] = None, @@ -52,7 +52,7 @@ class MessageQueue: async def process_queue(self, handler): """Process the message queue.""" - while not self.queue.empty(): + while True: message = await self.queue.get() await self._process_message(message, handler) self.queue.task_done() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py index 5c23de4a18d..a39c75dfe9e 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py @@ -8,10 +8,10 @@ from typing import Optional import uvicorn from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from openbb_core.provider.utils.websockets.database import Database +from openbb_core.provider.utils.websockets.helpers import get_logger, parse_kwargs from starlette.websockets import WebSocketState -from openbb_websockets.helpers import get_logger, parse_kwargs, setup_database - connected_clients: set = set() kwargs = parse_kwargs() @@ -25,6 +25,7 @@ TABLE_NAME = kwargs.pop("table_name", None) or "records" SLEEP_TIME = kwargs.pop("sleep_time", None) or 0.25 AUTH_TOKEN = kwargs.pop("auth_token", None) +DATABASE = Database(results_file=RESULTS_FILE, table_name=TABLE_NAME) app = FastAPI() @@ -151,14 +152,14 @@ class BroadcastServer: # pylint: disable=too-many-instance-attributes def _encrypt_value(self, value: str) -> str: """Encrypt the value for storage.""" # pylint: disable=import-outside-toplevel - from openbb_websockets.helpers import encrypt_value + from openbb_core.provider.utils.websockets.helpers import encrypt_value return encrypt_value(self._key, self._iv, value) def _decrypt_value(self, value: str) -> str: """Decrypt the value for use.""" # pylint: disable=import-outside-toplevel - from openbb_websockets.helpers import decrypt_value + from openbb_core.provider.utils.websockets.helpers import decrypt_value return decrypt_value(self._key, self._iv, value) @@ -294,5 +295,5 @@ if __name__ == "__main__": # pylint: disable=import-outside-toplevel from openbb_core.provider.utils.helpers import run_async - run_async(setup_database, RESULTS_FILE, TABLE_NAME) + run_async(DATABASE._setup_database) main() diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index 2c669035d56..1b14d3dce88 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -99,12 +99,14 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes import os import tempfile import threading - from aiosqlite import DatabaseError from queue import Queue 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 + from openbb_core.provider.utils.websockets.database import Database + from openbb_core.provider.utils.websockets.helpers import ( + encrypt_value, + get_logger, + ) self.name = name self.module = module.replace(".py", "") # type: ignore @@ -156,8 +158,15 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes atexit.register(self._atexit) try: - self._setup_database() - except DatabaseError as e: + self.database = Database( + results_file=self.results_file, + table_name=self.table_name, + limit=self._limit, + keep_results=self.save_results, + logger=self.logger, + data_model=self.data_model, + ) + except Exception as e: # pylint: disable=broad-except msg = f"Unexpected error setting up the SQLite database and table -> {e.__class___.__name__}: {e}" self.logger.error(msg) self._exception = OpenBBError(msg) @@ -178,14 +187,6 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes if os.path.exists(self.results_file) and not self.save_results: # type: ignore os.remove(self.results_file) # type: ignore - def _setup_database(self) -> None: - """Set up the SQLite database and table.""" - # pylint: disable=import-outside-toplevel - from openbb_core.provider.utils.helpers import run_async # noqa - from openbb_websockets.helpers import setup_database - - 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.""" # pylint: disable=import-outside-toplevel @@ -193,7 +194,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes import queue import sys from openbb_core.provider.utils.errors import UnauthorizedError - from openbb_websockets.helpers import clean_message + from openbb_core.provider.utils.websockets.helpers import clean_message from pydantic import ValidationError while not self._stop_log_thread_event.is_set(): @@ -274,7 +275,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes # pylint: disable=import-outside-toplevel import queue # noqa import sys - from openbb_websockets.helpers import clean_message + from openbb_core.provider.utils.websockets.helpers import clean_message while not self._stop_broadcasting_event.is_set(): try: @@ -321,7 +322,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes import threading import time from openbb_core.app.model.abstract.error import OpenBBError - from openbb_websockets.helpers import decrypt_value + from openbb_core.provider.utils.websockets.helpers import decrypt_value if self.is_running: self.logger.info("Provider connection already running.") @@ -495,59 +496,22 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes return False @property - def results(self) -> Union[list[dict], list["Data"], None]: + def results(self) -> list: """Retrieve the deserialized results from the results file.""" # pylint: disable=import-outside-toplevel - import json # noqa - import sqlite3 from openbb_core.app.model.abstract.error import OpenBBError - output: list = [] - file_path = self.results_path - if file_path.exists(): - with sqlite3.connect(file_path) as conn: - try: - cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa - for row in cursor: - _, message = row - if self.data_model: - message = json.loads(message) - if isinstance(message, (str, bytes)): - output.append( - self.data_model.model_validate_json(message) - ) - elif isinstance(message, dict): - output.append(self.data_model(**message)) - else: - output.append(json.loads(message)) - except Exception as e: - raise OpenBBError(f"Error retrieving results: {e}") from e - if output: - return output - - self.logger.info("No results found in %s", self.results_file) - - return None + try: + return self.database.fetch_all() + except Exception as e: # pylint: disable=broad-except + raise OpenBBError(f"Error retrieving results: {e}") from e @results.deleter def results(self): """Clear results stored from the WebSocket stream.""" # pylint: disable=import-outside-toplevel - import sqlite3 # noqa - try: - with sqlite3.connect(self.results_path) as conn: - conn.execute("PRAGMA journal_mode=WAL;") - conn.execute(f"DELETE FROM {self.table_name}") # noqa - conn.commit() - - self._setup_database() - - self.logger.info( - "Results cleared from table %s in %s", - self.table_name, - self.results_file, - ) + self.database.clear_results() except Exception as e: # pylint: disable=broad-except msg = f"Error clearing results: {e.__class__.__name__}: {e}" self.logger.error(msg) @@ -608,7 +572,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes import psutil import queue from openbb_platform_api.utils.api import check_port - from openbb_websockets.helpers import decrypt_value + from openbb_core.provider.utils.websockets.helpers import decrypt_value if ( self._broadcast_process is not None diff --git a/openbb_platform/extensions/websockets/openbb_websockets/listen.py b/openbb_platform/extensions/websockets/openbb_websockets/listen.py index 75e20bed126..329dc745828 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/listen.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/listen.py @@ -25,7 +25,10 @@ class Listener: import websockets from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.utils.errors import UnauthorizedError - from openbb_websockets.helpers import clean_message, get_logger + from openbb_core.provider.utils.websockets.helpers import ( + clean_message, + get_logger, + ) from websockets.exceptions import InvalidStatusCode kwargs = kwargs or {} 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 c28715185a9..67b06be549f 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -2,18 +2,17 @@ import asyncio import json -import os import signal import sys import websockets import websockets.exceptions +from openbb_core.provider.utils.websockets.database import Database from openbb_core.provider.utils.websockets.helpers import ( get_logger, handle_termination_signal, handle_validation_error, parse_kwargs, - write_to_db, ) from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_fmp.models.websocket_connection import FmpWebSocketData @@ -24,7 +23,13 @@ kwargs = parse_kwargs() queue = MessageQueue() command_queue = MessageQueue() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) -kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) + +DATABASE = Database( + results_file=kwargs["results_file"], + table_name=kwargs["table_name"], + limit=kwargs.get("limit"), + logger=logger, +) async def login(websocket): @@ -125,12 +130,7 @@ async def process_message(message): except ValidationError: raise e from e if result: - await write_to_db( - result, - kwargs["results_file"], - kwargs["table_name"], - kwargs.get("limit"), - ) + await DATABASE._write_to_db(result) # pylint: disable=protected-access async def connect_and_stream(): 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 470cc375b4b..803787aaeaa 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -9,12 +9,12 @@ import sys from typing import Any from openbb_core.provider.utils.helpers import run_async +from openbb_core.provider.utils.websockets.database import Database from openbb_core.provider.utils.websockets.helpers import ( get_logger, handle_termination_signal, handle_validation_error, parse_kwargs, - write_to_db, ) from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_intrinio.models.websocket_connection import IntrinioWebSocketData @@ -26,8 +26,15 @@ kwargs = parse_kwargs() command_queue = MessageQueue(logger=logger) CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) +DATABASE = Database( + results_file=kwargs["results_file"], + table_name=kwargs["table_name"], + limit=kwargs.get("limit"), + logger=logger, +) -async def process_message(message): + +def process_message(message): """Process the message and write to the database.""" result: Any = None try: @@ -43,17 +50,12 @@ async def process_message(message): except ValidationError: raise e from e if result: - await write_to_db( - result, - kwargs["results_file"], - kwargs["table_name"], - kwargs["limit"], - ) + DATABASE.write_to_db(result) def on_message(message, backlog): """Process the message and write to the database.""" - run_async(process_message, message) + process_message(message) options = { @@ -73,6 +75,7 @@ client = IntrinioRealtimeClient( async def subscribe(symbol, event): """Subscribe or unsubscribe to a symbol.""" ticker = symbol.split(",") if isinstance(symbol, str) else symbol + ticker = ["lobby"] if "*" in ticker or "LOBBY" in ticker else ticker try: if event == "subscribe": client.join(ticker) @@ -123,7 +126,7 @@ if __name__ == "__main__": try: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - loop.set_exception_handler(lambda loop, context: None) + # loop.set_exception_handler(lambda loop, context: None) for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, handle_termination_signal, logger) 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 a5fe307d682..b6345d32ea8 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -8,12 +8,12 @@ import sys import websockets import websockets.exceptions +from openbb_core.provider.utils.websockets.database import Database from openbb_core.provider.utils.websockets.helpers import ( get_logger, handle_termination_signal, handle_validation_error, parse_kwargs, - write_to_db, ) from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_polygon.models.websocket_connection import ( @@ -32,6 +32,13 @@ FEED = kwargs.pop("feed", None) ASSET_TYPE = kwargs.pop("asset_type", None) kwargs["results_file"] = os.path.abspath(kwargs.get("results_file")) +DATABASE = Database( + results_file=kwargs["results_file"], + table_name=kwargs["table_name"], + limit=kwargs.get("limit"), + logger=logger, +) + async def handle_symbol(symbol): """Handle the symbol and map it to the correct format.""" @@ -181,12 +188,7 @@ async def process_message(message): raise e from e if result: - await write_to_db( - result, - kwargs["results_file"], - kwargs["table_name"], - kwargs.get("limit"), - ) + await DATABASE._write_to_db(result) else: logger.info("PROVIDER INFO: %s", msg) 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 e977d2d3110..4a6a17bf511 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -8,12 +8,12 @@ import sys import websockets from openbb_core.provider.utils.errors import UnauthorizedError +from openbb_core.provider.utils.websockets.database import Database from openbb_core.provider.utils.websockets.helpers import ( get_logger, handle_termination_signal, handle_validation_error, parse_kwargs, - write_to_db, ) from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_tiingo.models.websocket_connection import TiingoWebSocketData @@ -75,6 +75,13 @@ kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) +DATABASE = Database( + results_file=kwargs["results_file"], + table_name=kwargs["table_name"], + limit=kwargs.get("limit"), + logger=logger, +) + # Subscribe and unsubscribe events are handled in a separate connection using the subscription_id set by the login event. async def update_symbols(symbol, event): @@ -189,12 +196,7 @@ async def process_message(message): # pylint: disable=too-many-branches raise e from e if result: - await write_to_db( - result, - kwargs["results_file"], - kwargs["table_name"], - kwargs.get("limit"), - ) + await DATABASE._write_to_db(result) # pylint: disable=protected-access return