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 6f169f30ec6..294a02b6879 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -756,7 +756,7 @@ class DatabaseWriter: chunk_size = 20000 minutes = self.export_interval or 5 - + latest_date = None if not self._export_running or not self.export_thread: return try: @@ -768,25 +768,31 @@ class DatabaseWriter: """ # noqa async with self.database.get_connection("read") as conn: - cursor = await conn.execute(latest_query) - latest_date = await cursor.fetchone() - if not latest_date: + try: + async with conn.execute(latest_query) as cursor: + latest_date = await cursor.fetchone() + if not latest_date: + return + + except asyncio.InvalidStateError as e: + self.database.logger.error(f"Database connection state error: {e}") + self._export_running = False return - latest_timestamp = to_datetime(latest_date[0]) + latest_timestamp = to_datetime(latest_date[0]) - # Round down to nearest interval - cutoff_time = latest_timestamp - timedelta( - minutes=latest_timestamp.minute % minutes, - seconds=latest_timestamp.second, - microseconds=latest_timestamp.microsecond, - ) + # Round down to nearest interval + cutoff_time = latest_timestamp - timedelta( + minutes=latest_timestamp.minute % minutes, + seconds=latest_timestamp.second, + microseconds=latest_timestamp.microsecond, + ) - # If we have processed data before, use that as reference - if self._last_processed_timestamp: - start_time = self._last_processed_timestamp - else: - start_time = cutoff_time - timedelta(minutes=minutes) + # If we have processed data before, use that as reference + if self._last_processed_timestamp: + start_time = self._last_processed_timestamp + else: + start_time = cutoff_time - timedelta(minutes=minutes) cutoff_time = self._last_processed_timestamp + timedelta(minutes=minutes) start_time = self._last_processed_timestamp 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 05eac814a9e..8133c769936 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py @@ -189,6 +189,10 @@ class IntrinioWebSocketFetcher( broadcast_host=query.broadcast_host, broadcast_port=query.broadcast_port, auth_token=query.auth_token, + export_directory=query.export_directory, + export_interval=query.export_interval, + prune_interval=query.prune_interval, + verbose=query.verbose, **kwargs, ) 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 4ddfad2fb41..917ab79cdf1 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -143,7 +143,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/models/websocket_connection.py b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py index cb3f078030b..ccfae237c35 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -1225,6 +1225,7 @@ class PolygonWebSocketFetcher( broadcast_host=query.broadcast_host, broadcast_port=query.broadcast_port, auth_token=query.auth_token, + verbose=query.verbose, **kwargs, ) 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 cf37ba5687f..ab6ef1459d1 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -32,14 +32,14 @@ Optional Keyword Arguments import asyncio import json -import orjson import os import signal import sys import time - -import websockets from concurrent.futures import ThreadPoolExecutor + +import orjson +import websockets from openbb_core.provider.utils.websockets.database import Database, DatabaseWriter from openbb_core.provider.utils.websockets.helpers import ( get_logger, @@ -67,10 +67,10 @@ URL_MAP = { } logger = get_logger("openbb.websocket.polygon") -process_queue = MessageQueue(logger=logger, backoff_factor=0) -input_queue = MessageQueue(logger=logger, backoff_factor=0) -command_queue = MessageQueue(logger=logger, backoff_factor=0) -db_queue = MessageQueue(logger=logger, backoff_factor=0, max_size=100000) +process_queue = MessageQueue(logger=logger, backoff_factor=0, max_size=1000000) +input_queue = MessageQueue(logger=logger, backoff_factor=0, max_size=1000000) +command_queue = MessageQueue(logger=logger, backoff_factor=0, max_size=1000000) +db_queue = MessageQueue(logger=logger, backoff_factor=0, max_size=1000000) kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) FEED = kwargs.pop("feed", None) @@ -78,7 +78,7 @@ ASSET_TYPE = kwargs.pop("asset_type", "crypto") kwargs["results_file"] = os.path.abspath(kwargs.get("results_file")) URL = URL_MAP.get(ASSET_TYPE) SUBSCRIBED_SYMBOLS: set = set() -message_thread_pool = ThreadPoolExecutor(max_workers=8) +message_thread_pool = ThreadPoolExecutor(max_workers=16) if not kwargs.get("api_key"): raise ValueError("No API key provided.") @@ -93,7 +93,6 @@ DATABASE = DatabaseWriter( limit=kwargs.get("limit"), logger=logger, ), - batch_size=1, queue=db_queue, ) @@ -281,10 +280,11 @@ async def connect_and_stream(): conn_kwargs.update( { - "ping_interval": 20, - "ping_timeout": 30, + "ping_interval": None, + "ping_timeout": None, "close_timeout": 1, - "max_queue": 1000, + "max_size": 2**63, + "max_queue": None, } ) @@ -293,6 +293,7 @@ async def connect_and_stream(): try: async for message in websocket: input_queue.queue.put_nowait(message) + except Exception as e: raise e from e @@ -337,17 +338,14 @@ async def connect_and_stream(): await DATABASE.start_writer() - for _ in range(30): - processor_task = asyncio.create_task( - input_queue.process_queue( - lambda message: process_input_messages(message) - ) - ) - tasks.add(processor_task) - handler_task = asyncio.create_task( - process_queue.process_queue(lambda message: process_message(message)) - ) - tasks.add(handler_task) + processor_task = asyncio.create_task( + input_queue.process_queue(lambda message: process_input_messages(message)) + ) + tasks.add(processor_task) + handler_task = asyncio.create_task( + process_queue.process_queue(lambda message: process_message(message)) + ) + tasks.add(handler_task) async for websocket in connect(URL, **conn_kwargs): try: @@ -374,6 +372,7 @@ async def connect_and_stream(): receiver_task = asyncio.create_task( message_receiver(websocket), name="receiver" ) + tasks.add(receiver_task) await asyncio.gather(receiver_task) # Attempt to reopen the connection