diff --git a/README.md b/README.md index d9812c6..5c6dd2f 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,30 @@ if __name__ == "__main__": asyncio.run(main()) ``` +### JetStream acknowledgement progress + +JetStream brokers expose acknowledgement progress for long-running tasks. Call it +before the consumer's acknowledgement timeout expires to tell JetStream that the +message is still being processed. + +```python +from taskiq import Context, TaskiqDepends + + +@broker.task(ack_type="manual") +async def long_running_task( + context: Context = TaskiqDepends(), +) -> None: + while still_working: + await do_some_work() + await context.ack_progress() + + await context.ack() +``` + +`ack_progress()` can be called repeatedly. It does not finally acknowledge the +message; call `ack()` when processing has completed successfully. + ## NatsBroker configuration Here's the constructor parameters: diff --git a/taskiq_nats/broker.py b/taskiq_nats/broker.py index 62b272d..79d65fb 100644 --- a/taskiq_nats/broker.py +++ b/taskiq_nats/broker.py @@ -196,6 +196,7 @@ async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]: yield AckableMessage( data=message.data, ack=message.ack, + ack_progress=message.in_progress, ) async def _startup_consumer(self) -> None: @@ -237,6 +238,7 @@ async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]: yield AckableMessage( data=nats_message.data, ack=nats_message.ack, + ack_progress=nats_message.in_progress, ) except NatsTimeoutError: continue diff --git a/tests/test_jetstream.py b/tests/test_jetstream.py index 49a4471..f14407e 100644 --- a/tests/test_jetstream.py +++ b/tests/test_jetstream.py @@ -1,5 +1,7 @@ import asyncio import uuid +from collections.abc import AsyncGenerator +from unittest.mock import AsyncMock, MagicMock from taskiq import AckableMessage, BrokerMessage @@ -7,6 +9,46 @@ from tests.utils import read_message +async def test_push_based_broker_exposes_ack_progress( + nats_urls: list[str], +) -> None: + nats_message = MagicMock() + nats_message.data = b"message" + nats_message.ack = AsyncMock() + nats_message.in_progress = AsyncMock() + + async def message_stream() -> AsyncGenerator[MagicMock, None]: + yield nats_message + + broker = PushBasedJetStreamBroker(servers=nats_urls) + broker.consumer = MagicMock(messages=message_stream()) + + message = await anext(broker.listen()) + + assert message.ack_progress is not None + await message.ack_progress() + nats_message.in_progress.assert_awaited_once_with() + + +async def test_pull_based_broker_exposes_ack_progress( + nats_urls: list[str], +) -> None: + nats_message = MagicMock() + nats_message.data = b"message" + nats_message.ack = AsyncMock() + nats_message.in_progress = AsyncMock() + + broker = PullBasedJetStreamBroker(servers=nats_urls) + broker.consumer = MagicMock() + broker.consumer.fetch = AsyncMock(return_value=[nats_message]) + + message = await anext(broker.listen()) + + assert message.ack_progress is not None + await message.ack_progress() + nats_message.in_progress.assert_awaited_once_with() + + async def test_push_based_broker_success( # (too many await) nats_urls: list[str], nats_subject: str, @@ -36,6 +78,8 @@ async def test_push_based_broker_success( # (too many await) ackable_msg = await asyncio.wait_for(read_message(broker), 0.5) assert isinstance(ackable_msg, AckableMessage) assert ackable_msg.data == sent_message.message + assert ackable_msg.ack_progress is not None + await ackable_msg.ack_progress() ack = ackable_msg.ack() if ack is not None: await ack @@ -76,6 +120,8 @@ async def test_pull_based_broker_success( ackable_msg = await asyncio.wait_for(read_message(broker), 0.5) assert isinstance(ackable_msg, AckableMessage) assert ackable_msg.data == sent_message.message + assert ackable_msg.ack_progress is not None + await ackable_msg.ack_progress() ack = ackable_msg.ack() if ack is not None: await ack