"""Authenticated workbench snapshot and resumable Server-Sent Events.""" from __future__ import annotations import asyncio import json from datetime import datetime, timezone from fastapi import APIRouter, Depends, Header, HTTPException, Query from fastapi.responses import StreamingResponse from psycopg import AsyncConnection from psycopg.conninfo import conninfo_to_dict from app.api.v1.endpoints.workbench_runtime import workbench_runtime from app.auth.permissions import OPERATIONAL_TASK_VIEW, permissions_for_context, require_permission from app.auth.project_dependencies import ( ProjectContext, get_metadata_repository, get_project_business_routing, get_project_pg_connection, ) from app.core.config import settings from app.domain.schemas.workbench_events import WorkbenchEvent, WorkbenchSyncSnapshot from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository from app.infra.db.postgresql.workbench_events import WorkbenchEventRepository from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository from app.infra.db.project_routing import ActiveProjectRouting from app.services.operational_tasks import OperationalTaskService from app.infra.db.postgresql.workbench_event_hub import workbench_event_hubs router = APIRouter() def _service(ctx: ProjectContext, conn: AsyncConnection) -> OperationalTaskService: return OperationalTaskService( OperationalTaskRepository(conn, ctx.project_id), ctx, permissions_for_context(ctx), ) @router.get("/workbench/sync-snapshot", response_model=WorkbenchSyncSnapshot) async def workbench_sync_snapshot( ctx: ProjectContext = Depends(require_permission(OPERATIONAL_TASK_VIEW)), conn: AsyncConnection = Depends(get_project_pg_connection), ): async with conn.transaction(isolation_level="repeatable read", read_only=True): cursor = await WorkbenchEventRepository(conn, ctx.project_id).current_cursor() runtime = await workbench_runtime(ctx=ctx, conn=conn) tasks, _ = await _service(ctx, conn).list( task_type=None, status=None, limit=50, offset=0 ) notifications, _ = await _service(ctx, conn).list_notifications( unread_only=False, limit=100, offset=0 ) return WorkbenchSyncSnapshot( schema="workbench-sync-snapshot@1", event_cursor=cursor, runtime=runtime, tasks=tasks, notifications=notifications, generated_at=datetime.now(timezone.utc), ) @router.get("/workbench/events") async def workbench_events( after: int | None = Query(default=None, ge=0), last_event_id: str | None = Header(default=None, alias="Last-Event-ID"), ctx: ProjectContext = Depends(require_permission(OPERATIONAL_TASK_VIEW)), routing: ActiveProjectRouting = Depends(get_project_business_routing), metadata_repo: MetadataRepository = Depends(get_metadata_repository), ): if not settings.WORKBENCH_EVENTS_ENABLED: raise HTTPException(status_code=503, detail="Workbench events are disabled") cursor = after if cursor is None and last_event_id: try: cursor = int(last_event_id) except ValueError as exc: raise HTTPException(status_code=400, detail="Invalid Last-Event-ID") from exc cursor = cursor or 0 current_identity = _database_identity(routing.business_dsn) try: active_routings = await metadata_repo.list_active_project_db_routings("biz_data") finally: # FastAPI keeps yield dependencies alive for a StreamingResponse. The # stream only needs immutable routing values, so release the shared # metadata checkout before the five-minute SSE body begins. await metadata_repo.session.close() if any( item.project_id != ctx.project_id and _database_identity(item.dsn) == current_identity for item in active_routings ): raise HTTPException( status_code=503, detail="Workbench events require one dedicated business database per project", ) async def stream(): started = asyncio.get_running_loop().time() last_keepalive = started async with workbench_event_hubs.subscribe(ctx.project_id, routing.business_dsn) as hub: oldest = await hub.oldest_cursor() latest = await hub.current_cursor() if cursor > 0 and ( oldest is None or cursor < oldest - 1 or cursor > latest ): payload = {"schema": "workbench-event@1", "cursor": cursor} yield f"event: reset_required\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n" return current = cursor wake_generation = hub.generation while asyncio.get_running_loop().time() - started < 300: rows = await hub.read_after(current) for row in rows: event = WorkbenchEvent( schema="workbench-event@1", **row, ) current = event.event_seq data = event.model_dump_json(by_alias=True) yield f"id: {current}\nevent: workbench_event\ndata: {data}\n\n" if rows: continue now = asyncio.get_running_loop().time() if now - last_keepalive >= 15: yield f": keepalive {int(now)}\n\n" last_keepalive = now # LISTEN is only a wake-up path. The timeout also reconciles # against the durable cursor if a notification was missed. wake_generation = await hub.wait_for_wake(wake_generation, timeout=1) yield "event: reconnect\ndata: {}\n\n" return StreamingResponse( stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no", "Content-Encoding": "identity", }, ) def _database_identity(dsn: str) -> tuple[str, str, str]: values = conninfo_to_dict(dsn) return ( values.get("host") or values.get("hostaddr") or "localhost", values.get("port") or "5432", values.get("dbname") or values.get("user") or "", )