Skip to content

Commit 21aae32

Browse files
use fastapi lifespan (#183)
1 parent b232949 commit 21aae32

File tree

1 file changed

+22
-12
lines changed

1 file changed

+22
-12
lines changed

stac_fastapi/pgstac/app.py

+22-12
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@
66
"""
77

88
import os
9+
from contextlib import asynccontextmanager
910

1011
from brotli_asgi import BrotliMiddleware
12+
from fastapi import FastAPI
1113
from fastapi.responses import ORJSONResponse
1214
from stac_fastapi.api.app import StacApi
1315
from stac_fastapi.api.middleware import CORSMiddleware, ProxyHeaderMiddleware
@@ -18,6 +20,7 @@
1820
create_post_request_model,
1921
create_request_model,
2022
)
23+
from stac_fastapi.api.openapi import update_openapi
2124
from stac_fastapi.extensions.core import (
2225
FieldsExtension,
2326
FilterExtension,
@@ -101,7 +104,26 @@
101104
post_request_model = create_post_request_model(extensions, base_model=PgstacSearch)
102105
get_request_model = create_get_request_model(extensions)
103106

107+
108+
@asynccontextmanager
109+
async def lifespan(app: FastAPI):
110+
"""FastAPI Lifespan."""
111+
await connect_to_db(app)
112+
yield
113+
await close_db_connection(app)
114+
115+
116+
fastapp = FastAPI(
117+
openapi_url=settings.openapi_url,
118+
docs_url=settings.docs_url,
119+
redoc_url=None,
120+
root_path=getattr(settings, "root_path", None),
121+
lifespan=lifespan,
122+
)
123+
124+
104125
api = StacApi(
126+
app=update_openapi(fastapp),
105127
settings=settings,
106128
extensions=extensions + [collection_search_extension]
107129
if collection_search_extension
@@ -125,18 +147,6 @@
125147
app = api.app
126148

127149

128-
@app.on_event("startup")
129-
async def startup_event():
130-
"""Connect to database on startup."""
131-
await connect_to_db(app)
132-
133-
134-
@app.on_event("shutdown")
135-
async def shutdown_event():
136-
"""Close database connection."""
137-
await close_db_connection(app)
138-
139-
140150
def run():
141151
"""Run app from command line using uvicorn if available."""
142152
try:

0 commit comments

Comments
 (0)