109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
"""WebSocket Connection Manager"""
|
|
|
|
from typing import Dict, Set, Optional
|
|
from fastapi import WebSocket
|
|
import redis.asyncio as redis
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, Set[WebSocket]] = {} # user_id -> connections
|
|
self.channel_subscriptions: Dict[str, Set[WebSocket]] = {}
|
|
self.websocket_channels: Dict[WebSocket, Set[str]] = {}
|
|
self.redis_client: Optional[redis.Redis] = None
|
|
|
|
async def connect(self, websocket: WebSocket, user_id: str):
|
|
await websocket.accept()
|
|
if user_id not in self.active_connections:
|
|
self.active_connections[user_id] = set()
|
|
self.active_connections[user_id].add(websocket)
|
|
|
|
if self.redis_client is None:
|
|
redis_url = settings.REDIS_URL
|
|
if redis_url.startswith("redis://"):
|
|
self.redis_client = redis.from_url(redis_url, decode_responses=True)
|
|
else:
|
|
self.redis_client = redis.Redis(
|
|
host=settings.REDIS_SERVER,
|
|
port=settings.REDIS_PORT,
|
|
db=settings.REDIS_DB,
|
|
decode_responses=True,
|
|
)
|
|
|
|
def disconnect(self, websocket: WebSocket, user_id: str):
|
|
if user_id in self.active_connections:
|
|
self.active_connections[user_id].discard(websocket)
|
|
if not self.active_connections[user_id]:
|
|
del self.active_connections[user_id]
|
|
self.unsubscribe_all(websocket)
|
|
|
|
def subscribe(self, websocket: WebSocket, channels: list[str]):
|
|
normalized_channels = {
|
|
str(channel).strip()
|
|
for channel in channels
|
|
if str(channel).strip()
|
|
}
|
|
if not normalized_channels:
|
|
return
|
|
|
|
socket_channels = self.websocket_channels.setdefault(websocket, set())
|
|
for channel in normalized_channels:
|
|
self.channel_subscriptions.setdefault(channel, set()).add(websocket)
|
|
socket_channels.add(channel)
|
|
|
|
def unsubscribe(self, websocket: WebSocket, channels: list[str]):
|
|
for channel in {str(channel).strip() for channel in channels if str(channel).strip()}:
|
|
subscribers = self.channel_subscriptions.get(channel)
|
|
if subscribers is not None:
|
|
subscribers.discard(websocket)
|
|
if not subscribers:
|
|
del self.channel_subscriptions[channel]
|
|
socket_channels = self.websocket_channels.get(websocket)
|
|
if socket_channels is not None:
|
|
socket_channels.discard(channel)
|
|
if not socket_channels:
|
|
del self.websocket_channels[websocket]
|
|
|
|
def unsubscribe_all(self, websocket: WebSocket):
|
|
channels = list(self.websocket_channels.get(websocket, set()))
|
|
if channels:
|
|
self.unsubscribe(websocket, channels)
|
|
|
|
async def send_personal_message(self, message: dict, user_id: str):
|
|
if user_id in self.active_connections:
|
|
for connection in self.active_connections[user_id]:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception:
|
|
pass
|
|
|
|
async def broadcast(self, message: dict, channel: str = "all"):
|
|
if channel == "all":
|
|
for user_id in self.active_connections:
|
|
await self.send_personal_message(message, user_id)
|
|
else:
|
|
for connection in list(self.channel_subscriptions.get(channel, set())):
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception:
|
|
self.unsubscribe_all(connection)
|
|
|
|
async def close_all(self):
|
|
for user_id in self.active_connections:
|
|
for connection in self.active_connections[user_id]:
|
|
await connection.close()
|
|
self.active_connections.clear()
|
|
self.channel_subscriptions.clear()
|
|
self.websocket_channels.clear()
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
async def get_websocket_manager() -> ConnectionManager:
|
|
return manager
|