diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/broadcast.py b/openbb_platform/core/openbb_core/provider/utils/websockets/broadcast.py new file mode 100644 index 00000000000..4f9c5d71796 --- /dev/null +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/broadcast.py @@ -0,0 +1,378 @@ +"""Broadcast server for streaming results to connected clients via WebSocket.""" + +import asyncio +import json +import logging +import os +import sys +from pathlib import Path +from typing import Optional + +import uvicorn +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from openbb_core.provider.utils.websockets.database import CHECK_FOR, Database +from openbb_core.provider.utils.websockets.helpers import parse_kwargs +from starlette.websockets import WebSocketState + +kwargs = parse_kwargs() + +HOST = kwargs.pop("host", None) or "localhost" +PORT = kwargs.pop("port", None) or 6666 +PORT = int(PORT) + +RESULTS_FILE = kwargs.pop("results_file", None) +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) + +SQL = kwargs.pop("sql", None) +SQL_CONNECT_KWARGS = kwargs.pop("sql_connect_kwargs", None) or {} + +app = FastAPI() + + +async def read_stdin(broadcast_server): + """Read from stdin.""" + while True: + line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) + sys.stdin.flush() + sys.stdout.flush() + + if not line: + break + + try: + command = ( + json.loads(line.strip()) + if line.strip().startswith("{") or line.strip().startswith("[") + else line.strip() + ) + await broadcast_server.websocket.send_json(json.dumps(command)) + except json.JSONDecodeError: + broadcast_server.logger.error("Invalid JSON received from stdin") + + +@app.websocket("/") +async def websocket_endpoint( # noqa: PLR0915 + websocket: WebSocket, + auth_token: Optional[str] = None, + replay: bool = False, +): + """Connect to the broadcast server.""" + headers = dict(websocket.headers) + sql = None + + if headers.get("sql"): + sql = headers.pop("sql", None) + + broadcast_server = BroadcastServer( + RESULTS_FILE, + TABLE_NAME, + SLEEP_TIME, + str(AUTH_TOKEN), + sql=sql, + ) + broadcast_server.replay = replay + auth_token = str(auth_token) + + if sql and ( + any(x in sql for x in CHECK_FOR) + or (broadcast_server.table_name not in sql and "message" not in sql) + ): + await websocket.accept() + await websocket.send_text("Connection refused because of invalid SQL.") + broadcast_server.logger.info("Invalid SQL query passed. -> %s", sql) + await websocket.close(code=1008, reason="Invalid parameter values.") + return + + if ( + broadcast_server.auth_token is not None + and auth_token + != broadcast_server._decrypt_value( # pylint: disable=protected-access + broadcast_server.auth_token + ) + ): + await websocket.accept() + await websocket.send_text( + "UnauthorizedError: Invalid authentication token. Could not connect to the broadcast." + ) + broadcast_server.logger.error( + "Invalid authentication token passed by a client connecting." + ) + await websocket.close(code=1008, reason="Invalid authentication token") + return + + await websocket.accept() + + if RESULTS_FILE is None: + raise ValueError("Results file path is required for WebSocket server.") + + broadcast_server.websocket = websocket + + stream_task = asyncio.create_task(broadcast_server.stream_results()) + stdin_task = asyncio.create_task(read_stdin(broadcast_server)) + try: + await websocket.receive_text() + + except WebSocketDisconnect: + pass + except Exception as e: # pylint: disable=broad-except + msg = f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" + broadcast_server.logger.error(msg) + finally: + stream_task.cancel() + stdin_task.cancel() + try: + await stream_task + await stdin_task + except asyncio.CancelledError: + broadcast_server.logger.info("A listener task was cancelled.") + for handler in broadcast_server.logger.handlers: + handler.flush() + except Exception as e: # pylint: disable=broad-except + msg = ( + f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" + ) + broadcast_server.logger.error(msg) + for handler in broadcast_server.logger.handlers: + handler.flush() + if websocket.client_state != WebSocketState.DISCONNECTED: + try: + await websocket.close() + except RuntimeError as e: + msg = f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" + broadcast_server.logger.error(msg) + for handler in broadcast_server.logger.handlers: + handler.flush() + + +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. + It is responsible for reading the results database and sending new messages to the connected client(s). + """ + + def __init__( + self, + results_file, + table_name, + sleep_time: float = 0.25, + auth_token: Optional[str] = None, + sql_connect_kwargs: Optional[dict] = None, + sql: Optional[str] = None, + logger: Optional[logging.Logger] = None, + **kwargs, + ): + """Initialize the BroadcastServer instance.""" + self.results_file = results_file + self.table_name = table_name + self.logger = logger if logger else logging.getLogger("uvicorn.error") + self.sleep_time = sleep_time + self._app = app + self._key = os.urandom(32) + self._iv = os.urandom(16) + self.auth_token = self._encrypt_value(auth_token) if auth_token else None + self.websocket = None + self.kwargs = kwargs + self.sql_connect_kwargs = ( + sql_connect_kwargs if sql_connect_kwargs is not None else {} + ) + self.database = Database( + results_file=results_file, + table_name=table_name, + logger=self.logger, + **sql_connect_kwargs or {}, + ) + self.sql = sql + + def _encrypt_value(self, value: str) -> str: + """Encrypt the value for storage.""" + # pylint: disable=import-outside-toplevel + 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_core.provider.utils.websockets.helpers import decrypt_value + + return decrypt_value(self._key, self._iv, value) + + async def stream_results( # noqa: PLR0915 # pylint: disable=too-many-branches + self, + sql: str = None, + replay: bool = False, + ): + """Continuously read the database and send new messages as JSON via WebSocket.""" + # pylint: disable=import-outside-toplevel + import aiosqlite + from openbb_core.app.model.abstract.error import OpenBBError + + file_path = Path(self.results_file).absolute() + last_id = 0 + + if not file_path.exists(): + self.logger.error("Results file not found: %s", str(file_path)) + return + + query = f"SELECT MAX(id) FROM {self.table_name}" # noqa:S608 + last_id = ( + 0 + if hasattr(self, "replay") and self.replay is True or replay is True + else self.database.query(query)[0] + ) + if sql and self.sql is None: + self.sql = sql + elif self.sql is not None and sql is None: + sql = self.sql + + if sql and sql.lower().startswith("json_extract"): + sql = f"SELECT * FROM {self.table_name} WHERE {sql}" # noqa:S608 + + if sql and ( + any(x.lower() in sql.lower() for x in CHECK_FOR) + or (self.table_name not in sql and "message" not in sql) + ): + await self.websocket.accept() + await self.websocket.send_text("Invalid SQL query passed.") + await self.websocket.close(code=1008, reason="Invalid query") + self.logger.error( + "Invalid query passed to the stream_results method: %s", sql + ) + return + + try: # pylint: disable=too-many-nested-blocks + while True: + try: + if file_path.exists(): + query = ( + sql.replace(";", "") + + f" {'AND' if 'WHERE' in sql else 'WHERE'} id > {last_id}" + if sql is not None + else f"SELECT * FROM {self.table_name} WHERE id > {last_id}" # noqa:S608 + ) + + if ":" not in self.results_file: + results_file = ( + "file:" + + ( + self.results_file + if self.results_file.startswith("/") + else "/" + self.results_file + ) + + "?mode=ro" + ) + else: + results_file = ( + self.results_file + + f"{'&mode=ro' if '?' in self.results_file else '?mode=ro'}" + ) + conn_kwargs = self.sql_connect_kwargs.copy() + conn_kwargs["uri"] = True + conn_kwargs["check_same_thread"] = False + + async with aiosqlite.connect( + results_file, **conn_kwargs + ) as conn, conn.execute(query) as cursor: + async for row in cursor: + last_id = row[0] if row[0] > last_id else last_id + await self.websocket.send_json( + json.dumps(json.loads(row[1])) + ) + if self.replay is True: + await asyncio.sleep(self.sleep_time / 10) + else: + self.logger.error("Results file not found: %s", str(file_path)) + break + + await asyncio.sleep(self.sleep_time) + except KeyboardInterrupt: + self.logger.info("\nResults stream cancelled.") + break + except aiosqlite.OperationalError as e: + if "no such table" in str(e): + self.logger.error( + "Results file was removed by the parent process." + ) + break + raise OpenBBError(e) from e + except asyncio.CancelledError: + break + except WebSocketDisconnect: + pass + except Exception as e: # pylint: disable=broad-except + msg = f"{e.__class__.__name__ if hasattr(e, "__class__") else e} -> {e}" + self.logger.error(msg) + return + + def start_app(self, host: str = "127.0.0.1", port: int = 6666): + """Start the FastAPI app with Uvicorn.""" + uvicorn.run( + self._app, + host=host, + port=port, + **kwargs, + ) + + +def create_broadcast_server( + results_file: str, + table_name: str, + sleep_time: float = 0.25, + auth_token: Optional[str] = None, + sql_connect_kwargs: Optional[dict] = None, + sql: Optional[str] = None, + **kwargs, +): + """Create a new BroadcastServer instance.""" + return BroadcastServer( + results_file, + table_name, + sleep_time, + auth_token, + sql_connect_kwargs, + sql, + **kwargs, + ) + + +def main(): + """Run the main function.""" + broadcast_server = create_broadcast_server( + RESULTS_FILE, + TABLE_NAME, + SLEEP_TIME, + str(AUTH_TOKEN), + SQL_CONNECT_KWARGS, + SQL, + ) + + try: + broadcast_server.start_app( + host=HOST, + port=PORT, + **kwargs, + ) + except TypeError as e: + msg = f"Invalid keyword argument passed to unvicorn. -> {e.args[0]}" + broadcast_server.logger.error(msg) + for handler in broadcast_server.logger.handlers: + handler.flush() + except KeyboardInterrupt: + broadcast_server.logger.info("Broadcast server terminated.") + for handler in broadcast_server.logger.handlers: + handler.flush() + finally: + sys.exit(0) + + +if __name__ == "__main__": + if not RESULTS_FILE: + raise ValueError("Results file path is required for Broadcast server.") + + try: + main() + except KeyboardInterrupt: + sys.exit(0) diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/client.py b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py index 7db087ac055..e08af867fe5 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/client.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py @@ -28,7 +28,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes The symbol(s) requested to subscribe on start. Enter multiple symbols separated by commas, without spaces. Where supported by the provider, * represents all symbols within the feed. limit : Optional[int] - The limit of records to hold in memory. Once the limit is reached, the oldest records are removed. + The limit of records to store. Once the limit is reached, a one-in-one-out policy is used. + A limit of None is the most efficient setting, but requires adequate disk storage to handle high volume. Default is 5000. 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. @@ -40,10 +41,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes table_name : Optional[str] SQL table name to store serialized data messages. By default, 'records'. save_results : bool - Whether to persist the results after the main Python session ends. Default is False. + Whether to persist the results after exiting. Default is False. data_model : Optional[BaseModel] Pydantic data model to validate the results before storing them in the database. - Also used to deserialize the results from the database. auth_token : Optional[str] Used to limit access to the broadcast stream. When provided, listeners should supply as a URL parameter. Example: 'ws://127.0.0.1:6666/?auth_token=SOME_TOKEN'. @@ -53,8 +53,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes A pre-configured logger to use for this instance. 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, - {'api_key': 'MY_API_KEY', 'connect_kwargs': {'key': 'value'}}. + To pass items to 'websocket.connect()', include them in the 'kwargs' dictionary as a nested dictionary, + with key, 'connect_kwargs'. + {'api_key': 'MY_API_KEY', 'connect_kwargs': {'key': 'value'}}. Properties ---------- @@ -68,6 +69,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes Check if the broadcast server process is running. broadcast_address : str URI address for connecting to the broadcast stream. + num_results : int + Number of results stored in the database. results : list[Data] All stored results from the provider connection. Clear the results by deleting the property. e.g., del client.results @@ -88,8 +91,10 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes Stop the broadcast server and disconnect all listening clients. send_message Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets. + get_latest_results + Get the latest results from the database, optionally filter by symbol. query_database - Run a SELECT query to the database. Returns a list of deserialized results, or validated models. + Run a SELECT query to the database. Returns a list of deserialized results. """ def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals @@ -309,9 +314,6 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes output = "INFO: " + f"Stream results from {address}" self._broadcast_address = address - if output and "Started server process" in output: - output = None - if output and "Waiting for application startup." in output: output = None @@ -325,16 +327,16 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes output = output.replace("INFO:", "BROADCAST INFO:") + "\n" output = output[0] if isinstance(output, tuple) else output output = clean_message(output) - if ( - output.startswith("BROADCAST ERROR:") - or "unexpected error" in output.lower() - ): - self._psutil_broadcast_process.kill() - self._broadcast_process.wait() - self._broadcast_thread.join() - sys.stdout.write(output + "\n") - sys.stdout.flush() - break + # if ( + # output.startswith("BROADCAST ERROR:") + # or "unexpected error" in output.lower() + # ): + # self._psutil_broadcast_process.kill() + # self._broadcast_process.wait() + # self._broadcast_thread.join() + # sys.stdout.write(output + "\n") + # sys.stdout.flush() + # continue sys.stdout.write(output + "\n") sys.stdout.flush() except queue.Empty: @@ -544,6 +546,13 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes return self._psutil_broadcast_process.is_running() return False + @property + def num_results(self) -> int: + """Get the number of results stored in the database.""" + return self.query_database(f"SELECT COUNT(*) FROM {self.table_name};")[ # noqa + 0 + ] + @property def results(self) -> list: """ @@ -616,10 +625,16 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes else None ) + def get_latest_results( + self, symbol: Optional[str] = None, limit: Optional[int] = 100 + ) -> list: + """Get the latest results from the database, optionally filter by symbol.""" + return self.database.get_latest_results(symbol=symbol, limit=limit) + def query_database( self, sql: Optional[str] = None, - limit: Optional[int] = 25, + limit: Optional[int] = 100, ) -> list: """ Make a SELECT query to the database for results. @@ -700,7 +715,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes command = [ sys.executable, "-m", - "openbb_websockets.broadcast", + "openbb_core.provider.utils.websockets.broadcast", f"host={host}", f"port={open_port}", f"results_file={self.results_file}", 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 f4df4b1344c..ef58e03f56a 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -1,6 +1,6 @@ """Database module for serialized websockets results.""" -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Iterable, Optional from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.utils.helpers import run_async @@ -11,6 +11,50 @@ if TYPE_CHECKING: from pydantic import BaseModel +CHECK_FOR = ( + "DATABASE", + "TABLE", + "BACKUP", + "DELETE", + "UPDATE", + "INSERT", + "CREATE", + "MODIFY", + "PRAGMA", + "ALTER", + "DROP", + "RENAME", + "REPLACE", + "TRUNCATE", + "VACUUM", + "ATTACH", + "DETACH", + "REINDEX", + "MOVE", + "1=1", + "=''", + '=""', + '"=""', + "= ''", + '= ""', + "or ''", + 'or ""', + "OR ''", + 'OR ""', + "AND ''", + 'AND ""', + "and ''", + 'and ""', + "('')", + '("")', + "('',)", + '("",)', + " ''", + ' ""', + "' '", + '" "', +) + class Database: """ @@ -52,12 +96,25 @@ class Database: ------- write_to_db(message) -> None Write the WebSocket message to the SQLite database. + Synchronous wrapper for _write_to_db. fetch_all(limit: Optional[int] = None) -> list Read the WebSocket message from the SQLite database. - query(sql: str) -> list - Run a SELECT query to the database. + Synchronous wrapper for _fetch_all. + get_latest_results(symbol: Optional[str] = None, limit: Optional[int] = None) -> list + Get the latest records from the database. Optionally filter by symbol. + Synchronous wrapper for _get_latest_results. + query(sql: str, parameters: Optional[Iterable[Any]]) -> list + Run a SELECT query to the database. Table name cannot be anything other than the originally assigned name. + For convenience, the query can start after WHERE, or provide a full query string to extract a specific array. + There are only two columns in the table: "id" - auto-increment index - and "message" - serialized JSON data row. + Synchronous wrapper for _query_db. clear_results() -> None Clear all results from the SQLite database. + + Raises + ------ + OpenBBError + All exceptions are raised as OpenBBError. """ def __init__( # pylint: disable=too-many-positional-arguments @@ -74,11 +131,11 @@ class Database: # pylint: disable=import-outside-toplevel import tempfile # noqa from pathlib import Path + from aiosqlite import ProgrammingError from openbb_core.provider.utils.websockets.helpers import get_logger self.results_file = None self.table_exists = False - self._exception = None self.logger = logger if logger else get_logger("openbb.websocket.database") if not results_file: @@ -86,7 +143,7 @@ class Database: temp_file_path = temp_file.name self.results_path = Path(temp_file_path).absolute() self.results_file = temp_file_path - elif results_file and "://" in results_file: + if ":" in results_file: self.results_file = results_file self.results_path = results_file # type: ignore kwargs["uri"] = True @@ -94,6 +151,13 @@ class Database: self.results_path = Path(results_file).absolute() self.results_file = results_file + if ( + " " in table_name + or table_name.isupper() + or any(x.lower() in table_name.lower() for x in CHECK_FOR) + ): + raise OpenBBError(ProgrammingError(f"Invalid table name, {table_name}.")) + self.table_name = table_name if table_name else "records" self.limit = limit self.loop = loop @@ -124,7 +188,6 @@ class Database: + f" -> {e}" ) self.logger.error(msg) - self._exception = e raise OpenBBError(msg) from e async with aiosqlite.connect( @@ -135,19 +198,19 @@ class Database: CREATE TABLE IF NOT EXISTS {self.table_name} ( id INTEGER PRIMARY KEY AUTOINCREMENT, message TEXT NOT NULL - ) + ); """ ) - await cursor.close() await conn.commit() + await cursor.close() self.table_exists = True + except Exception as e: msg = ( "Unexpected error while creating SQLite database ->" f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" ) self.logger.error(msg) - self._exception = e # type: ignore raise OpenBBError(msg) from e async def _write_to_db(self, message) -> None: @@ -163,23 +226,15 @@ class Database: if not isinstance(message, str): message = json.dumps(message) + if not self.table_exists: + raise aiosqlite.OperationalError( + "Attempt to write to non-existent table." + ) + async with aiosqlite.connect( self.results_file, loop=self.loop, **self.kwargs ) as conn: - cursor = await conn.execute("PRAGMA journal_mode=WAL;") - await cursor.close() - if not self.table_exists: - cursor = await conn.execute( - f""" - CREATE TABLE IF NOT EXISTS {self.table_name} ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - message TEXT NOT NULL - ) - """ - ) - await conn.commit() - await cursor.close() - self.table_exists = True + await conn.execute("PRAGMA journal_mode=WAL;") cursor = await conn.execute( f""" INSERT INTO {self.table_name} (message) @@ -187,35 +242,39 @@ class Database: """, # noqa (message,), ) - await conn.commit() - await cursor.close() + self._at_limit = False - if self.limit is not None: - count = await conn.execute( - f"SELECT COUNT(*) FROM {self.table_name}" # noqa - ) - current_count = int((await count.fetchone())[0]) - await count.close() - limit = 0 if int(self.limit) < 0 else int(self.limit) + if self.limit is not None and not self._at_limit: + limit = max(0, int(self.limit)) - if current_count > limit and limit != 0: - del_cursor = await conn.execute( - f""" - DELETE FROM {self.table_name} - WHERE id IN ( - SELECT id FROM {self.table_name} - ORDER BY id DESC - LIMIT -1 OFFSET ? - ) - """, # noqa: S608 - (limit,), + if limit > 0: + count_cursor = await conn.execute( + f"SELECT COUNT(*) FROM {self.table_name}" # noqa ) - await del_cursor.close() + count = await count_cursor.fetchone() - await conn.commit() + if count[0] > limit: + self._at_limit = True + + await count_cursor.close() + + if self._at_limit: + await conn.execute( + f""" + DELETE FROM {self.table_name} + WHERE id = ( + SELECT id FROM {self.table_name} + ORDER BY id ASC + LIMIT 1 + ) + """, # noqa + ) + + await cursor.close() + await conn.commit() except Exception as e: # pylint: disable=broad-except - raise e from e + raise OpenBBError(e) from e def write_to_db(self, message) -> None: """Write the WebSocket message to the SQLite database.""" @@ -227,7 +286,6 @@ class Database: f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" ) self.logger.error(msg) - self._exception = e # type: ignore raise OpenBBError(msg) from e async def _fetch_all(self, limit: Optional[int] = None) -> list: @@ -238,10 +296,29 @@ class Database: try: rows: list = [] conn_kwargs = self.kwargs.copy() + if not conn_kwargs.get("check_same_thread"): conn_kwargs["check_same_thread"] = False + + if ":" not in self.results_file: + results_file = ( + "file:" + + ( + self.results_file + if self.results_file.startswith("/") + else "/" + self.results_file + ) + + "?mode=ro" + ) + else: + results_file = ( + self.results_file + + f"{'&mode=ro' if '?' in self.results_file else '?mode=ro'}" + ) + conn_kwargs["uri"] = True + async with aiosqlite.connect( - self.results_file, + results_file, loop=self.loop, **conn_kwargs, ) as conn: @@ -249,36 +326,35 @@ class Database: f"SELECT message FROM {self.table_name} ORDER BY id DESC" # noqa ) if limit is not None: - query += f" LIMIT {limit}" - async with conn.execute(query) as cursor: + query += " LIMIT ?" + params = (limit,) + else: + params = () + async with conn.execute(query, params) as cursor: async for row in cursor: - rows.append(self.deserialize_row(row)) + rows.append(await self.deserialize_row(row)) return rows except Exception as e: - raise e from e + raise OpenBBError(e) from e - def deserialize_row(self, row: str) -> Union["BaseModel", Any]: + async def deserialize_row(self, row: str) -> dict: """Deserialize a row from the SQLite database.""" # pylint: disable=import-outside-toplevel import json try: return ( - self.data_model.model_validate_json(row[0]) - if self.data_model is not None - else ( - json.loads(row[0]) - if ( - ( - isinstance(row[0], str) - and (row[0].startswith("{") or row[0].startswith("[")) - ) - or isinstance(row[0], bytes) + json.loads(row[0]) + if ( + ( + isinstance(row[0], str) + and (row[0].startswith("{") or row[0].startswith("[")) ) - else row[0] + or isinstance(row[0], bytes) ) + else row[0] ) except Exception as e: msg = ( @@ -286,8 +362,7 @@ class Database: f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" ) self.logger.error(msg) - self._exception = e # type: ignore - raise OpenBBError(msg) from e + raise OpenBBError(e) from e def fetch_all(self, limit: Optional[int] = None) -> list: """Fetch all the results from the SQLite database.""" @@ -299,15 +374,50 @@ class Database: f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" ) self.logger.error(msg) - self._exception = e # type: ignore - raise OpenBBError(msg) from e + raise OpenBBError(e) from e - async def _query_db(self, sql) -> list: + async def _get_latest_results( + self, symbol: Optional[str] = None, limit: Optional[int] = None + ) -> list: + """Get the latest records from the database. Optionally filter by symbol.""" + if symbol: + symbols = symbol.split(",") + sym_str = "(" + for sym in symbols: + sym_str += f"'{sym.upper()}'" + ("," if sym != symbols[-1] else "") + sym_str += ")" + query = f"json_extract (message, '$.symbol') IN {sym_str}" + else: + query = f"SELECT message FROM {self.table_name}" # noqa + + query += " ORDER BY json_extract (message, '$.date') DESC" + + if limit is not None: + query += f" LIMIT {limit};" + + return await self._query_db(query) + + def get_latest_results( + self, symbol: Optional[str] = None, limit: Optional[int] = None + ) -> list: + """Get the latest records from the database. Optionally filter by symbol.""" + try: + return run_async(self._get_latest_results, symbol, limit) + except Exception as e: + msg = ( + "Unexpected error while getting latest records ->" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e} -> {e.args}" + ) + self.logger.error(msg) + + async def _query_db(self, sql, parameters: Optional[Iterable[Any]] = None) -> list: """Query the SQLite database.""" # pylint: disable=import-outside-toplevel import aiosqlite # noqa import json + if not sql or sql in ("", "''"): + raise OpenBBError("Empty query not allowed.") query = ( sql if sql.startswith("SELECT") @@ -316,46 +426,58 @@ class Database: if not query.endswith(";"): query += ";" - if not query.startswith("SELECT"): - raise OpenBBError( - "Operation not allowed. Only 'SELECT' operations allowed." - ) + if ( + not query.startswith("SELECT") + or any(x.lower() in query.lower() for x in CHECK_FOR) + or (self.table_name not in query and "message" not in query) + ): + raise OpenBBError(f"Invalid operation: {sql}.") + rows: list = [] try: conn_kwargs = self.kwargs.copy() if not conn_kwargs.get("check_same_thread"): conn_kwargs["check_same_thread"] = False + + if ":" not in self.results_file: + results_file = ( + "file:" + + ( + self.results_file + if self.results_file.startswith("/") + else "/" + self.results_file + ) + + "?mode=ro" + ) + else: + results_file = ( + self.results_file + + f"{'&mode=ro' if '?' in self.results_file else '?mode=ro'}" + ) + conn_kwargs["uri"] = True + async with aiosqlite.connect( - self.results_file, + results_file, loop=self.loop, **conn_kwargs, - ) as conn, conn.execute(query) as cursor: + ) as conn, conn.execute(query, parameters) as cursor: async for row in cursor: rows.append( - self.data_model.model_validate_json(row[0]) - if self.data_model is not None - and query.startswith( - f"SELECT message FROM {self.table_name}" # noqa - ) - else ( - json.loads(row[0]) - if ( - ( - isinstance(row[0], str) - and ( - row[0].startswith("{") or row[0].startswith("[") - ) - ) - or isinstance(row[0], bytes) + json.loads(row[0]) + if ( + ( + isinstance(row[0], str) + and (row[0].startswith("{") or row[0].startswith("[")) ) - else row[0] + or isinstance(row[0], bytes) ) + else row[0] ) return rows except Exception as e: # pylint: disable=broad-except - raise e from e + raise OpenBBError(e) from e - def query(self, sql: str) -> list: + def query(self, sql: str, parameters: Optional[Iterable[Any]] = None) -> list: """ Run a SELECT query to the database. @@ -379,11 +501,10 @@ class Database: >>> """ try: - return run_async(self._query_db, sql) + return run_async(self._query_db, sql, parameters) except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error while querying SQLite database -> {e.__class__.__name__}: {e}" + msg = f"{e.__class__.__name__ if hasattr(e, '__class__') else e}: {e.args}" self.logger.error(msg) - raise OpenBBError(e) from e async def _clear_results(self): """Clear the results from the SQLite database.""" @@ -415,7 +536,7 @@ class Database: self.results_file, ) except Exception as e: # pylint: disable=broad-except - raise e from e + raise OpenBBError(e) from e def clear_results(self) -> None: """Clear all results from the SQLite database.""" 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 d54d590e278..4ecddf97963 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 @@ -11,7 +11,7 @@ class MessageQueue: self, max_size: int = 10000, max_retries=5, - backoff_factor=0.5, + backoff_factor=0.75, logger: Optional[logging.Logger] = None, ): """Initialize the MessageQueue.""" diff --git a/openbb_platform/extensions/websockets/openbb_websockets/listen.py b/openbb_platform/extensions/websockets/openbb_websockets/listen.py index 329dc745828..da5f5cf2198 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/listen.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/listen.py @@ -53,11 +53,14 @@ class Listener: async with websockets.connect(url, **kwargs) as websocket: self.websocket = websocket url = clean_message(url) - msg = f"\nListening for messages from {clean_message(url)}" + msg = f"\nConnecting to {clean_message(url)} ..." self.logger.info(msg) for handler in self.logger.handlers: handler.flush() async for message in websocket: + if "invalid SQL" in message: + raise websockets.exceptions.WebSocketException(message) + if ( isinstance(message, str) and "Invalid authentication token" in message @@ -79,6 +82,9 @@ class Listener: msg = f"The process hosting {clean_message(url)} was terminated." self.logger.error(msg) break + except websockets.exceptions.WebSocketException as error: + self.logger.error(error) + break except websockets.exceptions.InvalidURI as error: msg = f"Invalid URI -> {error}" self.logger.error(msg) @@ -97,8 +103,12 @@ class Listener: msg = f"An error occurred while attempting to connect to: {clean_message(url)} -> {error}" self.logger.error(msg) break + except Exception as error: # pylint: disable=broad-except - msg = f"Unexpected error -> {error.__class__.__name__}: {error}" + msg = ( + "Unexpected error -> " + f"{error.__class__.__name__ if hasattr(error, '__class__') else error}: {error.args}" + ) self.logger.error(msg) raise OpenBBError(error) from error finally: 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 ef1664d8e20..139eb41f1a3 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -131,12 +131,15 @@ async def read_stdin_and_update_symbols(): if not line: break - line = json.loads(line.strip()) + if "qsize" in line: + logger.info(f"PROVIDER INFO: Queue size: {queue.queue.qsize()}") + else: + line = json.loads(line.strip()) - if line: - symbol = line.get("symbol") - event = line.get("event") - await update_symbols(symbol, event) + if line: + symbol = line.get("symbol") + event = line.get("event") + await update_symbols(symbol, event) async def process_message(message): # pylint: disable=too-many-branches @@ -211,9 +214,17 @@ async def process_message(message): # pylint: disable=too-many-branches async def connect_and_stream(): """Connect to the WebSocket and stream data to file.""" + tasks: set = set() + handler_task = asyncio.create_task( queue.process_queue(lambda message: process_message(message)) ) + tasks.add(handler_task) + for i in range(0, 15): + new_task = asyncio.create_task( + queue.process_queue(lambda message: process_message(message)) + ) + tasks.add(new_task) stdin_task = asyncio.create_task(read_stdin_and_update_symbols()) ticker: list = [] @@ -265,10 +276,11 @@ async def connect_and_stream(): sys.exit(1) finally: - handler_task.cancel() - await handler_task - stdin_task.cancel() - await stdin_task + tasks.add(stdin_task) + for task in tasks: + task.cancel() + await task + asyncio.gather(*tasks, return_exceptions=True) sys.exit(0) @@ -290,5 +302,5 @@ if __name__ == "__main__": logger.error(ERR) finally: - loop.stop() + loop.close() sys.exit(0)