完善牧安云哨-后端

bbit_ai/app/routers/WS.py

@@ -1,6 +1,8 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, Query
 from starlette.websockets import WebSocket, WebSocketDisconnect
 
+from config.security import get_user_id_from_token_from_ws
+from db.postgres import get_dept_id_by_user_id
 from db.postgres.ws_manager import ConnectionManager
 
 ws_manager = ConnectionManager()
@@ -10,8 +12,13 @@ iot_ws_router = APIRouter()
 
 
 @iot_ws_router.websocket("/device-status")
-async def websocket_device_status(websocket: WebSocket):
-    await ws_manager.connect(websocket)
+async def websocket_device_status(
+    websocket: WebSocket,
+    token: str = Query(...),
+):
+    user_id = get_user_id_from_token_from_ws(token)
+    dept_id = get_dept_id_by_user_id(user_id)  # 查数据库或缓存
+    await ws_manager.connect(websocket, user_id, dept_id, 0)
     print("[WS] client connected")
 
     try:
@@ -21,3 +28,22 @@ async def websocket_device_status(websocket: WebSocket):
     except WebSocketDisconnect:
         await ws_manager.disconnect(websocket)
         print("[WS] client disconnected")
+
+
+@iot_ws_router.websocket("/sentinel_record")
+async def websocket_sentinel_record(
+    websocket: WebSocket,
+    token: str = Query(...),
+):
+    user_id = get_user_id_from_token_from_ws(token)
+    dept_id = get_dept_id_by_user_id(user_id)  # 查数据库或缓存
+    print("user_id:", user_id)
+    print("dept_id:", dept_id)
+    print("已接入")
+    await ws_manager.connect(websocket, user_id, dept_id, 1)
+
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        await ws_manager.disconnect(websocket)