mqtt-ingester/api/main.py

"""FastAPI service exposing telemetry data from TimescaleDB over HTTPS.

Requires mutual TLS: clients must present a certificate signed by the CA.
Requests are served from a pool of database connections, so each request
reuses an established session instead of paying connect/auth costs.
"""

import os
from contextlib import asynccontextmanager
from typing import Any, Final, Optional

import psycopg
from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.responses import JSONResponse
from psycopg import sql
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from shared.access import COOKIE_NAME, Identity, hash_token, load_identity, topic_conditions

_SELECT_SQL: Final = "SELECT time, topic, device_id, payload FROM telemetry"

# Connections use row_factory=dict_row, so each row is a dict.
_DictConnection = psycopg.Connection[dict[str, Any]]

pool: Optional[ConnectionPool[_DictConnection]] = None


def _connect_kwargs() -> dict[str, Any]:
    return {
        "host": os.getenv("DB_HOST", "timescaledb"),
        "port": int(os.getenv("DB_PORT", "5432")),
        "dbname": os.getenv("DB_NAME", "telemetry"),
        "user": os.getenv("DB_USER", "ingester"),
        "password": os.getenv("DB_PASSWORD", "changeme"),
        "row_factory": dict_row,
    }


@asynccontextmanager
async def lifespan(app: FastAPI):
    global pool
    # Subscripting solves the pool's invariant connection type; row_factory
    # in kwargs makes every pooled connection return dict rows.
    server_pool = ConnectionPool[_DictConnection](
        kwargs=_connect_kwargs(),
        min_size=int(os.getenv("DB_POOL_MIN_SIZE", "2")),
        max_size=int(os.getenv("DB_POOL_MAX_SIZE", "10")),
    )
    server_pool.open(wait=False)
    pool = server_pool
    try:
        yield
    finally:
        pool = None
        server_pool.close()


app = FastAPI(title="MQTT Telemetry API", lifespan=lifespan)


def get_pool() -> ConnectionPool[_DictConnection]:
    if pool is None:
        raise RuntimeError("Database pool not initialized")
    return pool


def _request_identity(request: Request) -> Optional[Identity]:
    """Resolve the request principal: mTLS cert, session cookie, or None."""
    # uvicorn populates the client cert into the transport when mTLS is enabled.
    cert = getattr(request.scope.get("transport"), "get_extra_info", lambda _: None)(
        "peercert"
    )
    if cert is not None:
        return Identity.service(cert)
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return None
    conn = get_pool().getconn()
    try:
        return load_identity(conn, hash_token(token))
    finally:
        get_pool().putconn(conn)


@app.middleware("http")
async def require_client_cert(request: Request, call_next):
    # Return a response directly: HTTPException raised from middleware is
    # not handled by the (inner) exception handlers and would surface
    # as a 500.
    identity = _request_identity(request)
    if identity is None:
        return JSONResponse(status_code=401, content={"detail": "Authentication required"})
    request.state.identity = identity
    return await call_next(request)


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/whoami")
def whoami(request: Request) -> dict[str, Any]:
    """Show the authenticated principal and effective access rules.

    Useful to verify which auth path a request took (client certificate
    vs session cookie) and which topics the effective rules allow.
    """
    identity: Identity = request.state.identity
    return {
        "username": identity.username,
        "role": identity.role,
        "rules": {"allow": list(identity.rules.allow), "deny": list(identity.rules.deny)},
    }


def _telemetry_query(
    identity: Identity,
    topic: Optional[str],
    device_id: Optional[str],
    limit: int,
) -> tuple[sql.Composed, list[Any]]:
    """Build the SELECT for the requested filters as composed SQL.

    The statement is composed with psycopg.sql, and all values are bound
    as parameters, so nothing is interpolated into the query string.
    Access rules are folded into the WHERE clause (deny wins over allow),
    so filtering happens in the database and LIMIT stays correct.
    """
    conditions, param_groups = topic_conditions(identity)
    params: list[Any] = []
    if topic:
        conditions.append(sql.SQL("topic = %s"))
        params.append(topic)
    if device_id:
        conditions.append(sql.SQL("device_id = %s"))
        params.append(device_id)
    params.extend(p for group in param_groups for p in group)
    query = sql.SQL(_SELECT_SQL)
    if conditions:
        query += sql.SQL(" WHERE ") + sql.SQL(" AND ").join(conditions)
    query += sql.SQL(" ORDER BY time DESC LIMIT %s")
    params.append(limit)
    return query, params


@app.get("/telemetry")
def list_telemetry(
    request: Request,
    topic: Optional[str] = Query(default=None),
    device_id: Optional[str] = Query(default=None),
    limit: int = Query(default=100, le=1000),
) -> list[dict[str, Any]]:
    identity: Identity = request.state.identity
    query, params = _telemetry_query(identity, topic, device_id, limit)
    with get_pool().connection() as conn, conn.cursor() as cur:
        cur.execute(query, params)
        return cur.fetchall()


@app.get("/telemetry/latest")
def latest_telemetry(
    request: Request,
    topic: Optional[str] = Query(default=None),
    device_id: Optional[str] = Query(default=None),
) -> Optional[dict[str, Any]]:
    identity: Identity = request.state.identity
    query, params = _telemetry_query(identity, topic, device_id, 1)
    with get_pool().connection() as conn, conn.cursor() as cur:
        cur.execute(query, params)
        row = cur.fetchone()
        if row is None:
            raise HTTPException(status_code=404, detail="No data found")
        return row