Add auth Prometheus metrics

src/infrastructure/metrics.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import time
+from collections import defaultdict
+from datetime import timedelta
+from typing import Iterable
+
+from fastapi import Request, Response
+from sqlalchemy import distinct, func, select
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.responses import Response as StarletteResponse
+
+from src.infrastructure.database.context import async_session_maker
+from src.infrastructure.database.models.sessions import Session
+from src.infrastructure.database.models.user import UserModel
+
+
+SERVICE_NAME = 'auth'
+WINDOWS: tuple[tuple[str, timedelta | None], ...] = (
+    ('all', None),
+    ('30d', timedelta(days=30)),
+    ('7d', timedelta(days=7)),
+    ('24h', timedelta(days=1)),
+    ('1h', timedelta(hours=1)),
+)
+LATENCY_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0)
+
+_http_requests: dict[tuple[str, str, str], int] = defaultdict(int)
+_http_latency_sum: dict[tuple[str, str], float] = defaultdict(float)
+_http_latency_count: dict[tuple[str, str], int] = defaultdict(int)
+_http_latency_buckets: dict[tuple[str, str, float], int] = defaultdict(int)
+
+
+def _format_labels(labels: dict[str, str]) -> str:
+    return ','.join(f'{key}="{value}"' for key, value in labels.items())
+
+
+def _sample(name: str, value: int | float, labels: dict[str, str] | None = None) -> str:
+    if not labels:
+        return f'{name} {value}'
+    return f'{name}{{{_format_labels(labels)}}} {value}'
+
+
+def _route_template(request: Request) -> str:
+    route = request.scope.get('route')
+    path = getattr(route, 'path', None)
+    if path:
+        return str(path)
+    return request.url.path
+
+
+class PrometheusMetricsMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> StarletteResponse:
+        if request.url.path == '/metrics':
+            return await call_next(request)
+
+        started = time.perf_counter()
+        status_code = '500'
+        route = request.url.path
+
+        try:
+            response = await call_next(request)
+            status_code = str(response.status_code)
+            return response
+        finally:
+            elapsed = time.perf_counter() - started
+            route = _route_template(request) or route
+            method = request.method
+
+            _http_requests[(method, route, status_code)] += 1
+            _http_latency_sum[(method, route)] += elapsed
+            _http_latency_count[(method, route)] += 1
+            for bucket in LATENCY_BUCKETS:
+                if elapsed <= bucket:
+                    _http_latency_buckets[(method, route, bucket)] += 1
+
+
+async def _count_scalar(statement) -> int:
+    async with async_session_maker() as session:
+        result = await session.execute(statement)
+        return int(result.scalar_one() or 0)
+
+
+async def _collect_user_metrics() -> list[str]:
+    lines: list[str] = []
+
+    users_total = await _count_scalar(select(func.count()).select_from(UserModel))
+    lines.append(_sample('auth_users_total', users_total, {'service': SERVICE_NAME}))
+
+    active_sessions = await _count_scalar(
+        select(func.count())
+        .select_from(Session)
+        .where(Session.revoked_at.is_(None))
+    )
+    lines.append(_sample('auth_sessions_active_total', active_sessions, {'service': SERVICE_NAME}))
+
+    for window_name, window_delta in WINDOWS:
+        created_conditions = []
+        active_conditions = [Session.revoked_at.is_(None)]
+        if window_delta is not None:
+            created_conditions.append(UserModel.created_at >= func.now() - window_delta)
+            active_conditions.append(Session.last_seen_at >= func.now() - window_delta)
+
+        new_users = await _count_scalar(
+            select(func.count())
+            .select_from(UserModel)
+            .where(*created_conditions)
+        )
+        active_users = await _count_scalar(
+            select(func.count(distinct(Session.user_id)))
+            .select_from(Session)
+            .where(*active_conditions)
+        )
+
+        labels = {'service': SERVICE_NAME, 'window': window_name}
+        lines.append(_sample('auth_user_registrations_total', new_users, labels))
+        lines.append(_sample('auth_active_unique_users_total', active_users, labels))
+
+    return lines
+
+
+def _collect_http_metrics() -> Iterable[str]:
+    yield '# TYPE auth_http_requests_total counter'
+    for (method, route, status_code), value in sorted(_http_requests.items()):
+        yield _sample(
+            'auth_http_requests_total',
+            value,
+            {
+                'service': SERVICE_NAME,
+                'method': method,
+                'route': route,
+                'status_code': status_code,
+            },
+        )
+
+    yield '# TYPE auth_http_request_duration_seconds histogram'
+    route_keys = sorted(_http_latency_count)
+    for method, route in route_keys:
+        cumulative = 0
+        for bucket in LATENCY_BUCKETS:
+            cumulative = _http_latency_buckets.get((method, route, bucket), cumulative)
+            yield _sample(
+                'auth_http_request_duration_seconds_bucket',
+                cumulative,
+                {
+                    'service': SERVICE_NAME,
+                    'method': method,
+                    'route': route,
+                    'le': str(bucket),
+                },
+            )
+        yield _sample(
+            'auth_http_request_duration_seconds_bucket',
+            _http_latency_count[(method, route)],
+            {
+                'service': SERVICE_NAME,
+                'method': method,
+                'route': route,
+                'le': '+Inf',
+            },
+        )
+        yield _sample(
+            'auth_http_request_duration_seconds_sum',
+            _http_latency_sum[(method, route)],
+            {'service': SERVICE_NAME, 'method': method, 'route': route},
+        )
+        yield _sample(
+            'auth_http_request_duration_seconds_count',
+            _http_latency_count[(method, route)],
+            {'service': SERVICE_NAME, 'method': method, 'route': route},
+        )
+
+
+async def metrics_response() -> Response:
+    lines = [
+        '# HELP auth_users_total Total registered users in auth database.',
+        '# TYPE auth_users_total gauge',
+        '# HELP auth_user_registrations_total Registered users counted from the database for the selected rolling window.',
+        '# TYPE auth_user_registrations_total gauge',
+        '# HELP auth_active_unique_users_total Unique users with non-revoked sessions seen in the selected rolling window.',
+        '# TYPE auth_active_unique_users_total gauge',
+        '# HELP auth_sessions_active_total Non-revoked auth sessions.',
+        '# TYPE auth_sessions_active_total gauge',
+    ]
+
+    try:
+        lines.extend(await _collect_user_metrics())
+        lines.append(_sample('auth_metrics_scrape_success', 1, {'service': SERVICE_NAME}))
+    except Exception:
+        lines.append(_sample('auth_metrics_scrape_success', 0, {'service': SERVICE_NAME}))
+
+    lines.extend(_collect_http_metrics())
+    return Response(
+        content='\n'.join(lines) + '\n',
+        media_type='text/plain; version=0.0.4; charset=utf-8',
+    )

src/main.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from contextlib import asynccontextmanager
 import secrets
 from typing import AsyncGenerator
-from fastapi import Depends, FastAPI, status
+from fastapi import Depends, FastAPI, Response, status
 from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import HTMLResponse
@@ -14,6 +14,7 @@ from src.infrastructure.config.settings import get_settings
 from src.infrastructure.vault import JwtKeyStore, start_jwt_keys_scheduler
 from src.infrastructure.utils import generate_instance_id
 from src.infrastructure.logger import logger
+from src.infrastructure.metrics import PrometheusMetricsMiddleware, metrics_response
 from src.infrastructure.config import settings
 from src.presentation.dependencies import get_rabbit
 from src.presentation.handler import application_exception_handler
@@ -157,6 +158,8 @@ app: FastAPI = FastAPI(
 app.add_exception_handler(ApplicationException, application_exception_handler)
 app.add_exception_handler(Exception, unhandled_exception_handler)
 
+app.add_middleware(PrometheusMetricsMiddleware)
+
 app.include_router(v1_router)
 
 app.add_middleware(TraceIDMiddleware, logger=logger)
@@ -200,6 +203,11 @@ async def custom_redoc_html(_credentials: HTTPBasicCredentials = Depends(verify_
     )
 
 
+@app.get('/metrics', include_in_schema=False)
+async def metrics() -> Response:
+    return await metrics_response()
+
+
 @app.post('/ping')
 async def ping() -> dict[str, str]:
     return {