refactor to use registered task

Author: Andrei Neagu

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

@@ -14,7 +14,7 @@
     DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
 from ...long_running_tasks.models import TaskGet
-from ...long_running_tasks.task import TaskContext, TaskProtocol
+from ...long_running_tasks.task import RegisteredTaskName, TaskContext
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
@@ -45,7 +45,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,
@@ -56,7 +56,7 @@ async def start_long_running_task(
     task_id = None
     try:
         task_id = long_running_manager.tasks_manager.start_task(
-            task_,
+            registerd_task_name,
             fire_and_forget=fire_and_forget,
             task_context=task_context,
             task_name=task_name,

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 = (
+        "notask 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}'"
 

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

@@ -6,7 +6,7 @@
 import urllib.parse
 from collections import deque
 from contextlib import suppress
-from typing import Any, Final, Protocol, TypeAlias
+from typing import Any, ClassVar, Final, Protocol, TypeAlias
 from uuid import uuid4
 
 from models_library.api_schemas_long_running_tasks.base import TaskProgress
@@ -21,6 +21,7 @@
     TaskExceptionError,
     TaskNotCompletedError,
     TaskNotFoundError,
+    TaskNotRegisteredError,
 )
 from .models import TaskId, TaskStatus, TrackedTask
 
@@ -35,6 +36,7 @@
     seconds=1
 ).total_seconds()
 
+RegisteredTaskName: TypeAlias = str
 Namespace: TypeAlias = str
 TrackedTaskGroupDict: TypeAlias = dict[TaskId, TrackedTask]
 TaskContext: TypeAlias = dict[str, Any]
@@ -49,6 +51,19 @@ async def __call__(
     def __name__(self) -> str: ...
 
 
+class TaskRegistry:
+    REGISTERED_TASKS: ClassVar[dict[RegisteredTaskName, TaskProtocol]] = {}
+
+    @classmethod
+    def register(cls, task: TaskProtocol) -> None:
+        cls.REGISTERED_TASKS[task.__name__] = task
+
+    @classmethod
+    def unregister(cls, task: TaskProtocol) -> None:
+        if task.__name__ in cls.REGISTERED_TASKS:
+            del cls.REGISTERED_TASKS[task.__name__]
+
+
 async def _await_task(task: asyncio.Task) -> None:
     await task
 
@@ -318,7 +333,7 @@ def _get_task_id(self, task_name: str, *, is_unique: bool) -> TaskId:
 
     def start_task(
         self,
-        task: TaskProtocol,
+        registered_task_name: RegisteredTaskName,
         *,
         unique: bool = False,
         task_context: TaskContext | None = None,
@@ -351,6 +366,10 @@ def start_task(
         Returns:
             TaskId: the task unique identifier
         """
+        if registered_task_name not in TaskRegistry.REGISTERED_TASKS:
+            raise TaskNotRegisteredError(task_name=registered_task_name)
+
+        task = TaskRegistry.REGISTERED_TASKS[registered_task_name]
 
         # NOTE: If not task name is given, it will be composed of the handler's module and it's name
         # to keep the urls shorter and more meaningful.

packages/service-library/tests/aiohttp/long_running_tasks/conftest.py

@@ -19,15 +19,15 @@
     TaskProgress,
     TaskStatus,
 )
-from servicelib.long_running_tasks.task import TaskContext
+from servicelib.long_running_tasks.task import TaskContext, TaskRegistry
 from tenacity.asyncio import AsyncRetrying
 from tenacity.retry import retry_if_exception_type
 from tenacity.stop import stop_after_delay
 from tenacity.wait import wait_fixed
 
 
 async def _string_list_task(
-    task_progress: TaskProgress,
+    progress: TaskProgress,
     num_strings: int,
     sleep_time: float,
     fail: bool,
@@ -36,7 +36,7 @@ async def _string_list_task(
     for index in range(num_strings):
         generated_strings.append(f"{index}")
         await asyncio.sleep(sleep_time)
-        task_progress.update(message="generated item", percent=index / num_strings)
+        progress.update(message="generated item", percent=index / num_strings)
         if fail:
             msg = "We were asked to fail!!"
             raise RuntimeError(msg)
@@ -47,6 +47,9 @@ async def _string_list_task(
     )
 
 
+TaskRegistry.register(_string_list_task)
+
+
 @pytest.fixture
 def task_context(faker: Faker) -> TaskContext:
     return {"user_id": faker.pyint(), "product": faker.pystr()}
@@ -73,7 +76,7 @@ async def generate_list_strings(request: web.Request) -> web.Response:
         query_params = parse_request_query_parameters_as(_LongTaskQueryParams, request)
         return await long_running_tasks.server.start_long_running_task(
             request,
-            _string_list_task,
+            _string_list_task.__name__,
             num_strings=query_params.num_strings,
             sleep_time=query_params.sleep_time,
             fail=query_params.fail,

packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks.py

@@ -97,7 +97,7 @@ async def test_workflow(
     # now get the result
     result_url = client.app.router["get_task_result"].url_for(task_id=task_id)
     result = await client.get(f"{result_url}")
-    task_result, error = await assert_status(result, status.HTTP_201_CREATED)
+    task_result, error = await assert_status(result, status.HTTP_200_OK)
     assert task_result
     assert not error
     assert task_result == [f"{x}" for x in range(10)]

packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks_with_task_context.py

@@ -151,7 +151,7 @@ async def test_get_task_result(
     await assert_status(resp, status.HTTP_404_NOT_FOUND)
     # calling with context should find the task
     resp = await client_with_task_context.get(f"{result_url.with_query(task_context)}")
-    await assert_status(resp, status.HTTP_201_CREATED)
+    await assert_status(resp, status.HTTP_200_OK)
 
 
 async def test_cancel_task(

packages/service-library/tests/fastapi/long_running_tasks/test_long_running_tasks.py

@@ -12,7 +12,7 @@
 import asyncio
 import json
 from collections.abc import AsyncIterator, Awaitable, Callable
-from typing import Final
+from typing import Annotated, Final
 
 import pytest
 from asgi_lifespan import LifespanManager
@@ -31,7 +31,7 @@
     TaskProgress,
     TaskStatus,
 )
-from servicelib.long_running_tasks.task import TaskContext
+from servicelib.long_running_tasks.task import TaskContext, TaskRegistry
 from tenacity.asyncio import AsyncRetrying
 from tenacity.retry import retry_if_exception_type
 from tenacity.stop import stop_after_delay
@@ -42,7 +42,7 @@
 
 
 async def _string_list_task(
-    task_progress: TaskProgress,
+    progress: TaskProgress,
     num_strings: int,
     sleep_time: float,
     fail: bool,
@@ -51,14 +51,17 @@ async def _string_list_task(
     for index in range(num_strings):
         generated_strings.append(f"{index}")
         await asyncio.sleep(sleep_time)
-        task_progress.update(message="generated item", percent=index / num_strings)
+        progress.update(message="generated item", percent=index / num_strings)
         if fail:
             msg = "We were asked to fail!!"
             raise RuntimeError(msg)
 
     return generated_strings
 
 
+TaskRegistry.register(_string_list_task)
+
+
 @pytest.fixture
 def server_routes() -> APIRouter:
     routes = APIRouter()
@@ -69,18 +72,18 @@ def server_routes() -> APIRouter:
     async def create_string_list_task(
         num_strings: int,
         sleep_time: float,
+        long_running_manager: Annotated[
+            FastAPILongRunningManager, Depends(get_long_running_manager)
+        ],
+        *,
         fail: bool = False,
-        long_running_manager: FastAPILongRunningManager = Depends(
-            get_long_running_manager
-        ),
     ) -> TaskId:
-        task_id = long_running_manager.tasks_manager.start_task(
-            _string_list_task,
+        return long_running_manager.tasks_manager.start_task(
+            _string_list_task.__name__,
             num_strings=num_strings,
             sleep_time=sleep_time,
             fail=fail,
         )
-        return task_id
 
     return routes
 

packages/service-library/tests/fastapi/long_running_tasks/test_long_running_tasks_context_manager.py

@@ -25,6 +25,7 @@
     TaskId,
     TaskProgress,
 )
+from servicelib.long_running_tasks.task import TaskRegistry
 
 TASK_SLEEP_INTERVAL: Final[PositiveFloat] = 0.1
 
@@ -38,17 +39,25 @@ async def _assert_task_removed(
     assert result.status_code == status.HTTP_404_NOT_FOUND
 
 
-async def a_test_task(task_progress: TaskProgress) -> int:
+async def a_test_task(progress: TaskProgress) -> int:
+    _ = progress
     await asyncio.sleep(TASK_SLEEP_INTERVAL)
     return 42
 
 
-async def a_failing_test_task(task_progress: TaskProgress) -> None:
+TaskRegistry.register(a_test_task)
+
+
+async def a_failing_test_task(progress: TaskProgress) -> None:
+    _ = progress
     await asyncio.sleep(TASK_SLEEP_INTERVAL)
     msg = "I am failing as requested"
     raise RuntimeError(msg)
 
 
+TaskRegistry.register(a_failing_test_task)
+
+
 @pytest.fixture
 def user_routes() -> APIRouter:
     router = APIRouter()
@@ -59,15 +68,17 @@ async def create_task_user_defined_route(
             FastAPILongRunningManager, Depends(get_long_running_manager)
         ],
     ) -> TaskId:
-        return long_running_manager.tasks_manager.start_task(task=a_test_task)
+        return long_running_manager.tasks_manager.start_task(a_test_task.__name__)
 
     @router.get("/api/failing", status_code=status.HTTP_200_OK)
     async def create_task_which_fails(
         long_running_manager: Annotated[
             FastAPILongRunningManager, Depends(get_long_running_manager)
         ],
     ) -> TaskId:
-        return long_running_manager.tasks_manager.start_task(task=a_failing_test_task)
+        return long_running_manager.tasks_manager.start_task(
+            a_failing_test_task.__name__
+        )
 
     return router