fix: use lifespan context manager in Starlette and FastAPI (#368)

The Starlette and FastAPI integrations now use a lifespan context manager instead of the deprecated startup/shutdown hooks

Author: cofin

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
       - id: unasyncd
         additional_dependencies: ["ruff"]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.9.2"
+    rev: "v0.9.3"
     hooks:
       # Run the linter.
     - id: ruff
@@ -32,7 +32,7 @@ repos:
     - id: ruff-format
       types_or: [ python, pyi ]
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.3.0
+    rev: v2.4.0
     hooks:
       - id: codespell
         exclude: "uv.lock|examples/us_state_lookup.json"

advanced_alchemy/config/common.py

@@ -202,8 +202,15 @@ def __post_init__(self) -> None:
 
             event.listen(Session, "before_flush", touch_updated_timestamp)
 
-    def __hash__(self) -> int:
-        return hash((self.__class__.__qualname__, self.bind_key))
+    def __hash__(self) -> int:  # pragma: no cover
+        return hash(
+            (
+                self.__class__.__qualname__,
+                self.connection_string,
+                self.engine_config.__class__.__qualname__,
+                self.bind_key,
+            )
+        )
 
     def __eq__(self, other: object) -> bool:
         return self.__hash__() == other.__hash__()
@@ -250,7 +257,7 @@ def get_engine(self) -> EngineT:
             del engine_config["json_serializer"]
             return self.create_engine_callable(self.connection_string, **engine_config)
 
-    def create_session_maker(self) -> Callable[[], SessionT]:
+    def create_session_maker(self) -> Callable[[], SessionT]:  # pragma: no cover
         """Get a session maker. If none exists yet, create one.
 
         Returns:

advanced_alchemy/extensions/starlette/config.py

@@ -140,8 +140,6 @@ def init_app(self, app: Starlette) -> None:
         )
 
         app.add_middleware(BaseHTTPMiddleware, dispatch=self.middleware_dispatch)
-        app.add_event_handler("startup", self.on_startup)  # pyright: ignore[reportUnknownMemberType]
-        app.add_event_handler("shutdown", self.on_shutdown)  # pyright: ignore[reportUnknownMemberType]
 
     async def on_startup(self) -> None:
         """Initialize the Starlette application with this configuration."""
@@ -288,8 +286,6 @@ def init_app(self, app: Starlette) -> None:
         )
         _ = self.create_session_maker()
         app.add_middleware(BaseHTTPMiddleware, dispatch=self.middleware_dispatch)
-        app.add_event_handler("startup", self.on_startup)  # pyright: ignore[reportUnknownMemberType]
-        app.add_event_handler("shutdown", self.on_shutdown)  # pyright: ignore[reportUnknownMemberType]
 
     async def on_startup(self) -> None:
         """Initialize the Starlette application with this configuration."""
@@ -381,5 +377,5 @@ async def on_shutdown(self) -> None:  # pragma: no cover
         if self.app is not None:
             with contextlib.suppress(AttributeError, KeyError):
                 delattr(self.app.state, self.engine_key)
-                delattr(self.app.state, self.session_key)
                 delattr(self.app.state, self.session_maker_key)
+                delattr(self.app.state, self.session_key)

advanced_alchemy/extensions/starlette/extension.py

@@ -3,7 +3,16 @@
 
 import contextlib
 from contextlib import asynccontextmanager, contextmanager
-from typing import TYPE_CHECKING, AsyncGenerator, Callable, Generator, Sequence, Union, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Callable,
+    Generator,
+    Sequence,
+    Union,
+    overload,
+)
 
 from starlette.applications import Starlette
 from starlette.requests import Request  # noqa: TC002
@@ -68,8 +77,33 @@ def init_app(self, app: Starlette) -> None:
 
         app.state.advanced_alchemy = self
 
+        original_lifespan = app.router.lifespan_context
+
+        @asynccontextmanager
+        async def wrapped_lifespan(app: Starlette) -> AsyncGenerator[Any, None]:  # pragma: no cover
+            async with self.lifespan(app), original_lifespan(app) as state:
+                yield state
+
+        app.router.lifespan_context = wrapped_lifespan
+
+    @asynccontextmanager
+    async def lifespan(self, app: Starlette) -> AsyncGenerator[Any, None]:  # pragma: no cover
+        """Context manager for lifespan events.
+
+        Args:
+            app: The starlette application.
+
+        Yields:
+            None
+        """
+        await self.on_startup()
+        try:
+            yield
+        finally:
+            await self.on_shutdown()
+
     @property
-    def app(self) -> Starlette:
+    def app(self) -> Starlette:  # pragma: no cover
         """Returns the Starlette application instance.
 
         Raises:
@@ -79,7 +113,7 @@ def app(self) -> Starlette:
         Returns:
             starlette.applications.Starlette: The Starlette application instance.
         """
-        if self._app is None:
+        if self._app is None:  # pragma: no cover
             msg = "Application not initialized. Did you forget to call init_app?"
             raise ImproperConfigurationError(msg)
 
@@ -119,36 +153,38 @@ def get_config(self, key: str | None = None) -> Union[SQLAlchemyAsyncConfig, SQL
         if key == "default" and len(self.config) == 1:
             key = self.config[0].bind_key or "default"
         config = self._mapped_configs.get(key)
-        if config is None:
+        if config is None:  # pragma: no cover
             msg = f"Config with key {key} not found"
             raise ImproperConfigurationError(msg)
         return config
 
     def get_async_config(self, key: str | None = None) -> SQLAlchemyAsyncConfig:
         """Get the async config for the given key."""
         config = self.get_config(key)
-        if not isinstance(config, SQLAlchemyAsyncConfig):
+        if not isinstance(config, SQLAlchemyAsyncConfig):  # pragma: no cover
             msg = "Expected an async config, but got a sync config"
             raise ImproperConfigurationError(msg)
         return config
 
     def get_sync_config(self, key: str | None = None) -> SQLAlchemySyncConfig:
         """Get the sync config for the given key."""
         config = self.get_config(key)
-        if not isinstance(config, SQLAlchemySyncConfig):
+        if not isinstance(config, SQLAlchemySyncConfig):  # pragma: no cover
             msg = "Expected a sync config, but got an async config"
             raise ImproperConfigurationError(msg)
         return config
 
     @asynccontextmanager
-    async def with_async_session(self, key: str | None = None) -> AsyncGenerator[AsyncSession, None]:
+    async def with_async_session(
+        self, key: str | None = None
+    ) -> AsyncGenerator[AsyncSession, None]:  # pragma: no cover
         """Context manager for getting an async session."""
         config = self.get_async_config(key)
         async with config.get_session() as session:
             yield session
 
     @contextmanager
-    def with_sync_session(self, key: str | None = None) -> Generator[Session, None]:
+    def with_sync_session(self, key: str | None = None) -> Generator[Session, None]:  # pragma: no cover
         """Context manager for getting a sync session."""
         config = self.get_sync_config(key)
         with config.get_session() as session:
@@ -164,7 +200,8 @@ def _get_session_from_request(request: Request, config: SQLAlchemySyncConfig) ->
 
     @staticmethod
     def _get_session_from_request(
-        request: Request, config: SQLAlchemyAsyncConfig | SQLAlchemySyncConfig
+        request: Request,
+        config: SQLAlchemyAsyncConfig | SQLAlchemySyncConfig,  # pragma: no cover
     ) -> Session | AsyncSession:  # pragma: no cover
         """Get the session for the given key."""
         session = getattr(request.state, config.session_key, None)
@@ -173,22 +210,24 @@ def _get_session_from_request(
             setattr(request.state, config.session_key, session)
         return session
 
-    def get_session(self, request: Request, key: str | None = None) -> Session | AsyncSession:
+    def get_session(self, request: Request, key: str | None = None) -> Session | AsyncSession:  # pragma: no cover
         """Get the session for the given key."""
         config = self.get_config(key)
         return self._get_session_from_request(request, config)
 
-    def get_async_session(self, request: Request, key: str | None = None) -> AsyncSession:
+    def get_async_session(self, request: Request, key: str | None = None) -> AsyncSession:  # pragma: no cover
         """Get the async session for the given key."""
         config = self.get_async_config(key)
         return self._get_session_from_request(request, config)
 
-    def get_sync_session(self, request: Request, key: str | None = None) -> Session:
+    def get_sync_session(self, request: Request, key: str | None = None) -> Session:  # pragma: no cover
         """Get the sync session for the given key."""
         config = self.get_sync_config(key)
         return self._get_session_from_request(request, config)
 
-    def provide_session(self, key: str | None = None) -> Callable[[Request], Session | AsyncSession]:
+    def provide_session(
+        self, key: str | None = None
+    ) -> Callable[[Request], Session | AsyncSession]:  # pragma: no cover
         """Get the session for the given key."""
         config = self.get_config(key)
 
@@ -197,7 +236,7 @@ def _get_session(request: Request) -> Session | AsyncSession:
 
         return _get_session
 
-    def provide_async_session(self, key: str | None = None) -> Callable[[Request], AsyncSession]:
+    def provide_async_session(self, key: str | None = None) -> Callable[[Request], AsyncSession]:  # pragma: no cover
         """Get the async session for the given key."""
         config = self.get_async_config(key)
 
@@ -206,7 +245,7 @@ def _get_session(request: Request) -> AsyncSession:
 
         return _get_session
 
-    def provide_sync_session(self, key: str | None = None) -> Callable[[Request], Session]:
+    def provide_sync_session(self, key: str | None = None) -> Callable[[Request], Session]:  # pragma: no cover
         """Get the sync session for the given key."""
         config = self.get_sync_config(key)
 
@@ -220,12 +259,12 @@ def get_engine(self, key: str | None = None) -> Engine | AsyncEngine:  # pragma:
         config = self.get_config(key)
         return config.get_engine()
 
-    def get_async_engine(self, key: str | None = None) -> AsyncEngine:
+    def get_async_engine(self, key: str | None = None) -> AsyncEngine:  # pragma: no cover
         """Get the async engine for the given key."""
         config = self.get_async_config(key)
         return config.get_engine()
 
-    def get_sync_engine(self, key: str | None = None) -> Engine:
+    def get_sync_engine(self, key: str | None = None) -> Engine:  # pragma: no cover
         """Get the sync engine for the given key."""
         config = self.get_sync_config(key)
         return config.get_engine()

tests/unit/test_extensions/test_fastapi.py

@@ -1,5 +1,6 @@
 import sys
-from typing import TYPE_CHECKING, Callable, Generator, Literal, Union, cast
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, AsyncGenerator, Callable, Generator, Literal, Union, cast
 from unittest.mock import MagicMock
 
 import pytest
@@ -356,3 +357,41 @@ def handler(
         )
         assert mock.call_args_list[0].kwargs["session"] is not mock.call_args_list[1].kwargs["session"]
         assert mock.call_args_list[0].kwargs["engine"] is not mock.call_args_list[1].kwargs["engine"]
+
+
+async def test_lifespan_startup_shutdown_called_fastapi(mocker: MockerFixture, app: FastAPI, config: AnyConfig) -> None:
+    mock_startup = mocker.patch.object(AdvancedAlchemy, "on_startup")
+    mock_shutdown = mocker.patch.object(AdvancedAlchemy, "on_shutdown")
+    _alchemy = AdvancedAlchemy(config, app=app)
+
+    with TestClient(app=app) as _client:  # TestClient context manager triggers lifespan events
+        pass  # App starts up and shuts down within this context
+
+    mock_startup.assert_called_once()
+    mock_shutdown.assert_called_once()
+
+
+async def test_lifespan_with_custom_lifespan_fastapi(mocker: MockerFixture, app: FastAPI, config: AnyConfig) -> None:
+    mock_aa_startup = mocker.patch.object(AdvancedAlchemy, "on_startup")
+    mock_aa_shutdown = mocker.patch.object(AdvancedAlchemy, "on_shutdown")
+    mock_custom_startup = mocker.MagicMock()
+    mock_custom_shutdown = mocker.MagicMock()
+
+    @asynccontextmanager
+    async def custom_lifespan(app_in: FastAPI) -> AsyncGenerator[None, None]:
+        mock_custom_startup()
+        yield
+        mock_custom_shutdown()
+
+    app.router.lifespan_context = custom_lifespan  # type: ignore[assignment] # Set a custom lifespan on the app
+    _alchemy = AdvancedAlchemy(config, app=app)
+
+    with TestClient(app=app) as _client:  # TestClient context manager triggers lifespan events
+        pass  # App starts up and shuts down within this context
+
+    mock_aa_startup.assert_called_once()
+    mock_aa_shutdown.assert_called_once()
+    mock_custom_startup.assert_called_once()
+    mock_custom_shutdown.assert_called_once()
+
+    # Optionally assert the order of calls if needed, e.g., using mocker.call_order

tests/unit/test_extensions/test_starlette.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import sys
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Callable, Generator, Union, cast
 from unittest.mock import MagicMock
 
@@ -476,3 +478,43 @@ async def handler(request: Request) -> Response:
         client.get("/")
 
         assert alchemy_1.get_sync_engine() is not alchemy_1.get_async_engine("other")
+
+
+async def test_lifespan_startup_shutdown_called_starlette(
+    mocker: MockerFixture, app: Starlette, config: AnyConfig
+) -> None:
+    mock_startup = mocker.patch.object(AdvancedAlchemy, "on_startup")
+    mock_shutdown = mocker.patch.object(AdvancedAlchemy, "on_shutdown")
+    _alchemy = AdvancedAlchemy(config, app=app)
+
+    with TestClient(app=app) as _client:  # TestClient context manager triggers lifespan events
+        pass  # App starts up and shuts down within this context
+
+    mock_startup.assert_called_once()
+    mock_shutdown.assert_called_once()
+
+
+async def test_lifespan_with_custom_lifespan_starlette(
+    mocker: MockerFixture, app: Starlette, config: AnyConfig
+) -> None:
+    mock_aa_startup = mocker.patch.object(AdvancedAlchemy, "on_startup")
+    mock_aa_shutdown = mocker.patch.object(AdvancedAlchemy, "on_shutdown")
+    mock_custom_startup = mocker.MagicMock()
+    mock_custom_shutdown = mocker.MagicMock()
+
+    @asynccontextmanager
+    async def custom_lifespan(app_in: Starlette) -> AsyncGenerator[None, None]:
+        mock_custom_startup()
+        yield
+        mock_custom_shutdown()
+
+    app.router.lifespan_context = custom_lifespan  # type: ignore[assignment] # Set a custom lifespan on the app
+    _alchemy = AdvancedAlchemy(config, app=app)
+
+    with TestClient(app=app) as _client:  # TestClient context manager triggers lifespan events
+        pass  # App starts up and shuts down within this context
+
+    mock_aa_startup.assert_called_once()
+    mock_aa_shutdown.assert_called_once()
+    mock_custom_startup.assert_called_once()
+    mock_custom_shutdown.assert_called_once()