Bryan Van Deusen
218 lines
6.9 KiB
Python
218 lines
6.9 KiB
Python
"""WebSocket API for real-time updates."""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Set
|
|
from quart import Blueprint, websocket, current_app
|
|
import redis.asyncio as aioredis
|
|
|
|
from app.config import get_settings
|
|
from app.events import EVENTS_CHANNEL
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
bp = Blueprint("websocket", __name__)
|
|
|
|
# Connected WebSocket clients
|
|
connected_clients: Set = set()
|
|
|
|
|
|
class WebSocketManager:
|
|
"""Manages WebSocket connections and broadcasts."""
|
|
|
|
_instance = None
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._instance.clients = set()
|
|
return cls._instance
|
|
|
|
async def connect(self, ws):
|
|
"""Register a new WebSocket connection."""
|
|
self.clients.add(ws)
|
|
logger.info(f"WebSocket connected. Total clients: {len(self.clients)}")
|
|
|
|
async def disconnect(self, ws):
|
|
"""Remove a WebSocket connection."""
|
|
self.clients.discard(ws)
|
|
logger.info(f"WebSocket disconnected. Total clients: {len(self.clients)}")
|
|
|
|
async def broadcast(self, event_type: str, data: dict):
|
|
"""Broadcast an event to all connected clients."""
|
|
if not self.clients:
|
|
return
|
|
|
|
message = json.dumps({
|
|
"type": event_type,
|
|
"data": data,
|
|
})
|
|
|
|
disconnected = set()
|
|
for client in self.clients:
|
|
try:
|
|
await client.send(message)
|
|
except Exception as e:
|
|
logger.debug(f"Failed to send to client: {e}")
|
|
disconnected.add(client)
|
|
|
|
# Clean up disconnected clients
|
|
self.clients -= disconnected
|
|
|
|
async def send_to_client(self, ws, event_type: str, data: dict):
|
|
"""Send an event to a specific client."""
|
|
message = json.dumps({
|
|
"type": event_type,
|
|
"data": data,
|
|
})
|
|
try:
|
|
await ws.send(message)
|
|
except Exception as e:
|
|
logger.debug(f"Failed to send to client: {e}")
|
|
|
|
|
|
# Global manager instance
|
|
ws_manager = WebSocketManager()
|
|
|
|
|
|
@bp.websocket("/events")
|
|
async def events():
|
|
"""WebSocket endpoint for real-time events.
|
|
|
|
Events sent from server:
|
|
- download.started: { download_id, source_id, url }
|
|
- download.progress: { download_id, file_count }
|
|
- download.completed: { download_id, file_count, duration }
|
|
- download.failed: { download_id, error_type, error_message }
|
|
- source.updated: { source_id, name, last_check }
|
|
- credential.expiring: { platform, expires_in_hours }
|
|
"""
|
|
await ws_manager.connect(websocket._get_current_object())
|
|
|
|
try:
|
|
# Send initial connection confirmation
|
|
await websocket.send(json.dumps({
|
|
"type": "connected",
|
|
"data": {"message": "Connected to Gallery Subscriber"},
|
|
}))
|
|
|
|
# Keep connection alive and handle incoming messages
|
|
while True:
|
|
try:
|
|
# Wait for messages (ping/pong or commands)
|
|
message = await asyncio.wait_for(websocket.receive(), timeout=30)
|
|
|
|
# Handle ping
|
|
if message == "ping":
|
|
await websocket.send("pong")
|
|
else:
|
|
# Parse JSON commands if needed
|
|
try:
|
|
data = json.loads(message)
|
|
await handle_client_message(websocket._get_current_object(), data)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
except asyncio.TimeoutError:
|
|
# Send keepalive ping
|
|
try:
|
|
await websocket.send(json.dumps({"type": "ping"}))
|
|
except Exception:
|
|
break
|
|
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
await ws_manager.disconnect(websocket._get_current_object())
|
|
|
|
|
|
async def handle_client_message(ws, data: dict):
|
|
"""Handle messages from WebSocket clients."""
|
|
msg_type = data.get("type")
|
|
|
|
if msg_type == "subscribe":
|
|
# Client wants to subscribe to specific events
|
|
# For now, all clients receive all events
|
|
await ws_manager.send_to_client(ws, "subscribed", {"status": "ok"})
|
|
|
|
elif msg_type == "ping":
|
|
await ws_manager.send_to_client(ws, "pong", {})
|
|
|
|
|
|
# Helper functions for broadcasting events from other parts of the app
|
|
|
|
async def broadcast_download_started(download_id: int, source_id: int, url: str):
|
|
"""Broadcast that a download has started."""
|
|
await ws_manager.broadcast("download.started", {
|
|
"download_id": download_id,
|
|
"source_id": source_id,
|
|
"url": url,
|
|
})
|
|
|
|
|
|
async def broadcast_download_completed(download_id: int, file_count: int, duration: float):
|
|
"""Broadcast that a download has completed."""
|
|
await ws_manager.broadcast("download.completed", {
|
|
"download_id": download_id,
|
|
"file_count": file_count,
|
|
"duration_seconds": duration,
|
|
})
|
|
|
|
|
|
async def broadcast_download_failed(download_id: int, error_type: str, error_message: str):
|
|
"""Broadcast that a download has failed."""
|
|
await ws_manager.broadcast("download.failed", {
|
|
"download_id": download_id,
|
|
"error_type": error_type,
|
|
"error_message": error_message,
|
|
})
|
|
|
|
|
|
async def broadcast_source_updated(source_id: int, name: str, last_check: str):
|
|
"""Broadcast that a source has been updated."""
|
|
await ws_manager.broadcast("source.updated", {
|
|
"source_id": source_id,
|
|
"name": name,
|
|
"last_check": last_check,
|
|
})
|
|
|
|
|
|
# Redis Pub/Sub integration for cross-process events (from Celery tasks)
|
|
|
|
async def start_redis_subscriber():
|
|
"""Start a background task that subscribes to Redis events and broadcasts them.
|
|
|
|
This should be called when the Quart app starts.
|
|
"""
|
|
settings = get_settings()
|
|
|
|
try:
|
|
redis_client = aioredis.from_url(settings.redis_url)
|
|
pubsub = redis_client.pubsub()
|
|
await pubsub.subscribe(EVENTS_CHANNEL)
|
|
|
|
logger.info(f"Started Redis subscriber on channel: {EVENTS_CHANNEL}")
|
|
|
|
async for message in pubsub.listen():
|
|
if message["type"] == "message":
|
|
try:
|
|
data = json.loads(message["data"])
|
|
event_type = data.get("type")
|
|
event_data = data.get("data", {})
|
|
|
|
# Broadcast to all WebSocket clients
|
|
await ws_manager.broadcast(event_type, event_data)
|
|
logger.debug(f"Broadcast event from Redis: {event_type}")
|
|
except json.JSONDecodeError:
|
|
logger.warning(f"Invalid JSON in Redis message: {message['data']}")
|
|
except Exception as e:
|
|
logger.error(f"Error broadcasting Redis event: {e}")
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info("Redis subscriber cancelled")
|
|
await pubsub.unsubscribe(EVENTS_CHANNEL)
|
|
await redis_client.close()
|
|
except Exception as e:
|
|
logger.error(f"Redis subscriber error: {e}")
|