From 4d586f437194e4b876bac06a520fcbb85902040e Mon Sep 17 00:00:00 2001 From: florian Date: Fri, 10 Oct 2025 11:41:36 +0200 Subject: [PATCH] Added /health check for Prometheus and added failure capturing incase something went wrong, previous approach would not increase the counter --- src/main.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/src/main.py b/src/main.py index 2bcf3ea..1b8044d 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, JSONResponse from fastapi.responses import JSONResponse from fastapi.security.api_key import APIKeyHeader from starlette.exceptions import HTTPException as StarletteHTTPException @@ -53,12 +53,36 @@ 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() +async def prometheus_middleware(request: Request, call_next): + status = 500 + try: + response = await call_next(request) + status = response.status_code + except Exception: + raise + finally: + REQUEST_COUNTER.labels(request.method, request.url.path, status).inc() return response +@api.get("/health", tags=["Health"]) +async def return_health(db=Depends(get_db)): + try: + cursor = db.cursor() + cursor.execute("SELECT 1") + cursor.fetchone() + db_status = "ok" + except Exception: + logger.error(f"Health check DB failed: {e}") + db_status = "error" + + return JSONResponse( + status_code=200 if db_status == "ok" else 500, + content={"status": db_status, "message": "Service is running"} + ) + + def verify_api_key_dependency(db=Depends(get_db), api_key: str = Depends(api_key_header)) -> int: cursor = db.cursor() cursor.execute("SELECT user_id, api_key FROM users WHERE status = 'active'")