bdlg2023-sms-back/app/main.py

"""Main FastAPI app instance declaration."""

import typing
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.openapi.utils import get_openapi

from app.api.api import api_router
from app.core import config
from app.core.events import broadcast


@asynccontextmanager
async def lifespan(app: FastAPI):
    await broadcast.connect()
    yield
    await broadcast.disconnect()


app = FastAPI(
    title=config.settings.PROJECT_NAME,
    version=config.settings.VERSION,
    description=config.settings.DESCRIPTION,
    openapi_url="/openapi.json",
    docs_url="/",
    lifespan=lifespan,
)
app.include_router(api_router)

# Sets all CORS enabled origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=[str(origin) for origin in config.settings.BACKEND_CORS_ORIGINS],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Guards against HTTP Host Header attacks
if config.settings.ENVIRONMENT != "PYTEST":
    app.add_middleware(TrustedHostMiddleware, allowed_hosts=config.settings.ALLOWED_HOSTS)


def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema

    openapi_schema = get_openapi(
        title=config.settings.PROJECT_NAME,
        version=config.settings.VERSION,
        description=config.settings.DESCRIPTION,
        routes=app.routes,
    )

    one_of_schemas = []

    for route in app.routes:
        if hasattr(route, "openapi_extra") and route.openapi_extra:
            event_name = route.openapi_extra.get("sse_event")
            if not event_name:
                continue

            method = list(route.methods)[0] if route.methods else "GET"
            payload_schema = {}

            # Determine payload schema for DELETE (Path params)
            model = getattr(route, "response_model", None)
            if method == "DELETE" and model is None:
                props = {}
                for param in getattr(route.dependant, "path_params", []):
                    props[param.name] = {"type": "string"}
                payload_schema = {"type": "object", "properties": props}

            # Determine payload schema for POST/PUT (Response Model)
            else:
                origin = typing.get_origin(model)
                if origin is list:
                    item_model = typing.get_args(model)[0]
                    if hasattr(item_model, "__name__"):
                        model_name = item_model.__name__
                        payload_schema = {
                            "type": "array",
                            "items": {"$ref": f"#/components/schemas/{model_name}"},
                        }
                elif hasattr(model, "__name__"):
                    model_name = model.__name__
                    payload_schema = {"$ref": f"#/components/schemas/{model_name}"}

            # Build the discriminated union object with traceability in the description
            event_schema = {
                "type": "object",
                "description": f"Triggered by `{method}` `{route.path}`",
                "properties": {
                    "event": {"type": "string", "enum": [event_name]},
                    "data": payload_schema,
                },
                "required": ["event", "data"],
            }
            one_of_schemas.append(event_schema)

    if one_of_schemas:
        if "components" not in openapi_schema:
            openapi_schema["components"] = {"schemas": {}}
        elif "schemas" not in openapi_schema["components"]:
            openapi_schema["components"]["schemas"] = {}

        openapi_schema["components"]["schemas"]["SseEventPayload"] = {
            "title": "SseEventPayload",
            "description": "Discriminated union of all possible SSE events and their payloads.",
            "oneOf": one_of_schemas,
        }

        for path, path_item in openapi_schema["paths"].items():
            if path.endswith("/stream") and "get" in path_item:
                path_item["get"]["responses"]["200"]["content"] = {
                    "text/event-stream": {
                        "schema": {"$ref": "#/components/schemas/SseEventPayload"}
                    }
                }

    app.openapi_schema = openapi_schema
    return app.openapi_schema


# 5. Override the default OpenAPI method
app.openapi = custom_openapi