From 87ee7fbd1a54c63e04309523c8e7991457cf708c Mon Sep 17 00:00:00 2001 From: MoonTestUse1 Date: Sun, 5 Jan 2025 03:01:18 +0600 Subject: [PATCH] add websockets suppor98 --- backend/app/main.py | 9 +- backend/app/websockets/notifications.py | 120 +++++++++++++----------- 2 files changed, 73 insertions(+), 56 deletions(-) diff --git a/backend/app/main.py b/backend/app/main.py index 87cb2fb..995b754 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -37,8 +37,13 @@ app.include_router(admin.router, prefix="/api/admin", tags=["admin"]) app.include_router(statistics.router, prefix="/api/statistics", tags=["statistics"]) # Добавляем WebSocket маршруты -app.websocket("/api/ws/admin")(notification_manager.admin_endpoint) -app.websocket("/api/ws/employee/{employee_id}")(notification_manager.employee_endpoint) +@app.websocket("/api/ws/admin") +async def admin_websocket(websocket): + await notification_manager.admin_endpoint(websocket) + +@app.websocket("/api/ws/employee/{employee_id}") +async def employee_websocket(websocket, employee_id: int): + await notification_manager.employee_endpoint(websocket, employee_id) @app.on_event("startup") async def startup_event(): diff --git a/backend/app/websockets/notifications.py b/backend/app/websockets/notifications.py index 24c900b..8a60a6a 100644 --- a/backend/app/websockets/notifications.py +++ b/backend/app/websockets/notifications.py @@ -1,81 +1,93 @@ +"""WebSocket notifications manager""" from fastapi import WebSocket -from typing import Dict, List -import json +from typing import Dict, List, Set import logging +import json logger = logging.getLogger(__name__) class NotificationManager: + """Менеджер WebSocket подключений и уведомлений""" def __init__(self): - self.active_connections: Dict[str, List[WebSocket]] = { - "admin": [], # Подключения админов - "employee": [] # Подключения сотрудников - } - - async def connect(self, websocket: WebSocket, client_type: str): + self.admin_connections: Set[WebSocket] = set() + self.employee_connections: Dict[int, WebSocket] = {} + + async def connect(self, websocket: WebSocket, client_type: str, employee_id: int = None): + """Установка WebSocket соединения""" await websocket.accept() - self.active_connections[client_type].append(websocket) - logger.info(f"New {client_type} connection. Total connections: {len(self.active_connections[client_type])}") + if client_type == "admin": + self.admin_connections.add(websocket) + logger.info("Admin connected to WebSocket") + elif client_type == "employee" and employee_id: + self.employee_connections[employee_id] = websocket + logger.info(f"Employee {employee_id} connected to WebSocket") - def disconnect(self, websocket: WebSocket, client_type: str): - if websocket in self.active_connections[client_type]: - self.active_connections[client_type].remove(websocket) - logger.info(f"{client_type} disconnected. Remaining connections: {len(self.active_connections[client_type])}") + def disconnect(self, websocket: WebSocket, client_type: str, employee_id: int = None): + """Закрытие WebSocket соединения""" + if client_type == "admin": + self.admin_connections.discard(websocket) + logger.info("Admin disconnected from WebSocket") + elif client_type == "employee" and employee_id: + self.employee_connections.pop(employee_id, None) + logger.info(f"Employee {employee_id} disconnected from WebSocket") async def broadcast_to_admins(self, message: dict): - """Отправка сообщения всем подключенным админам""" - logger.info(f"Broadcasting to admins: {message}") - logger.info(f"Number of admin connections: {len(self.active_connections['admin'])}") - - if not self.active_connections["admin"]: - logger.warning("No admin connections available") - return - - disconnected = [] - - for connection in self.active_connections["admin"]: + """Отправка сообщения всем админам""" + disconnected = set() + for websocket in self.admin_connections: try: - await connection.send_json(message) - logger.info("Message sent successfully to admin") + await websocket.send_json(message) except Exception as e: logger.error(f"Error sending message to admin: {e}") - disconnected.append(connection) - continue + disconnected.add(websocket) # Удаляем отключенные соединения - for connection in disconnected: - self.disconnect(connection, "admin") + for websocket in disconnected: + self.disconnect(websocket, "admin") - async def broadcast_to_employees(self, employee_id: int, message: dict): + async def send_to_employee(self, employee_id: int, message: dict): """Отправка сообщения конкретному сотруднику""" - logger.info(f"Broadcasting to employee {employee_id}: {message}") - logger.info(f"Number of employee connections: {len(self.active_connections['employee'])}") - - if not self.active_connections["employee"]: - logger.warning("No employee connections available") - return - - disconnected = [] - - for connection in self.active_connections["employee"]: + websocket = self.employee_connections.get(employee_id) + if websocket: try: - await connection.send_json(message) - logger.info("Message sent successfully to employee") + await websocket.send_json(message) except Exception as e: - logger.error(f"Error sending message to employee: {e}") - disconnected.append(connection) - continue - - # Удаляем отключенные соединения - for connection in disconnected: - self.disconnect(connection, "employee") + logger.error(f"Error sending message to employee {employee_id}: {e}") + self.disconnect(websocket, "employee", employee_id) async def handle_ping(self, websocket: WebSocket): - """Обработка ping сообщений""" + """Обработка ping-сообщений""" try: await websocket.send_json({"type": "pong"}) - logger.debug("Sent pong response") except Exception as e: - logger.error(f"Error sending pong: {e}") + logger.error(f"Error handling ping: {e}") + async def admin_endpoint(self, websocket: WebSocket): + """WebSocket endpoint для админов""" + await self.connect(websocket, "admin") + try: + while True: + data = await websocket.receive_json() + if data.get("type") == "ping": + await self.handle_ping(websocket) + except Exception as e: + logger.error(f"Error in admin websocket: {e}") + finally: + self.disconnect(websocket, "admin") + + async def employee_endpoint(self, websocket: WebSocket, employee_id: int): + """WebSocket endpoint для сотрудников""" + await self.connect(websocket, "employee", employee_id) + try: + while True: + data = await websocket.receive_json() + if data.get("type") == "ping": + await self.handle_ping(websocket) + except Exception as e: + logger.error(f"Error in employee websocket: {e}") + finally: + self.disconnect(websocket, "employee", employee_id) + + +# Создаем глобальный экземпляр менеджера уведомлений notification_manager = NotificationManager() \ No newline at end of file