Enhance health check endpoint to verify database connectivity

- Introduced PgStacApi class to extend StacApi with enhanced health checks.
- Updated /_mgmt/ping endpoint to check database readiness.
- Modified tests to validate new response format and database status.

Author: emmanuelmathot

Makefile

@@ -32,6 +32,10 @@ docker-shell:
 test:
 	$(runtests) /bin/bash -c 'export && python -m pytest /app/tests/api/test_api.py --log-cli-level $(LOG_LEVEL)'
 
+.PHONY: test-mgmt
+test-mgmt:
+	$(runtests) /bin/bash -c 'export && python -m pytest /app/tests/resources/test_mgmt.py --log-cli-level $(LOG_LEVEL)'
+
 .PHONY: run-database
 run-database:
 	docker compose run --rm database

pr_description.md

@@ -0,0 +1,52 @@
+# Enhance /_mgmt/ping endpoint to check if pgstac is ready
+
+## Overview
+This PR enhances the `/_mgmt/ping` health check endpoint to actually verify database connectivity before returning a positive response. The current implementation always returns `{"message": "PONG"}` regardless of whether the database is actually accessible, which can lead to misleading health checks in production environments.
+
+## Changes Made
+- Created a `PgStacApi` class in `stac_fastapi/pgstac/app.py` that extends the base `StacApi` class
+- Overrode the `add_health_check` method to implement database connectivity checks
+- Modified the ping endpoint to:
+  - Test database connectivity by attempting to connect to the read pool
+  - Verify pgstac is properly set up by querying the `pgstac.migrations` table
+  - Return appropriate status codes and error messages when the database is unavailable
+- Updated the test in `tests/resources/test_mgmt.py` to verify the new response format
+
+## Implementation Details
+The enhanced endpoint now:
+- Returns `{"message": "PONG", "database": "OK"}` with status 200 when the database is healthy
+- Returns a 503 Service Unavailable with descriptive error message when the database cannot be reached or pgstac is not properly set up
+
+## Why This Is Important
+- Provides more accurate health/readiness checks in containerized environments
+- Better integration with container orchestration systems like Kubernetes
+- Faster detection of database connectivity issues
+- Helps operators quickly identify when database connectivity is the root cause of issues
+
+## Testing
+The implementation can be tested by:
+
+1. Running the API with a working database connection:
+```bash
+# Start with a working database
+docker-compose up -d
+# The endpoint should return status 200
+curl -v http://localhost:8080/_mgmt/ping
+```
+
+2. Testing with a non-functioning database:
+```bash
+# Stop the database
+docker-compose stop pgstac
+# The endpoint should now return a 503 error
+curl -v http://localhost:8080/_mgmt/ping
+```
+
+## Alternatives Considered
+I considered two alternative approaches:
+
+1. **Using a simple connection check without querying any tables** - This would only verify the database is running but not that pgstac is properly set up.
+
+2. **Implementing a more extensive set of health checks** - While more comprehensive, this would add complexity and potential performance overhead to the health check endpoint.
+
+The current implementation provides a good balance between thoroughness and simplicity.

stac_fastapi/pgstac/app.py

@@ -9,8 +9,8 @@
 from contextlib import asynccontextmanager
 
 from brotli_asgi import BrotliMiddleware
-from fastapi import FastAPI
-from fastapi.responses import ORJSONResponse
+from fastapi import FastAPI, Request, APIRouter
+from fastapi.responses import ORJSONResponse, JSONResponse
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.middleware import CORSMiddleware, ProxyHeaderMiddleware
 from stac_fastapi.api.models import (
@@ -160,7 +160,39 @@ async def lifespan(app: FastAPI):
     await close_db_connection(app)
 
 
-api = StacApi(
+class PgStacApi(StacApi):
+    """PgStac API with enhanced health checks."""
+
+    def add_health_check(self):
+        """Add a health check with pgstac database readiness check."""
+        mgmt_router = APIRouter(prefix=self.app.state.router_prefix)
+
+        @mgmt_router.get("/_mgmt/ping")
+        async def ping(request: Request):
+            """Liveliness/readiness probe that checks database connection."""
+            try:
+                # Test read connection
+                async with request.app.state.get_connection(request, "r") as conn:
+                    # Execute a simple query to verify pgstac is ready
+                    # Check if we can query the migrations table which should exist in pgstac
+                    result = await conn.fetchval("SELECT 1 FROM pgstac.migrations LIMIT 1")
+                    if result is not None:
+                        return {"message": "PONG", "database": "OK"}
+                    else:
+                        return JSONResponse(
+                            status_code=503,
+                            content={"message": "Database tables not found", "database": "ERROR"}
+                        )
+            except Exception as e:
+                return JSONResponse(
+                    status_code=503,
+                    content={"message": f"Database connection failed: {str(e)}", "database": "ERROR"}
+                )
+
+        self.app.include_router(mgmt_router, tags=["Liveliness/Readiness"])
+
+
+api = PgStacApi(
     app=FastAPI(
         openapi_url=settings.openapi_url,
         docs_url=settings.docs_url,

tests/api/test_api.py

@@ -10,7 +10,7 @@
 from pypgstac.db import PgstacDB
 from pypgstac.load import Loader
 from pystac import Collection, Extent, Item, SpatialExtent, TemporalExtent
-from stac_fastapi.api.app import StacApi
+from stac_fastapi.pgstac.app import PgStacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
 from stac_fastapi.extensions.core import (
     CollectionSearchExtension,
@@ -748,7 +748,7 @@ async def get_collection(
         ]
     )
 
-    api = StacApi(
+    api = PgStacApi(
         client=Client(pgstac_search_model=post_request_model),
         settings=settings,
         extensions=extensions,
@@ -806,7 +806,7 @@ async def test_no_extension(
     )
     extensions = []
     post_request_model = create_post_request_model(extensions, base_model=PgstacSearch)
-    api = StacApi(
+    api = PgStacApi(
         client=CoreCrudClient(pgstac_search_model=post_request_model),
         settings=settings,
         extensions=extensions,

tests/conftest.py

@@ -15,7 +15,7 @@
 from pypgstac.db import PgstacDB
 from pypgstac.migrate import Migrate
 from pytest_postgresql.janitor import DatabaseJanitor
-from stac_fastapi.api.app import StacApi
+from stac_fastapi.pgstac.app import PgStacApi
 from stac_fastapi.api.models import (
     ItemCollectionUri,
     create_get_request_model,
@@ -181,7 +181,7 @@ def api_client(request):
         search_extensions, base_model=PgstacSearch
     )
 
-    api = StacApi(
+    api = PgStacApi(
         settings=api_settings,
         extensions=application_extensions,
         client=CoreCrudClient(pgstac_search_model=search_post_request_model),
@@ -296,7 +296,7 @@ def api_client_no_ext():
     api_settings = Settings(
         testing=True,
     )
-    return StacApi(
+    return PgStacApi(
         settings=api_settings,
         extensions=[
             TransactionExtension(client=TransactionsClient(), settings=api_settings)

tests/resources/test_mgmt.py

@@ -6,4 +6,7 @@ async def test_ping_no_param(app_client):
     """
     res = await app_client.get("/_mgmt/ping")
     assert res.status_code == 200
-    assert res.json() == {"message": "PONG"}
+    response_json = res.json()
+    assert response_json["message"] == "PONG"
+    assert "database" in response_json
+    assert response_json["database"] == "OK"