refactor: storage rpc routes

Author: giancarloromeo

packages/celery-library/src/celery_library/rpc/__init__.py

@@ -7,7 +7,6 @@
     TransferrableCeleryError,
     decode_celery_transferrable_error,
 )
-from fastapi import FastAPI
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
     AsyncJobId,
@@ -22,21 +21,22 @@
     JobSchedulerError,
 )
 from servicelib.celery.models import TaskState
+from servicelib.celery.task_manager import TaskManager
 from servicelib.logging_utils import log_catch
 from servicelib.rabbitmq import RPCRouter
 
-from ...modules.celery import get_task_manager_from_app
-
 _logger = logging.getLogger(__name__)
 router = RPCRouter()
 
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
-async def cancel(app: FastAPI, job_id: AsyncJobId, job_id_data: AsyncJobNameData):
-    assert app  # nosec
+async def cancel(
+    task_manager: TaskManager, job_id: AsyncJobId, job_id_data: AsyncJobNameData
+):
+    assert task_manager  # nosec
     assert job_id_data  # nosec
     try:
-        await get_task_manager_from_app(app).cancel_task(
+        await task_manager.cancel_task(
             task_context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
@@ -46,13 +46,13 @@ async def cancel(app: FastAPI, job_id: AsyncJobId, job_id_data: AsyncJobNameData
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
 async def status(
-    app: FastAPI, job_id: AsyncJobId, job_id_data: AsyncJobNameData
+    task_manager: TaskManager, job_id: AsyncJobId, job_id_data: AsyncJobNameData
 ) -> AsyncJobStatus:
-    assert app  # nosec
+    assert task_manager  # nosec
     assert job_id_data  # nosec
 
     try:
-        task_status = await get_task_manager_from_app(app).get_task_status(
+        task_status = await task_manager.get_task_status(
             task_context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
@@ -75,20 +75,20 @@ async def status(
     )
 )
 async def result(
-    app: FastAPI, job_id: AsyncJobId, job_id_data: AsyncJobNameData
+    task_manager: TaskManager, job_id: AsyncJobId, job_id_data: AsyncJobNameData
 ) -> AsyncJobResult:
-    assert app  # nosec
+    assert task_manager  # nosec
     assert job_id  # nosec
     assert job_id_data  # nosec
 
     try:
-        _status = await get_task_manager_from_app(app).get_task_status(
+        _status = await task_manager.get_task_status(
             task_context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
         if not _status.is_done:
             raise JobNotDoneError(job_id=job_id)
-        _result = await get_task_manager_from_app(app).get_task_result(
+        _result = await task_manager.get_task_result(
             task_context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
@@ -122,12 +122,12 @@ async def result(
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
 async def list_jobs(
-    app: FastAPI, filter_: str, job_id_data: AsyncJobNameData
+    task_manager: TaskManager, filter_: str, job_id_data: AsyncJobNameData
 ) -> list[AsyncJobGet]:
     _ = filter_
-    assert app  # nosec
+    assert task_manager  # nosec
     try:
-        tasks = await get_task_manager_from_app(app).list_tasks(
+        tasks = await task_manager.list_tasks(
             task_context=job_id_data.model_dump(),
         )
     except CeleryError as exc:

services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py renamed to packages/celery-library/src/celery_library/rpc/_async_jobs.py

@@ -1,16 +1,15 @@
 import logging
 from pathlib import Path
 
-from fastapi import FastAPI
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
     AsyncJobNameData,
 )
 from models_library.projects_nodes_io import LocationID
 from servicelib.celery.models import TaskMetadata
+from servicelib.celery.task_manager import TaskManager
 from servicelib.rabbitmq import RPCRouter
 
-from ...modules.celery import get_task_manager_from_app
 from .._worker_tasks._paths import compute_path_size as remote_compute_path_size
 from .._worker_tasks._paths import delete_paths as remote_delete_paths
 
@@ -20,13 +19,13 @@
 
 @router.expose(reraise_if_error_type=None)
 async def compute_path_size(
-    app: FastAPI,
+    task_manager: TaskManager,
     job_id_data: AsyncJobNameData,
     location_id: LocationID,
     path: Path,
 ) -> AsyncJobGet:
     task_name = remote_compute_path_size.__name__
-    task_uuid = await get_task_manager_from_app(app).submit_task(
+    task_uuid = await task_manager.submit_task(
         task_metadata=TaskMetadata(
             name=task_name,
         ),
@@ -41,13 +40,13 @@ async def compute_path_size(
 
 @router.expose(reraise_if_error_type=None)
 async def delete_paths(
-    app: FastAPI,
+    task_manager: TaskManager,
     job_id_data: AsyncJobNameData,
     location_id: LocationID,
     paths: set[Path],
 ) -> AsyncJobGet:
     task_name = remote_delete_paths.__name__
-    task_uuid = await get_task_manager_from_app(app).submit_task(
+    task_uuid = await task_manager.submit_task(
         task_metadata=TaskMetadata(
             name=task_name,
         ),

services/storage/src/simcore_service_storage/api/rpc/_paths.py

@@ -1,27 +1,26 @@
-from fastapi import FastAPI
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
     AsyncJobNameData,
 )
 from models_library.api_schemas_storage.storage_schemas import FoldersBody
 from models_library.api_schemas_webserver.storage import PathToExport
 from servicelib.celery.models import TaskMetadata, TasksQueue
+from servicelib.celery.task_manager import TaskManager
 from servicelib.rabbitmq import RPCRouter
 
-from ...modules.celery import get_task_manager_from_app
 from .._worker_tasks._simcore_s3 import deep_copy_files_from_project, export_data
 
 router = RPCRouter()
 
 
 @router.expose(reraise_if_error_type=None)
 async def copy_folders_from_project(
-    app: FastAPI,
+    task_manager: TaskManager,
     job_id_data: AsyncJobNameData,
     body: FoldersBody,
 ) -> AsyncJobGet:
     task_name = deep_copy_files_from_project.__name__
-    task_uuid = await get_task_manager_from_app(app).submit_task(
+    task_uuid = await task_manager.submit_task(
         task_metadata=TaskMetadata(
             name=task_name,
         ),
@@ -35,10 +34,12 @@ async def copy_folders_from_project(
 
 @router.expose()
 async def start_export_data(
-    app: FastAPI, job_id_data: AsyncJobNameData, paths_to_export: list[PathToExport]
+    task_manager: TaskManager,
+    job_id_data: AsyncJobNameData,
+    paths_to_export: list[PathToExport],
 ) -> AsyncJobGet:
     task_name = export_data.__name__
-    task_uuid = await get_task_manager_from_app(app).submit_task(
+    task_uuid = await task_manager.submit_task(
         task_metadata=TaskMetadata(
             name=task_name,
             ephemeral=False,

services/storage/src/simcore_service_storage/api/rpc/_simcore_s3.py

@@ -1,12 +1,14 @@
 import logging
 
+from celery_library.rpc import _async_jobs
 from fastapi import FastAPI
 from models_library.api_schemas_storage import STORAGE_RPC_NAMESPACE
 from servicelib.logging_utils import log_context
 from servicelib.rabbitmq import RPCRouter
+from simcore_service_storage.modules.celery import get_task_manager_from_app
 
 from ...modules.rabbitmq import get_rabbitmq_rpc_server
-from . import _async_jobs, _paths, _simcore_s3
+from . import _paths, _simcore_s3
 
 _logger = logging.getLogger(__name__)
 
@@ -18,15 +20,18 @@
 ]
 
 
-def setup_rpc_api_routes(app: FastAPI) -> None:
+def setup_rpc_routes(app: FastAPI) -> None:
     async def startup() -> None:
         with log_context(
             _logger,
             logging.INFO,
             msg="Storage startup RPC API Routes",
         ):
             rpc_server = get_rabbitmq_rpc_server(app)
+            task_manager = get_task_manager_from_app(app)
             for router in ROUTERS:
-                await rpc_server.register_router(router, STORAGE_RPC_NAMESPACE, app)
+                await rpc_server.register_router(
+                    router, STORAGE_RPC_NAMESPACE, task_manager=task_manager
+                )
 
     app.add_event_handler("startup", startup)

services/storage/src/simcore_service_storage/api/rpc/routes.py

@@ -32,7 +32,7 @@
     APP_WORKER_STARTED_BANNER_MSG,
 )
 from ..api.rest.routes import setup_rest_api_routes
-from ..api.rpc.routes import setup_rpc_api_routes
+from ..api.rpc.routes import setup_rpc_routes
 from ..dsm import setup_dsm
 from ..dsm_cleaner import setup_dsm_cleaner
 from ..exceptions.handlers import set_exception_handlers
@@ -92,10 +92,11 @@ def create_app(settings: ApplicationSettings) -> FastAPI:  # noqa: C901
 
     if not settings.STORAGE_WORKER_MODE:
         setup_rabbitmq(app)
-        setup_rpc_api_routes(app)
 
         assert settings.STORAGE_CELERY  # nosec
         setup_task_manager(app, celery_settings=settings.STORAGE_CELERY)
+
+        setup_rpc_routes(app)
     setup_rest_api_long_running_tasks_for_uploads(app)
     setup_rest_api_routes(app, API_VTAG)
     set_exception_handlers(app)

services/storage/src/simcore_service_storage/core/application.py

@@ -13,7 +13,7 @@
 
 def setup_task_manager(app: FastAPI, celery_settings: CelerySettings) -> None:
     async def on_startup() -> None:
-        app.state.celery_client = await create_task_manager(
+        app.state.task_manager = await create_task_manager(
             create_app(celery_settings), celery_settings
         )
 
@@ -24,7 +24,7 @@ async def on_startup() -> None:
 
 
 def get_task_manager_from_app(app: FastAPI) -> CeleryTaskManager:
-    assert hasattr(app.state, "celery_client")  # nosec
-    celery_client = app.state.celery_client
-    assert isinstance(celery_client, CeleryTaskManager)  # nosec
-    return celery_client
+    assert hasattr(app.state, "task_manager")  # nosec
+    task_manager = app.state.task_manager
+    assert isinstance(task_manager, CeleryTaskManager)  # nosec
+    return task_manager