完善牧安云哨-后端

bbit_ai/app/db/postgres/ws_manager.py

@@ -1,25 +1,50 @@
 import asyncio
 from typing import List
+from uuid import UUID
 
 from fastapi import WebSocket
 
 
 class ConnectionManager:
     def __init__(self):
-        self.active_connections: List[WebSocket] = []
+        self.active_connections: List[dict] = []  # 保存 websocket 和用户信息
         self.lock = asyncio.Lock()
 
-    async def connect(self, websocket: WebSocket):
+    # proj_id:0：在线状态 1：畜牧车辆进入
+    async def connect(
+        self, websocket: WebSocket, user_id: UUID, dept_id: str, proj_id: int
+    ):
         await websocket.accept()
         async with self.lock:
-            self.active_connections.append(websocket)
+            self.active_connections.append(
+                {
+                    "ws": websocket,
+                    "user_id": user_id,
+                    "dept_id": dept_id,
+                    "proj_id": proj_id,
+                }
+            )
 
     async def disconnect(self, websocket: WebSocket):
         async with self.lock:
-            if websocket in self.active_connections:
-                self.active_connections.remove(websocket)
+            self.active_connections = [
+                conn for conn in self.active_connections if conn["ws"] != websocket
+            ]
 
-    async def broadcast(self, message: dict):
+    async def noticeOnlineStatus(self, message: dict):
         async with self.lock:
-            for ws in self.active_connections:
-                await ws.send_json(message)
+            for conn in self.active_connections:
+                if conn["proj_id"] == 0:
+                    await conn["ws"].send_json(message)
+
+    async def noticeSentinel(
+        self, message: dict, target_departments: List[UUID] = None
+    ):
+        """
+        target_departments: 指定哪些部门能收到消息
+        """
+        async with self.lock:
+            for conn in self.active_connections:
+                if target_departments:
+                    if conn["proj_id"] == 1 and conn["dept_id"] in target_departments:
+                        await conn["ws"].send_json(message)