""" Router WebSocket — suivi temps réel du pipeline de traitement d'images. Endpoints : - WS /ws/pipeline/{image_id}?token= → événements d'un pipeline - WS /ws/admin/monitor?token= → monitoring admin global """ import json import logging from typing import Any from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.database import AsyncSessionLocal, get_db from app.dependencies.auth import hash_api_key from app.metrics import hub_active_websockets from app.models.client import APIClient from app.models.image import Image, ProcessingStatus logger = logging.getLogger(__name__) router = APIRouter(tags=["WebSocket"]) # ───────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────── async def _authenticate_ws(websocket: WebSocket, db: AsyncSession) -> APIClient | None: """ Authentifie une connexion WebSocket via le query param `token`. Retourne le client ou None si invalide. """ token = websocket.query_params.get("token") if not token: return None # Vérification Master Key from app.config import settings if settings.ADMIN_API_KEY and token == settings.ADMIN_API_KEY: return APIClient( id="admin-master", name="Imago Master Admin", scopes=["admin", "images:read", "images:write", "ai:use"], plan="premium", ) key_hash = hash_api_key(token) result = await db.execute( select(APIClient).where(APIClient.api_key_hash == key_hash) ) client = result.scalar_one_or_none() if client and client.is_active: return client return None async def _get_image(image_id: int, db: AsyncSession) -> Image | None: """Charge une image depuis la BDD.""" result = await db.execute(select(Image).where(Image.id == image_id)) return result.scalar_one_or_none() # ───────────────────────────────────────────────────────────── # WS /ws/pipeline/{image_id} — suivi d'un pipeline # ───────────────────────────────────────────────────────────── from fastapi import Depends @router.websocket("/ws/pipeline/{image_id}") async def ws_pipeline( websocket: WebSocket, image_id: int, db: AsyncSession = Depends(get_db) ) -> None: """ WebSocket de suivi temps réel d'un pipeline image. - Authentification via query param `token` - Vérifie que l'image appartient au client - Envoie le buffer de reconnexion puis les événements live - Ferme après pipeline.done ou pipeline.error """ # ── Authentification ────────────────────────────────────── client = await _authenticate_ws(websocket, db) if client is None: await websocket.close(code=4001, reason="Token manquant ou invalide") return # ── Vérification propriété de l'image ───────────────────── image = await _get_image(image_id, db) if image is None: await websocket.accept() await websocket.close(code=4004, reason="Image introuvable") return # Admin peut voir toutes les images, sinon vérifier ownership if not client.has_scope("admin") and image.client_id != client.id: await websocket.close(code=4003, reason="Accès interdit") return # ── Accepter la connexion ───────────────────────────────── await websocket.accept() hub_active_websockets.inc() try: # ── Image déjà terminée → message synthétique ───────── if image.processing_status == ProcessingStatus.DONE: await websocket.send_json({ "event": "pipeline.done", "image_id": image_id, "status": "done", "synthetic": True, }) return if image.processing_status == ProcessingStatus.ERROR: await websocket.send_json({ "event": "pipeline.error", "image_id": image_id, "error": image.processing_error or "Erreur inconnue", "synthetic": True, }) return # ── Récupérer le buffer de reconnexion depuis Redis ─── redis = getattr(websocket.app.state, "redis", None) if redis is not None: try: buffer_key = f"pipeline:buffer:{image_id}" buffered = await redis.lrange(buffer_key, 0, -1) for raw_event in buffered: try: event_data = json.loads(raw_event) await websocket.send_json(event_data) except (json.JSONDecodeError, Exception): pass except Exception as e: logger.warning("ws.buffer_read_error", extra={"error": str(e)}) # ── S'abonner au channel Redis et écouter les événements if redis is not None: pubsub = redis.pubsub() try: await pubsub.subscribe(f"pipeline:{image_id}") async for message in pubsub.listen(): if message["type"] != "message": continue try: data = json.loads(message["data"]) except (json.JSONDecodeError, TypeError): continue await websocket.send_json(data) # Fermer après pipeline.done ou pipeline.error event_type = data.get("event", "") if event_type in ("pipeline.done", "pipeline.error"): break finally: await pubsub.unsubscribe(f"pipeline:{image_id}") await pubsub.close() else: # Pas de Redis — envoyer un message d'info et fermer await websocket.send_json({ "event": "error", "message": "Redis indisponible — utilisez le polling GET /images/{id}/status", }) except WebSocketDisconnect: logger.info("ws.client_disconnected", extra={ "image_id": image_id, "client_id": client.id, }) except Exception as e: logger.error("ws.unexpected_error", extra={ "image_id": image_id, "error": str(e), }) finally: hub_active_websockets.dec() # ───────────────────────────────────────────────────────────── # WS /ws/admin/monitor — monitoring admin global # ───────────────────────────────────────────────────────────── @router.websocket("/ws/admin/monitor") async def ws_admin_monitor( websocket: WebSocket, db: AsyncSession = Depends(get_db) ) -> None: """ WebSocket admin pour surveiller tous les pipelines en temps réel. Nécessite le scope `admin`. Pousse un événement à chaque démarrage ou fin de pipeline sur n'importe quelle image. """ # ── Authentification ────────────────────────────────────── client = await _authenticate_ws(websocket, db) if client is None: await websocket.close(code=4001, reason="Token manquant ou invalide") return if not client.has_scope("admin"): await websocket.close(code=4003, reason="Scope admin requis") return # ── Accepter la connexion ───────────────────────────────── await websocket.accept() hub_active_websockets.inc() try: redis = getattr(websocket.app.state, "redis", None) if redis is None: await websocket.send_json({ "event": "error", "message": "Redis indisponible", }) return pubsub = redis.pubsub() try: await pubsub.subscribe("pipeline:admin") async for message in pubsub.listen(): if message["type"] != "message": continue try: data = json.loads(message["data"]) except (json.JSONDecodeError, TypeError): continue await websocket.send_json(data) finally: await pubsub.unsubscribe("pipeline:admin") await pubsub.close() except WebSocketDisconnect: logger.info("ws.admin_disconnected", extra={"client_id": client.id}) except Exception as e: logger.error("ws.admin_error", extra={"error": str(e)}) finally: hub_active_websockets.dec()