♻️ preparing TasksManager's interface to be extracted into a common interface (ITISFoundation#7884)

Co-authored-by: Andrei Neagu <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

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

@@ -2,7 +2,8 @@
 from datetime import datetime
 from typing import Any
 
-from pydantic import BaseModel, field_validator
+from common_library.exclude import Unset
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from .base import TaskId, TaskProgress
 
@@ -20,15 +21,30 @@ class TaskResult(BaseModel):
 
 class TaskBase(BaseModel):
     task_id: TaskId
-    task_name: str
+    task_name: str | Unset = Unset.VALUE
+
+    @model_validator(mode="after")
+    def try_populate_task_name_from_task_id(self) -> "TaskBase":
+        # NOTE: currently this model is used to validate tasks coming from
+        # the celery backend and form long_running_tasks
+        # 1. if a task comes from Celery, it will keep it's given name
+        # 2. if a task comes from long_running_tasks, it will extract it form
+        #   the task_id, which looks like "{PREFIX}.{TASK_NAME}.UNIQUE|{UUID}"
+
+        if self.task_id and self.task_name == Unset.VALUE:
+            parts = self.task_id.split(".")
+            if len(parts) > 1:
+                self.task_name = urllib.parse.unquote(parts[1])
+
+        if self.task_name == Unset.VALUE:
+            self.task_name = self.task_id
+
+        return self
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class TaskGet(TaskBase):
     status_href: str
     result_href: str
     abort_href: str
-
-    @field_validator("task_name")
-    @classmethod
-    def unquote_str(cls, v) -> str:
-        return urllib.parse.unquote(v)

packages/models-library/tests/test_api_schemas_long_running_tasks_tasks.py

@@ -0,0 +1,56 @@
+import pytest
+from models_library.api_schemas_long_running_tasks.tasks import TaskGet
+from pydantic import TypeAdapter
+
+
+def _get_data_without_task_name(task_id: str) -> dict:
+    return {
+        "task_id": task_id,
+        "status_href": "",
+        "result_href": "",
+        "abort_href": "",
+    }
+
+
+@pytest.mark.parametrize(
+    "data, expected_task_name",
+    [
+        (_get_data_without_task_name("a.b.c.d"), "b"),
+        (_get_data_without_task_name("a.b.c"), "b"),
+        (_get_data_without_task_name("a.b"), "b"),
+        (_get_data_without_task_name("a"), "a"),
+    ],
+)
+def test_try_extract_task_name(data: dict, expected_task_name: str) -> None:
+    task_get = TaskGet(**data)
+    assert task_get.task_name == expected_task_name
+
+    task_get = TypeAdapter(TaskGet).validate_python(data)
+    assert task_get.task_name == expected_task_name
+
+
+def _get_data_with_task_name(task_id: str, task_name: str) -> dict:
+    return {
+        "task_id": task_id,
+        "task_name": task_name,
+        "status_href": "",
+        "result_href": "",
+        "abort_href": "",
+    }
+
+
+@pytest.mark.parametrize(
+    "data, expected_task_name",
+    [
+        (_get_data_with_task_name("a.b.c.d", "a_name"), "a_name"),
+        (_get_data_with_task_name("a.b.c", "a_name"), "a_name"),
+        (_get_data_with_task_name("a.b", "a_name"), "a_name"),
+        (_get_data_with_task_name("a", "a_name"), "a_name"),
+    ],
+)
+def test_task_name_is_provided(data: dict, expected_task_name: str) -> None:
+    task_get = TaskGet(**data)
+    assert task_get.task_name == expected_task_name
+
+    task_get = TypeAdapter(TaskGet).validate_python(data)
+    assert task_get.task_name == expected_task_name

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

@@ -4,7 +4,7 @@
 from pydantic import BaseModel
 from servicelib.aiohttp import status
 
-from ...long_running_tasks import http_endpoint_responses
+from ...long_running_tasks import lrt_api
 from ...long_running_tasks.models import TaskGet, TaskId, TaskStatus
 from ..requests_validation import parse_request_path_parameters_as
 from ..rest_responses import create_data_response
@@ -24,12 +24,11 @@ async def list_tasks(request: web.Request) -> web.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(
+            for t in lrt_api.list_tasks(
                 long_running_manager.tasks_manager,
                 long_running_manager.get_task_context(request),
             )
@@ -42,7 +41,7 @@ async def get_task_status(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
     long_running_manager = get_long_running_manager(request.app)
 
-    task_status: TaskStatus = http_endpoint_responses.get_task_status(
+    task_status: TaskStatus = lrt_api.get_task_status(
         long_running_manager.tasks_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,
@@ -56,7 +55,7 @@ async def get_task_result(request: web.Request) -> web.Response | Any:
     long_running_manager = get_long_running_manager(request.app)
 
     # NOTE: this might raise an exception that will be catched by the _error_handlers
-    return await http_endpoint_responses.get_task_result(
+    return await lrt_api.get_task_result(
         long_running_manager.tasks_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,
@@ -68,7 +67,7 @@ async def cancel_and_delete_task(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
     long_running_manager = get_long_running_manager(request.app)
 
-    await http_endpoint_responses.remove_task(
+    await lrt_api.remove_task(
         long_running_manager.tasks_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,

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

@@ -9,12 +9,13 @@
 from pydantic import AnyHttpUrl, TypeAdapter
 
 from ...aiohttp import status
+from ...long_running_tasks import lrt_api
 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 ...long_running_tasks.task import RegisteredTaskName, TaskContext
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
@@ -25,11 +26,11 @@
 from ._manager import AiohttpLongRunningManager, get_long_running_manager
 
 
-def no_ops_decorator(handler: Handler):
+def _no_ops_decorator(handler: Handler):
     return handler
 
 
-def no_task_context_decorator(handler: Handler):
+def _no_task_context_decorator(handler: Handler):
     @wraps(handler)
     async def _wrap(request: web.Request):
         request[RQT_LONG_RUNNING_TASKS_CONTEXT_KEY] = {}
@@ -45,7 +46,7 @@ def _create_task_name_from_request(request: web.Request) -> str:
 async def start_long_running_task(
     # NOTE: positional argument are suffixed with "_" to avoid name conflicts with "task_kwargs" keys
     request_: web.Request,
-    task_: TaskProtocol,
+    registerd_task_name: RegisteredTaskName,
     *,
     fire_and_forget: bool = False,
     task_context: TaskContext,
@@ -55,9 +56,9 @@ async def start_long_running_task(
     task_name = _create_task_name_from_request(request_)
     task_id = None
     try:
-        task_id = start_task(
+        task_id = await lrt_api.start_task(
             long_running_manager.tasks_manager,
-            task_,
+            registerd_task_name,
             fire_and_forget=fire_and_forget,
             task_context=task_context,
             task_name=task_name,
@@ -78,7 +79,6 @@ async def start_long_running_task(
         )
         task_get = TaskGet(
             task_id=task_id,
-            task_name=task_name,
             status_href=f"{status_url}",
             result_href=f"{result_url}",
             abort_href=f"{abort_url}",
@@ -121,8 +121,8 @@ def setup(
     app: web.Application,
     *,
     router_prefix: str,
-    handler_check_decorator: Callable = no_ops_decorator,
-    task_request_context_decorator: Callable = no_task_context_decorator,
+    handler_check_decorator: Callable = _no_ops_decorator,
+    task_request_context_decorator: Callable = _no_task_context_decorator,
     stale_task_check_interval: datetime.timedelta = DEFAULT_STALE_TASK_CHECK_INTERVAL,
     stale_task_detect_timeout: datetime.timedelta = DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 ) -> None:

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

@@ -2,7 +2,7 @@
 
 from fastapi import APIRouter, Depends, Request, status
 
-from ...long_running_tasks import http_endpoint_responses
+from ...long_running_tasks import lrt_api
 from ...long_running_tasks.models import TaskGet, TaskId, TaskResult, TaskStatus
 from ..requests_decorators import cancel_on_disconnect
 from ._dependencies import get_long_running_manager
@@ -23,14 +23,13 @@ async def list_tasks(
     return [
         TaskGet(
             task_id=t.task_id,
-            task_name=t.task_name,
             status_href=str(request.url_for("get_task_status", task_id=t.task_id)),
             result_href=str(request.url_for("get_task_result", task_id=t.task_id)),
             abort_href=str(
                 request.url_for("cancel_and_delete_task", task_id=t.task_id)
             ),
         )
-        for t in http_endpoint_responses.list_tasks(
+        for t in lrt_api.list_tasks(
             long_running_manager.tasks_manager, task_context=None
         )
     ]
@@ -52,7 +51,7 @@ async def get_task_status(
     ],
 ) -> TaskStatus:
     assert request  # nosec
-    return http_endpoint_responses.get_task_status(
+    return lrt_api.get_task_status(
         long_running_manager.tasks_manager, task_context=None, task_id=task_id
     )
 
@@ -75,7 +74,7 @@ async def get_task_result(
     ],
 ) -> TaskResult | Any:
     assert request  # nosec
-    return await http_endpoint_responses.get_task_result(
+    return await lrt_api.get_task_result(
         long_running_manager.tasks_manager, task_context=None, task_id=task_id
     )
 
@@ -98,6 +97,6 @@ async def cancel_and_delete_task(
     ],
 ) -> None:
     assert request  # nosec
-    await http_endpoint_responses.remove_task(
+    await lrt_api.remove_task(
         long_running_manager.tasks_manager, task_context=None, task_id=task_id
     )

packages/service-library/src/servicelib/logging_errors.py

@@ -74,7 +74,7 @@ def create_troubleshootting_log_kwargs(
             ...
         except MyException as exc
             _logger.exception(
-                **create_troubleshotting_log_kwargs(
+                **create_troubleshootting_log_kwargs(
                     user_error_msg=frontend_msg,
                     error=exc,
                     error_context={

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

@@ -5,6 +5,13 @@ class BaseLongRunningError(OsparcErrorMixin, Exception):
     """base exception for this module"""
 
 
+class TaskNotRegisteredError(BaseLongRunningError):
+    msg_template: str = (
+        "no task with task_name='{task_name}' was found in the task registry. "
+        "Make sure it's registered before starting it."
+    )
+
+
 class TaskAlreadyRunningError(BaseLongRunningError):
     msg_template: str = "{task_name} must be unique, found: '{managed_task}'"