fix: job missing error

Author: giancarloromeo

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

@@ -1,6 +1,8 @@
 import base64
 import pickle
 
+from common_library.errors_classes import OsparcErrorMixin
+
 
 class TransferrableCeleryError(Exception):
     def __repr__(self) -> str:
@@ -22,3 +24,7 @@ def decode_celery_transferrable_error(error: TransferrableCeleryError) -> Except
     assert isinstance(error, TransferrableCeleryError)  # nosec
     result: Exception = pickle.loads(base64.b64decode(error.args[0]))  # noqa: S301
     return result
+
+
+class TaskNotFoundError(OsparcErrorMixin, Exception):
+    msg_template = "Task with id '{task_id}' not found"

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

@@ -13,6 +13,7 @@
 from models_library.api_schemas_rpc_async_jobs.exceptions import (
     JobAbortedError,
     JobError,
+    JobMissingError,
     JobNotDoneError,
     JobSchedulerError,
 )
@@ -22,6 +23,7 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ..errors import (
+    TaskNotFoundError,
     TransferrableCeleryError,
     decode_celery_transferrable_error,
 )
@@ -30,7 +32,7 @@
 router = RPCRouter()
 
 
-@router.expose(reraise_if_error_type=(JobSchedulerError,))
+@router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def cancel(
     task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
 ):
@@ -42,11 +44,13 @@ async def cancel(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(exc=f"{exc}") from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
 
-@router.expose(reraise_if_error_type=(JobSchedulerError,))
+@router.expose(reraise_if_error_type=(JobSchedulerError, JobMissingError))
 async def status(
     task_manager: TaskManager, job_id: AsyncJobId, job_filter: AsyncJobFilter
 ) -> AsyncJobStatus:
@@ -59,6 +63,8 @@ async def status(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(exc=f"{exc}") from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
@@ -71,9 +77,10 @@ async def status(
 
 @router.expose(
     reraise_if_error_type=(
+        JobAbortedError,
         JobError,
+        JobMissingError,
         JobNotDoneError,
-        JobAbortedError,
         JobSchedulerError,
     )
 )
@@ -97,6 +104,8 @@ async def result(
             task_filter=task_filter,
             task_uuid=job_id,
         )
+    except TaskNotFoundError as exc:
+        raise JobMissingError(exc=f"{exc}") from exc
     except CeleryError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 

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

@@ -21,6 +21,7 @@
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
+from .errors import TaskNotFoundError
 from .utils import build_task_id
 
 _logger = logging.getLogger(__name__)
@@ -74,6 +75,9 @@ async def cancel_task(self, task_filter: TaskFilter, task_uuid: TaskUUID) -> Non
             msg=f"task cancellation: {task_filter=} {task_uuid=}",
         ):
             task_id = build_task_id(task_filter, task_uuid)
+            if not await self.exists_task(task_id):
+                raise TaskNotFoundError(task_id=task_id)
+
             await self._task_info_store.remove_task(task_id)
             await self._forget_task(task_id)
 
@@ -93,6 +97,9 @@ async def get_task_result(
             msg=f"Get task result: {task_filter=} {task_uuid=}",
         ):
             task_id = build_task_id(task_filter, task_uuid)
+            if not await self.exists_task(task_id):
+                raise TaskNotFoundError(task_id=task_id)
+
             async_result = self._celery_app.AsyncResult(task_id)
             result = async_result.result
             if async_result.ready():
@@ -103,10 +110,9 @@ async def get_task_result(
             return result
 
     async def _get_task_progress_report(
-        self, task_filter: TaskFilter, task_uuid: TaskUUID, task_state: TaskState
+        self, task_id: TaskID, task_state: TaskState
     ) -> ProgressReport:
         if task_state in (TaskState.STARTED, TaskState.RETRY):
-            task_id = build_task_id(task_filter, task_uuid)
             progress = await self._task_info_store.get_task_progress(task_id)
             if progress is not None:
                 return progress
@@ -135,14 +141,14 @@ async def get_task_status(
         ):
             task_id = build_task_id(task_filter, task_uuid)
             if not await self.exists_task(task_id):
-                task_state = TaskState.ABORTED
-            else:
-                task_state = await self._get_task_celery_state(task_id)
+                raise TaskNotFoundError(task_id=task_id)
+
+            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_filter, task_uuid, task_state
+                    task_id, task_state
                 ),
             )
 

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

@@ -20,8 +20,8 @@
     AsyncJobGet,
 )
 from models_library.api_schemas_rpc_async_jobs.exceptions import (
-    JobAbortedError,
     JobError,
+    JobMissingError,
 )
 from models_library.products import ProductName
 from models_library.rabbitmq_basic_types import RPCNamespace
@@ -308,12 +308,6 @@ async def test_async_jobs_cancel(
         job_filter=job_filter,
     )
 
-    await _wait_for_job(
-        async_jobs_rabbitmq_rpc_client,
-        async_job_get=async_job_get,
-        job_filter=job_filter,
-    )
-
     jobs = await async_jobs.list_jobs(
         async_jobs_rabbitmq_rpc_client,
         rpc_namespace=ASYNC_JOBS_RPC_NAMESPACE,
@@ -322,7 +316,15 @@ async def test_async_jobs_cancel(
     )
     assert async_job_get.job_id not in [job.job_id for job in jobs]
 
-    with pytest.raises(JobAbortedError):
+    with pytest.raises(JobMissingError):
+        await async_jobs.status(
+            async_jobs_rabbitmq_rpc_client,
+            rpc_namespace=ASYNC_JOBS_RPC_NAMESPACE,
+            job_id=async_job_get.job_id,
+            job_filter=job_filter,
+        )
+
+    with pytest.raises(JobMissingError):
         await async_jobs.result(
             async_jobs_rabbitmq_rpc_client,
             rpc_namespace=ASYNC_JOBS_RPC_NAMESPACE,