Files

85 lines
2.5 KiB
Python

from __future__ import annotations
import asyncio
import json
from typing import List
import contextlib
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from app.streaming.binance_hub import hub as binance_hub
from app.streaming.alpha_hub import alpha_hub
router = APIRouter(prefix="/stream", tags=["Stream"])
@router.websocket("/klines")
async def stream_klines_ws(
websocket: WebSocket,
symbols: str = Query("BTCUSDT,XAUUSD"),
timeframe: str = Query("1m"),
):
await websocket.accept()
syms: List[str] = [s.strip().upper().replace("/", "") for s in symbols.split(",") if s.strip()]
async def forward_alpha(sym: str):
queue, unsubscribe = await alpha_hub.subscribe(sym, timeframe="1m")
try:
while True:
evt = await queue.get()
if evt is None:
break
try:
await websocket.send_text(json.dumps(evt))
except WebSocketDisconnect:
break
except Exception:
await asyncio.sleep(0)
finally:
try:
await unsubscribe()
except Exception:
pass
async def forward_binance(sym: str):
queue, unsubscribe = await binance_hub.subscribe(sym, timeframe="1m")
try:
while True:
evt = await queue.get()
if evt is None:
break
# evt already normalized with iso timestamps
try:
await websocket.send_text(json.dumps(evt))
except WebSocketDisconnect:
break
except Exception:
await asyncio.sleep(0)
finally:
try:
await unsubscribe()
except Exception:
pass
tasks: List[asyncio.Task] = []
try:
for s in syms:
if s == "XAUUSD" or s.startswith("XAU"):
tasks.append(asyncio.create_task(forward_alpha("XAUUSD")))
else:
tasks.append(asyncio.create_task(forward_binance(s)))
# Wait for disconnect
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
except WebSocketDisconnect:
pass
finally:
for t in tasks:
t.cancel()
with contextlib.suppress(Exception):
await t
with contextlib.suppress(Exception):
await websocket.close()