introduce TaslkFilterBase model

Author: bisgaard-itis

packages/celery-library/src/celery_library/backends/_redis.py

@@ -7,7 +7,7 @@
 from pydantic import ValidationError
 from servicelib.celery.models import (
     Task,
-    TaskContext,
+    TaskFilterBase,
     TaskID,
     TaskMetadata,
     TaskUUID,
@@ -82,7 +82,7 @@ async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
             )
             return None
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
+    async def list_tasks(self, task_context: TaskFilterBase) -> list[Task]:
         search_key = (
             _CELERY_TASK_INFO_PREFIX
             + build_task_id_prefix(task_context)

packages/celery-library/src/celery_library/task_manager.py

@@ -11,7 +11,7 @@
 from models_library.progress_bar import ProgressReport
 from servicelib.celery.models import (
     Task,
-    TaskContext,
+    TaskFilterBase,
     TaskID,
     TaskInfoStore,
     TaskMetadata,
@@ -41,7 +41,7 @@ async def submit_task(
         self,
         task_metadata: TaskMetadata,
         *,
-        task_context: TaskContext,
+        task_context: TaskFilterBase,
         **task_params,
     ) -> TaskUUID:
         with log_context(
@@ -72,7 +72,9 @@ async def submit_task(
     def _abort_task(self, task_id: TaskID) -> None:
         AbortableAsyncResult(task_id, app=self._celery_app).abort()
 
-    async def cancel_task(self, task_context: TaskContext, task_uuid: TaskUUID) -> None:
+    async def cancel_task(
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
+    ) -> None:
         with log_context(
             _logger,
             logging.DEBUG,
@@ -88,7 +90,7 @@ def _forget_task(self, task_id: TaskID) -> None:
         AbortableAsyncResult(task_id, app=self._celery_app).forget()
 
     async def get_task_result(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
     ) -> Any:
         with log_context(
             _logger,
@@ -106,7 +108,7 @@ async def get_task_result(
             return result
 
     async def _get_task_progress_report(
-        self, task_context: TaskContext, task_uuid: TaskUUID, task_state: TaskState
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID, task_state: TaskState
     ) -> ProgressReport:
         if task_state in (TaskState.STARTED, TaskState.RETRY, TaskState.ABORTED):
             task_id = build_task_id(task_context, task_uuid)
@@ -131,7 +133,7 @@ def _get_task_celery_state(self, task_id: TaskID) -> TaskState:
         return TaskState(self._celery_app.AsyncResult(task_id).state)
 
     async def get_task_status(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
     ) -> TaskStatus:
         with log_context(
             _logger,
@@ -148,7 +150,7 @@ async def get_task_status(
                 ),
             )
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
+    async def list_tasks(self, task_context: TaskFilterBase) -> list[Task]:
         with log_context(
             _logger,
             logging.DEBUG,

packages/celery-library/src/celery_library/utils.py

@@ -2,20 +2,19 @@
 
 from celery import Celery  # type: ignore[import-untyped]
 from servicelib.celery.app_server import BaseAppServer
-from servicelib.celery.models import TaskContext, TaskID, TaskUUID
+from servicelib.celery.models import TaskFilterBase, TaskID, TaskUUID
 
 _APP_SERVER_KEY = "app_server"
 
 _TASK_ID_KEY_DELIMITATOR: Final[str] = ":"
 
 
-def build_task_id_prefix(task_context: TaskContext) -> str:
-    return _TASK_ID_KEY_DELIMITATOR.join(
-        [f"{task_context[key]}" for key in sorted(task_context)]
-    )
+def build_task_id_prefix(task_context: TaskFilterBase) -> str:
+    _dict = task_context.model_dump()
+    return _TASK_ID_KEY_DELIMITATOR.join([f"{_dict[key]}" for key in sorted(_dict)])
 
 
-def build_task_id(task_context: TaskContext, task_uuid: TaskUUID) -> TaskID:
+def build_task_id(task_context: TaskFilterBase, task_uuid: TaskUUID) -> TaskID:
     return _TASK_ID_KEY_DELIMITATOR.join(
         [build_task_id_prefix(task_context), f"{task_uuid}"]
     )

packages/celery-library/tests/unit/test_tasks.py

@@ -20,7 +20,7 @@
 from common_library.errors_classes import OsparcErrorMixin
 from models_library.progress_bar import ProgressReport
 from servicelib.celery.models import (
-    TaskContext,
+    TaskFilterBase,
     TaskID,
     TaskMetadata,
     TaskState,
@@ -93,7 +93,7 @@ def _(celery_app: Celery) -> None:
 async def test_submitting_task_calling_async_function_results_with_success_state(
     celery_task_manager: CeleryTaskManager,
 ):
-    task_context = TaskContext(user_id=42)
+    task_context = TaskFilterBase(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -123,7 +123,7 @@ async def test_submitting_task_calling_async_function_results_with_success_state
 async def test_submitting_task_with_failure_results_with_error(
     celery_task_manager: CeleryTaskManager,
 ):
-    task_context = TaskContext(user_id=42)
+    task_context = TaskFilterBase(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -151,7 +151,7 @@ async def test_submitting_task_with_failure_results_with_error(
 async def test_cancelling_a_running_task_aborts_and_deletes(
     celery_task_manager: CeleryTaskManager,
 ):
-    task_context = TaskContext(user_id=42)
+    task_context = TaskFilterBase(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -185,7 +185,7 @@ async def test_cancelling_a_running_task_aborts_and_deletes(
 async def test_listing_task_uuids_contains_submitted_task(
     celery_task_manager: CeleryTaskManager,
 ):
-    task_context = TaskContext(user_id=42)
+    task_context = TaskFilterBase(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(

packages/service-library/src/servicelib/celery/models.py

@@ -1,19 +1,24 @@
 import datetime
 from enum import StrEnum
-from typing import Annotated, Any, Protocol, TypeAlias
+from typing import Annotated, Protocol, TypeAlias
 from uuid import UUID
 
 from models_library.progress_bar import ProgressReport
-from pydantic import BaseModel, StringConstraints
+from pydantic import BaseModel, ConfigDict, StringConstraints
 
-TaskContext: TypeAlias = dict[str, Any]
 TaskID: TypeAlias = str
 TaskName: TypeAlias = Annotated[
     str, StringConstraints(strip_whitespace=True, min_length=1)
 ]
 TaskUUID: TypeAlias = UUID
 
 
+class TaskFilterBase(BaseModel):
+    __root__: dict[str, str]
+
+    model_config = ConfigDict(extra="forbid")
+
+
 class TaskState(StrEnum):
     PENDING = "PENDING"
     STARTED = "STARTED"
@@ -56,7 +61,7 @@ async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None: ...
 
     async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None: ...
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]: ...
+    async def list_tasks(self, task_context: TaskFilterBase) -> list[Task]: ...
 
     async def remove_task(self, task_id: TaskID) -> None: ...
 

packages/service-library/src/servicelib/celery/task_manager.py

@@ -4,7 +4,7 @@
 
 from ..celery.models import (
     Task,
-    TaskContext,
+    TaskFilterBase,
     TaskID,
     TaskMetadata,
     TaskStatus,
@@ -14,22 +14,22 @@
 
 class TaskManager(Protocol):
     async def submit_task(
-        self, task_metadata: TaskMetadata, *, task_context: TaskContext, **task_param
+        self, task_metadata: TaskMetadata, *, task_context: TaskFilterBase, **task_param
     ) -> TaskUUID: ...
 
     async def cancel_task(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
     ) -> None: ...
 
     async def get_task_result(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
     ) -> Any: ...
 
     async def get_task_status(
-        self, task_context: TaskContext, task_uuid: TaskUUID
+        self, task_context: TaskFilterBase, task_uuid: TaskUUID
     ) -> TaskStatus: ...
 
-    async def list_tasks(self, task_context: TaskContext) -> list[Task]: ...
+    async def list_tasks(self, task_context: TaskFilterBase) -> list[Task]: ...
 
     async def set_task_progress(
         self, task_id: TaskID, report: ProgressReport