File size: 2,437 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import asyncio
import json
import logging
from datetime import datetime
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.services.market_data import market_data_service
from app.services.order_manager import order_manager
logger = logging.getLogger(__name__)
ws_router = APIRouter()
class ConnectionManager:
def __init__(self):
self.active_connections: list[WebSocket] = []
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket):
if websocket in self.active_connections:
self.active_connections.remove(websocket)
async def broadcast(self, message: str):
dead = []
for connection in self.active_connections:
try:
await connection.send_text(message)
except Exception:
dead.append(connection)
for conn in dead:
self.disconnect(conn)
manager = ConnectionManager()
def json_serializer(obj):
if isinstance(obj, datetime):
return obj.isoformat()
raise TypeError(f"Type {type(obj)} not serializable")
@ws_router.websocket("/ws/market")
async def market_ws(websocket: WebSocket):
await manager.connect(websocket)
queue = market_data_service.subscribe()
try:
while True:
data = await queue.get()
msg = json.dumps({
"type": "market_data",
"data": data.model_dump(),
}, default=json_serializer)
await websocket.send_text(msg)
except WebSocketDisconnect:
pass
except Exception as e:
logger.error("WebSocket error: %s", e)
finally:
market_data_service.unsubscribe(queue)
manager.disconnect(websocket)
@ws_router.websocket("/ws/account")
async def account_ws(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
account = order_manager.account
msg = json.dumps({
"type": "account_update",
"data": account.model_dump(),
}, default=json_serializer)
await websocket.send_text(msg)
await asyncio.sleep(2)
except WebSocketDisconnect:
pass
except Exception as e:
logger.error("Account WS error: %s", e)
finally:
manager.disconnect(websocket)
|