Add health checks

Author: chrisburr

diracx-routers/pyproject.toml

@@ -41,10 +41,11 @@ types = [
 ]
 
 [project.entry-points."diracx.services"]
-jobs = "diracx.routers.jobs:router"
-config = "diracx.routers.configuration:router"
-auth = "diracx.routers.auth:router"
 ".well-known" = "diracx.routers.auth.well_known:router"
+auth = "diracx.routers.auth:router"
+config = "diracx.routers.configuration:router"
+health = "diracx.routers.health:router"
+jobs = "diracx.routers.jobs:router"
 
 [project.entry-points."diracx.access_policies"]
 WMSAccessPolicy = "diracx.routers.jobs.access_policies:WMSAccessPolicy"

diracx-routers/src/diracx/routers/fastapi_classes.py

@@ -82,9 +82,10 @@ def __init__(
         *,
         dependencies=None,
         require_auth: bool = True,
+        include_in_schema: bool = True,
         path_root: str = "/api",
     ):
-        super().__init__(dependencies=dependencies)
+        super().__init__(dependencies=dependencies, include_in_schema=include_in_schema)
         self.diracx_require_auth = require_auth
         self.diracx_path_root = path_root
 

diracx-routers/src/diracx/routers/health/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__all__ = ["router"]
+
+from ..fastapi_classes import DiracxRouter
+from .probes import router as probes_router
+
+router = DiracxRouter(require_auth=False, include_in_schema=False)
+router.include_router(probes_router)

diracx-routers/src/diracx/routers/health/probes.py

@@ -0,0 +1,51 @@
+"""Health probes for use in Kubernetes and other orchestration systems."""
+
+from __future__ import annotations
+
+__all__ = ["router"]
+
+from fastapi import HTTPException
+from starlette.responses import JSONResponse
+
+from ..dependencies import AuthDB, Config
+from ..fastapi_classes import DiracxRouter
+
+router = DiracxRouter(require_auth=False)
+
+
+@router.get("/live", include_in_schema=False)
+async def liveness():
+    """Returns a simple status to indicate the app is running."""
+    return JSONResponse(content={"status": "live"})
+
+
+@router.get("/ready", include_in_schema=False)
+async def readiness(config: Config, auth_db: AuthDB):
+    """Readiness endpoint.
+
+    Checks if at least the configuration is loaded and the AuthDB database
+    connection is available.
+    """
+    if not any(vo_registry.Users for vo_registry in config.Registry.values()):
+        raise HTTPException(status_code=503, detail="No users in registry")
+    try:
+        await auth_db.ping()
+    except Exception as e:
+        raise HTTPException(status_code=503, detail="AuthDB ping failed") from e
+    return JSONResponse(content={"status": "ready"})
+
+
+@router.get("/startup", include_in_schema=False)
+async def startup_check(config: Config, auth_db: AuthDB):
+    """Startup endpoint.
+
+    Checks if at least the configuration is loaded and the AuthDB database
+    connection is available.
+    """
+    if not any(vo_registry.Users for vo_registry in config.Registry.values()):
+        raise HTTPException(status_code=503, detail="No users in registry")
+    try:
+        await auth_db.ping()
+    except Exception as e:
+        raise HTTPException(status_code=503, detail="AuthDB ping failed") from e
+    return JSONResponse(content={"status": "startup complete"})