diff --git a/requirements.txt b/requirements.txt index e514438..177a42f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,6 +9,7 @@ h11==0.16.0 idna==3.10 mysql-connector-python==9.4.0 pika==1.3.2 +prometheus_client==0.23.1 pycparser==2.23 pydantic==2.12.0 pydantic_core==2.41.1 diff --git a/src/main.py b/src/main.py index 347dafc..2169745 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,4 @@ -from fastapi import FastAPI, Depends, HTTPException, Request +from fastapi import FastAPI, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from fastapi.security.api_key import APIKeyHeader from starlette.exceptions import HTTPException as StarletteHTTPException @@ -11,6 +11,8 @@ from rabbitmq_handler import send_message_to_rmq, create_connection, close_conne import uvicorn from uvicorn_logging_config import LOGGING_CONFIG from contextlib import asynccontextmanager +from metrics_server import REQUEST_COUNTER +import asyncio logger = setup_logger(__name__) @@ -50,6 +52,12 @@ api = FastAPI( lifespan=lifespan ) +@api.middleware("http") +async def prometheus_middleware(request, call_next): + response = await call_next(request) + REQUEST_COUNTER.labels(request.method, request.url.path, response.status_code).inc() + return response + def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str: cursor = db.cursor() cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'") @@ -84,11 +92,14 @@ def receive_notifications( logger.info("Successfully delivered message to RMQ") return {"status": "queued"} +async def start_servers(): + config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8101, log_config=LOGGING_CONFIG, log_level="info") + config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level="info") + + server_main = uvicorn.Server(config_main) + server_metrics = uvicorn.Server(config_metrics) + + await asyncio.gather(server_main.serve(), server_metrics.serve()) + if __name__ == "__main__": - uvicorn.run( - "main:api", - host="0.0.0.0", - port=8101, - log_config=LOGGING_CONFIG, - log_level="info" - ) + asyncio.run(start_servers()) diff --git a/src/metrics_server.py b/src/metrics_server.py new file mode 100644 index 0000000..c3fdc9b --- /dev/null +++ b/src/metrics_server.py @@ -0,0 +1,10 @@ +from fastapi import FastAPI, Response +from prometheus_client import generate_latest, CONTENT_TYPE_LATEST, Counter + +metrics_api = FastAPI(title="Metrics Server", description="Prometheus metrics endpoint") + +REQUEST_COUNTER = Counter("http_requests_total", "Total HTTP Requests", ["method", "endpoint", "status"]) + +@metrics_api.get("/metrics") +async def metrics(): + return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)