and some more..

This commit is contained in:
Danglewood
2024-12-09 19:31:18 -08:00
parent 02d92cddf5
commit 7935d6802b
3 changed files with 29 additions and 24 deletions

View File

@@ -237,6 +237,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
if "SymbolError" in output:
err = ValueError(output)
self._exception = err
sys.stdout.write(output + "\n")
sys.stdout.flush()
continue
# Other errors are logged to stdout and the process is killed.
# If the exception is raised by the parent thread, it will be treated as an unexpected error.
@@ -262,7 +264,8 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
elif output.startswith("INFO:"):
output = output.replace("INFO:", "PROVIDER INFO:")
self.logger.info(output)
sys.stdout.write(output + "\n")
sys.stdout.flush()
except queue.Empty:
continue
@@ -394,9 +397,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes
time.sleep(2)
if self._exception is not None:
with self._exception as exc:
self._exception = None
raise OpenBBError(exc)
exc = getattr(self, "_exception", None)
self._exception = None
raise OpenBBError(exc)
if not self.is_running:
self.logger.error("The provider server failed to start.")
@@ -722,9 +725,10 @@ def non_blocking_websocket(client, output_queue, provider_message_queue) -> None
output_queue.put(output.strip())
except Exception as e:
client.logger.error(
msg = (
f"Unexpected error in non_blocking_websocket: {e.__class__.__name__} -> {e}"
)
client.logger.error(msg)
raise e from e
finally:
client._process.stdout.close()
@@ -794,9 +798,10 @@ def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> Non
if output:
output_queue.put(output.strip())
except Exception as e: # pylint: disable=broad-except
client.logger.error(
err = (
f"Unexpected error in non_blocking_broadcast: {e.__class__.__name__} -> {e}"
)
client.logger.error(err)
finally:
client._broadcast_process.stdout.close()
client._broadcast_process.wait()

View File

@@ -122,8 +122,8 @@ def validate_date(cls, v):
v = v / 1e9 # Convert nanoseconds to seconds
dt = datetime.fromtimestamp(v, tz=timezone("UTC"))
dt = dt.astimezone(timezone("America/New_York"))
return dt
return dt
return v
class PolygonWebSocketQueryParams(WebSocketQueryParams):
@@ -353,9 +353,9 @@ class PolygonCryptoTradeWebSocketData(WebSocketData):
"""Validate the conditions."""
if v is None or isinstance(v, list) and v[0] == 0:
return None
if isinstance(v, list) and v[0] == 1:
elif isinstance(v, list) and v[0] == 1:
return "sellside"
if isinstance(v, list) and v[0] == 2:
elif isinstance(v, list) and v[0] == 2:
return "buyside"
return str(v)
@@ -1200,7 +1200,7 @@ class PolygonWebSocketData(Data):
or MODEL_MAP.get(data.get("type", ""))
)
if not model:
return super().__new__(cls) # pylint: disable=E1120
return super().__new__(cls)
return model.model_validate(data) # type: ignore

View File

@@ -1,4 +1,4 @@
"""Polygon WebSocket server."""
"""Polygon WebSocket client."""
import asyncio
import json
@@ -30,7 +30,7 @@ kwargs = parse_kwargs()
CONNECT_KWARGS = kwargs.pop("connect_kwargs", {})
FEED = kwargs.pop("feed", None)
ASSET_TYPE = kwargs.pop("asset_type", None)
kwargs["results_file"] = os.path.abspath(kwargs["results_file"])
kwargs["results_file"] = os.path.abspath(kwargs.get("results_file"))
async def handle_symbol(symbol):
@@ -85,9 +85,8 @@ async def handle_symbol(symbol):
return ",".join(new_symbols)
async def login(websocket, api_key):
"""Login to the WebSocket."""
login_event = f'{{"action":"auth","params":"{api_key}"}}'
async def login(websocket):
login_event = f'{{"action":"auth","params":"{kwargs["api_key"]}"}}'
try:
await websocket.send(login_event)
res = await websocket.recv()
@@ -108,7 +107,7 @@ async def login(websocket, api_key):
logger.error(err)
sys.exit(1)
logger.info("PROVIDER INFO: %s", msg.get("message"))
except Exception as e: # pylint: disable=broad-except
except Exception as e:
logger.error("PROVIDER ERROR: %s -> %s", e.__class__.__name__, e.args[0])
sys.exit(1)
@@ -123,7 +122,7 @@ async def subscribe(websocket, symbol, event):
subscribe_event = f'{{"action":"{event}","params":"{ticker}"}}'
try:
await websocket.send(subscribe_event)
except Exception as e: # pylint: disable=broad-except
except Exception as e:
msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}"
logger.error(msg)
@@ -183,7 +182,7 @@ async def process_message(message):
if result:
await write_to_db(
result,
kwargs["results_path"],
kwargs["results_file"],
kwargs["table_name"],
kwargs.get("limit"),
)
@@ -191,7 +190,7 @@ async def process_message(message):
logger.info("PROVIDER INFO: %s", msg)
async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-statements
async def connect_and_stream():
"""Connect to the WebSocket and stream data to file."""
handler_task = asyncio.create_task(
@@ -207,7 +206,8 @@ async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-s
try:
async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket:
await login(websocket, kwargs["api_key"])
await login(websocket)
response = await websocket.recv()
messages = json.loads(response)
await process_message(messages)
@@ -262,7 +262,7 @@ async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-s
logger.error(msg)
sys.exit(1)
except Exception as e: # pylint: disable=broad-except
except Exception as e:
msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}"
logger.error(msg)
sys.exit(1)
@@ -293,8 +293,8 @@ if __name__ == "__main__":
logger.error("PROVIDER ERROR: WebSocket connection closed")
except Exception as e: # pylint: disable=broad-except
ERR = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}"
logger.error(ERR)
msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}"
logger.error(msg)
finally:
sys.exit(0)