BBIT-Kai
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
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()
|
|
|
|
|
|
iot_ws_router = APIRouter()
|
|
|
|
|
|
@iot_ws_router.websocket("/device-status")
|
|
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:
|
|
while True:
|
|
# 这里不需要接收任何消息
|
|
await websocket.receive_text()
|
|
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)
|