first commit

app/routers/websocket.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import asyncio
+import json
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from app.core.logger import setup_logger
+from app.services.realtime import realtime_service
+
+router = APIRouter(tags=["websocket"])
+logger = setup_logger("ws_router")
+
+_connected_clients: set[WebSocket] = set()
+
+
+async def _broadcast(data: dict) -> None:
+    message = json.dumps(data, ensure_ascii=False)
+    disconnected: list[WebSocket] = []
+    for client in _connected_clients:
+        try:
+            await client.send_text(message)
+        except Exception:
+            disconnected.append(client)
+    for client in disconnected:
+        _connected_clients.discard(client)
+
+
+@router.websocket("/ws/realtime")
+async def websocket_endpoint(websocket: WebSocket) -> None:
+    await websocket.accept()
+    _connected_clients.add(websocket)
+    logger.info("WebSocket 클라이언트 연결 (총 %d)", len(_connected_clients))
+
+    realtime_service.on("price", _broadcast)
+    realtime_service.on("orderbook", _broadcast)
+
+    try:
+        while True:
+            raw = await websocket.receive_text()
+            try:
+                msg = json.loads(raw)
+                action = msg.get("action", "")
+
+                if action == "subscribe":
+                    stock_code = msg.get("stock_code", "")
+                    data_type = msg.get("data_type", "price")
+                    await realtime_service.subscribe(stock_code, data_type)
+                    await websocket.send_text(
+                        json.dumps({"status": "subscribed", "stock_code": stock_code, "data_type": data_type})
+                    )
+
+                elif action == "unsubscribe":
+                    stock_code = msg.get("stock_code", "")
+                    data_type = msg.get("data_type", "price")
+                    await realtime_service.unsubscribe(stock_code, data_type)
+                    await websocket.send_text(
+                        json.dumps({"status": "unsubscribed", "stock_code": stock_code})
+                    )
+
+            except json.JSONDecodeError:
+                pass
+
+    except WebSocketDisconnect:
+        pass
+    finally:
+        _connected_clients.discard(websocket)
+        logger.info("WebSocket 클라이언트 해제 (총 %d)", len(_connected_clients))