fixed interfaces

Author: Andrei Neagu

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_routes.py

@@ -1,12 +1,16 @@
-from typing import Any
+from typing import Annotated, Any
 
 from aiohttp import web
-from pydantic import BaseModel
+from models_library.rest_base import RequestParameters
+from pydantic import BaseModel, Field
 
 from ...aiohttp import status
 from ...long_running_tasks import lrt_api
 from ...long_running_tasks.models import TaskGet, TaskId
-from ..requests_validation import parse_request_path_parameters_as
+from ..requests_validation import (
+    parse_request_path_parameters_as,
+    parse_request_query_parameters_as,
+)
 from ..rest_responses import create_data_response
 from ._manager import get_long_running_manager
 
@@ -65,16 +69,29 @@ async def get_task_result(request: web.Request) -> web.Response | Any:
     )
 
 
+class MyRequestQueryParams(RequestParameters):
+    wait_for_removal: Annotated[
+        bool,
+        Field(
+            description=(
+                "when True waits for the task to be removed "
+                "completly instead of returning immediately"
+            )
+        ),
+    ] = True
+
+
 @routes.delete("/{task_id}", name="remove_task")
 async def remove_task(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
+    query_params = parse_request_query_parameters_as(MyRequestQueryParams, request)
     long_running_manager = get_long_running_manager(request.app)
 
     await lrt_api.remove_task(
         long_running_manager.rpc_client,
         long_running_manager.lrt_namespace,
         long_running_manager.get_task_context(request),
         path_params.task_id,
-        wait_for_removal=False,  # frontend does not care about waiting for this
+        wait_for_removal=query_params.wait_for_removal,
     )
     return web.json_response(status=status.HTTP_204_NO_CONTENT)

packages/service-library/src/servicelib/fastapi/long_running_tasks/_routes.py

@@ -1,6 +1,6 @@
 from typing import Annotated, Any
 
-from fastapi import APIRouter, Depends, Request, status
+from fastapi import APIRouter, Depends, Query, Request, status
 
 from ...long_running_tasks import lrt_api
 from ...long_running_tasks.models import TaskGet, TaskId, TaskResult, TaskStatus
@@ -30,7 +30,7 @@ async def list_tasks(
         for t in await lrt_api.list_tasks(
             long_running_manager.rpc_client,
             long_running_manager.lrt_namespace,
-            task_context={},
+            long_running_manager.get_task_context(request),
         )
     ]
 
@@ -45,16 +45,16 @@ async def list_tasks(
 @cancel_on_disconnect
 async def get_task_status(
     request: Request,
-    task_id: TaskId,
     long_running_manager: Annotated[
         FastAPILongRunningManager, Depends(get_long_running_manager)
     ],
+    task_id: TaskId,
 ) -> TaskStatus:
     assert request  # nosec
     return await lrt_api.get_task_status(
         long_running_manager.rpc_client,
         long_running_manager.lrt_namespace,
-        task_context={},
+        long_running_manager.get_task_context(request),
         task_id=task_id,
     )
 
@@ -71,16 +71,16 @@ async def get_task_status(
 @cancel_on_disconnect
 async def get_task_result(
     request: Request,
-    task_id: TaskId,
     long_running_manager: Annotated[
         FastAPILongRunningManager, Depends(get_long_running_manager)
     ],
+    task_id: TaskId,
 ) -> TaskResult | Any:
     assert request  # nosec
     return await lrt_api.get_task_result(
         long_running_manager.rpc_client,
         long_running_manager.lrt_namespace,
-        task_context={},
+        long_running_manager.get_task_context(request),
         task_id=task_id,
     )
 
@@ -97,16 +97,26 @@ async def get_task_result(
 @cancel_on_disconnect
 async def remove_task(
     request: Request,
-    task_id: TaskId,
     long_running_manager: Annotated[
         FastAPILongRunningManager, Depends(get_long_running_manager)
     ],
+    task_id: TaskId,
+    *,
+    wait_for_removal: Annotated[
+        bool,
+        Query(
+            description=(
+                "when True waits for the task to be removed "
+                "completly instead of returning immediately"
+            ),
+        ),
+    ] = True,
 ) -> None:
     assert request  # nosec
     await lrt_api.remove_task(
         long_running_manager.rpc_client,
         long_running_manager.lrt_namespace,
-        task_context={},
+        long_running_manager.get_task_context(request),
         task_id=task_id,
-        wait_for_removal=True,  # only used by internal services, they will wait as before
+        wait_for_removal=wait_for_removal,
     )

packages/service-library/src/servicelib/long_running_tasks/_lrt_client.py

@@ -137,7 +137,6 @@ async def remove_task(
     *,
     task_context: TaskContext,
     task_id: TaskId,
-    reraise_errors: bool,
     wait_for_removal: bool,
     cancellation_timeout: timedelta | None,
 ) -> None:
@@ -158,7 +157,6 @@ async def remove_task(
         task_context=task_context,
         task_id=task_id,
         wait_for_removal=wait_for_removal,
-        reraise_errors=reraise_errors,
         timeout_s=timeout_s,
     )
     assert result is None  # nosec

packages/service-library/src/servicelib/long_running_tasks/_lrt_server.py

@@ -1,10 +1,11 @@
 import logging
+from contextlib import suppress
 from typing import TYPE_CHECKING, Any
 
 from ..logging_errors import create_troubleshootting_log_kwargs
 from ..rabbitmq import RPCRouter
 from ._serialization import string_to_object
-from .errors import BaseLongRunningError
+from .errors import BaseLongRunningError, TaskNotFoundError
 from .models import (
     ErrorResponse,
     RegisteredTaskName,
@@ -110,12 +111,12 @@ async def get_task_result(
         raise ValueError(msg)
     finally:
         # Ensure the task is removed regardless of the result
-        await long_running_manager.tasks_manager.remove_task(
-            task_id,
-            with_task_context=task_context,
-            wait_for_removal=True,
-            reraise_errors=False,
-        )
+        with suppress(TaskNotFoundError):
+            await long_running_manager.tasks_manager.remove_task(
+                task_id,
+                with_task_context=task_context,
+                wait_for_removal=True,
+            )
 
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
@@ -125,11 +126,9 @@ async def remove_task(
     task_context: TaskContext,
     task_id: TaskId,
     wait_for_removal: bool,
-    reraise_errors: bool,
 ) -> None:
     await long_running_manager.tasks_manager.remove_task(
         task_id,
         with_task_context=task_context,
         wait_for_removal=wait_for_removal,
-        reraise_errors=reraise_errors,
     )

packages/service-library/src/servicelib/long_running_tasks/lrt_api.py

@@ -125,7 +125,6 @@ async def remove_task(
         lrt_namespace,
         task_id=task_id,
         task_context=task_context,
-        reraise_errors=True,
         wait_for_removal=wait_for_removal,
         cancellation_timeout=cancellation_timeout,
     )

packages/service-library/src/servicelib/long_running_tasks/task.py

@@ -202,13 +202,13 @@ async def setup(self) -> None:
     async def teardown(self) -> None:
         # ensure all created tasks are cancelled
         for tracked_task in await self._tasks_data.list_tasks_data():
-            await self.remove_task(
-                tracked_task.task_id,
-                tracked_task.task_context,
-                # when closing we do not care about pending errors
-                wait_for_removal=True,
-                reraise_errors=False,
-            )
+            # when closing we do not care about pending errors
+            with suppress(TaskNotFoundError):
+                await self.remove_task(
+                    tracked_task.task_id,
+                    tracked_task.task_context,
+                    wait_for_removal=True,
+                )
 
         for task in self._created_tasks.values():
             _logger.warning(
@@ -278,10 +278,7 @@ async def _stale_tasks_monitor(self) -> None:
                     ).model_dump_json(),
                 )
                 await self.remove_task(
-                    task_id,
-                    with_task_context=task_context,
-                    wait_for_removal=True,
-                    reraise_errors=False,
+                    task_id, with_task_context=task_context, wait_for_removal=True
                 )
 
     async def _cancelled_tasks_removal(self) -> None:
@@ -438,15 +435,12 @@ async def remove_task(
         with_task_context: TaskContext,
         *,
         wait_for_removal: bool,
-        reraise_errors: bool = True,
     ) -> None:
-        """cancels and removes task"""
-        try:
-            tracked_task = await self._get_tracked_task(task_id, with_task_context)
-        except TaskNotFoundError:
-            if reraise_errors:
-                raise
-            return
+        """
+        cancels and removes task
+        raises TaskNotFoundError if the task cannot be found
+        """
+        tracked_task = await self._get_tracked_task(task_id, with_task_context)
 
         await self._tasks_data.mark_task_for_removal(
             tracked_task.task_id, tracked_task.task_context

packages/service-library/tests/long_running_tasks/test_long_running_tasks_task.py

@@ -498,13 +498,6 @@ async def test_remove_unknown_task(
             "invalid_id", with_task_context=empty_context, wait_for_removal=True
         )
 
-    await long_running_manager.tasks_manager.remove_task(
-        "invalid_id",
-        with_task_context=empty_context,
-        wait_for_removal=True,
-        reraise_errors=False,
-    )
-
 
 async def test__cancelled_tasks_worker_equivalent_of_cancellation_from_a_different_process(
     long_running_manager: BaseLongRunningManager, empty_context: TaskContext