xiaoxia-saas/apps/api/app/middleware/prometheus_metrics.py

"""
Prometheus metrics middleware and metrics endpoint for the API.

Exposes:
- HTTP request duration histogram (by method, path, status)
- HTTP requests total counter (by method, path, status)
- Requests in progress gauge
- Custom application metrics
"""

import re
import time
from collections.abc import Callable

from fastapi import Request, Response
from prometheus_client import (
    CONTENT_TYPE_LATEST,
    REGISTRY,
    Counter,
    Gauge,
    Histogram,
    generate_latest,
)
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import PlainTextResponse

# Buckets for HTTP request duration (seconds)
HTTP_DURATION_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0)

# --- Metrics Definitions ---

HTTP_REQUESTS_TOTAL = Counter(
    "http_requests_total",
    "Total HTTP requests",
    ["method", "endpoint", "status"],
    registry=REGISTRY,
)

HTTP_REQUEST_DURATION_SECONDS = Histogram(
    "http_request_duration_seconds",
    "HTTP request duration in seconds",
    ["method", "endpoint", "status"],
    buckets=HTTP_DURATION_BUCKETS,
    registry=REGISTRY,
)

HTTP_REQUESTS_IN_PROGRESS = Gauge(
    "http_requests_in_progress",
    "Number of HTTP requests currently being processed",
    ["method"],
    registry=REGISTRY,
)

# --- Application Metrics ---

APP_INFO = Gauge(
    "app_info",
    "Application information",
    ["version", "environment"],
    registry=REGISTRY,
)


def _normalize_path(path: str) -> str:
    """
    Normalize request path to reduce cardinality.
    Replace UUIDs and numeric IDs with placeholders.
    """
    # Replace UUIDs
    path = re.sub(
        r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",
        ":uuid",
        path,
    )
    # Replace numeric IDs in path segments
    path = re.sub(r"/\d+(/|$)", "/:id\\1", path)
    return path


class PrometheusMetricsMiddleware(BaseHTTPMiddleware):
    """Middleware that records Prometheus metrics for every HTTP request."""

    def __init__(self, app, exclude_paths: list[str] | None = None):
        super().__init__(app)
        self.exclude_paths = set(exclude_paths or ["/metrics", "/health", "/ready", "/startup"])

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        path = request.url.path
        if path in self.exclude_paths:
            return await call_next(request)

        method = request.method
        normalized_path = _normalize_path(path)

        HTTP_REQUESTS_IN_PROGRESS.labels(method=method).inc()
        start_time = time.perf_counter()

        try:
            response = await call_next(request)
            status = str(response.status_code)
            return response
        except Exception:
            status = "500"
            raise
        finally:
            duration = time.perf_counter() - start_time

            HTTP_REQUESTS_TOTAL.labels(
                method=method,
                endpoint=normalized_path,
                status=status,
            ).inc()

            HTTP_REQUEST_DURATION_SECONDS.labels(
                method=method,
                endpoint=normalized_path,
                status=status,
            ).observe(duration)

            HTTP_REQUESTS_IN_PROGRESS.labels(method=method).dec()


async def metrics_endpoint(request: Request) -> PlainTextResponse:
    """FastAPI endpoint that returns Prometheus metrics in text format.

    需要 Bearer Token 认证，Token 通过 METRICS_AUTH_TOKEN 环境变量配置。
    """
    import os

    # Bearer Token 认证
    auth_token = os.getenv("METRICS_AUTH_TOKEN", "")
    if auth_token:
        auth_header = request.headers.get("Authorization", "")
        if not auth_header.startswith("Bearer ") or auth_header[7:] != auth_token:
            return PlainTextResponse(content="Unauthorized", status_code=401)

    version = os.getenv("APP_VERSION", "unknown")
    environment = os.getenv("APP_ENV", "unknown")
    APP_INFO.labels(version=version, environment=environment).set(1)

    metrics_output = generate_latest(REGISTRY)
    return PlainTextResponse(
        content=metrics_output,
        media_type=CONTENT_TYPE_LATEST,
    )