refactor: move task metadata

Author: giancarloromeo

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

@@ -82,10 +82,10 @@ async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
             )
             return None
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
+    async def list_tasks(self, context: TaskContext) -> list[Task]:
         search_key = (
             _CELERY_TASK_INFO_PREFIX
-            + build_task_id_prefix(task_context)
+            + build_task_id_prefix(context)
             + _CELERY_TASK_ID_KEY_SEPARATOR
         )
         search_key_len = len(search_key)

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

@@ -38,7 +38,7 @@ async def cancel(
     assert job_id_data  # nosec
     try:
         await task_manager.cancel_task(
-            task_context=job_id_data.model_dump(),
+            context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
     except CeleryError as exc:
@@ -54,7 +54,7 @@ async def status(
 
     try:
         task_status = await task_manager.get_task_status(
-            task_context=job_id_data.model_dump(),
+            context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
     except CeleryError as exc:
@@ -84,13 +84,13 @@ async def result(
 
     try:
         _status = await task_manager.get_task_status(
-            task_context=job_id_data.model_dump(),
+            context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
         if not _status.is_done:
             raise JobNotDoneError(job_id=job_id)
         _result = await task_manager.get_task_result(
-            task_context=job_id_data.model_dump(),
+            context=job_id_data.model_dump(),
             task_uuid=job_id,
         )
     except CeleryError as exc:
@@ -129,7 +129,7 @@ async def list_jobs(
     assert task_manager  # nosec
     try:
         tasks = await task_manager.list_tasks(
-            task_context=job_id_data.model_dump(),
+            context=job_id_data.model_dump(),
         )
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc

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

@@ -10,11 +10,14 @@
 from common_library.async_tools import make_async
 from models_library.progress_bar import ProgressReport
 from servicelib.celery.models import (
+    TASK_QUEUE_DEFAULT,
     Task,
     TaskContext,
     TaskID,
     TaskInfoStore,
     TaskMetadata,
+    TaskName,
+    TaskQueue,
     TaskState,
     TaskStatus,
     TaskUUID,
@@ -39,63 +42,69 @@ class CeleryTaskManager:
 
     async def send_task(
         self,
-        task_metadata: TaskMetadata,
+        name: TaskName,
+        context: TaskContext,
         *,
-        task_context: TaskContext,
-        **task_params,
+        is_ephemeral: bool = False,
+        queue: TaskQueue = TASK_QUEUE_DEFAULT,
+        **params,
     ) -> TaskUUID:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Send {task_metadata.name=}: {task_context=} {task_params=}",
+            msg=f"Send {name=}: {context=} {params=}",
         ):
             task_uuid = uuid4()
-            task_id = build_task_id(task_context, task_uuid)
+            task_id = build_task_id(context, task_uuid)
             self._celery_app.send_task(
-                task_metadata.name,
+                name,
                 task_id=task_id,
-                kwargs={"task_id": task_id} | task_params,
-                queue=task_metadata.queue,
+                kwargs={"task_id": task_id} | params,
+                queue=queue,
             )
 
             expiry = (
                 self._celery_settings.CELERY_EPHEMERAL_RESULT_EXPIRES
-                if task_metadata.ephemeral
+                if is_ephemeral
                 else self._celery_settings.CELERY_RESULT_EXPIRES
             )
             await self._task_info_store.create_task(
-                task_id, task_metadata, expiry=expiry
+                task_id,
+                TaskMetadata(
+                    name=name,
+                    ephemeral=is_ephemeral,
+                    queue=queue,
+                ),
+                expiry=expiry,
             )
             return task_uuid
 
     @make_async()
     def _abort_task(self, task_id: TaskID) -> None:
         AbortableAsyncResult(task_id, app=self._celery_app).abort()
 
-    async def cancel_task(self, task_context: TaskContext, task_uuid: TaskUUID) -> None:
+    async def cancel_task(self, context: TaskContext, task_uuid: TaskUUID) -> None:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"task cancellation: {task_context=} {task_uuid=}",
+            msg=f"task cancellation: {context=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_context, task_uuid)
-            if not (await self.get_task_status(task_context, task_uuid)).is_done:
+            task_id = build_task_id(context, task_uuid)
+            if not (await self.get_task_status(context, task_uuid)).is_done:
                 await self._abort_task(task_id)
             await self._task_info_store.remove_task(task_id)
 
     @make_async()
     def _forget_task(self, task_id: TaskID) -> None:
         AbortableAsyncResult(task_id, app=self._celery_app).forget()
 
-    async def get_task_result(
-        self, task_context: TaskContext, task_uuid: TaskUUID
-    ) -> Any:
+    async def get_task_result(self, context: TaskContext, task_uuid: TaskUUID) -> Any:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Get task result: {task_context=} {task_uuid=}",
+            msg=f"Get task result: {context=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_context, task_uuid)
+            task_id = build_task_id(context, task_uuid)
             async_result = self._celery_app.AsyncResult(task_id)
             result = async_result.result
             if async_result.ready():
@@ -105,15 +114,15 @@ async def get_task_result(
                     await self._task_info_store.remove_task(task_id)
             return result
 
-    async def _get_task_progress_report(
-        self, task_context: TaskContext, task_uuid: TaskUUID, task_state: TaskState
+    async def _get_progress_report(
+        self, context: TaskContext, task_uuid: TaskUUID, state: TaskState
     ) -> ProgressReport:
-        if task_state in (TaskState.STARTED, TaskState.RETRY, TaskState.ABORTED):
-            task_id = build_task_id(task_context, task_uuid)
+        if state in (TaskState.STARTED, TaskState.RETRY, TaskState.ABORTED):
+            task_id = build_task_id(context, task_uuid)
             progress = await self._task_info_store.get_task_progress(task_id)
             if progress is not None:
                 return progress
-        if task_state in (
+        if state in (
             TaskState.SUCCESS,
             TaskState.FAILURE,
         ):
@@ -131,30 +140,30 @@ def _get_task_celery_state(self, task_id: TaskID) -> TaskState:
         return TaskState(self._celery_app.AsyncResult(task_id).state)
 
     async def get_task_status(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, context: TaskContext, task_uuid: TaskUUID
     ) -> TaskStatus:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Getting task status: {task_context=} {task_uuid=}",
+            msg=f"Getting task status: {context=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_context, task_uuid)
+            task_id = build_task_id(context, task_uuid)
             task_state = await self._get_task_celery_state(task_id)
             return TaskStatus(
                 task_uuid=task_uuid,
                 task_state=task_state,
-                progress_report=await self._get_task_progress_report(
-                    task_context, task_uuid, task_state
+                progress_report=await self._get_progress_report(
+                    context, task_uuid, task_state
                 ),
             )
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
+    async def list_tasks(self, context: TaskContext) -> list[Task]:
         with log_context(
             _logger,
             logging.DEBUG,
-            msg=f"Listing tasks: {task_context=}",
+            msg=f"Listing tasks: {context=}",
         ):
-            return await self._task_info_store.list_tasks(task_context)
+            return await self._task_info_store.list_tasks(context)
 
     async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
         await self._task_info_store.set_task_progress(

packages/celery-library/tests/conftest.py

@@ -120,7 +120,7 @@ def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
 
 
 @pytest.fixture
-async def celery_task_manager(
+async def task_manager(
     celery_app: Celery,
     celery_settings: CelerySettings,
     with_celery_worker: TestWorkController,

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

@@ -27,7 +27,7 @@
 from models_library.rabbitmq_basic_types import RPCNamespace
 from models_library.users import UserID
 from pydantic import TypeAdapter
-from servicelib.celery.models import TaskID, TaskMetadata
+from servicelib.celery.models import TaskID
 from servicelib.celery.task_manager import TaskManager
 from servicelib.rabbitmq import RabbitMQRPCClient, RPCRouter
 from servicelib.rabbitmq.rpc_interfaces.async_jobs import async_jobs
@@ -83,7 +83,9 @@ async def rpc_sync_job(
 ) -> AsyncJobGet:
     task_name = sync_job.__name__
     task_uuid = await task_manager.send_task(
-        TaskMetadata(name=task_name), task_context=job_id_data.model_dump(), **kwargs
+        name=task_name,
+        context=job_id_data.model_dump(),
+        **kwargs,
     )
 
     return AsyncJobGet(job_id=task_uuid, job_name=task_name)
@@ -95,7 +97,9 @@ async def rpc_async_job(
 ) -> AsyncJobGet:
     task_name = async_job.__name__
     task_uuid = await task_manager.send_task(
-        TaskMetadata(name=task_name), task_context=job_id_data.model_dump(), **kwargs
+        name=task_name,
+        context=job_id_data.model_dump(),
+        **kwargs,
     )
 
     return AsyncJobGet(job_id=task_uuid, job_name=task_name)
@@ -139,13 +143,13 @@ async def async_job(task: Task, task_id: TaskID, action: Action, payload: Any) -
 
 @pytest.fixture
 async def register_rpc_routes(
-    async_jobs_rabbitmq_rpc_client: RabbitMQRPCClient, celery_task_manager: TaskManager
+    async_jobs_rabbitmq_rpc_client: RabbitMQRPCClient, task_manager: TaskManager
 ) -> None:
     await async_jobs_rabbitmq_rpc_client.register_router(
-        _async_jobs.router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=celery_task_manager
+        _async_jobs.router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=task_manager
     )
     await async_jobs_rabbitmq_rpc_client.register_router(
-        router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=celery_task_manager
+        router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=task_manager
     )