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 7291d1cd24c..0aef53649ac 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py @@ -14,7 +14,7 @@ from openbb_websockets.models import ( WebSocketData, WebSocketQueryParams, ) -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator URL_MAP = { "stock": "wss://api.tiingo.com/iex", @@ -89,25 +89,22 @@ class TiingoWebSocketQueryParams(WebSocketQueryParams): } symbol: str = Field( - description=QUERY_DESCRIPTIONS.get("symbol", "") + "Use '*' for all symbols.", + description=QUERY_DESCRIPTIONS.get("symbol", "") + " Use '*' for all symbols.", ) asset_type: Literal["stock", "fx", "crypto"] = Field( default="crypto", - description="The asset type for the feed.", + description="The asset type for the feed. Choices are 'stock', 'fx', or 'crypto'.", ) feed: Literal["trade", "trade_and_quote"] = Field( default="trade_and_quote", - description="The type of data feed to subscribe to. FX only supports quote.", + description="The type of data feed to subscribe to. FX only supports quote." + + " Choices are 'trade' or 'trade_and_quote'.", ) class TiingoWebSocketData(WebSocketData): """Tiingo WebSocket data model.""" - timestamp: Optional[int] = Field( - default=None, - description="Nanoseconds since POSIX time UTC.", - ) type: Literal["quote", "trade", "break"] = Field( description="The type of data.", ) @@ -175,25 +172,35 @@ class TiingoWebSocketData(WebSocketData): "quote" if v == "Q" else "trade" if v == "T" else "break" if v == "B" else v ) - @field_validator("date", mode="before", check_fields=False) + @field_validator("date", "timestamp", mode="before", check_fields=False) def _validate_date(cls, v): """Validate the date.""" # pylint: disable=import-outside-toplevel + from pandas import to_datetime from pytz import timezone if isinstance(v, str): - return datetime.fromisoformat(v) - try: - return datetime.fromtimestamp(v / 1000) - except Exception: - if isinstance(v, (int, float)): - # Check if the timestamp is in nanoseconds and convert to seconds - if v > 1e12: - v = v / 1e9 # Convert nanoseconds to seconds - dt = datetime.fromtimestamp(v) - dt = timezone("America/New_York").localize(dt) - return dt - return v + dt = to_datetime(v, utc=True).tz_convert(timezone("America/New_York")) + else: + try: + dt = datetime.fromtimestamp(v / 1000) + except Exception: + if isinstance(v, (int, float)): + # Check if the timestamp is in nanoseconds and convert to seconds + if v > 1e12: + v = v / 1e9 # Convert nanoseconds to seconds + dt = datetime.fromtimestamp(v) + else: + dt = v + + return dt + + @model_validator(mode="before") + @classmethod + def _validate_model(cls, values): + """Validate the model.""" + + return values class TiingoWebSocketConnection(WebSocketConnection): @@ -244,6 +251,7 @@ class TiingoWebSocketFetcher( "url": url, "api_key": api_key, "threshold_level": threshold_level, + "connect_kwargs": query.connect_kwargs, } client = WebSocketClient( @@ -264,13 +272,18 @@ class TiingoWebSocketFetcher( try: client.connect() - # Unhandled exceptions are caught and raised as OpenBBError - except Exception as e: # pylint: disable=broad-except - client.disconnect() - raise OpenBBError(e) from e + except OpenBBError as e: + if client.is_running: + client.disconnect() + raise e from e - # Wait for the connection to be established before returning. - await sleep(2) + await sleep(1) + + if client._exception: + exc = getattr(client, "_exception", None) + client._exception = None + client._atexit() + raise OpenBBError(exc) if client.is_running: return client 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 5c31928890b..cf256c146a0 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -13,9 +13,11 @@ from openbb_websockets.helpers import ( MessageQueue, get_logger, handle_termination_signal, + handle_validation_error, parse_kwargs, write_to_db, ) +from pydantic import ValidationError # These are the data array definitions. IEX_FIELDS = [ @@ -66,10 +68,11 @@ CRYPTO_QUOTE_FIELDS = [ "ask_size", "ask_price", ] -subscription_id = None +SUBSCRIPTION_ID = "" queue = MessageQueue() logger = get_logger("openbb.websocket.tiingo") kwargs = parse_kwargs() +CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) # Subscribe and unsubscribe events are handled in a separate connection using the subscription_id set by the login event. @@ -77,7 +80,7 @@ async def update_symbols(symbol, event): """Update the symbols to subscribe to.""" url = kwargs["url"] - if not subscription_id: + if not SUBSCRIPTION_ID: logger.error( "PROVIDER ERROR: Must be assigned a subscription ID to update symbols. Try logging in." ) @@ -87,7 +90,7 @@ async def update_symbols(symbol, event): "eventName": event, "authorization": kwargs["api_key"], "eventData": { - "subscriptionId": subscription_id, + "subscriptionId": SUBSCRIPTION_ID, "tickers": symbol, }, } @@ -96,13 +99,10 @@ async def update_symbols(symbol, event): await websocket.send(json.dumps(update_event)) response = await websocket.recv() message = json.loads(response) - if message.get("response", {}).get("code") != 200: - logger.error(f"PROVIDER ERROR: {message}") - else: - msg = ( - f"PROVIDER INFO: {message.get('response', {}).get('message')}. " - f"Subscribed to symbols: {message.get('data', {}).get('tickers')}" - ) + if "tickers" in message.get("data", {}): + tickers = message["data"]["tickers"] + threshold_level = message["data"].get("thresholdLevel") + msg = f"PROVIDER INFO: Subscribed to {tickers} with threshold level {threshold_level}" logger.info(msg) @@ -124,10 +124,11 @@ async def read_stdin_and_update_symbols(): async def process_message(message, results_path, table_name, limit): - result = {} - data_message = {} - message = json.loads(message) - msg = "" + """Process the message and write to the database.""" + result: dict = {} + data_message: dict = {} + message = message if isinstance(message, (dict, list)) else json.loads(message) + msg: str = "" if message.get("messageType") == "E": response = message.get("response", {}) msg = f"PROVIDER ERROR: {response.get('code')}: {response.get('message')}" @@ -147,20 +148,21 @@ 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 + global SUBSCRIPTION_ID + SUBSCRIPTION_ID = message["data"]["subscriptionId"] - subscription_id = message["data"]["subscriptionId"] + if "tickers" in response.get("data", {}): + tickers = message["data"]["tickers"] + threshold_level = message["data"].get("thresholdLevel") + msg = f"PROVIDER INFO: Subscribed to {tickers} with threshold level {threshold_level}" + logger.info(msg) - if "tickers" in message.get("data", {}): - tickers = message["data"]["tickers"] - threshold_level = message["data"].get("thresholdLevel") - msg = f"PROVIDER INFO: Subscribed to {tickers} with threshold level {threshold_level}" - logger.info(msg) elif message.get("messageType") == "A": data = message.get("data", []) service = message.get("service") if service == "iex": data_message = {IEX_FIELDS[i]: data[i] for i in range(len(data))} + _ = data_message.pop("timestamp", None) elif service == "fx": data_message = {FX_FIELDS[i]: data[i] for i in range(len(data))} elif service == "crypto_data": @@ -179,10 +181,12 @@ async def process_message(message, results_path, table_name, limit): result = TiingoWebSocketData.model_validate(data_message).model_dump_json( exclude_none=True, exclude_unset=True ) - except Exception as e: - msg = f"PROVIDER ERROR: Error validating data: {e}" - logger.error(msg) - return + except ValidationError as e: + try: + handle_validation_error(logger, e) + except ValidationError: + raise e from e + if result: await write_to_db(result, results_path, table_name, limit) return @@ -207,33 +211,45 @@ async def connect_and_stream( subscribe_event = { "eventName": "subscribe", "authorization": api_key, - "eventData": {"thresholdLevel": threshold_level, "tickers": ticker}, + "eventData": { + "thresholdLevel": threshold_level, + "tickers": ticker, + }, } + connect_kwargs = CONNECT_KWARGS.copy() + if "ping_timeout" not in connect_kwargs: + connect_kwargs["ping_timeout"] = None + if "close_timeout" not in connect_kwargs: + connect_kwargs["close_timeout"] = None + try: - async with websockets.connect( - url, ping_interval=20, ping_timeout=20, max_queue=1000 - ) as websocket: + async with websockets.connect(url, **connect_kwargs) as websocket: logger.info("PROVIDER INFO: WebSocket connection established.") await websocket.send(json.dumps(subscribe_event)) while True: message = await websocket.recv() await queue.enqueue(message) + except UnauthorizedError as e: + logger.error(str(e)) + sys.exit(1) + except websockets.ConnectionClosed as e: msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e.reason}" logger.info(msg) # 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 ) except websockets.WebSocketException as e: - logger.error(e) + logger.error(str(e)) sys.exit(1) except Exception as e: - msg = f"PROVIDER ERROR: Unexpected error -> {e}" + msg = f"Unexpected error -> {e.__class__.__name__}: {e.__str__()}" logger.error(msg) sys.exit(1) @@ -268,10 +284,10 @@ if __name__ == "__main__": loop.run_forever() except (KeyboardInterrupt, websockets.ConnectionClosed): - logger.error("PROVIDER ERROR: WebSocket connection closed") + logger.error("PROVIDER ERROR: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.args[0]}" + msg = f"Unexpected error -> {e.__class__.__name__}: {e.__str__()}" logger.error(msg) finally: