🎨🐛 Fix filtering bug in celery tasks (#8355)

Author: bisgaard-itis

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

@@ -1,24 +1,22 @@
 import contextlib
 import logging
 from datetime import timedelta
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 from models_library.progress_bar import ProgressReport
 from pydantic import ValidationError
 from servicelib.celery.models import (
     Task,
     TaskFilter,
     TaskID,
+    TaskInfoStore,
     TaskMetadata,
-    TaskUUID,
+    Wildcard,
 )
 from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
-from ..utils import build_task_id_prefix
-
 _CELERY_TASK_INFO_PREFIX: Final[str] = "celery-task-info-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
-_CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
 _CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 1000
 _CELERY_TASK_METADATA_KEY: Final[str] = "metadata"
 _CELERY_TASK_PROGRESS_KEY: Final[str] = "progress"
@@ -88,17 +86,14 @@ async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
             return None
 
     async def list_tasks(self, task_filter: TaskFilter) -> list[Task]:
-        search_key = (
-            _CELERY_TASK_INFO_PREFIX
-            + build_task_id_prefix(task_filter)
-            + _CELERY_TASK_ID_KEY_SEPARATOR
+        search_key = _CELERY_TASK_INFO_PREFIX + task_filter.create_task_id(
+            task_uuid=Wildcard()
         )
-        search_key_len = len(search_key)
 
         keys: list[str] = []
         pipeline = self._redis_client_sdk.redis.pipeline()
         async for key in self._redis_client_sdk.redis.scan_iter(
-            match=search_key + "*", count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
+            match=search_key, count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
         ):
             # fake redis (tests) returns bytes, real redis returns str
             _key = (
@@ -120,7 +115,7 @@ async def list_tasks(self, task_filter: TaskFilter) -> list[Task]:
                 task_metadata = TaskMetadata.model_validate_json(raw_metadata)
                 tasks.append(
                     Task(
-                        uuid=TaskUUID(key[search_key_len:]),
+                        uuid=TaskFilter.get_task_uuid(key),
                         metadata=task_metadata,
                     )
                 )
@@ -143,3 +138,7 @@ async def task_exists(self, task_id: TaskID) -> bool:
         n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
         assert isinstance(n, int)  # nosec
         return n > 0
+
+
+if TYPE_CHECKING:
+    _: type[TaskInfoStore] = RedisTaskInfoStore

packages/celery-library/src/celery_library/rpc/_async_jobs.py

@@ -134,9 +134,8 @@ async def result(
 
 @router.expose(reraise_if_error_type=(JobSchedulerError,))
 async def list_jobs(
-    task_manager: TaskManager, filter_: str, job_filter: AsyncJobFilter
+    task_manager: TaskManager, job_filter: AsyncJobFilter
 ) -> list[AsyncJobGet]:
-    _ = filter_
     assert task_manager  # nosec
     task_filter = TaskFilter.model_validate(job_filter.model_dump())
     try:

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

@@ -22,7 +22,6 @@
 from settings_library.celery import CelerySettings
 
 from .errors import TaskNotFoundError
-from .utils import build_task_id
 
 _logger = logging.getLogger(__name__)
 
@@ -50,7 +49,7 @@ async def submit_task(
             msg=f"Submit {task_metadata.name=}: {task_filter=} {task_params=}",
         ):
             task_uuid = uuid4()
-            task_id = build_task_id(task_filter, task_uuid)
+            task_id = task_filter.create_task_id(task_uuid=task_uuid)
             self._celery_app.send_task(
                 task_metadata.name,
                 task_id=task_id,
@@ -74,7 +73,7 @@ async def cancel_task(self, task_filter: TaskFilter, task_uuid: TaskUUID) -> Non
             logging.DEBUG,
             msg=f"task cancellation: {task_filter=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_filter, task_uuid)
+            task_id = task_filter.create_task_id(task_uuid=task_uuid)
             if not await self.task_exists(task_id):
                 raise TaskNotFoundError(task_id=task_id)
 
@@ -96,7 +95,7 @@ async def get_task_result(
             logging.DEBUG,
             msg=f"Get task result: {task_filter=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_filter, task_uuid)
+            task_id = task_filter.create_task_id(task_uuid=task_uuid)
             if not await self.task_exists(task_id):
                 raise TaskNotFoundError(task_id=task_id)
 
@@ -139,7 +138,7 @@ async def get_task_status(
             logging.DEBUG,
             msg=f"Getting task status: {task_filter=} {task_uuid=}",
         ):
-            task_id = build_task_id(task_filter, task_uuid)
+            task_id = task_filter.create_task_id(task_uuid=task_uuid)
             if not await self.task_exists(task_id):
                 raise TaskNotFoundError(task_id=task_id)
 

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

@@ -1,26 +1,8 @@
-from typing import Final
-
 from celery import Celery  # type: ignore[import-untyped]
 from servicelib.celery.app_server import BaseAppServer
-from servicelib.celery.models import TaskFilter, TaskID, TaskUUID
 
 _APP_SERVER_KEY = "app_server"
 
-_TASK_ID_KEY_DELIMITATOR: Final[str] = ":"
-
-
-def build_task_id_prefix(task_filter: TaskFilter) -> str:
-    filter_dict = task_filter.model_dump()
-    return _TASK_ID_KEY_DELIMITATOR.join(
-        [f"{filter_dict[key]}" for key in sorted(filter_dict)]
-    )
-
-
-def build_task_id(task_filter: TaskFilter, task_uuid: TaskUUID) -> TaskID:
-    return _TASK_ID_KEY_DELIMITATOR.join(
-        [build_task_id_prefix(task_filter), f"{task_uuid}"]
-    )
-
 
 def get_app_server(app: Celery) -> BaseAppServer:
     app_server = app.conf[_APP_SERVER_KEY]

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

@@ -258,7 +258,6 @@ async def test_async_jobs_workflow(
     jobs = await async_jobs.list_jobs(
         async_jobs_rabbitmq_rpc_client,
         rpc_namespace=ASYNC_JOBS_RPC_NAMESPACE,
-        filter_="",  # currently not used
         job_filter=job_filter,
     )
     assert len(jobs) > 0
@@ -311,7 +310,6 @@ async def test_async_jobs_cancel(
     jobs = await async_jobs.list_jobs(
         async_jobs_rabbitmq_rpc_client,
         rpc_namespace=ASYNC_JOBS_RPC_NAMESPACE,
-        filter_="",  # currently not used
         job_filter=job_filter,
     )
     assert async_job_get.job_id not in [job.job_id for job in jobs]

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

@@ -18,16 +18,21 @@
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.utils import get_app_server
 from common_library.errors_classes import OsparcErrorMixin
+from faker import Faker
 from models_library.progress_bar import ProgressReport
 from servicelib.celery.models import (
     TaskFilter,
     TaskID,
     TaskMetadata,
     TaskState,
+    TaskUUID,
+    Wildcard,
 )
 from servicelib.logging_utils import log_context
 from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
+_faker = Faker()
+
 _logger = logging.getLogger(__name__)
 
 pytest_simcore_core_services_selection = ["redis"]
@@ -199,5 +204,53 @@ async def test_listing_task_uuids_contains_submitted_task(
             tasks = await celery_task_manager.list_tasks(task_filter)
             assert any(task.uuid == task_uuid for task in tasks)
 
-    tasks = await celery_task_manager.list_tasks(task_filter)
-    assert any(task.uuid == task_uuid for task in tasks)
+        tasks = await celery_task_manager.list_tasks(task_filter)
+        assert any(task.uuid == task_uuid for task in tasks)
+
+
+async def test_filtering_listing_tasks(
+    celery_task_manager: CeleryTaskManager,
+):
+    class MyFilter(TaskFilter):
+        user_id: int
+        product_name: str | Wildcard
+        client_app: str | Wildcard
+
+    user_id = 42
+    expected_task_uuids: set[TaskUUID] = set()
+
+    for _ in range(5):
+        task_filter = MyFilter(
+            user_id=user_id,
+            product_name=_faker.word(),
+            client_app=_faker.word(),
+        )
+        task_uuid = await celery_task_manager.submit_task(
+            TaskMetadata(
+                name=dreamer_task.__name__,
+            ),
+            task_filter=task_filter,
+        )
+        expected_task_uuids.add(task_uuid)
+
+    for _ in range(3):
+        task_filter = MyFilter(
+            user_id=_faker.pyint(min_value=100, max_value=200),
+            product_name=_faker.word(),
+            client_app=_faker.word(),
+        )
+        await celery_task_manager.submit_task(
+            TaskMetadata(
+                name=dreamer_task.__name__,
+            ),
+            task_filter=task_filter,
+        )
+
+    search_filter = MyFilter(
+        user_id=user_id,
+        product_name=Wildcard(),
+        client_app=Wildcard(),
+    )
+    tasks = await celery_task_manager.list_tasks(search_filter)
+    assert expected_task_uuids == {task.uuid for task in tasks}
+    await asyncio.sleep(5 * 60)

packages/models-library/src/models_library/api_schemas_rpc_async_jobs/async_jobs.py

@@ -67,7 +67,7 @@ class AsyncJobFilter(AsyncJobFilterBase):
 
     product_name: ProductName
     user_id: UserID
-    client_name: Annotated[
+    client_name: Annotated[  # this is the name of the app which *submits* the async job. It is mainly used for filtering purposes
         str,
         StringConstraints(min_length=1, pattern=r"^[^\s]+$"),
     ]

packages/pytest-simcore/src/pytest_simcore/helpers/storage_rpc_server.py

@@ -13,8 +13,6 @@
     AsyncJobGet,
 )
 from models_library.api_schemas_webserver.storage import PathToExport
-from models_library.products import ProductName
-from models_library.users import UserID
 from pydantic import TypeAdapter, validate_call
 from pytest_mock import MockType
 from servicelib.rabbitmq._client_rpc import RabbitMQRPCClient
@@ -27,14 +25,12 @@ async def start_export_data(
         self,
         rabbitmq_rpc_client: RabbitMQRPCClient | MockType,
         *,
-        user_id: UserID,
-        product_name: ProductName,
         paths_to_export: list[PathToExport],
         export_as: Literal["path", "download_link"],
+        job_filter: AsyncJobFilter,
     ) -> tuple[AsyncJobGet, AsyncJobFilter]:
         assert rabbitmq_rpc_client
-        assert user_id
-        assert product_name
+        assert job_filter
         assert paths_to_export
         assert export_as
 

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

@@ -1,20 +1,103 @@
 import datetime
 from enum import StrEnum
-from typing import Annotated, Final, Protocol, TypeAlias
+from typing import Annotated, Any, Final, Protocol, Self, TypeAlias, TypeVar
 from uuid import UUID
 
 from models_library.progress_bar import ProgressReport
-from pydantic import BaseModel, ConfigDict, StringConstraints
+from pydantic import BaseModel, ConfigDict, StringConstraints, model_validator
 from pydantic.config import JsonDict
 
+ModelType = TypeVar("ModelType", bound=BaseModel)
+
 TaskID: TypeAlias = str
 TaskName: TypeAlias = Annotated[
     str, StringConstraints(strip_whitespace=True, min_length=1)
 ]
 TaskUUID: TypeAlias = UUID
+_TASK_ID_KEY_DELIMITATOR: Final[str] = ":"
+_WILDCARD: Final[str] = "*"
+_FORBIDDEN_CHARS = (_WILDCARD, _TASK_ID_KEY_DELIMITATOR, "=")
+
+
+class Wildcard:
+    def __str__(self) -> str:
+        return _WILDCARD
+
+
+class TaskFilter(BaseModel):
+    """
+    Class for associating metadata with a celery task. The implementation is very flexible and allows "clients" to define their own metadata.
+    The class exposes a filtering mechanism to list tasks using wildcards.
+
+    Example usage:
+        class MyTaskFilter(TaskFilter):
+            user_id: int | Wildcard
+            product_name: int | Wildcard
+            client_name: str
+
+        Listing tasks using the filter `MyTaskFilter(user_id=123, product_name=Wildcard(), client_name="my-app")` will return all tasks with
+        user_id 123, any product_name submitted from my-app.
+
+    If the metadata schema is known, the class allows deserializing the metadata (recreate_as_model). I.e. one can recover the metadata from the task:
+        metadata -> task_uuid -> metadata
+
+    """
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+    @model_validator(mode="after")
+    def _check_valid_filters(self) -> Self:
+        for key, value in self.model_dump().items():
+            # forbidden keys
+            if any(x in key for x in _FORBIDDEN_CHARS):
+                raise ValueError(f"Invalid filter key: '{key}'")
+            # forbidden values
+            if not isinstance(value, Wildcard) and any(
+                x in f"{value}" for x in _FORBIDDEN_CHARS
+            ):
+                raise ValueError(f"Invalid filter value for key '{key}': '{value}'")
+        return self
+
+    def _build_task_id_prefix(self) -> str:
+        filter_dict = self.model_dump()
+        return _TASK_ID_KEY_DELIMITATOR.join(
+            [f"{key}={filter_dict[key]}" for key in sorted(filter_dict)]
+        )
 
+    def create_task_id(self, task_uuid: TaskUUID | Wildcard) -> TaskID:
+        return _TASK_ID_KEY_DELIMITATOR.join(
+            [
+                self._build_task_id_prefix(),
+                f"task_uuid={task_uuid}",
+            ]
+        )
 
-class TaskFilter(BaseModel): ...
+    @classmethod
+    def recreate_as_model(cls, task_id: TaskID, schema: type[ModelType]) -> ModelType:
+        filter_dict = cls._recreate_data(task_id)
+        return schema.model_validate(filter_dict)
+
+    @classmethod
+    def _recreate_data(cls, task_id: TaskID) -> dict[str, Any]:
+        """Recreates the filter data from a task_id string
+        WARNING: does not validate types. For that use `recreate_model` instead
+        """
+        try:
+            parts = task_id.split(_TASK_ID_KEY_DELIMITATOR)
+            return {
+                key: value
+                for part in parts[:-1]
+                if (key := part.split("=")[0]) and (value := part.split("=")[1])
+            }
+        except (IndexError, ValueError) as err:
+            raise ValueError(f"Invalid task_id format: {task_id}") from err
+
+    @classmethod
+    def get_task_uuid(cls, task_id: TaskID) -> TaskUUID:
+        try:
+            return UUID(task_id.split(_TASK_ID_KEY_DELIMITATOR)[-1].split("=")[1])
+        except (IndexError, ValueError) as err:
+            raise ValueError(f"Invalid task_id format: {task_id}") from err
 
 
 class TaskState(StrEnum):

packages/service-library/src/servicelib/rabbitmq/rpc_interfaces/async_jobs/async_jobs.py

@@ -92,13 +92,11 @@ async def list_jobs(
     rabbitmq_rpc_client: RabbitMQRPCClient,
     *,
     rpc_namespace: RPCNamespace,
-    filter_: str,
     job_filter: AsyncJobFilter,
 ) -> list[AsyncJobGet]:
     _result: list[AsyncJobGet] = await rabbitmq_rpc_client.request(
         rpc_namespace,
         TypeAdapter(RPCMethodName).validate_python("list_jobs"),
-        filter_=filter_,
         job_filter=job_filter,
         timeout_s=_DEFAULT_TIMEOUT_S,
     )