Add per-task ack type support (#636)

Author: vvanglro

docs/guide/cli.md

@@ -58,6 +58,21 @@ This can be configured using `--ack-type` parameter. For example:
 taskiq worker --ack-type when_executed mybroker:broker
 ```
 
+It can also be overridden per task using the `ack_type` task label:
+
+```python
+@broker.task(ack_type="when_received")
+async def best_effort_task() -> None:
+    ...
+
+
+@broker.task(ack_type="when_saved")
+async def critical_task() -> None:
+    ...
+```
+
+If a task has no `ack_type` label, the worker-level `--ack-type` value is used.
+
 ### Type casts
 
 One of features taskiq have is automatic type casts. For example you have a type-hinted task like this:

taskiq/__init__.py

@@ -9,7 +9,7 @@
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.abc.schedule_source import ScheduleSource
-from taskiq.acks import AckableMessage
+from taskiq.acks import AckableMessage, AcknowledgeType
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.brokers.shared_broker import async_shared_broker
 from taskiq.brokers.zmq_broker import ZeroMQBroker
@@ -41,6 +41,7 @@
 
 __all__ = [
     "AckableMessage",
+    "AcknowledgeType",
     "AsyncBroker",
     "AsyncResultBackend",
     "AsyncTaskiqDecoratedTask",

taskiq/acks.py

@@ -1,5 +1,6 @@
 import enum
 from collections.abc import Awaitable, Callable
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -20,6 +21,20 @@ class AcknowledgeType(str, enum.Enum):
     WHEN_SAVED = "when_saved"
 
 
+def parse_acknowledge_type(value: Any) -> AcknowledgeType:
+    """Parse acknowledge type from a task label value."""
+    if isinstance(value, AcknowledgeType):
+        return value
+    if isinstance(value, str):
+        try:
+            return AcknowledgeType(value.lower())
+        except ValueError as exc:
+            raise ValueError(f"Unknown acknowledge type value: {value}.") from exc
+    raise ValueError(
+        f"Unsupported acknowledge type: {value} use str or AcknowledgeType",
+    )
+
+
 class AckableMessage(BaseModel):
     """
     Message that can be acknowledged.

taskiq/receiver/receiver.py

@@ -15,7 +15,7 @@
 
 from taskiq.abc.broker import AckableMessage, AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
-from taskiq.acks import AcknowledgeType
+from taskiq.acks import AcknowledgeType, parse_acknowledge_type
 from taskiq.context import Context
 from taskiq.exceptions import NoResultError
 from taskiq.message import TaskiqMessage
@@ -148,13 +148,14 @@ async def callback(  # noqa: C901, PLR0912
                     ),
                 )
 
+        ack_time = self._get_ack_time(taskiq_msg)
         logger.info(
             "Executing task %s with ID: %s",
             taskiq_msg.task_name,
             taskiq_msg.task_id,
         )
 
-        if self.ack_time == AcknowledgeType.WHEN_RECEIVED and isinstance(
+        if ack_time == AcknowledgeType.WHEN_RECEIVED and isinstance(
             message,
             AckableMessage,
         ):
@@ -165,7 +166,7 @@ async def callback(  # noqa: C901, PLR0912
             message=taskiq_msg,
         )
 
-        if self.ack_time == AcknowledgeType.WHEN_EXECUTED and isinstance(
+        if ack_time == AcknowledgeType.WHEN_EXECUTED and isinstance(
             message,
             AckableMessage,
         ):
@@ -192,12 +193,28 @@ async def callback(  # noqa: C901, PLR0912
             if raise_err:
                 raise exc
 
-        if self.ack_time == AcknowledgeType.WHEN_SAVED and isinstance(
+        if ack_time == AcknowledgeType.WHEN_SAVED and isinstance(
             message,
             AckableMessage,
         ):
             await maybe_awaitable(message.ack())
 
+    def _get_ack_time(self, message: TaskiqMessage) -> AcknowledgeType:
+        """
+        Get acknowledge time for a task.
+
+        Task-level `ack_type` label overrides worker-level configuration.
+        """
+        ack_type = message.labels.get("ack_type")
+        if ack_type is None:
+            return self.ack_time
+        try:
+            return parse_acknowledge_type(ack_type)
+        except ValueError as exc:
+            raise ValueError(
+                f"Invalid ack_type label {ack_type!r} for task {message.task_name}.",
+            ) from exc
+
     async def run_task(  # noqa: C901, PLR0912, PLR0915
         self,
         target: Callable[..., Any],

tests/receiver/test_receiver.py

@@ -13,6 +13,8 @@
 
 from taskiq.abc.broker import AckableMessage, AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.abc.result_backend import AsyncResultBackend
+from taskiq.acks import AcknowledgeType
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.exceptions import NoResultError, TaskiqResultTimeoutError
 from taskiq.message import TaskiqMessage
@@ -26,6 +28,7 @@ def get_receiver(
     no_parse: bool = False,
     max_async_tasks: int | None = None,
     max_async_tasks_jitter: int = 0,
+    ack_type: AcknowledgeType | None = None,
 ) -> Receiver:
     """
     Returns receiver with custom broker and args.
@@ -44,9 +47,50 @@ def get_receiver(
         validate_params=not no_parse,
         max_async_tasks=max_async_tasks,
         max_async_tasks_jitter=max_async_tasks_jitter,
+        ack_type=ack_type,
     )
 
 
+class _EventResultBackend(AsyncResultBackend[Any]):
+    def __init__(self, events: list[str]) -> None:
+        self.events = events
+        self.results: dict[str, TaskiqResult[Any]] = {}
+
+    async def set_result(self, task_id: str, result: TaskiqResult[Any]) -> None:
+        self.events.append("save")
+        self.results[task_id] = result
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        return task_id in self.results
+
+    async def get_result(
+        self,
+        task_id: str,
+        with_logs: bool = False,
+    ) -> TaskiqResult[Any]:
+        return self.results[task_id]
+
+
+class _EventMiddleware(TaskiqMiddleware):
+    def __init__(self, events: list[str]) -> None:
+        super().__init__()
+        self.events = events
+
+    def post_execute(
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult[Any],
+    ) -> None:
+        self.events.append("post_execute")
+
+    def post_save(
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult[Any],
+    ) -> None:
+        self.events.append("post_save")
+
+
 async def test_run_task_successful_async() -> None:
     """Tests that run_task can run async tasks."""
 
@@ -327,6 +371,192 @@ async def ack_callback() -> None:
     assert acked
 
 
+async def test_task_ack_type_when_received_overrides_worker_ack_type() -> None:
+    """Task ack_type label overrides worker-level ack type."""
+    events: list[str] = []
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(
+            _EventResultBackend(events),
+        )
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task
+    async def my_task() -> int:
+        events.append("task")
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_SAVED)
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={"ack_type": "when_received"},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["ack", "task", "post_execute", "save", "post_save"]
+
+
+async def test_task_ack_type_when_executed_overrides_worker_ack_type() -> None:
+    """Task ack_type label can delay a worker-level when_received ack."""
+    events: list[str] = []
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(
+            _EventResultBackend(events),
+        )
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task(ack_type="when_executed")
+    async def my_task() -> int:
+        events.append("task")
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_RECEIVED)
+    broker_message = broker.formatter.dumps(my_task.kicker()._prepare_message())
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["task", "ack", "post_execute", "save", "post_save"]
+
+
+async def test_task_ack_type_when_saved_overrides_worker_ack_type() -> None:
+    """Task ack_type label can use when_saved with an early-ack worker."""
+    events: list[str] = []
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(
+            _EventResultBackend(events),
+        )
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task
+    async def my_task() -> int:
+        events.append("task")
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_RECEIVED)
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={"ack_type": "when_saved"},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["task", "post_execute", "save", "post_save", "ack"]
+
+
+async def test_worker_ack_type_is_used_without_task_ack_type() -> None:
+    """Worker-level ack_type is used when the task has no ack_type label."""
+    events: list[str] = []
+    broker = (
+        InMemoryBroker()
+        .with_result_backend(
+            _EventResultBackend(events),
+        )
+        .with_middlewares(_EventMiddleware(events))
+    )
+
+    @broker.task
+    async def my_task() -> int:
+        events.append("task")
+        return 1
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_RECEIVED)
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    await receiver.callback(
+        AckableMessage(
+            data=broker_message.message,
+            ack=ack_callback,
+        ),
+    )
+
+    assert events == ["ack", "task", "post_execute", "save", "post_save"]
+
+
+async def test_invalid_task_ack_type_raises_before_execution() -> None:
+    """Invalid task ack_type fails before task execution and acknowledgement."""
+    events: list[str] = []
+    broker = InMemoryBroker()
+
+    @broker.task
+    async def my_task() -> None:
+        events.append("task")
+
+    def ack_callback() -> None:
+        events.append("ack")
+
+    receiver = get_receiver(broker, ack_type=AcknowledgeType.WHEN_RECEIVED)
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={"ack_type": "when_save"},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    with pytest.raises(ValueError, match="Invalid ack_type label"):
+        await receiver.callback(
+            AckableMessage(
+                data=broker_message.message,
+                ack=ack_callback,
+            ),
+        )
+
+    assert events == []
+
+
 async def test_callback_wrong_format() -> None:
     """Test that wrong format of a message won't throw an error."""
     receiver = get_receiver()