refactor valkey into objects with health check

_On_Ice/ws_extended_fund_rate_all.py

@@ -0,0 +1,132 @@
+import asyncio
+import json
+import logging
+import socket
+import traceback
+from datetime import datetime, timezone
+from typing import AsyncContextManager
+import math
+
+import numpy as np
+import pandas as pd
+import requests.packages.urllib3.util.connection as urllib3_cn  # type: ignore
+from sqlalchemy import text
+import websockets
+from sqlalchemy.ext.asyncio import create_async_engine
+import valkey
+import os
+from dotenv import load_dotenv
+import modules.utils as utils
+
+### Allow only ipv4 ###
+def allowed_gai_family():
+    return socket.AF_INET
+urllib3_cn.allowed_gai_family = allowed_gai_family
+
+### Database ###
+USE_DB: bool = False
+USE_VK: bool = True
+# VK_FUND_RATE = 'fund_rate_extended'
+VK_FUND_RATE_ALL = 'fund_rate_extended_all'
+
+CON: AsyncContextManager
+VAL_KEY: valkey.Valkey
+
+### Logging ###
+load_dotenv()
+LOG_FILEPATH: str = f'{os.getenv("LOGS_PATH")}/Fund_Rate_Extended_FR_ALL.log'
+
+### Globals ###
+WSS_URL = "wss://api.starknet.extended.exchange/stream.extended.exchange/v1/funding/"
+LOCAL_FUNDING_RATES = []
+
+def time_round_down(dt, interval_mins=5) -> int: # returns timestamp in seconds
+    interval_secs = interval_mins * 60
+    seconds = dt.timestamp()
+    rounded_seconds = math.floor(seconds / interval_secs) * interval_secs
+    
+    return rounded_seconds
+
+### Websocket ###
+async def ws_stream():
+    global LOCAL_FUNDING_RATES
+
+    async for websocket in websockets.connect(WSS_URL):
+        logging.info(f"Connected to {WSS_URL}")
+        try:
+            async for message in websocket:
+                ts_arrival = round(datetime.now().timestamp()*1000)
+                if isinstance(message, str):
+                    try: 
+                        data = json.loads(message)
+                        if data.get('data', None) is not None:
+                            # print(f'FR: {data}')
+                            # fr_next_update_ts = (time_round_down(dt=datetime.now(timezone.utc), interval_mins=60)+(60*60))*1000
+                            fr_update = {
+                                'sequence_id': data['seq'],
+                                'timestamp_arrival': ts_arrival,
+                                'timestamp_msg': data['ts'],
+                                'symbol': data['data']['m'],
+                                'funding_rate': float(data['data']['f']),
+                                'funding_rate_updated_ts_ms': data['data']['T'],
+                                # 'next_funding_time_ts_ms': fr_next_update_ts,
+                            }
+                            LOCAL_FUNDING_RATES = utils.upsert_list_of_dicts_by_id(LOCAL_FUNDING_RATES, fr_update, id='symbol', seq_check_field=None)
+                            VAL_KEY.set(VK_FUND_RATE_ALL, json.dumps(LOCAL_FUNDING_RATES))
+                            # print(f'VK_SAVED: {data}')
+                            continue
+                        else:
+                            logging.info(f'Initial or unexpected data struct, skipping: {data}')
+                            continue
+                    except (json.JSONDecodeError, ValueError):
+                        logging.warning(f'Message not in JSON format, skipping: {message}')
+                        continue
+                else:
+                    raise ValueError(f'Type: {type(data)}  not expected: {message}')
+        except websockets.ConnectionClosed as e:
+            logging.error(f'Connection closed: {e}')
+            logging.error(traceback.format_exc())
+            continue
+        except Exception as e:
+            logging.error(f'Connection closed: {e}')
+            logging.error(traceback.format_exc())
+
+### Startup ###
+async def main():
+    global VAL_KEY
+    global CON
+    
+    if USE_VK:
+        VAL_KEY = valkey.Valkey(host='localhost', port=6379, db=0)
+    else:
+        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')
+        async with engine.connect() as CON:
+            # await create_rtds_btcusd_table(CON=CON)
+            await ws_stream()
+    else:
+        logging.warning("DATABASE NOT BEING USED, NO DATA WILL BE RECORDED")
+        raise NotImplementedError('DB not implemented')
+
+
+if __name__ == '__main__':
+    START_TIME = round(datetime.now().timestamp()*1000)
+
+    logging.info(f'Log FilePath: {LOG_FILEPATH}')
+
+    logging.basicConfig(
+        force=True,
+        filename=LOG_FILEPATH,
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        filemode='w'
+    )
+    logging.info(f"STARTED: {START_TIME}")
+    
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        logging.info("Stream stopped")