114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
"""WebSocket connection manager with channel subscriptions."""
|
|
|
|
from typing import Dict, Optional, Set
|
|
|
|
from fastapi import WebSocket
|
|
import redis.asyncio as redis
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manage user connections and channel subscriptions."""
|
|
|
|
def __init__(self) -> None:
|
|
self.user_connections: Dict[str, Set[WebSocket]] = {}
|
|
self.socket_users: Dict[WebSocket, str] = {}
|
|
self.channel_connections: Dict[str, Set[WebSocket]] = {}
|
|
self.socket_channels: Dict[WebSocket, Set[str]] = {}
|
|
self.redis_client: Optional[redis.Redis] = None
|
|
|
|
@property
|
|
def active_connections(self) -> Dict[str, Set[WebSocket]]:
|
|
"""Compatibility alias for existing callers."""
|
|
return self.user_connections
|
|
|
|
async def connect(self, websocket: WebSocket, user_id: str) -> None:
|
|
await websocket.accept()
|
|
self.user_connections.setdefault(user_id, set()).add(websocket)
|
|
self.socket_users[websocket] = user_id
|
|
self.socket_channels.setdefault(websocket, set())
|
|
|
|
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) -> None:
|
|
self.unsubscribe(websocket, list(self.socket_channels.get(websocket, set())))
|
|
|
|
if user_id in self.user_connections:
|
|
self.user_connections[user_id].discard(websocket)
|
|
if not self.user_connections[user_id]:
|
|
del self.user_connections[user_id]
|
|
|
|
self.socket_users.pop(websocket, None)
|
|
self.socket_channels.pop(websocket, None)
|
|
|
|
def subscribe(self, websocket: WebSocket, channels: list[str]) -> list[str]:
|
|
subscribed_channels: list[str] = []
|
|
socket_channel_set = self.socket_channels.setdefault(websocket, set())
|
|
|
|
for channel in channels:
|
|
normalized_channel = channel.strip()
|
|
if not normalized_channel:
|
|
continue
|
|
self.channel_connections.setdefault(normalized_channel, set()).add(websocket)
|
|
socket_channel_set.add(normalized_channel)
|
|
subscribed_channels.append(normalized_channel)
|
|
|
|
return subscribed_channels
|
|
|
|
def unsubscribe(self, websocket: WebSocket, channels: list[str]) -> None:
|
|
socket_channel_set = self.socket_channels.setdefault(websocket, set())
|
|
|
|
for channel in channels:
|
|
normalized_channel = channel.strip()
|
|
if not normalized_channel:
|
|
continue
|
|
if normalized_channel in self.channel_connections:
|
|
self.channel_connections[normalized_channel].discard(websocket)
|
|
if not self.channel_connections[normalized_channel]:
|
|
del self.channel_connections[normalized_channel]
|
|
socket_channel_set.discard(normalized_channel)
|
|
|
|
async def send_personal_message(self, message: dict, user_id: str) -> None:
|
|
for connection in list(self.user_connections.get(user_id, set())):
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception:
|
|
self.disconnect(connection, user_id)
|
|
|
|
async def broadcast(self, message: dict, channel: str = "all") -> None:
|
|
if channel == "all":
|
|
targets = list(self.socket_users.keys())
|
|
else:
|
|
targets = list(self.channel_connections.get(channel, set()))
|
|
|
|
for connection in targets:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception:
|
|
user_id = self.socket_users.get(connection)
|
|
if user_id is not None:
|
|
self.disconnect(connection, user_id)
|
|
|
|
async def close_all(self) -> None:
|
|
for websocket, user_id in list(self.socket_users.items()):
|
|
await websocket.close()
|
|
self.disconnect(websocket, user_id)
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
async def get_websocket_manager() -> ConnectionManager:
|
|
return manager
|