bdlg2023-sms-back/app/api/SSEBroadcasterRoute.py

import asyncio
import json

from fastapi import Request, Response
from fastapi.routing import APIRoute

from app.core.events import broadcast


class SSEBroadcasterRoute(APIRoute):
    def get_route_handler(self):
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            # 1. Execute the route and let FastAPI serialize the response
            response = await original_route_handler(request)

            openapi_extra = self.openapi_extra or {}
            event_name = openapi_extra.get("sse_event")

            # 2. Check if an event is declared and the request was successful
            if event_name and response.status_code in (200, 201):
                project_id = request.path_params.get("project_id")
                client_id = request.headers.get("X-Client-ID", "unknown")

                if request.method == "DELETE":
                    # For DELETE, the response body is often empty. We broadcast the path parameters.
                    item_json = json.dumps(request.path_params)
                else:
                    # For POST/PUT, response.body is already serialized as a JSON string by FastAPI
                    item_json = response.body.decode("utf-8")

                # Assemble the raw payload string
                payload_str = (
                    f'{{"client_id": "{client_id}", "event": "{event_name}", "item": {item_json}}}'
                )

                channel = f"project_{project_id}"
                asyncio.create_task(broadcast.publish(channel, payload_str))

            return response

        return custom_route_handler