diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 77865d3f..dee944c3 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -7,16 +7,165 @@ import ssl import certifi from .models import * -from websockets.client import connect, WebSocketClientProtocol -from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from picows import WSFrame, WSListener, WSMsgType, WSCloseCode, WSTransport, ws_connect from ..logging import get_logger import logging from ..exceptions import AuthError +import traceback env_key = "POLYGON_API_KEY" logger = get_logger("WebSocketClient") +class PolygonWSListener(WSListener): + def __init__(self, client): + self.client = client + self.transport = None + self.processor = None + self.reconnects = 0 + self._full_msg = bytearray() + self._full_msg_type = None + + async def _delay_subscription(self): + # Wait for 3 seconds after authentication before subscribing + await asyncio.sleep(3) + if self.client.schedule_resub: + self.client._handle_subscriptions() + + def on_ws_connected(self, transport: WSTransport): + """Called when WebSocket connection is established""" + self.transport = transport + logger.info("WebSocket connected") + # Send auth message + transport.send(WSMsgType.TEXT, self.client.json.dumps({"action": "auth", "params": self.client.api_key})) + + def on_ws_frame(self, transport: WSTransport, frame: WSFrame): + """Called when a WebSocket frame is received""" + # Handle control frames + if frame.msg_type == WSMsgType.CLOSE: + close_code = frame.get_close_code() + close_message = frame.get_close_message() + logger.info(f"WebSocket connection closed: code={close_code}, reason={close_message}") + transport.send_close(close_code) + return + elif frame.msg_type == WSMsgType.PING: + transport.send(WSMsgType.PONG, frame.get_payload_as_bytes()) + return + elif frame.msg_type == WSMsgType.PONG: + # Just ignore pong frames + return + + # Handle data frames (with fragmentation support) + if frame.fin: + if self._full_msg: + # Last fragment of a fragmented message + if frame.msg_type == 0: # Continuation frame + self._full_msg.extend(frame.get_payload_as_memoryview()) + # Process the message directly without creating a copy + message = self._full_msg + msg_type = self._full_msg_type + # We'll clear the buffer after processing + else: + # This shouldn't happen - fin=True but not a continuation frame + logger.warning(f"Unexpected frame type {frame.msg_type} with fin=True when processing fragmented message") + self._full_msg.clear() + self._full_msg_type = None + return + else: + # Single-frame message + if frame.msg_type == WSMsgType.TEXT: + message = frame.get_payload_as_memoryview() + msg_type = WSMsgType.TEXT + else: + logger.info(f"Ignoring unexpected frame type: {frame.msg_type}") + return + else: + # Fragment (not final) + if not self._full_msg: + # First fragment + self._full_msg_type = frame.msg_type + self._full_msg.extend(frame.get_payload_as_memoryview()) + return # Wait for more fragments + + # Process the complete message + try: + # Handle potential JSON parsing errors more gracefully + try: + msgJson = self.client.json.loads(message) + except json.JSONDecodeError as json_err: + error_pos = json_err.pos + logger.error(f"JSON decode error at position {error_pos}: {json_err}") + + # Try to recover by trimming the message if it appears to be truncated + if "unexpected end of data" in str(json_err): + # Find the last complete JSON object by looking for the last '}]' sequence + last_complete = message.rfind(b'}]' if isinstance(message, (bytes, bytearray, memoryview)) else '}]') + if last_complete > 0: + try: + # Try parsing up to the last complete object + fixed_msg = message[:last_complete+2] + msgJson = self.client.json.loads(fixed_msg) + logger.info(f"Recovered from truncated JSON by trimming to length {len(fixed_msg)}") + except json.JSONDecodeError: + raise json_err + else: + raise json_err + else: + raise json_err + + # Handle auth response + if msgJson and isinstance(msgJson, list) and len(msgJson) > 0 and "status" in msgJson[0]: + if msgJson[0]["status"] == "auth_failed": + logger.error(f"Authentication failed: {msgJson[0]['message']}") + transport.send_close(WSCloseCode.PROTOCOL_ERROR) + return + elif msgJson[0]["status"] == "connected": + logger.info("Authentication successful") + # Add a delay before subscribing to ensure auth is fully processed + asyncio.create_task(self._delay_subscription()) + return + + # Handle regular messages + if not self.client.raw: + has_status_messages = False # Initialize the variable + for m in msgJson: + # Check for subscription status messages + if "ev" in m and m["ev"] == "status": + has_status_messages = True + status_msg = m.get('message', '') + logger.info(f"Status message: {status_msg}") + + # Only parse if we have messages to process + if msgJson and (not has_status_messages or len(msgJson) > 1): + cmsg = parse(msgJson, logger) + else: + cmsg = [] + else: + cmsg = message + + if cmsg and len(cmsg) > 0 and self.client.processor: + # Create a task for async processing + asyncio.create_task(self.client.processor(cmsg)) + + except Exception as e: + # Log the error with minimal message details + logger.error(f"Error processing message: {e}") + traceback.print_exc() + finally: + # Clear the buffer if we were using a fragmented message + if self._full_msg and msg_type == self._full_msg_type: + self._full_msg.clear() + self._full_msg_type = None + + def on_ws_disconnected(self, transport): + """Called when WebSocket connection is closed""" + logger.info("WebSocket connection closed") + self.reconnects += 1 + self.client.scheduled_subs = set(self.client.subs) + self.client.subs = set() + self.client.schedule_resub = True + + class WebSocketClient: def __init__( self, @@ -62,7 +211,9 @@ def __init__( self.subscribed = False self.subs: Set[str] = set() self.max_reconnects = max_reconnects - self.websocket: Optional[WebSocketClientProtocol] = None + self.transport = None + self.listener = PolygonWSListener(self) + self.processor = None if subscriptions is None: subscriptions = [] self.scheduled_subs: Set[str] = set(subscriptions) @@ -72,7 +223,8 @@ def __init__( else: self.json = json - # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection + + async def connect( self, processor: Union[ @@ -89,73 +241,84 @@ async def connect( :param close_timeout: How long to wait for handshake when calling .close. :raises AuthError: If invalid API key is supplied. """ - reconnects = 0 - logger.debug("connect: %s", self.url) - # darwin needs some extra <3 - ssl_context = None - if self.url.startswith("wss://"): - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - ssl_context.load_verify_locations(certifi.where()) + logger.info(f"Connecting to: {self.url}") + self.processor = processor + + # For picows, we don't need to explicitly pass SSL context + # The library handles secure connections based on the URL scheme + - async for s in connect( - self.url, close_timeout=close_timeout, ssl=ssl_context, **kwargs - ): - self.websocket = s + + # Always resubscribe on reconnect + self.schedule_resub = True + + while True: try: - msg = await s.recv() - logger.debug("connected: %s", msg) - logger.debug("authing...") - await s.send( - self.json.dumps({"action": "auth", "params": self.api_key}) + # Connect using picows + _, client = await ws_connect( + lambda: self.listener, + self.url, +# max_frame_size=1024*1024, # from 10K to 1MB + **kwargs ) - auth_msg = await s.recv() - auth_msg_parsed = self.json.loads(auth_msg) - logger.debug("authed: %s", auth_msg) - if auth_msg_parsed[0]["status"] == "auth_failed": - raise AuthError(auth_msg_parsed[0]["message"]) - while True: - if self.schedule_resub: - logger.debug( - "reconciling: %s %s", self.subs, self.scheduled_subs - ) - new_subs = self.scheduled_subs.difference(self.subs) - await self._subscribe(new_subs) - old_subs = self.subs.difference(self.scheduled_subs) - await self._unsubscribe(old_subs) - self.subs = self.scheduled_subs - self.subs = set(self.scheduled_subs) - self.schedule_resub = False + self.transport = client.transport + + # Wait for disconnection + await self.transport.wait_disconnected() + + # Check if we should reconnect + if (self.max_reconnects is not None and + self.listener.reconnects > self.max_reconnects): + logger.info(f"Max reconnects ({self.max_reconnects}) reached") + break + + # Wait before reconnecting + await asyncio.sleep(1) + + except Exception as e: + logger.error(f"Connection error: {e}") + traceback.print_exc() + await asyncio.sleep(1) + + # Check if we should reconnect + self.listener.reconnects += 1 + if (self.max_reconnects is not None and + self.listener.reconnects > self.max_reconnects): + logger.info(f"Max reconnects ({self.max_reconnects}) reached") + break + - try: - cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( - await asyncio.wait_for(s.recv(), timeout=1) - ) - except asyncio.TimeoutError: - continue + + def _handle_subscriptions(self): + """Handle subscription reconciliation""" + if not self.transport: + return + + logger.info("Reconciling subscriptions") + + # Always send all subscriptions to ensure they're properly registered + # This is more reliable than just sending the difference + if self.scheduled_subs: + subs = ",".join(self.scheduled_subs) + logger.info(f"Subscribing to all: {subs}") + self.transport.send(WSMsgType.TEXT, + self.json.dumps({"action": "subscribe", "params": subs}) + ) - if not self.raw: - # we know cmsg is Data - msgJson = self.json.loads(cmsg) # type: ignore - for m in msgJson: - if m["ev"] == "status": - logger.debug("status: %s", m["message"]) - continue - cmsg = parse(msgJson, logger) + + # Handle unsubscriptions + old_subs = self.subs.difference(self.scheduled_subs) + if old_subs: + subs = ",".join(old_subs) + logger.info(f"Unsubscribing from: {subs}") + self.transport.send(WSMsgType.TEXT, + self.json.dumps({"action": "unsubscribe", "params": subs}) + ) + + self.subs = set(self.scheduled_subs) + self.schedule_resub = False + - if len(cmsg) > 0: - await processor(cmsg) # type: ignore - except ConnectionClosedOK as e: - logger.debug("connection closed (OK): %s", e) - return - except ConnectionClosedError as e: - logger.debug("connection closed (ERR): %s", e) - reconnects += 1 - self.scheduled_subs = set(self.subs) - self.subs = set() - self.schedule_resub = True - if self.max_reconnects is not None and reconnects > self.max_reconnects: - return - continue def run( self, @@ -180,20 +343,20 @@ async def handle_msg_wrapper(msgs): asyncio.run(self.connect(handle_msg_wrapper, close_timeout, **kwargs)) async def _subscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug("subbing: %s", subs) - await self.websocket.send( + logger.info(f"Subscribing to: {subs}") + self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) async def _unsubscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug("unsubbing: %s", subs) - await self.websocket.send( + logger.info(f"Unsubscribing from: {subs}") + self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "unsubscribe", "params": subs}) ) @@ -217,7 +380,7 @@ def subscribe(self, *subscriptions: str): topic, sym = self._parse_subscription(s) if topic == None: continue - logger.debug("sub desired: %s", s) + logger.info("Adding subscription: %s", s) self.scheduled_subs.add(s) # If user subs to X.*, remove other X.\w+ if sym == "*": @@ -237,7 +400,7 @@ def unsubscribe(self, *subscriptions: str): topic, sym = self._parse_subscription(s) if topic == None: continue - logger.debug("sub undesired: %s", s) + logger.info("Removing subscription: %s", s) self.scheduled_subs.discard(s) # If user unsubs to X.*, remove other X.\w+ @@ -261,8 +424,8 @@ async def close(self): """ logger.debug("closing") - if self.websocket: - await self.websocket.close() - self.websocket = None + if self.transport: + self.transport.send_close(WSCloseCode.GOING_AWAY) + self.transport = None else: - logger.warning("no websocket open to close") + logger.warning("no websocket connection open to close") diff --git a/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json b/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json deleted file mode 100644 index 0cad5983..00000000 --- a/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "results": [ - { - "ticker": "AAA", - "name": "AAF First Priority CLO Bond ETF", - "market": "stocks", - "locale": "us", - "primary_exchange": "ARCX", - "type": "ETF", - "active": true, - "currency_name": "usd", - "composite_figi": "BBG00X5FSP48", - "share_class_figi": "BBG00X5FSPZ4", - "last_updated_utc": "2022-04-27T00:00:00Z" - }, - { - "ticker": "AAAU", - "name": "Goldman Sachs Physical Gold ETF Shares", - "market": "stocks", - "locale": "us", - "primary_exchange": "BATS", - "type": "ETF", - "active": true, - "currency_name": "usd", - "cik": "0001708646", - "composite_figi": "BBG00LPXX872", - "share_class_figi": "BBG00LPXX8Z1", - "last_updated_utc": "2022-04-27T00:00:00Z" - } - ], - "status": "OK", - "request_id": "40d60d83fa0628503b4d13387b7bde2a", - "count": 2 -} \ No newline at end of file