diff --git a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py index dee00bd9f4f..5c23de4a18d 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py @@ -229,7 +229,7 @@ class BroadcastServer: # pylint: disable=too-many-instance-attributes disconnected_clients = set() for client in connected_clients.copy(): try: - await client.websocket.send_text(message) + await client.websocket.send_json(message) except WebSocketDisconnect: disconnected_clients.add(client) except Exception as e: # pylint: disable=broad-except diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index d73a4a32bbd..bfd9c6f0cc1 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -519,7 +519,7 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes elif isinstance(message, dict): output.append(self.data_model(**message)) else: - output.append(json.loads(json.loads(message))) + output.append(json.loads(message)) except Exception as e: raise OpenBBError(f"Error retrieving results: {e}") from e if output: diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py index c2888b7585f..914e8b70cfe 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py @@ -174,15 +174,18 @@ async def setup_database(results_path, table_name): if os.path.exists(results_path): try: await conn.execute("SELECT name FROM sqlite_master WHERE type='table';") - except aiosqlite.DatabaseError: - os.remove(results_path) - + except aiosqlite.DatabaseError as e: + raise OpenBBError( + "Unexpected error caused by an invalid SQLite database file." + "Please check the path, and inspect the file if it exists." + + f" -> {e}" + ) async with aiosqlite.connect(results_path) as conn: await conn.execute( f""" CREATE TABLE IF NOT EXISTS {table_name} ( id INTEGER PRIMARY KEY AUTOINCREMENT, - message TEXT + message JSON ) """ ) @@ -197,44 +200,46 @@ async def write_to_db(message, results_path, table_name, limit): conn = await aiosqlite.connect(results_path) - # Check if the table exists and create it if it doesn't - await conn.execute( - f""" - CREATE TABLE IF NOT EXISTS {table_name} ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - message TEXT - ) - """ - ) - await conn.commit() - - await conn.execute( - f"INSERT INTO {table_name} (message) VALUES (?)", # noqa - (json.dumps(message),), - ) - await conn.commit() - - records = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa - count = (await records.fetchone())[0] - count = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa - current_count = int((await count.fetchone())[0]) - limit = 0 if limit is None else int(limit) - - if current_count > limit and limit != 0: + try: + # Check if the table exists and create it if it doesn't await conn.execute( f""" - DELETE FROM {table_name} - WHERE id IN ( - SELECT id FROM {table_name} - ORDER BY id DESC - LIMIT -1 OFFSET ? + CREATE TABLE IF NOT EXISTS {table_name} ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message JSON ) - """, # noqa: S608 - (limit,), + """ ) + await conn.commit() - await conn.commit() - await conn.close() + await conn.execute( + f"INSERT INTO {table_name} (message) VALUES (?)", # noqa + (message,), + ) + await conn.commit() + + records = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa + count = (await records.fetchone())[0] + count = await conn.execute(f"SELECT COUNT(*) FROM {table_name}") # noqa + current_count = int((await count.fetchone())[0]) + limit = 0 if limit is None else int(limit) + + if current_count > limit and limit != 0: + await conn.execute( + f""" + DELETE FROM {table_name} + WHERE id IN ( + SELECT id FROM {table_name} + ORDER BY id DESC + LIMIT -1 OFFSET ? + ) + """, # noqa: S608 + (limit,), + ) + + await conn.commit() + finally: + await conn.close() class StdOutSink: @@ -271,7 +276,13 @@ class AuthTokenFilter(logging.Formatter): class MessageQueue: """Async message queue for the WebSocket connection.""" - def __init__(self, max_size: int = 1000, max_retries=5, backoff_factor=0.5): + def __init__( + self, + max_size: int = 5000, + max_retries=5, + backoff_factor=0.5, + logger: Optional[logging.Logger] = None, + ): """Initialize the MessageQueue.""" # pylint: disable=import-outside-toplevel from asyncio import Queue @@ -279,7 +290,7 @@ class MessageQueue: self.queue: Queue = Queue(maxsize=max_size) self.max_retries = max_retries self.backoff_factor = backoff_factor - self.logger = get_logger("openbb.websocket.queue") + self.logger = logger if logger else get_logger("openbb.websocket.queue") async def dequeue(self): """Dequeue a message.""" 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 b3b2c853337..f4b8b05f1a5 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py @@ -7,13 +7,14 @@ from typing import Any, Literal, Optional from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.abstract.fetcher import Fetcher +from openbb_intrinio.utils.references import TRADE_CONDITIONS, VENUES from openbb_websockets.client import WebSocketClient from openbb_websockets.models import ( WebSocketConnection, WebSocketData, WebSocketQueryParams, ) -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator class IntrinioWebSocketQueryParams(WebSocketQueryParams): @@ -78,7 +79,7 @@ class IntrinioWebSocketData(WebSocketData): default=None, description="The total volume of the trade or quote.", ) - condition: Optional[str] = Field( + conditions: Optional[str] = Field( default=None, description="The condition attached to the trade or quote.", ) @@ -103,11 +104,35 @@ class IntrinioWebSocketData(WebSocketData): return dt.astimezone(timezone("America/New_York")) - @field_validator("condition", mode="before", check_fields=False) + @field_validator("exchange", mode="before", check_fields=False) @classmethod - def _validate_condition(cls, v): - """Strip the empty spaces from the condition.""" - return str(v).strip().replace(" ", "") if v else None + def _validate_exchange(cls, v): + """Validate the exchange.""" + return VENUES.get(v, v) + + @model_validator(mode="before") + @classmethod + def _validate_conditions(cls, values): + """Validate the exchange.""" + new_values = values.copy() + conditions = new_values.pop("condition", None) + trade_type = new_values.get("type") + if trade_type == "trade": + + if not conditions: + return new_values + + new_conditions = [] + conditions = conditions.replace(" ", "") + for char in range(len(conditions)): + if trade_type == "trade": + new_conditions.append( + TRADE_CONDITIONS.get(conditions[char], conditions[char]) + ) + new_values["conditions"] = "; ".join(new_conditions) + else: + new_values["conditions"] = conditions if conditions else None + return new_values class IntrinioWebSocketConnection(WebSocketConnection): diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/utils/references.py b/openbb_platform/providers/intrinio/openbb_intrinio/utils/references.py index 23c15c14c6c..79131bf8f76 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/references.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/references.py @@ -36,10 +36,46 @@ VENUES = { "X": "NASDAQ OMX PSX, Inc. LLC", "Y": "Bats BYX Exchange, Inc.", "Z": "Bats BZX Exchange, Inc.", - "u": "Other OTC Markets", + "U": "Other OTC Markets", } +TRADE_CONDITIONS = { + "@": "Regular Sale", + "A": "Acquisition", + "B": "Bunched Trade", + "C": "Cash Sale", + "D": "Distribution", + "E": "Placeholder", + "F": "Intermarket Sweep", + "G": "Bunched Sold Trade", + "H": "Priced Variation Trade", + "I": "Odd Lot Trade", + "K": "Rule 155 Trade (AMEX)", + "L": "Sold Last", + "M": "Market Center Official Close", + "N": "Next Day", + "O": "Opening Prints", + "P": "Prior Reference Price", + "Q": "Market Center Official Open", + "R": "Seller", + "S": "Split Trade", + "T": "Form T", + "U": "Extended Trading Hours (Sold Out of Sequence)", + "V": "Contingent Trade", + "W": "Average Price Trade", + "X": "Cross/Periodic Auction Trade", + "Y": "Yellow Flag Regular Trade", + "Z": "Sold (Out of Sequence)", + "1": "Stopped Stock (Regular Trade)", + "4": "Derivatively Priced", + "5": "Re-Opening Prints", + "6": "Closing Prints", + "7": "Qualified Contingent Trade (QCT)", + "8": "Placeholder for 611 Exempt", + "9": "Corrected Consolidated Close", +} + ETF_EXCHANGES = Literal[ "xnas", "arcx", 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 62b94233129..9020afb8f87 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -22,7 +22,7 @@ from pydantic import ValidationError logger = get_logger("openbb.websocket.intrinio") kwargs = parse_kwargs() -command_queue = MessageQueue() +command_queue = MessageQueue(logger=logger) CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) @@ -34,10 +34,15 @@ async def process_message(message): if hasattr(message, "__dict__"): message = message.__dict__ if is_trade: - message["type"] = "trade" if is_trade else "quote" + message["type"] = "trade" try: - result = IntrinioWebSocketData.model_validate(message).model_dump_json() + result = IntrinioWebSocketData.model_validate(message) + result = ( + {} + if result.exchange == "!" or result.price == 0 + else result.model_dump_json() + ) except ValidationError as e: try: handle_validation_error(logger, e) @@ -45,7 +50,10 @@ async def process_message(message): raise e from e if result: await write_to_db( - result, kwargs["results_file"], kwargs["table_name"], kwargs["limit"] + result, + kwargs["results_file"], + kwargs["table_name"], + kwargs["limit"], ) 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 2a18dc2c517..97f6ea998f2 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -23,8 +23,8 @@ from openbb_websockets.helpers import ( from pydantic import ValidationError logger = get_logger("openbb.websocket.polygon") -queue = MessageQueue() -command_queue = MessageQueue() +queue = MessageQueue(logger=logger) +command_queue = MessageQueue(logger=logger) kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})