diff --git a/src/infrastructure/metrics.py b/src/infrastructure/metrics.py new file mode 100644 index 0000000..c1e09a1 --- /dev/null +++ b/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', + ) diff --git a/src/main.py b/src/main.py index 00f5264..d7e27de 100644 --- a/src/main.py +++ b/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 {