♻️ common http API interface for long_running_tasks (#7843)

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

api/specs/web-server/_long_running_tasks.py

@@ -65,6 +65,7 @@ def cancel_async_job(
 
 @router.get(
     "/tasks/{task_id}/result",
+    response_model=Any,
     name="get_task_result",
     description="Retrieves the result of a task",
     responses=_export_data_responses,

api/specs/web-server/_long_running_tasks_legacy.py

@@ -4,7 +4,7 @@
 # pylint: disable=too-many-arguments
 
 
-from typing import Annotated
+from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, status
 from models_library.generics import Envelope
@@ -54,6 +54,7 @@ def cancel_and_delete_task(
 @router.get(
     "/{task_id}/result",
     name="get_task_result",
+    response_model=Any,
     description="Retrieves the result of a task",
 )
 def get_task_result(

packages/models-library/src/models_library/api_schemas_long_running_tasks/tasks.py

@@ -18,9 +18,12 @@ class TaskResult(BaseModel):
     error: Any | None
 
 
-class TaskGet(BaseModel):
+class TaskBase(BaseModel):
     task_id: TaskId
     task_name: str
+
+
+class TaskGet(TaskBase):
     status_href: str
     result_href: str
     abort_href: str

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

@@ -3,9 +3,9 @@
 from pydantic import PositiveFloat
 
 MINUTE: Final[PositiveFloat] = 60
-APP_LONG_RUNNING_TASKS_MANAGER_KEY: Final[
-    str
-] = f"{__name__ }.long_running_tasks.tasks_manager"
-RQT_LONG_RUNNING_TASKS_CONTEXT_KEY: Final[
-    str
-] = f"{__name__}.long_running_tasks.context"
+APP_LONG_RUNNING_MANAGER_KEY: Final[str] = (
+    f"{__name__ }.long_running_tasks.tasks_manager"
+)
+RQT_LONG_RUNNING_TASKS_CONTEXT_KEY: Final[str] = (
+    f"{__name__}.long_running_tasks.context"
+)

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

@@ -0,0 +1,41 @@
+import datetime
+
+from aiohttp import web
+
+from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.task import TaskContext, TasksManager
+from ._constants import APP_LONG_RUNNING_MANAGER_KEY
+from ._request import get_task_context
+
+
+class AiohttpLongRunningManager(BaseLongRunningManager):
+    def __init__(
+        self,
+        app: web.Application,
+        stale_task_check_interval: datetime.timedelta,
+        stale_task_detect_timeout: datetime.timedelta,
+    ):
+        self._app = app
+        self._tasks_manager = TasksManager(
+            stale_task_check_interval=stale_task_check_interval,
+            stale_task_detect_timeout=stale_task_detect_timeout,
+        )
+
+    @property
+    def tasks_manager(self) -> TasksManager:
+        return self._tasks_manager
+
+    async def setup(self) -> None:
+        await self._tasks_manager.setup()
+
+    async def teardown(self) -> None:
+        await self._tasks_manager.teardown()
+
+    @staticmethod
+    def get_task_context(request: web.Request) -> TaskContext:
+        return get_task_context(request)
+
+
+def get_long_running_manager(app: web.Application) -> AiohttpLongRunningManager:
+    output: AiohttpLongRunningManager = app[APP_LONG_RUNNING_MANAGER_KEY]
+    return output

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

@@ -0,0 +1,10 @@
+from typing import Any
+
+from aiohttp import web
+
+from ._constants import RQT_LONG_RUNNING_TASKS_CONTEXT_KEY
+
+
+def get_task_context(request: web.Request) -> dict[str, Any]:
+    output: dict[str, Any] = request[RQT_LONG_RUNNING_TASKS_CONTEXT_KEY]
+    return output

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

@@ -1,18 +1,15 @@
-import logging
 from typing import Any
 
 from aiohttp import web
-from common_library.json_serialization import json_dumps
 from pydantic import BaseModel
 from servicelib.aiohttp import status
 
-from ...long_running_tasks.errors import TaskNotCompletedError, TaskNotFoundError
+from ...long_running_tasks import http_endpoint_responses
 from ...long_running_tasks.models import TaskGet, TaskId, TaskStatus
-from ...long_running_tasks.task import TrackedTask
 from ..requests_validation import parse_request_path_parameters_as
-from ._dependencies import get_task_context, get_tasks_manager
+from ..rest_responses import create_data_response
+from ._manager import get_long_running_manager
 
-_logger = logging.getLogger(__name__)
 routes = web.RouteTableDef()
 
 
@@ -22,79 +19,58 @@ class _PathParam(BaseModel):
 
 @routes.get("", name="list_tasks")
 async def list_tasks(request: web.Request) -> web.Response:
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
-    tracked_tasks: list[TrackedTask] = tasks_manager.list_tasks(
-        with_task_context=task_context
-    )
-
-    return web.json_response(
-        {
-            "data": [
-                TaskGet(
-                    task_id=t.task_id,
-                    task_name=t.task_name,
-                    status_href=f"{request.app.router['get_task_status'].url_for(task_id=t.task_id)}",
-                    result_href=f"{request.app.router['get_task_result'].url_for(task_id=t.task_id)}",
-                    abort_href=f"{request.app.router['cancel_and_delete_task'].url_for(task_id=t.task_id)}",
-                )
-                for t in tracked_tasks
-            ]
-        },
-        dumps=json_dumps,
+    long_running_manager = get_long_running_manager(request.app)
+    return create_data_response(
+        [
+            TaskGet(
+                task_id=t.task_id,
+                task_name=t.task_name,
+                status_href=f"{request.app.router['get_task_status'].url_for(task_id=t.task_id)}",
+                result_href=f"{request.app.router['get_task_result'].url_for(task_id=t.task_id)}",
+                abort_href=f"{request.app.router['cancel_and_delete_task'].url_for(task_id=t.task_id)}",
+            )
+            for t in http_endpoint_responses.list_tasks(
+                long_running_manager.tasks_manager,
+                long_running_manager.get_task_context(request),
+            )
+        ]
     )
 
 
 @routes.get("/{task_id}", name="get_task_status")
 async def get_task_status(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
 
-    task_status: TaskStatus = tasks_manager.get_task_status(
-        task_id=path_params.task_id, with_task_context=task_context
+    task_status: TaskStatus = http_endpoint_responses.get_task_status(
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
     )
-    return web.json_response({"data": task_status}, dumps=json_dumps)
+    return create_data_response(task_status)
 
 
 @routes.get("/{task_id}/result", name="get_task_result")
 async def get_task_result(request: web.Request) -> web.Response | Any:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
 
     # NOTE: this might raise an exception that will be catched by the _error_handlers
-    try:
-        task_result = tasks_manager.get_task_result(
-            task_id=path_params.task_id, with_task_context=task_context
-        )
-        # NOTE: this will fail if the task failed for some reason....
-        await tasks_manager.remove_task(
-            path_params.task_id, with_task_context=task_context, reraise_errors=False
-        )
-        return task_result
-    except (TaskNotFoundError, TaskNotCompletedError):
-        raise
-    except Exception:
-        # the task shall be removed in this case
-        await tasks_manager.remove_task(
-            path_params.task_id, with_task_context=task_context, reraise_errors=False
-        )
-        raise
+    return await http_endpoint_responses.get_task_result(
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
+    )
 
 
 @routes.delete("/{task_id}", name="cancel_and_delete_task")
 async def cancel_and_delete_task(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
-    await tasks_manager.remove_task(path_params.task_id, with_task_context=task_context)
-    return web.json_response(status=status.HTTP_204_NO_CONTENT)
+    long_running_manager = get_long_running_manager(request.app)
 
-
-__all__: tuple[str, ...] = (
-    "get_tasks_manager",
-    "TaskId",
-    "TaskGet",
-    "TaskStatus",
-)
+    await http_endpoint_responses.remove_task(
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
+    )
+    return web.json_response(status=status.HTTP_204_NO_CONTENT)

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

@@ -1,32 +1,28 @@
 import asyncio
-import logging
+import datetime
 from collections.abc import AsyncGenerator, Callable
 from functools import wraps
 from typing import Any
 
 from aiohttp import web
 from common_library.json_serialization import json_dumps
-from pydantic import AnyHttpUrl, PositiveFloat, TypeAdapter
+from pydantic import AnyHttpUrl, TypeAdapter
 
 from ...aiohttp import status
-from ...long_running_tasks.models import TaskGet
-from ...long_running_tasks.task import (
-    TaskContext,
-    TaskProtocol,
-    TasksManager,
-    start_task,
+from ...long_running_tasks.constants import (
+    DEFAULT_STALE_TASK_CHECK_INTERVAL,
+    DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
+from ...long_running_tasks.models import TaskGet
+from ...long_running_tasks.task import TaskContext, TaskProtocol, start_task
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
-    APP_LONG_RUNNING_TASKS_MANAGER_KEY,
-    MINUTE,
+    APP_LONG_RUNNING_MANAGER_KEY,
     RQT_LONG_RUNNING_TASKS_CONTEXT_KEY,
 )
-from ._dependencies import create_task_name_from_request, get_tasks_manager
 from ._error_handlers import base_long_running_error_handler
-
-_logger = logging.getLogger(__name__)
+from ._manager import AiohttpLongRunningManager, get_long_running_manager
 
 
 def no_ops_decorator(handler: Handler):
@@ -42,6 +38,10 @@ async def _wrap(request: web.Request):
     return _wrap
 
 
+def _create_task_name_from_request(request: web.Request) -> str:
+    return f"{request.method} {request.rel_url}"
+
+
 async def start_long_running_task(
     # NOTE: positional argument are suffixed with "_" to avoid name conflicts with "task_kwargs" keys
     request_: web.Request,
@@ -51,12 +51,12 @@ async def start_long_running_task(
     task_context: TaskContext,
     **task_kwargs: Any,
 ) -> web.Response:
-    task_manager = get_tasks_manager(request_.app)
-    task_name = create_task_name_from_request(request_)
+    long_running_manager = get_long_running_manager(request_.app)
+    task_name = _create_task_name_from_request(request_)
     task_id = None
     try:
         task_id = start_task(
-            task_manager,
+            long_running_manager.tasks_manager,
             task_,
             fire_and_forget=fire_and_forget,
             task_context=task_context,
@@ -91,8 +91,10 @@ async def start_long_running_task(
     except asyncio.CancelledError:
         # cancel the task, the client has disconnected
         if task_id:
-            task_manager = get_tasks_manager(request_.app)
-            await task_manager.cancel_task(task_id, with_task_context=None)
+            long_running_manager = get_long_running_manager(request_.app)
+            await long_running_manager.tasks_manager.cancel_task(
+                task_id, with_task_context=None
+            )
         raise
 
 
@@ -121,8 +123,8 @@ def setup(
     router_prefix: str,
     handler_check_decorator: Callable = no_ops_decorator,
     task_request_context_decorator: Callable = no_task_context_decorator,
-    stale_task_check_interval_s: PositiveFloat = 1 * MINUTE,
-    stale_task_detect_timeout_s: PositiveFloat = 5 * MINUTE,
+    stale_task_check_interval: datetime.timedelta = DEFAULT_STALE_TASK_CHECK_INTERVAL,
+    stale_task_detect_timeout: datetime.timedelta = DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 ) -> None:
     """
     - `router_prefix` APIs are mounted on `/...`, this
@@ -135,21 +137,24 @@ def setup(
     """
 
     async def on_cleanup_ctx(app: web.Application) -> AsyncGenerator[None, None]:
+        # add error handlers
+        app.middlewares.append(base_long_running_error_handler)
+
         # add components to state
-        app[APP_LONG_RUNNING_TASKS_MANAGER_KEY] = long_running_task_manager = (
-            TasksManager(
-                stale_task_check_interval_s=stale_task_check_interval_s,
-                stale_task_detect_timeout_s=stale_task_detect_timeout_s,
+        app[APP_LONG_RUNNING_MANAGER_KEY] = long_running_manager = (
+            AiohttpLongRunningManager(
+                app=app,
+                stale_task_check_interval=stale_task_check_interval,
+                stale_task_detect_timeout=stale_task_detect_timeout,
             )
         )
 
-        # add error handlers
-        app.middlewares.append(base_long_running_error_handler)
+        await long_running_manager.setup()
 
         yield
 
         # cleanup
-        await long_running_task_manager.close()
+        await long_running_manager.teardown()
 
     # add routing (done at setup-time)
     _wrap_and_add_routes(