Skip to content

Commit 662d9da

Browse files
vincentsaragohrodmngadomski
authored
update stac-fastapi requirements and add health-check (#235)
* update stac-fastapi requirements and add health-check * Update stac_fastapi/pgstac/core.py Co-authored-by: Henry Rodman <[email protected]> * update from review * Update stac_fastapi/pgstac/core.py Co-authored-by: Pete Gadomski <[email protected]> * update test --------- Co-authored-by: Henry Rodman <[email protected]> Co-authored-by: Pete Gadomski <[email protected]>
1 parent 4ce9075 commit 662d9da

File tree

6 files changed

+134
-8
lines changed

6 files changed

+134
-8
lines changed

CHANGES.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,11 @@
22

33
## [Unreleased]
44

5+
### Changed
6+
7+
- update `stac-fastapi-*` version requirements to `>=5.2,<6.0`
8+
- add pgstac health-check in `/_mgmt/health`
9+
510
## [5.0.2] - 2025-04-07
611

712
### Fixed

setup.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010
"orjson",
1111
"pydantic",
1212
"stac_pydantic==3.1.*",
13-
"stac-fastapi.api>=5.1,<6.0",
14-
"stac-fastapi.extensions>=5.1,<6.0",
15-
"stac-fastapi.types>=5.1,<6.0",
13+
"stac-fastapi.api>=5.2,<6.0",
14+
"stac-fastapi.extensions>=5.2,<6.0",
15+
"stac-fastapi.types>=5.2,<6.0",
1616
"asyncpg",
1717
"buildpg",
1818
"brotli_asgi",

stac_fastapi/pgstac/app.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@
1010

1111
from brotli_asgi import BrotliMiddleware
1212
from fastapi import FastAPI
13-
from fastapi.responses import ORJSONResponse
1413
from stac_fastapi.api.app import StacApi
1514
from stac_fastapi.api.middleware import CORSMiddleware, ProxyHeaderMiddleware
1615
from stac_fastapi.api.models import (
1716
EmptyRequest,
1817
ItemCollectionUri,
18+
JSONResponse,
1919
create_get_request_model,
2020
create_post_request_model,
2121
create_request_model,
@@ -40,7 +40,7 @@
4040
from starlette.middleware import Middleware
4141

4242
from stac_fastapi.pgstac.config import Settings
43-
from stac_fastapi.pgstac.core import CoreCrudClient
43+
from stac_fastapi.pgstac.core import CoreCrudClient, health_check
4444
from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
4545
from stac_fastapi.pgstac.extensions import QueryExtension
4646
from stac_fastapi.pgstac.extensions.filter import FiltersClient
@@ -54,7 +54,7 @@
5454
"transaction": TransactionExtension(
5555
client=TransactionsClient(),
5656
settings=settings,
57-
response_class=ORJSONResponse,
57+
response_class=JSONResponse,
5858
),
5959
"bulk_transactions": BulkTransactionExtension(client=BulkTransactionsClient()),
6060
}
@@ -174,7 +174,7 @@ async def lifespan(app: FastAPI):
174174
settings=settings,
175175
extensions=application_extensions,
176176
client=CoreCrudClient(pgstac_search_model=post_request_model),
177-
response_class=ORJSONResponse,
177+
response_class=JSONResponse,
178178
items_get_request_model=items_get_request_model,
179179
search_get_request_model=get_request_model,
180180
search_post_request_model=post_request_model,
@@ -188,6 +188,7 @@ async def lifespan(app: FastAPI):
188188
allow_methods=settings.cors_methods,
189189
),
190190
],
191+
health_check=health_check,
191192
)
192193
app = api.app
193194

stac_fastapi/pgstac/core.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,3 +605,49 @@ def _clean_search_args( # noqa: C901
605605
clean[k] = v
606606

607607
return clean
608+
609+
610+
async def health_check(request: Request) -> Union[Dict, JSONResponse]:
611+
"""PgSTAC HealthCheck."""
612+
resp = {
613+
"status": "UP",
614+
"lifespan": {
615+
"status": "UP",
616+
},
617+
}
618+
if not hasattr(request.app.state, "get_connection"):
619+
return JSONResponse(
620+
status_code=503,
621+
content={
622+
"status": "DOWN",
623+
"lifespan": {
624+
"status": "DOWN",
625+
"message": "application lifespan wasn't run",
626+
},
627+
"pgstac": {
628+
"status": "DOWN",
629+
"message": "Could not connect to database",
630+
},
631+
},
632+
)
633+
634+
try:
635+
async with request.app.state.get_connection(request, "r") as conn:
636+
q, p = render(
637+
"""SELECT pgstac.get_version();""",
638+
)
639+
version = await conn.fetchval(q, *p)
640+
except Exception as e:
641+
resp["status"] = "DOWN"
642+
resp["pgstac"] = {
643+
"status": "DOWN",
644+
"message": str(e),
645+
}
646+
return JSONResponse(status_code=503, content=resp)
647+
648+
resp["pgstac"] = {
649+
"status": "UP",
650+
"pgstac_version": version,
651+
}
652+
653+
return resp

tests/conftest.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
from stac_pydantic import Collection, Item
4343

4444
from stac_fastapi.pgstac.config import PostgresSettings, Settings
45-
from stac_fastapi.pgstac.core import CoreCrudClient
45+
from stac_fastapi.pgstac.core import CoreCrudClient, health_check
4646
from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
4747
from stac_fastapi.pgstac.extensions import QueryExtension
4848
from stac_fastapi.pgstac.extensions.filter import FiltersClient
@@ -191,6 +191,7 @@ def api_client(request):
191191
collections_get_request_model=collection_search_extension.GET,
192192
response_class=ORJSONResponse,
193193
router=APIRouter(prefix=prefix),
194+
health_check=health_check,
194195
)
195196

196197
return api
@@ -302,6 +303,7 @@ def api_client_no_ext():
302303
TransactionExtension(client=TransactionsClient(), settings=api_settings)
303304
],
304305
client=CoreCrudClient(),
306+
health_check=health_check,
305307
)
306308

307309

tests/resources/test_mgmt.py

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,11 @@
1+
from httpx import ASGITransport, AsyncClient
2+
from stac_fastapi.api.app import StacApi
3+
4+
from stac_fastapi.pgstac.config import PostgresSettings, Settings
5+
from stac_fastapi.pgstac.core import CoreCrudClient, health_check
6+
from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
7+
8+
19
async def test_ping_no_param(app_client):
210
"""
311
Test ping endpoint with a mocked client.
@@ -7,3 +15,67 @@ async def test_ping_no_param(app_client):
715
res = await app_client.get("/_mgmt/ping")
816
assert res.status_code == 200
917
assert res.json() == {"message": "PONG"}
18+
19+
20+
async def test_health(app_client):
21+
"""
22+
Test health endpoint
23+
24+
Args:
25+
app_client (TestClient): mocked client fixture
26+
27+
"""
28+
res = await app_client.get("/_mgmt/health")
29+
assert res.status_code == 200
30+
body = res.json()
31+
assert body["status"] == "UP"
32+
assert body["pgstac"]["status"] == "UP"
33+
assert body["pgstac"]["pgstac_version"]
34+
35+
36+
async def test_health_503(database):
37+
"""Test health endpoint error."""
38+
39+
# No lifespan so no `get_connection` is application state
40+
api = StacApi(
41+
settings=Settings(testing=True),
42+
extensions=[],
43+
client=CoreCrudClient(),
44+
health_check=health_check,
45+
)
46+
47+
async with AsyncClient(
48+
transport=ASGITransport(app=api.app), base_url="http://test"
49+
) as client:
50+
res = await client.get("/_mgmt/health")
51+
assert res.status_code == 503
52+
body = res.json()
53+
assert body["status"] == "DOWN"
54+
assert body["lifespan"]["status"] == "DOWN"
55+
assert body["lifespan"]["message"] == "application lifespan wasn't run"
56+
assert body["pgstac"]["status"] == "DOWN"
57+
assert body["pgstac"]["message"] == "Could not connect to database"
58+
59+
# No lifespan so no `get_connection` is application state
60+
postgres_settings = PostgresSettings(
61+
postgres_user=database.user,
62+
postgres_pass=database.password,
63+
postgres_host_reader=database.host,
64+
postgres_host_writer=database.host,
65+
postgres_port=database.port,
66+
postgres_dbname=database.dbname,
67+
)
68+
# Create connection pool but close it just after
69+
await connect_to_db(api.app, postgres_settings=postgres_settings)
70+
await close_db_connection(api.app)
71+
72+
async with AsyncClient(
73+
transport=ASGITransport(app=api.app), base_url="http://test"
74+
) as client:
75+
res = await client.get("/_mgmt/health")
76+
assert res.status_code == 503
77+
body = res.json()
78+
assert body["status"] == "DOWN"
79+
assert body["lifespan"]["status"] == "UP"
80+
assert body["pgstac"]["status"] == "DOWN"
81+
assert body["pgstac"]["message"] == "pool is closed"

0 commit comments

Comments
 (0)