diff --git a/CHANGES.md b/CHANGES.md index f3481d1..b48f44f 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/Makefile b/Makefile index 57b56b6..9a313fc 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/docs/api/stac_fastapi/pgstac/app.md b/docs/api/stac_fastapi/pgstac/app.md index f8129a2..d488586 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. diff --git a/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/app.py index fcc1288..6525aeb 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 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, diff --git a/tests/api/test_api.py b/tests/api/test_api.py index e4ce8d0..ac3c4e6 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.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, diff --git a/tests/conftest.py b/tests/conftest.py index 7944e8d..957e46b 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.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) diff --git a/tests/resources/test_mgmt.py b/tests/resources/test_mgmt.py index 9d2bc3d..0a80831 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"