🎨 TaskID -> TaskKey (#8409)

Author: bisgaard-itis

packages/celery-library/src/celery_library/backends/redis.py

@@ -10,8 +10,8 @@
     ExecutionMetadata,
     OwnerMetadata,
     Task,
-    TaskID,
     TaskInfoStore,
+    TaskKey,
 )
 from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
@@ -24,8 +24,8 @@
 _logger = logging.getLogger(__name__)
 
 
-def _build_key(task_id: TaskID) -> str:
-    return _CELERY_TASK_INFO_PREFIX + task_id
+def _build_key(task_key: TaskKey) -> str:
+    return _CELERY_TASK_INFO_PREFIX + task_key
 
 
 class RedisTaskInfoStore:
@@ -34,27 +34,27 @@ def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
 
     async def create_task(
         self,
-        task_id: TaskID,
+        task_key: TaskKey,
         execution_metadata: ExecutionMetadata,
         expiry: timedelta,
     ) -> None:
-        task_key = _build_key(task_id)
+        redis_key = _build_key(task_key)
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
-                name=task_key,
+                name=redis_key,
                 key=_CELERY_TASK_METADATA_KEY,
                 value=execution_metadata.model_dump_json(),
             )
         )
         await self._redis_client_sdk.redis.expire(
-            task_key,
+            redis_key,
             expiry,
         )
 
-    async def get_task_metadata(self, task_id: TaskID) -> ExecutionMetadata | None:
+    async def get_task_metadata(self, task_key: TaskKey) -> ExecutionMetadata | None:
         raw_result = await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hget(
-                _build_key(task_id), _CELERY_TASK_METADATA_KEY
+                _build_key(task_key), _CELERY_TASK_METADATA_KEY
             )
         )
         if not raw_result:
@@ -64,14 +64,16 @@ async def get_task_metadata(self, task_id: TaskID) -> ExecutionMetadata | None:
             return ExecutionMetadata.model_validate_json(raw_result)
         except ValidationError as exc:
             _logger.debug(
-                "Failed to deserialize task metadata for task %s: %s", task_id, f"{exc}"
+                "Failed to deserialize task metadata for task %s: %s",
+                task_key,
+                f"{exc}",
             )
             return None
 
-    async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
+    async def get_task_progress(self, task_key: TaskKey) -> ProgressReport | None:
         raw_result = await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hget(
-                _build_key(task_id), _CELERY_TASK_PROGRESS_KEY
+                _build_key(task_key), _CELERY_TASK_PROGRESS_KEY
             )
         )
         if not raw_result:
@@ -81,12 +83,14 @@ async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
             return ProgressReport.model_validate_json(raw_result)
         except ValidationError as exc:
             _logger.debug(
-                "Failed to deserialize task progress for task %s: %s", task_id, f"{exc}"
+                "Failed to deserialize task progress for task %s: %s",
+                task_key,
+                f"{exc}",
             )
             return None
 
     async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
-        search_key = _CELERY_TASK_INFO_PREFIX + owner_metadata.model_dump_task_id(
+        search_key = _CELERY_TASK_INFO_PREFIX + owner_metadata.model_dump_task_key(
             task_uuid=WILDCARD
         )
 
@@ -122,20 +126,22 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
 
         return tasks
 
-    async def remove_task(self, task_id: TaskID) -> None:
-        await self._redis_client_sdk.redis.delete(_build_key(task_id))
+    async def remove_task(self, task_key: TaskKey) -> None:
+        await self._redis_client_sdk.redis.delete(_build_key(task_key))
 
-    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+    async def set_task_progress(
+        self, task_key: TaskKey, report: ProgressReport
+    ) -> None:
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
-                name=_build_key(task_id),
+                name=_build_key(task_key),
                 key=_CELERY_TASK_PROGRESS_KEY,
                 value=report.model_dump_json(),
             )
         )
 
-    async def task_exists(self, task_id: TaskID) -> bool:
-        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
+    async def task_exists(self, task_key: TaskKey) -> bool:
+        n = await self._redis_client_sdk.redis.exists(_build_key(task_key))
         assert isinstance(n, int)  # nosec
         return n > 0
 

packages/celery-library/src/celery_library/errors.py

@@ -27,10 +27,8 @@ def decode_celery_transferrable_error(error: TransferrableCeleryError) -> Except
 
 
 class TaskSubmissionError(OsparcErrorMixin, Exception):
-    msg_template = (
-        "Unable to submit task {task_name} with id '{task_id}' and params {task_params}"
-    )
+    msg_template = "Unable to submit task {task_name} with key '{task_key}' and params {task_params}"
 
 
 class TaskNotFoundError(OsparcErrorMixin, Exception):
-    msg_template = "Task with id '{task_id}' was not found"
+    msg_template = "Task with uuid '{task_uuid}' and owner_metadata '{owner_metadata}' was not found"

packages/celery-library/src/celery_library/task.py

@@ -10,7 +10,7 @@
 from celery.exceptions import Ignore  # type: ignore[import-untyped]
 from common_library.async_tools import cancel_wait_task
 from pydantic import NonNegativeInt
-from servicelib.celery.models import TaskID
+from servicelib.celery.models import TaskKey
 
 from .errors import encode_celery_transferrable_error
 from .utils import get_app_server
@@ -47,7 +47,7 @@ def wrapper(task: Task, *args: P.args, **kwargs: P.kwargs) -> R:
             # NOTE: task.request is a thread local object, so we need to pass the id explicitly
             assert task.request.id is not None  # nosec
 
-            async def _run_task(task_id: TaskID) -> R:
+            async def _run_task(task_key: TaskKey) -> R:
                 try:
                     async with asyncio.TaskGroup() as tg:
                         async_io_task = tg.create_task(
@@ -57,7 +57,7 @@ async def _run_task(task_id: TaskID) -> R:
                         async def _abort_monitor():
                             while not async_io_task.done():
                                 if not await app_server.task_manager.task_exists(
-                                    task_id
+                                    task_key
                                 ):
                                     await cancel_wait_task(
                                         async_io_task,
@@ -140,7 +140,7 @@ def wrapper(task: Task, *args: P.args, **kwargs: P.kwargs) -> R:
 @overload
 def register_task(
     app: Celery,
-    fn: Callable[Concatenate[Task, TaskID, P], Coroutine[Any, Any, R]],
+    fn: Callable[Concatenate[Task, TaskKey, P], Coroutine[Any, Any, R]],
     task_name: str | None = None,
     timeout: timedelta | None = _DEFAULT_TASK_TIMEOUT,
     max_retries: NonNegativeInt = _DEFAULT_MAX_RETRIES,
@@ -164,7 +164,7 @@ def register_task(
 def register_task(  # type: ignore[misc]
     app: Celery,
     fn: (
-        Callable[Concatenate[Task, TaskID, P], Coroutine[Any, Any, R]]
+        Callable[Concatenate[Task, TaskKey, P], Coroutine[Any, Any, R]]
         | Callable[Concatenate[Task, P], R]
     ),
     task_name: str | None = None,

packages/celery-library/src/celery_library/task_manager.py

@@ -12,8 +12,8 @@
     ExecutionMetadata,
     OwnerMetadata,
     Task,
-    TaskID,
     TaskInfoStore,
+    TaskKey,
     TaskState,
     TaskStatus,
     TaskUUID,
@@ -50,7 +50,7 @@ async def submit_task(
             msg=f"Submit {execution_metadata.name=}: {owner_metadata=} {task_params=}",
         ):
             task_uuid = uuid4()
-            task_id = owner_metadata.model_dump_task_id(task_uuid=task_uuid)
+            task_key = owner_metadata.model_dump_task_key(task_uuid=task_uuid)
 
             expiry = (
                 self._celery_settings.CELERY_EPHEMERAL_RESULT_EXPIRES
@@ -60,26 +60,26 @@ async def submit_task(
 
             try:
                 await self._task_info_store.create_task(
-                    task_id, execution_metadata, expiry=expiry
+                    task_key, execution_metadata, expiry=expiry
                 )
                 self._celery_app.send_task(
                     execution_metadata.name,
-                    task_id=task_id,
-                    kwargs={"task_id": task_id} | task_params,
+                    task_id=task_key,
+                    kwargs={"task_key": task_key} | task_params,
                     queue=execution_metadata.queue.value,
                 )
             except CeleryError as exc:
                 try:
-                    await self._task_info_store.remove_task(task_id)
+                    await self._task_info_store.remove_task(task_key)
                 except CeleryError:
                     _logger.warning(
                         "Unable to cleanup task '%s' during error handling",
-                        task_id,
+                        task_key,
                         exc_info=True,
                     )
                 raise TaskSubmissionError(
                     task_name=execution_metadata.name,
-                    task_id=task_id,
+                    task_key=task_key,
                     task_params=task_params,
                 ) from exc
 
@@ -93,19 +93,21 @@ async def cancel_task(
             logging.DEBUG,
             msg=f"task cancellation: {owner_metadata=} {task_uuid=}",
         ):
-            task_id = owner_metadata.model_dump_task_id(task_uuid=task_uuid)
-            if not await self.task_exists(task_id):
-                raise TaskNotFoundError(task_id=task_id)
+            task_key = owner_metadata.model_dump_task_key(task_uuid=task_uuid)
+            if not await self.task_exists(task_key):
+                raise TaskNotFoundError(
+                    task_uuid=task_uuid, owner_metadata=owner_metadata
+                )
 
-            await self._task_info_store.remove_task(task_id)
-            await self._forget_task(task_id)
+            await self._task_info_store.remove_task(task_key)
+            await self._forget_task(task_key)
 
-    async def task_exists(self, task_id: TaskID) -> bool:
-        return await self._task_info_store.task_exists(task_id)
+    async def task_exists(self, task_key: TaskKey) -> bool:
+        return await self._task_info_store.task_exists(task_key)
 
     @make_async()
-    def _forget_task(self, task_id: TaskID) -> None:
-        self._celery_app.AsyncResult(task_id).forget()
+    def _forget_task(self, task_key: TaskKey) -> None:
+        self._celery_app.AsyncResult(task_key).forget()
 
     async def get_task_result(
         self, owner_metadata: OwnerMetadata, task_uuid: TaskUUID
@@ -115,24 +117,26 @@ async def get_task_result(
             logging.DEBUG,
             msg=f"Get task result: {owner_metadata=} {task_uuid=}",
         ):
-            task_id = owner_metadata.model_dump_task_id(task_uuid=task_uuid)
-            if not await self.task_exists(task_id):
-                raise TaskNotFoundError(task_id=task_id)
+            task_key = owner_metadata.model_dump_task_key(task_uuid=task_uuid)
+            if not await self.task_exists(task_key):
+                raise TaskNotFoundError(
+                    task_uuid=task_uuid, owner_metadata=owner_metadata
+                )
 
-            async_result = self._celery_app.AsyncResult(task_id)
+            async_result = self._celery_app.AsyncResult(task_key)
             result = async_result.result
             if async_result.ready():
-                task_metadata = await self._task_info_store.get_task_metadata(task_id)
+                task_metadata = await self._task_info_store.get_task_metadata(task_key)
                 if task_metadata is not None and task_metadata.ephemeral:
-                    await self._task_info_store.remove_task(task_id)
-                    await self._forget_task(task_id)
+                    await self._task_info_store.remove_task(task_key)
+                    await self._forget_task(task_key)
             return result
 
     async def _get_task_progress_report(
-        self, task_id: TaskID, task_state: TaskState
+        self, task_key: TaskKey, task_state: TaskState
     ) -> ProgressReport:
         if task_state in (TaskState.STARTED, TaskState.RETRY):
-            progress = await self._task_info_store.get_task_progress(task_id)
+            progress = await self._task_info_store.get_task_progress(task_key)
             if progress is not None:
                 return progress
 
@@ -147,8 +151,8 @@ async def _get_task_progress_report(
         )
 
     @make_async()
-    def _get_task_celery_state(self, task_id: TaskID) -> TaskState:
-        return TaskState(self._celery_app.AsyncResult(task_id).state)
+    def _get_task_celery_state(self, task_key: TaskKey) -> TaskState:
+        return TaskState(self._celery_app.AsyncResult(task_key).state)
 
     async def get_task_status(
         self, owner_metadata: OwnerMetadata, task_uuid: TaskUUID
@@ -158,16 +162,18 @@ async def get_task_status(
             logging.DEBUG,
             msg=f"Getting task status: {owner_metadata=} {task_uuid=}",
         ):
-            task_id = owner_metadata.model_dump_task_id(task_uuid=task_uuid)
-            if not await self.task_exists(task_id):
-                raise TaskNotFoundError(task_id=task_id)
+            task_key = owner_metadata.model_dump_task_key(task_uuid=task_uuid)
+            if not await self.task_exists(task_key):
+                raise TaskNotFoundError(
+                    task_uuid=task_uuid, owner_metadata=owner_metadata
+                )
 
-            task_state = await self._get_task_celery_state(task_id)
+            task_state = await self._get_task_celery_state(task_key)
             return TaskStatus(
                 task_uuid=task_uuid,
                 task_state=task_state,
                 progress_report=await self._get_task_progress_report(
-                    task_id, task_state
+                    task_key, task_state
                 ),
             )
 
@@ -179,9 +185,11 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         ):
             return await self._task_info_store.list_tasks(owner_metadata)
 
-    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+    async def set_task_progress(
+        self, task_key: TaskKey, report: ProgressReport
+    ) -> None:
         await self._task_info_store.set_task_progress(
-            task_id=task_id,
+            task_key=task_key,
             report=report,
         )
 

packages/celery-library/tests/unit/test_async_jobs.py

@@ -26,7 +26,7 @@
 from models_library.rabbitmq_basic_types import RPCNamespace
 from models_library.users import UserID
 from pydantic import TypeAdapter
-from servicelib.celery.models import ExecutionMetadata, OwnerMetadata, TaskID
+from servicelib.celery.models import ExecutionMetadata, OwnerMetadata, TaskKey
 from servicelib.celery.task_manager import TaskManager
 from servicelib.rabbitmq import RabbitMQRPCClient, RPCRouter
 from servicelib.rabbitmq.rpc_interfaces.async_jobs import async_jobs
@@ -121,15 +121,15 @@ async def _process_action(action: str, payload: Any) -> Any:
     return None
 
 
-def sync_job(task: Task, task_id: TaskID, action: Action, payload: Any) -> Any:
+def sync_job(task: Task, task_key: TaskKey, action: Action, payload: Any) -> Any:
     _ = task
-    _ = task_id
+    _ = task_key
     return asyncio.run(_process_action(action, payload))
 
 
-async def async_job(task: Task, task_id: TaskID, action: Action, payload: Any) -> Any:
+async def async_job(task: Task, task_key: TaskKey, action: Action, payload: Any) -> Any:
     _ = task
-    _ = task_id
+    _ = task_key
     return await _process_action(action, payload)