handle task not found errors properly in api-server

Author: bisgaard-itis

services/api-server/src/simcore_service_api_server/api/routes/tasks.py

@@ -1,6 +1,9 @@
 import logging
+from contextlib import contextmanager
 from typing import Annotated, Any
 
+from celery_library import task
+from celery_library.errors import TaskNotFoundError
 from common_library.error_codes import create_error_code
 from common_library.logging.logging_errors import create_troubleshooting_log_kwargs
 from fastapi import APIRouter, Depends, FastAPI, HTTPException, status
@@ -18,6 +21,7 @@
 from servicelib.celery.models import TaskState, TaskUUID
 from servicelib.fastapi.dependencies import get_app
 
+from ...exceptions.backend_errors import CeleryTaskNotFoundError
 from ...models.domain.celery_models import (
     ApiServerOwnerMetadata,
 )
@@ -42,6 +46,14 @@
 }
 
 
+@contextmanager
+def _exception_mapper(task_uuid: TaskUUID):
+    try:
+        yield
+    except TaskNotFoundError as exc:
+        raise CeleryTaskNotFoundError(task_uuid=task_uuid) from exc
+
+
 @router.get(
     "",
     response_model=ApiServerEnvelope[list[TaskGet]],
@@ -109,10 +121,11 @@ async def get_task_status(
         user_id=user_id,
         product_name=product_name,
     )
-    task_status = await task_manager.get_task_status(
-        owner_metadata=owner_metadata,
-        task_uuid=TaskUUID(f"{task_uuid}"),
-    )
+    with _exception_mapper(task_uuid=task_uuid):
+        task_status = await task_manager.get_task_status(
+            owner_metadata=owner_metadata,
+            task_uuid=TaskUUID(f"{task_uuid}"),
+        )
 
     return TaskStatus(
         task_progress=TaskProgress(
@@ -147,10 +160,11 @@ async def cancel_task(
         user_id=user_id,
         product_name=product_name,
     )
-    await task_manager.cancel_task(
-        owner_metadata=owner_metadata,
-        task_uuid=TaskUUID(f"{task_uuid}"),
-    )
+    with _exception_mapper(task_uuid=task_uuid):
+        await task_manager.cancel_task(
+            owner_metadata=owner_metadata,
+            task_uuid=TaskUUID(f"{task_uuid}"),
+        )
 
 
 @router.get(
@@ -183,37 +197,38 @@ async def get_task_result(
         product_name=product_name,
     )
 
-    task_status = await task_manager.get_task_status(
-        owner_metadata=owner_metadata,
-        task_uuid=TaskUUID(f"{task_uuid}"),
-    )
-
-    if not task_status.is_done:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail="Task result not available yet",
+    with _exception_mapper(task_uuid=task_uuid):
+        task_status = await task_manager.get_task_status(
+            owner_metadata=owner_metadata,
+            task_uuid=TaskUUID(f"{task_uuid}"),
         )
 
-    task_result = await task_manager.get_task_result(
-        owner_metadata=owner_metadata,
-        task_uuid=TaskUUID(f"{task_uuid}"),
-    )
-
-    if task_status.task_state == TaskState.FAILURE:
-        assert isinstance(task_result, Exception)
-        user_error_msg = f"The execution of task {task_uuid} failed"
-        support_id = create_error_code(task_result)
-        _logger.exception(
-            **create_troubleshooting_log_kwargs(
-                user_error_msg,
-                error=task_result,
-                error_code=support_id,
-                tip="Unexpected error in Celery",
+        if not task_status.is_done:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Task result not available yet",
             )
+
+        task_result = await task_manager.get_task_result(
+            owner_metadata=owner_metadata,
+            task_uuid=TaskUUID(f"{task_uuid}"),
         )
-        raise HTTPException(
-            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-            detail=user_error_msg,
-        )
+
+        if task_status.task_state == TaskState.FAILURE:
+            assert isinstance(task_result, Exception)
+            user_error_msg = f"The execution of task {task_uuid} failed"
+            support_id = create_error_code(task_result)
+            _logger.exception(
+                **create_troubleshooting_log_kwargs(
+                    user_error_msg,
+                    error=task_result,
+                    error_code=support_id,
+                    tip="Unexpected error in Celery",
+                )
+            )
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail=user_error_msg,
+            )
 
     return TaskResult(result=task_result, error=None)

services/api-server/src/simcore_service_api_server/exceptions/backend_errors.py

@@ -146,3 +146,8 @@ class LicensedItemCheckoutNotFoundError(BaseBackEndError):
 class JobAssetsMissingError(BaseBackEndError):
     msg_template = "Job assets missing for job {job_id}"
     status_code = status.HTTP_409_CONFLICT
+
+
+class CeleryTaskNotFoundError(BaseBackEndError):
+    msg_template = "Task {task_uuid} not found"
+    status_code = status.HTTP_404_NOT_FOUND

services/api-server/src/simcore_service_api_server/exceptions/handlers/__init__.py

@@ -45,16 +45,6 @@ def setup(app: FastAPI, *, is_debug: bool = False):
         ),
     )
 
-    app.add_exception_handler(
-        TaskNotFoundError,
-        make_handler_for_exception(
-            TaskNotFoundError,
-            status.HTTP_404_NOT_FOUND,
-            error_message=MSG_CLIENT_ERROR_USER_FRIENDLY_TEMPLATE,
-            add_exception_to_message=True,
-        ),
-    )
-
     app.add_exception_handler(
         CeleryError,
         make_handler_for_exception(