use task_id for coroutines

Author: sanderegg

services/storage/src/simcore_service_storage/api/_worker_tasks/_files.py

@@ -10,18 +10,21 @@
 
 from ...dsm import get_dsm_provider
 from ...models import FileMetaData
+from ...modules.celery.models import TaskId
 from ...modules.celery.utils import get_fastapi_app
 
 _logger = logging.getLogger(__name__)
 
 
 async def complete_upload_file(
     task: Task,
+    task_id: TaskId,
     user_id: UserID,
     location_id: LocationID,
     file_id: StorageFileID,
     body: FileUploadCompletionBody,
 ) -> FileMetaData:
+    assert task_id  # nosec
     with log_context(
         _logger,
         logging.INFO,

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

@@ -8,14 +8,16 @@
 from servicelib.logging_utils import log_context
 
 from ...dsm import get_dsm_provider
+from ...modules.celery.models import TaskId
 from ...modules.celery.utils import get_fastapi_app
 
 _logger = logging.getLogger(__name__)
 
 
 async def compute_path_size(
-    task: Task, user_id: UserID, location_id: LocationID, path: Path
+    task: Task, task_id: TaskId, user_id: UserID, location_id: LocationID, path: Path
 ) -> ByteSize:
+    assert task_id  # nosec
     with log_context(
         _logger,
         logging.INFO,

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

@@ -10,6 +10,7 @@
 from servicelib.progress_bar import ProgressBarData
 
 from ...dsm import get_dsm_provider
+from ...modules.celery.models import TaskId
 from ...modules.celery.utils import get_celery_worker, get_fastapi_app
 from ...simcore_s3_dsm import SimcoreS3DataManager
 
@@ -28,9 +29,8 @@ def _task_progress_cb(task: Task, task_id: str, report: ProgressReport) -> None:
 
 
 async def deep_copy_files_from_project(
-    task: Task, task_id: str, user_id: UserID, body: FoldersBody
+    task: Task, task_id: TaskId, user_id: UserID, body: FoldersBody
 ) -> dict[str, Any]:
-    # _logger.error("%s", f"{task=}, {task.request.id=}, {task_id=}")
     with log_context(
         _logger,
         logging.INFO,

services/storage/src/simcore_service_storage/modules/celery/_task.py

@@ -3,7 +3,7 @@
 import traceback
 from collections.abc import Callable, Coroutine
 from functools import wraps
-from typing import Any, Concatenate, ParamSpec, TypeAlias, TypeVar, overload
+from typing import Any, Concatenate, ParamSpec, TypeVar, overload
 
 from celery import (  # type: ignore[import-untyped]
     Celery,
@@ -13,7 +13,7 @@
 from celery.exceptions import Ignore  # type: ignore[import-untyped]
 
 from . import get_event_loop
-from .models import TaskError, TaskState
+from .models import TaskError, TaskId, TaskState
 from .utils import get_fastapi_app
 
 _logger = logging.getLogger(__name__)
@@ -52,8 +52,6 @@ def wrapper(task: Task, *args: Any, **kwargs: Any) -> Any:
 P = ParamSpec("P")
 R = TypeVar("R")
 
-TaskId: TypeAlias = str
-
 
 def _async_task_wrapper(
     app: Celery,
@@ -67,7 +65,7 @@ def decorator(
         @wraps(coro)
         def wrapper(task: Task, *args: P.args, **kwargs: P.kwargs) -> R:
             fastapi_app = get_fastapi_app(app)
-            _logger.debug("BEFORE task id: %s", task.request.id)
+            _logger.debug("task id: %s", task.request.id)
             # NOTE: task.request is a thread local object, so we need to pass the id explicitly
             assert task.request.id is not None  # nosec
             return asyncio.run_coroutine_threadsafe(

services/storage/src/simcore_service_storage/modules/celery/models.py

@@ -55,3 +55,6 @@ def _check_consistency(self) -> Self:
 class TaskError(BaseModel):
     exc_type: str
     exc_msg: str
+
+
+TaskId: TypeAlias = str