✨ Refactor lifespan management by introducing a context manager for logging and state tracking in FastAPI

Author: pcrespov

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

@@ -1,8 +1,11 @@
+import contextlib
+from collections.abc import Iterator
 from typing import Final
 
 from common_library.errors_classes import OsparcErrorMixin
 from fastapi import FastAPI
 from fastapi_lifespan_manager import State
+from servicelib.logging_utils import log_context
 
 
 class LifespanError(OsparcErrorMixin, RuntimeError): ...
@@ -54,3 +57,16 @@ def ensure_lifespan_called(state: State, lifespan_name: str) -> None:
     """
     if not is_lifespan_called(state, lifespan_name):
         raise LifespanExpectedCalledError(lifespan_name=lifespan_name)
+
+
+@contextlib.contextmanager
+def lifespan_context(
+    logger, level, lifespan_name: str, state: State
+) -> Iterator[State]:
+    """Helper context manager to log lifespan event and mark lifespan as called."""
+
+    with log_context(logger, level, lifespan_name):
+        # Check if lifespan has already been called
+        called_state = mark_lifespace_called(state, lifespan_name)
+
+        yield called_state

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

@@ -5,12 +5,12 @@
 
 from fastapi import FastAPI
 from fastapi_lifespan_manager import State
-from servicelib.logging_utils import log_catch, log_context
+from servicelib.logging_utils import log_catch
 from settings_library.postgres import PostgresSettings
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 from ..db_asyncpg_utils import create_async_engine_and_database_ready
-from .lifespan_utils import LifespanOnStartupError, mark_lifespace_called
+from .lifespan_utils import LifespanOnStartupError, lifespan_context
 
 _logger = logging.getLogger(__name__)
 
@@ -30,11 +30,10 @@ def create_postgres_database_input_state(settings: PostgresSettings) -> State:
 
 async def postgres_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[State]:
 
-    with log_context(_logger, logging.INFO, f"{__name__}"):
-
-        # Mark lifespan as called
-        called_state = mark_lifespace_called(state, "postgres_database_lifespan")
+    _lifespan_name = f"{__name__}.{postgres_database_lifespan.__name__}"
 
+    with lifespan_context(_logger, logging.INFO, _lifespan_name, state) as called_state:
+        # Validate input state
         settings = state[PostgresLifespanState.POSTGRES_SETTINGS]
 
         if settings is None or not isinstance(settings, PostgresSettings):

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

@@ -4,13 +4,12 @@
 from fastapi import FastAPI
 from fastapi_lifespan_manager import State
 from pydantic import BaseModel, ValidationError
-from servicelib.logging_utils import log_context
 from servicelib.rabbitmq import wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
 from .lifespan_utils import (
     LifespanOnStartupError,
-    mark_lifespace_called,
+    lifespan_context,
 )
 
 _logger = logging.getLogger(__name__)
@@ -33,10 +32,7 @@ async def rabbitmq_connectivity_lifespan(
     """
     _lifespan_name = f"{__name__}.{rabbitmq_connectivity_lifespan.__name__}"
 
-    with log_context(_logger, logging.INFO, _lifespan_name):
-
-        # Check if lifespan has already been called
-        called_state = mark_lifespace_called(state, _lifespan_name)
+    with lifespan_context(_logger, logging.INFO, _lifespan_name, state) as called_state:
 
         # Validate input state
         try:

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

@@ -5,12 +5,12 @@
 
 from fastapi import FastAPI
 from fastapi_lifespan_manager import State
-from pydantic import BaseModel, ConfigDict, StringConstraints, ValidationError
+from pydantic import BaseModel, StringConstraints, ValidationError
 from servicelib.logging_utils import log_catch, log_context
 from settings_library.redis import RedisDatabase, RedisSettings
 
 from ..redis import RedisClientSDK
-from .lifespan_utils import LifespanOnStartupError, mark_lifespace_called
+from .lifespan_utils import LifespanOnStartupError, lifespan_context
 
 _logger = logging.getLogger(__name__)
 
@@ -24,17 +24,11 @@ class RedisLifespanState(BaseModel):
     REDIS_CLIENT_NAME: Annotated[str, StringConstraints(min_length=3, max_length=32)]
     REDIS_CLIENT_DB: RedisDatabase
 
-    model_config = ConfigDict(
-        extra="allow",
-        arbitrary_types_allowed=True,  # RedisClientSDK has some arbitrary types and this class will never be serialized
-    )
-
 
 async def redis_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[State]:
-    with log_context(_logger, logging.INFO, f"{__name__}"):
+    _lifespan_name = f"{__name__}.{redis_database_lifespan.__name__}"
 
-        # Check if lifespan has already been called
-        called_state = mark_lifespace_called(state, "redis_database_lifespan")
+    with lifespan_context(_logger, logging.INFO, _lifespan_name, state) as called_state:
 
         # Validate input state
         try: