diff --git a/src/main.py b/src/main.py index 2169745..6fb1587 100644 --- a/src/main.py +++ b/src/main.py @@ -53,9 +53,15 @@ api = FastAPI( ) @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 def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str: @@ -78,6 +84,38 @@ async def custom_http_exception_handler(request,exc): content={"detail": exc.detail} ) +@api.get("/health", tags=["Health"]) +def return_health(request:Request, db=Depends(get_db)): + try: + cursor = db.cursor() + cursor.execute("SELECT 1") + cursor.fetchone() + db_status = "ok" + except Exception as e: + logger.error(f"Health check DB failed: {e}") + db_status = "error" + + try: + rmq_conn = getattr(request.app.state, "rmq_connection", None) + if not rmq_conn or not rmq_conn.is_open: + logger.error("Health check RMQ failed: connection closed or missing") + rmq_status = "error" + except Exception as e: + logger.error(f"Health check RMQ failed: {e}") + rmq_status = "error" + + overall_status = "ok" if db_status == "ok" and rmq_status == "ok" else "error" + status_code = 200 if overall_status == "ok" else 500 + + return JSONResponse( + status_code=status_code, + content={"status": overall_status, + "components": { + "database": db_status, + "rabbitmq": rmq_status + }, + "message": "Service is running" if overall_status == "ok" else "One or more checks failed"} + ) @api.post("/internal/receive-notifications") def receive_notifications(