feat: support delayed tasks

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
 

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"]):

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"

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

tests/test_broker_kick.py

@@ -1,12 +1,14 @@
+import asyncio
+
 import capo_sqs
 import pytest
 from taskiq import BrokerMessage
 
 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)