added extend symbol change for ws

ws_aster.py

@@ -5,7 +5,7 @@ import socket
 import traceback
 from datetime import datetime
 from typing import AsyncContextManager
-
+import time
 import numpy as np
 import pandas as pd
 import requests.packages.urllib3.util.connection as urllib3_cn  # type: ignore
@@ -26,35 +26,79 @@ urllib3_cn.allowed_gai_family = allowed_gai_family
 ### Database ###
 USE_DB: bool = True
 USE_VK: bool = True
+CON: AsyncContextManager
+VAL_KEY: valkey.Valkey
 VK_FUND_RATE = 'fund_rate_aster'
 VK_TICKER = 'fut_ticker_aster'
 VK_LAST_TRADE = 'fut_last_trade_aster'
-CON: AsyncContextManager | None = None
-VAL_KEY = None
 
 ### Logging ###
 load_dotenv()
-LOG_FILEPATH: str = os.getenv("LOGS_PATH") + '/Fund_Rate_Aster.log'
+LOG_FILEPATH: str = f'{os.getenv(key="LOGS_PATH")}/Fund_Rate_Aster.log'
 
 ### CONSTANTS ###
 SYMBOL: str = 'ETHUSDT'
+
 STREAM_MARKPRICE: str = f'{SYMBOL.lower()}@markPrice@1s'
 STREAM_BOOKTICKER: str = f'{SYMBOL.lower()}@bookTicker'
 STREAM_TRADES: str = f'{SYMBOL.lower()}@aggTrade'
 
 ### Globals ###
-WSS_URL = f"wss://fstream.asterdex.com/stream?streams={STREAM_MARKPRICE}/{STREAM_BOOKTICKER}/{STREAM_TRADES}"
+WSS_URL: str = f"wss://fstream.asterdex.com/stream?streams={STREAM_MARKPRICE}/{STREAM_BOOKTICKER}/{STREAM_TRADES}"
+ALLOW_SYMBOL_CHG: bool = False
+
+### Funcs ###
+async def subscribe_streams(websocket, streams: list[str]) -> None:
+    logging.info(f'Trying to sub: {streams}')
+    msg = {
+        "method": "SUBSCRIBE", 
+        "params": streams,
+        "id": int(round(number=datetime.now().timestamp()*1000)) 
+    }
+    await websocket.send(json.dumps(obj=msg))
+    logging.info(f'Success sub: {streams}')
+
+
+async def unsubscribe_streams(websocket, streams: list[str]) -> None:
+    logging.info(f'Trying to unsub: {streams}')
+    msg = {
+        "method": "UNSUBSCRIBE", 
+        "params": streams,
+        "id": int(round(number=datetime.now().timestamp()*1000)) 
+    }
+    await websocket.send(json.dumps(obj=msg))
+    logging.info(f'Success unsub: {streams}')
+
 
 ### Websocket ###
 async def ws_stream():
-    async for websocket in websockets.connect(WSS_URL):
-        logging.info(f"Connected to {WSS_URL}")
+    global SYMBOL
+    global STREAM_MARKPRICE
+    global STREAM_BOOKTICKER
+    global STREAM_TRADES
 
+    async for websocket in websockets.connect(WSS_URL, ping_interval=5):
+        logging.info(msg=f"Connected to {WSS_URL}")
         try:
             async for message in websocket:
+                ### Update Symbol if Algo Outputs Change ###
+                if ALLOW_SYMBOL_CHG:
+                    best_symbol_by_exchange: dict = json.loads(s=VAL_KEY.get(name='fr_algo_working_symbol'))  # ty:ignore[invalid-argument-type]
+                    best_symbol: str = f'{best_symbol_by_exchange['ASTER']['lh_asset']}{best_symbol_by_exchange['ASTER']['rh_asset']}'
+                    if best_symbol != SYMBOL:
+                        logging.info(f'Symbol Change: {SYMBOL} -> {best_symbol}')
+                        SYMBOL = best_symbol
+                        await unsubscribe_streams(websocket = websocket, streams=[STREAM_MARKPRICE,STREAM_BOOKTICKER,STREAM_TRADES])
+                        
+                        STREAM_MARKPRICE = f'{SYMBOL.lower()}@markPrice@1s'
+                        STREAM_BOOKTICKER = f'{SYMBOL.lower()}@bookTicker'
+                        STREAM_TRADES = f'{SYMBOL.lower()}@aggTrade'
+                        await subscribe_streams(websocket = websocket, streams=[STREAM_MARKPRICE,STREAM_BOOKTICKER,STREAM_TRADES])
+                        continue
+                
                 ts_arrival = round(datetime.now().timestamp()*1000)
                 if isinstance(message, str):
-                    try: 
+                    try:
                         data = json.loads(message)
                         channel = data.get('stream', None)
                         if channel is not None:
@@ -132,8 +176,8 @@ async def main():
     if USE_VK:
         VAL_KEY = valkey.Valkey(host='localhost', port=6379, db=0)
     else:
-        VAL_KEY = None
         logging.warning("VALKEY NOT BEING USED, NO DATA WILL BE PUBLISHED")
+        raise NotImplementedError('Cannot run without VK')
 
     if USE_DB:
         engine = create_async_engine('mysql+asyncmy://root:pwd@localhost/fund_rate')
@@ -141,9 +185,8 @@ async def main():
             await aster_db.create_fr_aster_mkt_trades(CON=CON)
             await ws_stream()
     else:
-        CON = None
         logging.warning("DATABASE NOT BEING USED, NO DATA WILL BE RECORDED")
-        await ws_stream()
+        raise NotImplementedError('Cannot run without DB')
 
 
 if __name__ == '__main__':