diff --git a/advanced_alchemy/config/common.py b/advanced_alchemy/config/common.py index 71eccee5..041ec90e 100644 --- a/advanced_alchemy/config/common.py +++ b/advanced_alchemy/config/common.py @@ -203,7 +203,14 @@ 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)) + 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__() diff --git a/advanced_alchemy/extensions/starlette/extension.py b/advanced_alchemy/extensions/starlette/extension.py index b0f7d712..6b78c7a6 100644 --- a/advanced_alchemy/extensions/starlette/extension.py +++ b/advanced_alchemy/extensions/starlette/extension.py @@ -67,6 +67,32 @@ def init_app(self, app: Starlette) -> None: config.init_app(app) app.state.advanced_alchemy = self + existing_lifespan = getattr(app, "lifespan", None) + + if existing_lifespan is not None: + + @asynccontextmanager + async def wrapped_lifespan(app: Starlette) -> AsyncGenerator[None, None]: + async with self.lifespan(app), existing_lifespan(app): # type: ignore[misc] + yield + + app.lifespan = wrapped_lifespan # type: ignore[attr-defined] + else: + app.lifespan = self.lifespan # type: ignore[attr-defined] + + @asynccontextmanager + async def lifespan(self, app: Starlette) -> AsyncGenerator[None, None]: + """Context manager for lifespan events. + + Args: + app: The starlette application. + + Yields: + None + """ + await self.startup() + yield + await self.shutdown() @property def app(self) -> Starlette: @@ -85,12 +111,12 @@ def app(self) -> Starlette: return self._app - async def on_startup(self) -> None: # pragma: no cover + async def startup(self) -> None: # pragma: no cover """Initializes the database.""" for config in self.config: await config.on_startup() - async def on_shutdown(self) -> None: # pragma: no cover + async def shutdown(self) -> None: # pragma: no cover """Handles the shutdown event by disposing of the SQLAlchemy engine. Ensures that all connections are properly closed during application shutdown.