From 7faa78c1a4ed0d8a2266e7e4e3fbbcd97e1e0ae8 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Wed, 16 Apr 2025 21:26:47 +0200 Subject: [PATCH 1/4] 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. --- Makefile | 4 +++ pr_description.md | 52 ++++++++++++++++++++++++++++++++++++ stac_fastapi/pgstac/app.py | 38 +++++++++++++++++++++++--- tests/api/test_api.py | 6 ++--- tests/conftest.py | 6 ++--- tests/resources/test_mgmt.py | 5 +++- 6 files changed, 101 insertions(+), 10 deletions(-) create mode 100644 pr_description.md diff --git a/Makefile b/Makefile index 57b56b6d..9a313fc2 100644 --- a/Makefile +++ b/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 diff --git a/pr_description.md b/pr_description.md new file mode 100644 index 00000000..b972989d --- /dev/null +++ b/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. diff --git a/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/app.py index fcc12889..824a6bfc 100644 --- a/stac_fastapi/pgstac/app.py +++ b/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, diff --git a/tests/api/test_api.py b/tests/api/test_api.py index e4ce8d06..e0ca8031 100644 --- a/tests/api/test_api.py +++ b/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, diff --git a/tests/conftest.py b/tests/conftest.py index 7944e8de..94119e60 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/resources/test_mgmt.py b/tests/resources/test_mgmt.py index 9d2bc3dc..0a80831b 100644 --- a/tests/resources/test_mgmt.py +++ b/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" From ec1326630a617e26af615e883d60e7c8584ed9f4 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Wed, 16 Apr 2025 21:29:08 +0200 Subject: [PATCH 2/4] Refactor import order and improve database connection error handling --- stac_fastapi/pgstac/app.py | 18 +++++++++++++----- tests/api/test_api.py | 2 +- tests/conftest.py | 2 +- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/app.py index 824a6bfc..6525aeb2 100644 --- a/stac_fastapi/pgstac/app.py +++ b/stac_fastapi/pgstac/app.py @@ -9,8 +9,8 @@ from contextlib import asynccontextmanager from brotli_asgi import BrotliMiddleware -from fastapi import FastAPI, Request, APIRouter -from fastapi.responses import ORJSONResponse, JSONResponse +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 ( @@ -175,18 +175,26 @@ async def ping(request: Request): 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") + 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"} + 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"} + content={ + "message": f"Database connection failed: {str(e)}", + "database": "ERROR", + }, ) self.app.include_router(mgmt_router, tags=["Liveliness/Readiness"]) diff --git a/tests/api/test_api.py b/tests/api/test_api.py index e0ca8031..ac3c4e61 100644 --- a/tests/api/test_api.py +++ b/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.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, @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index 94119e60..957e46ba 100644 --- a/tests/conftest.py +++ b/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.pgstac.app import PgStacApi 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 From de07a63e681af092c0842362afb7ddc4b3b4d9f0 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Wed, 16 Apr 2025 21:39:35 +0200 Subject: [PATCH 3/4] updated docs and changelog --- CHANGES.md | 4 ++ docs/api/stac_fastapi/pgstac/app.md | 69 ++++++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index f3481d1c..b48f44fa 100644 --- a/CHANGES.md +++ b/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 diff --git a/docs/api/stac_fastapi/pgstac/app.md b/docs/api/stac_fastapi/pgstac/app.md index f8129a24..d4885866 100644 --- a/docs/api/stac_fastapi/pgstac/app.md +++ b/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. From ec58cfc9095a55b22360d897e52bec8487207649 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Thu, 17 Apr 2025 08:34:55 +0200 Subject: [PATCH 4/4] Remove outdated PR description for the enhanced /_mgmt/ping endpoint --- pr_description.md | 52 ----------------------------------------------- 1 file changed, 52 deletions(-) delete mode 100644 pr_description.md diff --git a/pr_description.md b/pr_description.md deleted file mode 100644 index b972989d..00000000 --- a/pr_description.md +++ /dev/null @@ -1,52 +0,0 @@ -# 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.