using lifespan

Author: Andrei Neagu

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/catalog/__init__.py

@@ -1,7 +1,7 @@
 from ._public_client import CatalogPublicClient
-from ._setup import setup_catalog
+from ._setup import lifespan_catalog
 
 __all__: tuple[str, ...] = (
     "CatalogPublicClient",
-    "setup_catalog",
+    "lifespan_catalog",
 )

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/catalog/_setup.py

@@ -1,21 +1,21 @@
+from collections.abc import AsyncIterator
+
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 
 from ._public_client import CatalogPublicClient
 from ._thin_client import CatalogThinClient
 
 
-def setup_catalog(app: FastAPI) -> None:
-    async def _on_startup() -> None:
-        thin_client = CatalogThinClient(app)
-        thin_client.set_to_app_state(app)
-        thin_client.attach_lifespan_to(app)
+async def lifespan_catalog(app: FastAPI) -> AsyncIterator[State]:
+    thin_client = CatalogThinClient(app)
+    thin_client.set_to_app_state(app)
+    thin_client.attach_lifespan_to(app)
 
-        public_client = CatalogPublicClient(app)
-        public_client.set_to_app_state(app)
+    public_client = CatalogPublicClient(app)
+    public_client.set_to_app_state(app)
 
-    async def _on_shutdown() -> None:
-        CatalogPublicClient.pop_from_app_state(app)
-        CatalogThinClient.pop_from_app_state(app)
+    yield {}
 
-    app.add_event_handler("startup", _on_startup)
-    app.add_event_handler("shutdown", _on_shutdown)
+    CatalogPublicClient.pop_from_app_state(app)
+    CatalogThinClient.pop_from_app_state(app)