Merge pull request #217 from lsst-sqre/tickets/DM-41630

DM-41630: Convert to FastAPI lifespan functions

docs/user-guide/arq.rst

@@ -19,21 +19,26 @@ Quick start
 Dependency set up and configuration
 -----------------------------------
 
-In your application's FastAPI setup module, typically :file:`main.py`, you need to initialize `safir.dependencies.arq.ArqDependency` during the start up event:
+In your application's FastAPI setup module, typically :file:`main.py`, you need to initialize `safir.dependencies.arq.ArqDependency` during your lifespan function.
 
 .. code-block:: python
 
+    from collections.abc import AsyncIterator
+    from contextlib import asynccontextmanager
+
     from fastapi import Depends, FastAPI
     from safir.dependencies.arq import arq_dependency
 
-    app = FastAPI()
-
 
-    @app.on_event("startup")
-    async def startup() -> None:
+    @asynccontextmanager
+    def lifespan(app: FastAPI) -> AsyncIterator[None]:
         await arq_dependency.initialize(
             mode=config.arq_mode, redis_settings=config.arq_redis_settings
         )
+        yield
+
+
+    app = FastAPI(lifespan=lifespan)
 
 The ``mode`` parameter for `safir.dependencies.arq.ArqDependency.initialize` takes `ArqMode` enum values of either ``"production"`` or ``"test"``. The ``"production"`` mode configures a real arq_ queue backed by Redis, whereas ``"test"`` configures a mock version of the arq_ queue.
 

docs/user-guide/database.rst

@@ -144,30 +144,32 @@ For FastAPI applications, Safir provides a FastAPI dependency that creates a dat
 This uses the `SQLAlchemy async_scoped_session <https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html#using-asyncio-scoped-session>`__ to transparently manage a separate session per running task.
 
 To use the database session dependency, it must first be initialized during application startup.
-Generally this is done inside the application startup event:
+Generally this is done inside the application lifespan function.
+You must also close the dependency during application shutdown.
 
 .. code-block:: python
 
+   from collections.abc import AsyncIterator
+   from contextlib import asynccontextmanager
+
+   from fastapi import FastAPI
    from safir.dependencies.db_session import db_session_dependency
 
    from .config import config
 
 
-   @app.on_event("startup")
-   async def startup_event() -> None:
+   @asynccontextmanager
+   async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        await db_session_dependency.initialize(
            config.database_url, config.database_password
        )
+       yield
+       await db_session_dependency.aclose()
 
-As with some of the examples above, this assumes the application has a ``config`` object with the application settings, including the database URL and password.
-
-You must also close the dependency during application shutdown:
 
-.. code-block:: python
+   app = FastAPI(lifespan=lifespan)
 
-   @app.on_event("shutdown")
-   async def shutdown_event() -> None:
-       await db_session_dependency.aclose()
+As with some of the examples above, this assumes the application has a ``config`` object with the application settings, including the database URL and password.
 
 Then, any handler that needs a database session can depend on the `~safir.dependencies.db_session.db_session_dependency`:
 

docs/user-guide/http-client.rst

@@ -11,20 +11,24 @@ Setting up the httpx.AsyncClient
 The ``httpx.AsyncClient`` will be dyanmically created during application startup.
 Nothing further is needed apart from importing the dependency.
 However, it must be closed during application shutdown or a warning will be generated.
-
-To do this, add a shutdown hook to your application:
+This is normally done during the lifespan function for the FastAPI app.
 
 .. code-block:: python
 
-   from safir.dependencies.http_client import http_client_dependency
+   from collections.abc import AsyncIterator
+   from contextlib import asynccontextmanager
 
-   app = FastAPI()
+   from safir.dependencies.http_client import http_client_dependency
 
 
-   @app.on_event("shutdown")
-   async def shutdown_event() -> None:
+   @asynccontextmanager
+   async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+       yield
        await http_client_dependency.aclose()
 
+
+   app = FastAPI(lifespan=lifespan)
+
 You can add this line to an existing shutdown hook if you already have one.
 
 Using the httpx.AsyncClient

docs/user-guide/kubernetes.rst

@@ -24,12 +24,17 @@ For example:
 
 .. code-block:: python
 
+   from collections.abc import AsyncIterator
+   from contextlib import asynccontextmanager
+
+   from fastapi import FastAPI
    from safir.kubernetes import initialize_kubernetes
 
 
-   @app.on_event("startup")
-   async def startup_event() -> None:
+   @asynccontextmanager
+   async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        await initialize_kubernetes()
+       yield
 
 Testing with mock Kubernetes
 ============================

pyproject.toml

@@ -44,6 +44,7 @@ db = [
 dev = [
     "asgi-lifespan",
     "coverage[toml]",
+    "fastapi>=0.93.0",
     "flake8",
     "mypy",
     "pre-commit",

src/safir/dependencies/arq.py

@@ -43,22 +43,27 @@ async def initialize(
         --------
         .. code-block:: python
 
+           from collections.abc import AsyncIterator
+           from contextlib import asynccontextmanager
+
            from fastapi import Depends, FastAPI
            from safir.arq import ArqMode, ArqQueue
            from safir.dependencies.arq import arq_dependency
 
-           app = FastAPI()
-
 
-           @app.on_event("startup")
-           async def startup() -> None:
+           @asynccontextmanager
+           async def lifespan(app: FastAPI) -> AsyncIterator[None]:
                await arq_dependency.initialize(mode=ArqMode.test)
+               yield
+
+
+           app = FastAPI()
 
 
            @app.post("/")
            async def post_job(
                arq_queue: ArqQueue = Depends(arq_dependency),
-           ) -> Dict[str, Any]:
+           ) -> dict[str, Any]:
                job = await arq_queue.enqueue("test_task", "hello", an_int=42)
                return {"job_id": job.id}
         """

src/safir/dependencies/http_client.py

@@ -27,14 +27,24 @@ class HTTPClientDependency:
 
     Notes
     -----
-    The application must call ``http_client_dependency.aclose()`` as part of a
-    shutdown hook:
+    The application must call ``http_client_dependency.aclose()`` in the
+    application lifespan hook:
 
     .. code-block:: python
 
-       @app.on_event("shutdown")
-       async def shutdown_event() -> None:
+       from collections.abc import AsyncIterator
+       from contextlib import asynccontextmanager
+
+       from fastapi import FastAPI
+
+
+       @asynccontextmanager
+       async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+           yield
            await http_client_dependency.aclose()
+
+
+       app = FastAPI(lifespan=lifespan)
     """
 
     def __init__(self) -> None:

tests/dependencies/arq_test.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
 import pytest
@@ -14,10 +16,16 @@
 from safir.dependencies.arq import arq_dependency
 
 
+@asynccontextmanager
+async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
+    await arq_dependency.initialize(mode=ArqMode.test, redis_settings=None)
+    yield
+
+
 @pytest.mark.asyncio
 async def test_arq_dependency_mock() -> None:
     """Test the arq dependency entirely through the MockArqQueue."""
-    app = FastAPI()
+    app = FastAPI(lifespan=_lifespan)
 
     @app.post("/")
     async def post_job(
@@ -109,10 +117,6 @@ async def post_job_complete(
         except JobNotFound as e:
             raise HTTPException(status_code=404, detail=str(e)) from e
 
-    @app.on_event("startup")
-    async def startup() -> None:
-        await arq_dependency.initialize(mode=ArqMode.test, redis_settings=None)
-
     async with LifespanManager(app):
         async with AsyncClient(app=app, base_url="http://example.com") as c:
             r = await c.post("/")

tests/dependencies/http_client_test.py

@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
 import pytest
 import respx
 from asgi_lifespan import LifespanManager
@@ -16,9 +19,15 @@ def non_mocked_hosts() -> list[str]:
     return ["example.com"]
 
 
+@asynccontextmanager
+async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
+    yield
+    await http_client_dependency.aclose()
+
+
 @pytest.mark.asyncio
 async def test_http_client(respx_mock: respx.Router) -> None:
-    app = FastAPI()
+    app = FastAPI(lifespan=_lifespan)
     respx_mock.get("https://www.google.com").respond(200)
 
     @app.get("/")
@@ -29,10 +38,6 @@ async def handler(
         await http_client.get("https://www.google.com")
         return {}
 
-    @app.on_event("shutdown")
-    async def shutdown_event() -> None:
-        await http_client_dependency.aclose()
-
     async with LifespanManager(app):
         async with AsyncClient(app=app, base_url="http://example.com") as c:
             r = await c.get("/")