rapid interations of server side app and firefox extension
This commit is contained in:
Bryan Van Deusen
@@ -0,0 +1,217 @@
|
||||
"""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}")
|
||||
Reference in New Issue
Block a user