diff --git a/README.md b/README.md index 62a2d73..8583f83 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ How to run: ## Multiple queues -`SQSBroker` accepts a single queue or a list of them. The first queue is the default one, used whenever a task doesn't say otherwise. To send a task to a specific queue, set the `sqs_queue` label with that queue's name: +`SQSBroker` accepts a single queue or a list of them. The first queue is the default one, used whenever a task doesn't say otherwise. To send a task to a specific queue, set the `queue_name` label with that queue's name: ```python from taskiq_sqs import SQSBroker @@ -65,12 +65,29 @@ broker = SQSBroker( ], ) -@broker.task(sqs_queue="high-priority-queue") # "sqs_queue" is taskiq_sqs.broker.SQS_QUEUE_LABEL +@broker.task(queue_name="high-priority-queue") async def urgent_task() -> None: ... ``` -A worker started against this broker consumes from every configured queue at once. Passing a queue name through the `sqs_queue` label that isn't configured on the broker raises `UnknownQueueError`. +A worker started against this broker consumes from every configured queue at once. Passing a queue name through the `queue_name` label that isn't configured on the broker raises `UnknownQueueError`. + +## Delayed tasks + +Set the `delay` label to delay delivery of a task by that many seconds (0-900, SQS's own limit) + +```python +from taskiq_sqs import SQSBroker +from taskiq_sqs.types import SQSQueue + +broker = SQSBroker(queues=SQSQueue(name="my-queue")) + +@broker.task(delay=30) # "delay" is taskiq_sqs.constants.SQS_DELAY_SECONDS_LABEL +async def send_reminder() -> None: + ... +``` + +A value outside the 0-900 range (or not an integer) raises `InvalidDelaySecondsError` when the task is kicked. ## Offloading large messages to S3 diff --git a/src/taskiq_sqs/broker.py b/src/taskiq_sqs/broker.py index 89be6b2..637a573 100644 --- a/src/taskiq_sqs/broker.py +++ b/src/taskiq_sqs/broker.py @@ -10,14 +10,12 @@ from taskiq.message import BrokerMessage from taskiq_sqs import constants -from taskiq_sqs.exceptions import BrokerInitError, UnknownQueueError +from taskiq_sqs.exceptions import BrokerInitError, InvalidDelaySecondsError, UnknownQueueError from taskiq_sqs.types import SQSQueue logger = logging.getLogger(__name__) -SQS_QUEUE_LABEL = "sqs_queue" - _QueueItem = AckableMessage | BaseException @@ -34,7 +32,7 @@ def __init__( ) -> None: """Initialize the SQS broker. - :param queues: a single queue configuration, or a sequence of them for multiqueue support.ф + :param queues: a single queue configuration, or a sequence of them for multiqueue support. :param endpoint_url: the SQS endpoint URL. :param aws_region_name: the AWS region name. :param aws_access_key_id: the AWS access key ID. @@ -139,20 +137,25 @@ async def _build_kick_kwargs( :param message: BrokerMessage object. :param queue_url: URL of the queue the message will be sent to. """ - return { + kwargs: dict[str, Any] = { "queue_url": queue_url, "message_body": message.message.decode("utf-8"), } + if constants.SQS_DELAY_SECONDS_LABEL in message.labels: + kwargs["delay_seconds"] = self._validate_delay_seconds(message.labels[constants.SQS_DELAY_SECONDS_LABEL]) + return kwargs - async def kick(self, message: BrokerMessage) -> None: - """Kick tasks out from current program to configured SQS queue. - - The target queue is picked from the `sqs_queue` label (see `SQS_QUEUE_LABEL`), falling back to the first - configured queue when the label isn't set. + @staticmethod + def _validate_delay_seconds(delay_seconds: Any) -> int: + if isinstance(delay_seconds, bool) or not isinstance(delay_seconds, int): + raise InvalidDelaySecondsError(delay_seconds=delay_seconds, max_delay_seconds=constants.MAX_DELAY_SECONDS) + if delay_seconds < 0 or delay_seconds > constants.MAX_DELAY_SECONDS: + raise InvalidDelaySecondsError(delay_seconds=delay_seconds, max_delay_seconds=constants.MAX_DELAY_SECONDS) + return delay_seconds - :param message: BrokerMessage object. - """ - queue = self._resolve_queue(message.labels.get(SQS_QUEUE_LABEL)) + async def kick(self, message: BrokerMessage) -> None: + """Kick tasks out from current program to configured SQS queue.""" + queue = self._resolve_queue(message.labels.get(constants.SQS_QUEUE_LABEL)) queue_url = await self._get_queue_url(queue["name"]) kwargs = await self._build_kick_kwargs(message, queue_url) with self._handle_exceptions(queue["name"]): diff --git a/src/taskiq_sqs/constants.py b/src/taskiq_sqs/constants.py index 5c1f62e..0954db3 100644 --- a/src/taskiq_sqs/constants.py +++ b/src/taskiq_sqs/constants.py @@ -5,6 +5,10 @@ MAX_WAIT_TIME_SECONDS: Final[int] = 20 MAX_NUMBER_OF_MESSAGES: Final[int] = 10 +MAX_DELAY_SECONDS: Final[int] = 900 SQS_MAX_MESSAGE_SIZE_BYTES: Final[int] = 262_144 DEFAULT_S3_OFFLOAD_THRESHOLD_BYTES: Final[int] = 200_000 + +SQS_QUEUE_LABEL: Final[str] = "queue_name" +SQS_DELAY_SECONDS_LABEL: Final[str] = "delay" diff --git a/src/taskiq_sqs/exceptions.py b/src/taskiq_sqs/exceptions.py index 030350b..f4fe525 100644 --- a/src/taskiq_sqs/exceptions.py +++ b/src/taskiq_sqs/exceptions.py @@ -54,3 +54,11 @@ class UnknownQueueError(BaseTaskiqSQSError): __template__ = "Message references queue '{queue_name}' which is not configured on this broker" queue_name: str + + +class InvalidDelaySecondsError(BaseTaskiqSQSError): + """Error if a message's delay label is outside SQS's allowed range.""" + + __template__ = "delay must be an integer between 0 and {max_delay_seconds}, got {delay_seconds!r}" + delay_seconds: object + max_delay_seconds: int diff --git a/tests/test_broker_kick.py b/tests/test_broker_kick.py index 18dc790..618ae55 100644 --- a/tests/test_broker_kick.py +++ b/tests/test_broker_kick.py @@ -1,3 +1,5 @@ +import asyncio + import capo_sqs import pytest from taskiq import BrokerMessage @@ -5,8 +7,8 @@ from tests.conftest import _queue_name_from_url from taskiq_sqs import SQSBroker -from taskiq_sqs.broker import SQS_QUEUE_LABEL -from taskiq_sqs.exceptions import BrokerInitError, UnknownQueueError +from taskiq_sqs.constants import SQS_DELAY_SECONDS_LABEL, SQS_QUEUE_LABEL +from taskiq_sqs.exceptions import BrokerInitError, InvalidDelaySecondsError, UnknownQueueError async def test_when_kick_called__than_message_should_be_published_to_queue( @@ -66,3 +68,34 @@ async def test_when_kick_called_with_unknown_queue_label__then_should_raise_an_e with pytest.raises(UnknownQueueError): await multiqueue_sqs_broker.kick(broker_message) + + +async def test_when_kick_called_with_delay_label__then_message_is_delayed( + sqs_broker: SQSBroker, + sqs_client: capo_sqs.AsyncSQSClient, + sqs_queue: str, + broker_message: BrokerMessage, +) -> None: + broker_message.labels[SQS_DELAY_SECONDS_LABEL] = 1 + + await sqs_broker.kick(broker_message) + + immediate = await sqs_client.receive_message(queue_url=sqs_queue) + assert not immediate.get("messages") + + await asyncio.sleep(1.2) + + delayed = await sqs_client.receive_message(queue_url=sqs_queue) + assert len(delayed.get("messages", [])) == 1 + + +@pytest.mark.parametrize("delay_seconds", [-1, 901, "10", 10.5, True]) +async def test_when_kick_called_with_invalid_delay_label__then_should_raise_an_error( + sqs_broker: SQSBroker, + broker_message: BrokerMessage, + delay_seconds: object, +) -> None: + broker_message.labels[SQS_DELAY_SECONDS_LABEL] = delay_seconds + + with pytest.raises(InvalidDelaySecondsError): + await sqs_broker.kick(broker_message)