simplified interface

Author: Andrei Neagu

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

@@ -1,15 +1,13 @@
-from typing import Annotated, Any
+from typing import Any
 
 from aiohttp import web
-from models_library.rest_base import RequestParameters
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 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,
-    parse_request_query_parameters_as,
 )
 from ..rest_responses import create_data_response
 from ._manager import get_long_running_manager
@@ -69,29 +67,15 @@ async def get_task_result(request: web.Request) -> web.Response | Any:
     )
 
 
-class _RemoveTaskQueryParams(RequestParameters):
-    wait_for_removal: Annotated[
-        bool,
-        Field(
-            description=(
-                "when True waits for the task to be removed "
-                "completly instead of returning immediately"
-            )
-        ),
-    ] = False
-
-
 @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(_RemoveTaskQueryParams, 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=query_params.wait_for_removal,
     )
     return web.json_response(status=status.HTTP_204_NO_CONTENT)

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

@@ -108,7 +108,6 @@ async def start_long_running_task(
                 long_running_manager.lrt_namespace,
                 task_context,
                 task_id,
-                wait_for_removal=False,
             )
         raise
 

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, Query, Request, status
+from fastapi import APIRouter, Depends, Request, status
 
 from ...long_running_tasks import lrt_api
 from ...long_running_tasks.models import TaskGet, TaskId, TaskResult, TaskStatus
@@ -101,22 +101,11 @@ async def remove_task(
         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"
-            ),
-        ),
-    ] = False,
 ) -> None:
     assert request  # nosec
     await lrt_api.remove_task(
         long_running_manager.rpc_client,
         long_running_manager.lrt_namespace,
         long_running_manager.get_task_context(request),
         task_id=task_id,
-        wait_for_removal=wait_for_removal,
     )

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

@@ -118,7 +118,6 @@ async def remove_task(
     *,
     task_context: TaskContext,
     task_id: TaskId,
-    wait_for_removal: bool,
     cancellation_timeout: timedelta | None,
 ) -> None:
     timeout_s = (
@@ -127,17 +126,11 @@ async def remove_task(
         else int(cancellation_timeout.total_seconds())
     )
 
-    # NOTE: task always gets cancelled even if not waiting for it
-    # request will return immediatlye, no need to wait so much
-    if wait_for_removal is False:
-        timeout_s = _RPC_TIMEOUT_SHORT_REQUESTS
-
     result = await rabbitmq_rpc_client.request(
         get_rabbit_namespace(namespace),
         TypeAdapter(RPCMethodName).validate_python("remove_task"),
         task_context=task_context,
         task_id=task_id,
-        wait_for_removal=wait_for_removal,
         timeout_s=timeout_s,
     )
     assert result is None  # nosec

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

@@ -98,10 +98,7 @@ async def remove_task(
     *,
     task_context: TaskContext,
     task_id: TaskId,
-    wait_for_removal: bool,
 ) -> None:
     await long_running_manager.tasks_manager.remove_task(
-        task_id,
-        with_task_context=task_context,
-        wait_for_removal=wait_for_removal,
+        task_id, with_task_context=task_context, wait_for_removal=False
     )

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

@@ -104,7 +104,6 @@ async def remove_task(
     task_context: TaskContext,
     task_id: TaskId,
     *,
-    wait_for_removal: bool,
     cancellation_timeout: timedelta | None = None,
 ) -> None:
     """cancels and removes a task
@@ -116,6 +115,5 @@ async def remove_task(
         lrt_namespace,
         task_id=task_id,
         task_context=task_context,
-        wait_for_removal=wait_for_removal,
         cancellation_timeout=cancellation_timeout,
     )

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

@@ -1,3 +1,4 @@
+# pylint: disable=protected-access
 # pylint: disable=redefined-outer-name
 # pylint: disable=unused-argument
 
@@ -161,18 +162,19 @@ async def _assert_list_tasks_from_all_managers(
 
 
 async def _assert_task_is_no_longer_present(
-    rabbitmq_rpc_client: RabbitMQRPCClient,
     long_running_managers: list[LongRunningManager],
     task_context: TaskContext,
     task_id: TaskId,
 ) -> None:
-    with pytest.raises(TaskNotFoundError):
-        await lrt_api.get_task_status(
-            rabbitmq_rpc_client,
-            _get_long_running_manager(long_running_managers).lrt_namespace,
-            task_context,
-            task_id,
-        )
+    manager = _get_long_running_manager(long_running_managers)
+    async for attempt in AsyncRetrying(**_RETRY_PARAMS):
+        with attempt:  # noqa: SIM117
+            with pytest.raises(TaskNotFoundError):  # noqa: PT012
+                # use internals to detirmine when it's no longer here
+                await manager._tasks_manager._get_tracked_task(  # noqa: SLF001
+                    task_id, task_context
+                )
+                raise TryAgain
 
 
 _TASK_CONTEXT: Final[list[TaskContext | None]] = [{"a": "context"}, None]
@@ -234,7 +236,7 @@ async def test_workflow_with_result(
         assert result == to_return
 
         await _assert_task_is_no_longer_present(
-            rabbitmq_rpc_client, long_running_managers, saved_context, task_id
+            long_running_managers, saved_context, task_id
         )
 
 
@@ -289,13 +291,14 @@ async def test_workflow_raises_error(
             )
 
         await _assert_task_is_no_longer_present(
-            rabbitmq_rpc_client, long_running_managers, saved_context, task_id
+            long_running_managers, saved_context, task_id
         )
 
 
 @pytest.mark.parametrize("task_context", _TASK_CONTEXT)
 @pytest.mark.parametrize("is_unique", _IS_UNIQUE)
 async def test_remove_task(
+    disable_stale_tasks_monitor: None,
     long_running_managers: list[LongRunningManager],
     rabbitmq_rpc_client: RabbitMQRPCClient,
     is_unique: bool,
@@ -321,9 +324,8 @@ async def test_remove_task(
         _get_long_running_manager(long_running_managers).lrt_namespace,
         saved_context,
         task_id,
-        wait_for_removal=True,
     )
 
     await _assert_task_is_no_longer_present(
-        rabbitmq_rpc_client, long_running_managers, saved_context, task_id
+        long_running_managers, saved_context, task_id
     )