✨ Enhance lifespan management: add checks for multiple calls and improve error messaging

Author: pcrespov

packages/service-library/src/servicelib/fastapi/lifespan_utils.py

@@ -1,12 +1,44 @@
+from typing import Final
+
 from common_library.errors_classes import OsparcErrorMixin
+from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 
 
 class LifespanError(OsparcErrorMixin, RuntimeError): ...
 
 
 class LifespanOnStartupError(LifespanError):
-    msg_template = "Failed during startup of {module}"
+    msg_template = "Failed during startup of {lifespan_name}"
 
 
 class LifespanOnShutdownError(LifespanError):
-    msg_template = "Failed during shutdown of {module}"
+    msg_template = "Failed during shutdown of {lifespan_name}"
+
+
+class LifespanAlreadyCalledError(LifespanError):
+    msg_template = "The lifespan '{lifespan_name}' has already been called."
+
+
+_CALLED_LIFESPANS_KEY: Final[str] = "_CALLED_LIFESPANS"
+
+
+def is_lifespan_called(state: State, lifespan_name: str) -> bool:
+    called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
+    return lifespan_name in called_lifespans
+
+
+def record_lifespan_called_once(state: State, lifespan_name: str) -> State:
+    """Validates if a lifespan has already been called and records it in the state.
+    Raises LifespanAlreadyCalledError if the lifespan has already been called.
+    """
+    assert not isinstance(  # nosec
+        state, FastAPI
+    ), "TIP: lifespan func has (app, state) positional arguments"
+
+    if is_lifespan_called(state, lifespan_name):
+        raise LifespanAlreadyCalledError(lifespan_name=lifespan_name)
+
+    called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
+    called_lifespans.add(lifespan_name)
+    return {_CALLED_LIFESPANS_KEY: called_lifespans}

packages/service-library/src/servicelib/fastapi/postgres_lifespan.py

@@ -10,7 +10,7 @@
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 from ..db_asyncpg_utils import create_async_engine_and_database_ready
-from .lifespan_utils import LifespanOnStartupError
+from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
 
 _logger = logging.getLogger(__name__)
 
@@ -32,6 +32,9 @@ async def postgres_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[
 
     with log_context(_logger, logging.INFO, f"{__name__}"):
 
+        # Mark lifespan as called
+        called_state = record_lifespan_called_once(state, "postgres_database_lifespan")
+
         settings = state[PostgresLifespanState.POSTGRES_SETTINGS]
 
         if settings is None or not isinstance(settings, PostgresSettings):
@@ -48,6 +51,7 @@ async def postgres_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[
 
             yield {
                 PostgresLifespanState.POSTGRES_ASYNC_ENGINE: async_engine,
+                **called_state,
             }
 
         finally:

packages/service-library/src/servicelib/fastapi/redis_lifespan.py

@@ -10,7 +10,7 @@
 from settings_library.redis import RedisDatabase, RedisSettings
 
 from ..redis import RedisClientSDK
-from .lifespan_utils import LifespanOnStartupError
+from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
 
 _logger = logging.getLogger(__name__)
 
@@ -31,16 +31,17 @@ class RedisLifespanState(BaseModel):
 
 
 async def redis_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[State]:
-
     with log_context(_logger, logging.INFO, f"{__name__}"):
 
+        # Check if lifespan has already been called
+        called_state = record_lifespan_called_once(state, "redis_database_lifespan")
+
         # Validate input state
         try:
             redis_state = RedisLifespanState.model_validate(state)
             redis_dsn_with_secrets = redis_state.REDIS_SETTINGS.build_redis_dsn(
                 redis_state.REDIS_CLIENT_DB
             )
-
         except ValidationError as exc:
             raise RedisConfigurationError(validation_error=exc, state=state) from exc
 
@@ -50,16 +51,13 @@ async def redis_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[Sta
             logging.INFO,
             f"Creating redis client with name={redis_state.REDIS_CLIENT_NAME}",
         ):
-
             redis_client = RedisClientSDK(
                 redis_dsn_with_secrets,
                 client_name=redis_state.REDIS_CLIENT_NAME,
             )
 
         try:
-
-            yield {"REDIS_CLIENT_SDK": redis_client}
-
+            yield {"REDIS_CLIENT_SDK": redis_client, **called_state}
         finally:
             # Teardown client
             if redis_client:

packages/service-library/tests/fastapi/test_lifespan_utils.py

@@ -16,8 +16,10 @@
 from pytest_mock import MockerFixture
 from pytest_simcore.helpers.logging_tools import log_context
 from servicelib.fastapi.lifespan_utils import (
+    LifespanAlreadyCalledError,
     LifespanOnShutdownError,
     LifespanOnStartupError,
+    record_lifespan_called_once,
 )
 
 
@@ -186,7 +188,7 @@ async def lifespan_failing_on_startup(app: FastAPI) -> AsyncIterator[State]:
                 startup_step(_name)
             except RuntimeError as exc:
                 handle_error(_name, exc)
-                raise LifespanOnStartupError(module=_name) from exc
+                raise LifespanOnStartupError(lifespan_name=_name) from exc
             yield {}
             shutdown_step(_name)
 
@@ -201,7 +203,7 @@ async def lifespan_failing_on_shutdown(app: FastAPI) -> AsyncIterator[State]:
                 shutdown_step(_name)
             except RuntimeError as exc:
                 handle_error(_name, exc)
-                raise LifespanOnShutdownError(module=_name) from exc
+                raise LifespanOnShutdownError(lifespan_name=_name) from exc
 
     return {
         "startup_step": startup_step,
@@ -228,7 +230,7 @@ async def test_app_lifespan_with_error_on_startup(
     assert not failing_lifespan_manager["startup_step"].called
     assert not failing_lifespan_manager["shutdown_step"].called
     assert exception.error_context() == {
-        "module": "lifespan_failing_on_startup",
+        "lifespan_name": "lifespan_failing_on_startup",
         "message": "Failed during startup of lifespan_failing_on_startup",
         "code": "RuntimeError.LifespanError.LifespanOnStartupError",
     }
@@ -250,7 +252,35 @@ async def test_app_lifespan_with_error_on_shutdown(
     assert failing_lifespan_manager["startup_step"].called
     assert not failing_lifespan_manager["shutdown_step"].called
     assert exception.error_context() == {
-        "module": "lifespan_failing_on_shutdown",
+        "lifespan_name": "lifespan_failing_on_shutdown",
         "message": "Failed during shutdown of lifespan_failing_on_shutdown",
         "code": "RuntimeError.LifespanError.LifespanOnShutdownError",
     }
+
+
+async def test_lifespan_called_more_than_once(is_pdb_enabled: bool):
+    state = {}
+
+    app_lifespan = LifespanManager()
+
+    @app_lifespan.add
+    async def _one(_, state: State) -> AsyncIterator[State]:
+        called_state = record_lifespan_called_once(state, "test_lifespan_one")
+        yield {"other": 0, **called_state}
+
+    @app_lifespan.add
+    async def _two(_, state: State) -> AsyncIterator[State]:
+        called_state = record_lifespan_called_once(state, "test_lifespan_two")
+        yield {"something": 0, **called_state}
+
+    app_lifespan.add(_one)  # added "by mistake"
+
+    with pytest.raises(LifespanAlreadyCalledError) as err_info:
+        async with ASGILifespanManager(
+            FastAPI(lifespan=app_lifespan),
+            startup_timeout=None if is_pdb_enabled else 10,
+            shutdown_timeout=None if is_pdb_enabled else 10,
+        ):
+            ...
+
+    assert err_info.value.lifespan_name == "test_lifespan_one"