bdlg2023-sms-back/app/api/endpoints/project_stream.py

import json
import secrets
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse
from starlette import status

from app.api import deps
from app.core.events import redis_client
from app.models import OrgRole, User

router = APIRouter()

TICKET_TTL_SECONDS = 30


@router.post("/project/{project_id}/stream-ticket")
async def issue_stream_ticket(
    project_id: UUID,
    current_user: User = Depends(
        deps.require_org_role(OrgRole.ORG_ADMIN, OrgRole.RESPO_BENEVOLE, OrgRole.RESPO_COMMISSION)
    ),
):
    ticket = secrets.token_urlsafe(32)
    key = f"sse_ticket:{ticket}"
    await redis_client.set(key, str(project_id), ex=TICKET_TTL_SECONDS)
    return {"ticket": ticket}


async def _consume_ticket(ticket: str) -> str:
    key = f"sse_ticket:{ticket}"
    user_id = await redis_client.getdel(key)
    if not user_id:
        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid or expired ticket")
    return user_id


@router.get("/project/{project_id}/stream", summary="Real-Time SSE Stream")
async def project_stream(project_id: str, client_id: str, ticket: str, request: Request):
    """
    Connect to this endpoint to receive real-time updates.

    **Required Query Parameters:**
    - `client_id`: A unique UUID generated by the frontend to prevent echoing.
    - `ticket`: A unique ticket generated by the issue_stream_ticket endpoint.
    """
    ticket_project_id = await _consume_ticket(ticket)
    if ticket_project_id != project_id:
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, "Insufficient permissions for this organization"
        )

    async def event_generator():
        channel = f"project_{project_id}"
        pubsub = redis_client.pubsub()
        await pubsub.subscribe(channel)
        try:
            async for event in pubsub.listen():
                if event["type"] != "message":
                    continue
                if await request.is_disconnected():
                    break
                message = json.loads(event["data"])
                if message.get("client_id") == client_id:
                    continue
                yield f"data: {json.dumps({'event': message['event'], 'item': message['item']})}\n\n"
        finally:
            await pubsub.unsubscribe(channel)
            await pubsub.aclose()

    return StreamingResponse(event_generator(), media_type="text/event-stream")