improve

Author: sanderegg

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, ParamSpec, TypeVar
+from typing import Any, Concatenate, ParamSpec, TypeAlias, TypeVar, overload
 
 from celery import (  # type: ignore[import-untyped]
     Celery,
@@ -52,24 +52,58 @@ def wrapper(task: Task, *args: Any, **kwargs: Any) -> Any:
 P = ParamSpec("P")
 R = TypeVar("R")
 
+TaskId: TypeAlias = str
+
 
 def _async_task_wrapper(
     app: Celery,
-) -> Callable[[Callable[P, Coroutine[Any, Any, R]]], Callable[P, R]]:
-    def decorator(coro: Callable[P, Coroutine[Any, Any, R]]) -> Callable[P, R]:
+) -> Callable[
+    [Callable[Concatenate[Task, TaskId, P], Coroutine[Any, Any, R]]],
+    Callable[Concatenate[Task, P], R],
+]:
+    def decorator(
+        coro: Callable[Concatenate[Task, TaskId, P], Coroutine[Any, Any, R]],
+    ) -> Callable[Concatenate[Task, P], R]:
         @wraps(coro)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        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)
+            # 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(
-                coro(*args, **kwargs), get_event_loop(fastapi_app)
+                coro(task, task.request.id, *args, **kwargs),
+                get_event_loop(fastapi_app),
             ).result()
 
         return wrapper
 
     return decorator
 
 
-def define_task(app: Celery, fn: Callable, task_name: str | None = None):
+@overload
+def define_task(
+    app: Celery,
+    fn: Callable[Concatenate[Task, TaskId, P], Coroutine[Any, Any, R]],
+    task_name: str | None = None,
+) -> None: ...
+
+
+@overload
+def define_task(
+    app: Celery,
+    fn: Callable[Concatenate[Task, P], R],
+    task_name: str | None = None,
+) -> None: ...
+
+
+def define_task(
+    app: Celery,
+    fn: (
+        Callable[Concatenate[Task, TaskId, P], Coroutine[Any, Any, R]]
+        | Callable[Concatenate[Task, P], R]
+    ),
+    task_name: str | None = None,
+):
     wrapped_fn = error_handling(fn)
     if asyncio.iscoroutinefunction(fn):
         wrapped_fn = _async_task_wrapper(app)(fn)