next-tjwater-server/app/api/v1/endpoints/workbench_events.py

"""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():
        await conn.execute(
            "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY"
        )
        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 "",
    )