🎨 TaskID -> TaskKey

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

@@ -6,12 +6,12 @@
 from models_library.progress_bar import ProgressReport
 from pydantic import ValidationError
 from servicelib.celery.models import (
+    WILDCARD,
+    ExecutionMetadata,
+    OwnerMetadata,
     Task,
-    TaskFilter,
-    TaskID,
     TaskInfoStore,
-    TaskMetadata,
-    Wildcard,
+    TaskKey,
 )
 from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
@@ -24,7 +24,7 @@
 _logger = logging.getLogger(__name__)
 
 
-def _build_key(task_id: TaskID) -> str:
+def _build_key(task_id: TaskKey) -> str:
     return _CELERY_TASK_INFO_PREFIX + task_id
 
 
@@ -34,44 +34,46 @@ def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
 
     async def create_task(
         self,
-        task_id: TaskID,
-        task_metadata: TaskMetadata,
+        task_id: TaskKey,
+        execution_metadata: ExecutionMetadata,
         expiry: timedelta,
     ) -> None:
         task_key = _build_key(task_id)
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
                 name=task_key,
                 key=_CELERY_TASK_METADATA_KEY,
-                value=task_metadata.model_dump_json(),
+                value=execution_metadata.model_dump_json(),
             )
         )
         await self._redis_client_sdk.redis.expire(
             task_key,
             expiry,
         )
 
-    async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | 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:
             return None
 
         try:
-            return TaskMetadata.model_validate_json(raw_result)
+            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,13 +83,15 @@ 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, task_filter: TaskFilter) -> list[Task]:
-        search_key = _CELERY_TASK_INFO_PREFIX + task_filter.create_task_id(
-            task_uuid=Wildcard()
+    async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
+        search_key = _CELERY_TASK_INFO_PREFIX + owner_metadata.model_dump_task_key(
+            task_uuid=WILDCARD
         )
 
         keys: list[str] = []
@@ -112,30 +116,32 @@ async def list_tasks(self, task_filter: TaskFilter) -> list[Task]:
                 continue
 
             with contextlib.suppress(ValidationError):
-                task_metadata = TaskMetadata.model_validate_json(raw_metadata)
+                execution_metadata = ExecutionMetadata.model_validate_json(raw_metadata)
                 tasks.append(
                     Task(
-                        uuid=TaskFilter.get_task_uuid(key),
-                        metadata=task_metadata,
+                        uuid=OwnerMetadata.get_task_uuid(key),
+                        metadata=execution_metadata,
                     )
                 )
 
         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/rpc/_async_jobs.py

@@ -4,7 +4,6 @@
 
 from celery.exceptions import CeleryError  # type: ignore[import-untyped]
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
-    AsyncJobFilter,
     AsyncJobGet,
     AsyncJobId,
     AsyncJobResult,
@@ -17,7 +16,7 @@
     JobNotDoneError,
     JobSchedulerError,
 )
-from servicelib.celery.models import TaskFilter, TaskState
+from servicelib.celery.models import OwnerMetadata, TaskState
 from servicelib.celery.task_manager import TaskManager
 from servicelib.logging_utils import log_catch
 from servicelib.rabbitmq import RPCRouter
@@ -34,14 +33,13 @@
 
 @router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def cancel(
-    task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
+    task_manager: TaskManager, job_id: AsyncJobId, owner_metadata: OwnerMetadata
 ):
     assert task_manager  # nosec
-    assert job_filter  # nosec
-    task_filter = TaskFilter.model_validate(job_filter.model_dump())
+    assert owner_metadata  # nosec
     try:
         await task_manager.cancel_task(
-            task_filter=task_filter,
+            owner_metadata=owner_metadata,
             task_uuid=job_id,
         )
     except TaskNotFoundError as exc:
@@ -52,15 +50,14 @@ async def cancel(
 
 @router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def status(
-    task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
+    task_manager: TaskManager, job_id: AsyncJobId, owner_metadata: OwnerMetadata
 ) -> AsyncJobStatus:
     assert task_manager  # nosec
-    assert job_filter  # nosec
+    assert owner_metadata  # nosec
 
-    task_filter = TaskFilter.model_validate(job_filter.model_dump())
     try:
         task_status = await task_manager.get_task_status(
-            task_filter=task_filter,
+            owner_metadata=owner_metadata,
             task_uuid=job_id,
         )
     except TaskNotFoundError as exc:
@@ -85,23 +82,21 @@ async def status(
     )
 )
 async def result(
-    task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
+    task_manager: TaskManager, job_id: AsyncJobId, owner_metadata: OwnerMetadata
 ) -> AsyncJobResult:
     assert task_manager  # nosec
     assert job_id  # nosec
-    assert job_filter  # nosec
-
-    task_filter = TaskFilter.model_validate(job_filter.model_dump())
+    assert owner_metadata  # nosec
 
     try:
         _status = await task_manager.get_task_status(
-            task_filter=task_filter,
+            owner_metadata=owner_metadata,
             task_uuid=job_id,
         )
         if not _status.is_done:
             raise JobNotDoneError(job_id=job_id)
         _result = await task_manager.get_task_result(
-            task_filter=task_filter,
+            owner_metadata=owner_metadata,
             task_uuid=job_id,
         )
     except TaskNotFoundError as exc:
@@ -134,13 +129,12 @@ async def result(
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
 async def list_jobs(
-    task_manager: TaskManager, job_filter: AsyncJobFilter
+    task_manager: TaskManager, owner_metadata: OwnerMetadata
 ) -> list[AsyncJobGet]:
     assert task_manager  # nosec
-    task_filter = TaskFilter.model_validate(job_filter.model_dump())
     try:
         tasks = await task_manager.list_tasks(
-            task_filter=task_filter,
+            owner_metadata=owner_metadata,
         )
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc

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_id: TaskKey) -> R:
                 try:
                     async with asyncio.TaskGroup() as tg:
                         async_io_task = tg.create_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,