diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 66c8510d..6b061b33 100644 --- a/docs/guide/cli.md +++ b/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: diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 2414754f..1a764bc4 100644 --- a/taskiq/__init__.py +++ b/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", diff --git a/taskiq/acks.py b/taskiq/acks.py index d55f5b2f..f8179755 100644 --- a/taskiq/acks.py +++ b/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. diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 5c2a1468..8773c68e 100644 --- a/taskiq/receiver/receiver.py +++ b/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], diff --git a/tests/receiver/test_receiver.py b/tests/receiver/test_receiver.py index eeb29c11..695d5c1a 100644 --- a/tests/receiver/test_receiver.py +++ b/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() diff --git a/tests/test_acks.py b/tests/test_acks.py new file mode 100644 index 00000000..899e731a --- /dev/null +++ b/tests/test_acks.py @@ -0,0 +1,24 @@ +import pytest + +from taskiq.acks import AcknowledgeType, parse_acknowledge_type + + +def test_parse_acknowledge_type_from_enum() -> None: + assert ( + parse_acknowledge_type(AcknowledgeType.WHEN_RECEIVED) + == AcknowledgeType.WHEN_RECEIVED + ) + + +def test_parse_acknowledge_type_from_string() -> None: + assert parse_acknowledge_type("when_executed") == AcknowledgeType.WHEN_EXECUTED + + +def test_parse_acknowledge_type_unknown_string() -> None: + with pytest.raises(ValueError, match="Unknown acknowledge type value"): + parse_acknowledge_type("when_save") + + +def test_parse_acknowledge_type_unsupported_value() -> None: + with pytest.raises(ValueError, match="Unsupported acknowledge type"): + parse_acknowledge_type(1)