feat: add a metrics module (#67)

src/fastapi_toolsets/metrics/__init__.py

@@ -0,0 +1,10 @@
+"""Prometheus metrics integration for FastAPI applications."""
+
+from .handler import init_metrics
+from .registry import Metric, MetricsRegistry
+
+__all__ = [
+    "Metric",
+    "MetricsRegistry",
+    "init_metrics",
+]

src/fastapi_toolsets/metrics/handler.py

@@ -0,0 +1,73 @@
+"""Prometheus metrics endpoint for FastAPI applications."""
+
+import asyncio
+import os
+
+from fastapi import FastAPI
+from fastapi.responses import Response
+from prometheus_client import (
+    CONTENT_TYPE_LATEST,
+    CollectorRegistry,
+    generate_latest,
+    multiprocess,
+)
+
+from ..logger import get_logger
+from .registry import MetricsRegistry
+
+logger = get_logger()
+
+
+def _is_multiprocess() -> bool:
+    """Check if prometheus multi-process mode is enabled."""
+    return "PROMETHEUS_MULTIPROC_DIR" in os.environ
+
+
+def init_metrics(
+    app: FastAPI,
+    registry: MetricsRegistry,
+    *,
+    path: str = "/metrics",
+) -> FastAPI:
+    """Register a Prometheus ``/metrics`` endpoint on a FastAPI app.
+
+    Args:
+        app: FastAPI application instance.
+        registry: A :class:`MetricsRegistry` containing providers and collectors.
+        path: URL path for the metrics endpoint (default ``/metrics``).
+
+    Returns:
+        The same FastAPI instance (for chaining).
+
+    Example:
+        from fastapi import FastAPI
+        from fastapi_toolsets.metrics import MetricsRegistry, init_metrics
+
+        metrics = MetricsRegistry()
+        app = FastAPI()
+        init_metrics(app, registry=metrics)
+    """
+    for provider in registry.get_providers():
+        logger.debug("Initialising metric provider '%s'", provider.name)
+        provider.func()
+
+    collectors = registry.get_collectors()
+
+    @app.get(path, include_in_schema=False)
+    async def metrics_endpoint() -> Response:
+        for collector in collectors:
+            if asyncio.iscoroutinefunction(collector.func):
+                await collector.func()
+            else:
+                collector.func()
+
+        if _is_multiprocess():
+            prom_registry = CollectorRegistry()
+            multiprocess.MultiProcessCollector(prom_registry)
+            output = generate_latest(prom_registry)
+        else:
+            output = generate_latest()
+
+        return Response(content=output, media_type=CONTENT_TYPE_LATEST)
+
+    return app

src/fastapi_toolsets/metrics/registry.py

@@ -0,0 +1,122 @@
+"""Metrics registry with decorator-based registration."""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+from ..logger import get_logger
+
+logger = get_logger()
+
+
+@dataclass
+class Metric:
+    """A metric definition with metadata."""
+
+    name: str
+    func: Callable[..., Any]
+    collect: bool = field(default=False)
+
+
+class MetricsRegistry:
+    """Registry for managing Prometheus metric providers and collectors.
+
+    Example:
+        from prometheus_client import Counter, Gauge
+        from fastapi_toolsets.metrics import MetricsRegistry
+
+        metrics = MetricsRegistry()
+
+        @metrics.register
+        def http_requests():
+            return Counter("http_requests_total", "Total HTTP requests", ["method", "status"])
+
+        @metrics.register(name="db_pool")
+        def database_pool_size():
+            return Gauge("db_pool_size", "Database connection pool size")
+
+        @metrics.register(collect=True)
+        def collect_queue_depth(gauge=Gauge("queue_depth", "Current queue depth")):
+            gauge.set(get_current_queue_depth())
+    """
+
+    def __init__(self) -> None:
+        self._metrics: dict[str, Metric] = {}
+
+    def register(
+        self,
+        func: Callable[..., Any] | None = None,
+        *,
+        name: str | None = None,
+        collect: bool = False,
+    ) -> Callable[..., Any]:
+        """Register a metric provider or collector function.
+
+        Can be used as a decorator with or without arguments.
+
+        Args:
+            func: The metric function to register.
+            name: Metric name (defaults to function name).
+            collect: If ``True``, the function is called on every scrape.
+                If ``False`` (default), called once at init time.
+
+        Example:
+            @metrics.register
+            def my_counter():
+                return Counter("my_counter", "A counter")
+
+            @metrics.register(collect=True, name="queue")
+            def collect_queue_depth():
+                gauge.set(compute_depth())
+        """
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            metric_name = name or cast(Any, fn).__name__
+            self._metrics[metric_name] = Metric(
+                name=metric_name,
+                func=fn,
+                collect=collect,
+            )
+            return fn
+
+        if func is not None:
+            return decorator(func)
+        return decorator
+
+    def include_registry(self, registry: "MetricsRegistry") -> None:
+        """Include another :class:`MetricsRegistry` into this one.
+
+        Args:
+            registry: The registry to merge in.
+
+        Raises:
+            ValueError: If a metric name already exists in the current registry.
+
+        Example:
+            main = MetricsRegistry()
+            sub = MetricsRegistry()
+
+            @sub.register
+            def sub_metric():
+                return Counter("sub_total", "Sub counter")
+
+            main.include_registry(sub)
+        """
+        for metric_name, definition in registry._metrics.items():
+            if metric_name in self._metrics:
+                raise ValueError(
+                    f"Metric '{metric_name}' already exists in the current registry"
+                )
+            self._metrics[metric_name] = definition
+
+    def get_all(self) -> list[Metric]:
+        """Get all registered metric definitions."""
+        return list(self._metrics.values())
+
+    def get_providers(self) -> list[Metric]:
+        """Get metric providers (called once at init)."""
+        return [m for m in self._metrics.values() if not m.collect]
+
+    def get_collectors(self) -> list[Metric]:
+        """Get collectors (called on each scrape)."""
+        return [m for m in self._metrics.values() if m.collect]