Enhance health check endpoint to verify database connectivity

CHANGES.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Enhance `/_mgmt/ping` endpoint to check if pgstac database is ready before responding ([#230](https://github.com/stac-utils/stac-fastapi-pgstac/pull/230))
+
 ## [5.0.2] - 2025-04-07
 
 ### Fixed

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

docs/api/stac_fastapi/pgstac/app.md

@@ -1 +1,68 @@
-::: stac_fastapi.pgstac.app
+# stac_fastapi.pgstac.app
+
+## Overview
+
+The `stac_fastapi.pgstac.app` module contains the main application configuration for the FastAPI-based STAC API that uses PgSTAC as the backend. This module defines how the application is constructed, which extensions are enabled, and how the API endpoints are registered.
+
+## Key Components
+
+### PgStacApi Class
+
+```python
+class PgStacApi(StacApi)
+```
+
+Extended version of the `StacApi` class that provides PgSTAC-specific functionality.
+
+#### Methods
+
+##### add_health_check
+
+```python
+def add_health_check(self)
+```
+
+Adds a health check endpoint at `/_mgmt/ping` that verifies database connectivity.
+
+- The endpoint attempts to establish a database connection using the application's read connection pool.
+- It verifies PgSTAC is properly set up by querying the `pgstac.migrations` table.
+- Returns:
+  - Status 200 with `{"message": "PONG", "database": "OK"}` when the database is healthy.
+  - Status 503 with error details when the database cannot be reached or when PgSTAC is not properly set up.
+
+### Application Creation
+
+The module defines several key components for the FastAPI application:
+
+1. **Settings**: Configuration settings for the application.
+2. **Extensions**: Various STAC API extensions that are enabled.
+3. **Search Models**: Request/response models for search endpoints.
+4. **Database Connection**: Configuration for connecting to PostgreSQL with PgSTAC.
+
+### Lifespan
+
+```python
+@asynccontextmanager
+async def lifespan(app: FastAPI)
+```
+
+Manages the application lifespan:
+- Connects to the database during startup
+- Closes database connections during shutdown
+
+## Usage
+
+The module creates a FastAPI application with a PgSTAC backend:
+
+```python
+api = PgStacApi(
+    app=FastAPI(...),
+    settings=settings,
+    extensions=application_extensions,
+    client=CoreCrudClient(pgstac_search_model=post_request_model),
+    ...
+)
+app = api.app
+```
+
+The application can be run directly with uvicorn or used as a Lambda handler.

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 APIRouter, FastAPI, Request
+from fastapi.responses import JSONResponse, ORJSONResponse
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.middleware import CORSMiddleware, ProxyHeaderMiddleware
 from stac_fastapi.api.models import (
@@ -160,7 +160,47 @@ 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,6 @@
 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.api.models import create_get_request_model, create_post_request_model
 from stac_fastapi.extensions.core import (
     CollectionSearchExtension,
@@ -20,6 +19,7 @@
 from stac_fastapi.extensions.core.fields import FieldsConformanceClasses
 from stac_fastapi.types import stac as stac_types
 
+from stac_fastapi.pgstac.app import PgStacApi
 from stac_fastapi.pgstac.config import PostgresSettings
 from stac_fastapi.pgstac.core import CoreCrudClient, Settings
 from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
@@ -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,6 @@
 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.api.models import (
     ItemCollectionUri,
     create_get_request_model,
@@ -41,6 +40,7 @@
 from stac_fastapi.extensions.third_party import BulkTransactionExtension
 from stac_pydantic import Collection, Item
 
+from stac_fastapi.pgstac.app import PgStacApi
 from stac_fastapi.pgstac.config import PostgresSettings, Settings
 from stac_fastapi.pgstac.core import CoreCrudClient
 from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
@@ -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"