feat: add support for message expiration

README.md

@@ -108,6 +108,26 @@ async def process_order() -> None:
 - `deduplication_id` sets `MessageDeduplicationId`; if not set, the queue must have content-based deduplication enabled, or SQS rejects the message.
 - The `delay` label (see [Delayed tasks](#delayed-tasks)) is not supported on FIFO queues — SQS only allows delay to be configured on the queue itself, not per message — and raises `FifoDelayNotSupportedError` if used.
 
+## Message expiration
+
+Set the `expiry` label to a unix timestamp; if a worker receives the message after that time, it's deleted without being executed:
+
+```python
+import time
+from taskiq_sqs import SQSBroker
+from taskiq_sqs.types import SQSQueue
+
+broker = SQSBroker(queues=SQSQueue(name="my-queue"))
+
+@broker.task()
+async def process_event() -> None:
+    ...
+
+await process_event.kicker().with_labels(expiry=time.time() + 300).kiq()  # discarded if received after 5 minutes
+```
+
+Expiration is checked by the worker on receipt, not by SQS itself — a message can still sit in the queue past its expiry (e.g. while workers are busy or scaled to zero), it just won't run once picked up. `expiry` must be a non-negative number; anything else raises `InvalidExpiryError` when the task is kicked.
+
 ## Offloading large messages to S3
 
 SQS messages are limited to 256 KiB. `S3OffloadMiddleware` transparently uploads task payloads that exceed a configurable threshold to S3 before sending them to the queue, and replaces the message with a reference to the uploaded object. The worker downloads the original payload back from S3 before executing the task, and (by default) removes it from S3 afterwards.

src/taskiq_sqs/broker.py

@@ -1,6 +1,7 @@
 import asyncio
 import contextlib
 import logging
+import time
 from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence
 from typing import Any
 
@@ -10,15 +11,14 @@
 from taskiq.message import BrokerMessage
 
 from taskiq_sqs import constants
-from taskiq_sqs.exceptions import (
-    BrokerInitError,
-    FifoDelayNotSupportedError,
-    InvalidDelaySecondsError,
-    InvalidMessageDeduplicationIdError,
-    InvalidMessageGroupIdError,
-    UnknownQueueError,
+from taskiq_sqs.exceptions import BrokerInitError, FifoDelayNotSupportedError, UnknownQueueError
+from taskiq_sqs.types.message import (
+    validate_delay_seconds,
+    validate_expiry,
+    validate_message_deduplication_id,
+    validate_message_group_id,
 )
-from taskiq_sqs.types import SQSQueue
+from taskiq_sqs.types.queue import SQSQueue, validate_queue
 
 
 logger = logging.getLogger(__name__)
@@ -68,22 +68,7 @@ def _normalize_queues(queues: SQSQueue | Sequence[SQSQueue]) -> list[SQSQueue]:
             raise BrokerInitError(details="Queue names must be unique.")
 
         for queue in queue_list:
-            max_number_of_messages = queue.get("max_number_of_messages", 1)
-            if max_number_of_messages > constants.MAX_NUMBER_OF_MESSAGES or max_number_of_messages < 1:
-                raise BrokerInitError(
-                    details=f"MaxNumberOfMessages for queue '{queue['name']}' can be no greater than 10 or less than 1",
-                )
-            wait_time_seconds = queue.get("wait_time_seconds", 0)
-            if wait_time_seconds > constants.MAX_WAIT_TIME_SECONDS or wait_time_seconds < 0:
-                raise BrokerInitError(
-                    details=f"WaitTimeSeconds for queue '{queue['name']}' can be no greater than 20 or less than 0",
-                )
-            ends_with_fifo_suffix = queue["name"].endswith(".fifo")
-            if "is_fifo" in queue and queue["is_fifo"] != ends_with_fifo_suffix:
-                raise BrokerInitError(
-                    details=f"Queue '{queue['name']}' has is_fifo={queue['is_fifo']}, but SQS requires FIFO queue "
-                    "names to end in '.fifo' and standard queue names not to",
-                )
+            validate_queue(queue)
         return queue_list
 
     def _resolve_queue(self, queue_name: str | None) -> SQSQueue:
@@ -160,38 +145,20 @@ async def _build_kick_kwargs(
         if constants.SQS_DELAY_SECONDS_LABEL in message.labels:
             if is_fifo:
                 raise FifoDelayNotSupportedError(queue_name=queue["name"])
-            kwargs["delay_seconds"] = self._validate_delay_seconds(message.labels[constants.SQS_DELAY_SECONDS_LABEL])
+            kwargs["delay_seconds"] = validate_delay_seconds(message.labels[constants.SQS_DELAY_SECONDS_LABEL])
         if is_fifo:
             group_id = message.labels.get(constants.SQS_MESSAGE_GROUP_ID_LABEL, message.task_name)
-            kwargs["message_group_id"] = self._validate_message_group_id(group_id)
+            kwargs["message_group_id"] = validate_message_group_id(group_id)
             if constants.SQS_MESSAGE_DEDUPLICATION_ID_LABEL in message.labels:
                 deduplication_id = message.labels[constants.SQS_MESSAGE_DEDUPLICATION_ID_LABEL]
-                kwargs["message_deduplication_id"] = self._validate_message_deduplication_id(deduplication_id)
+                kwargs["message_deduplication_id"] = validate_message_deduplication_id(deduplication_id)
+        if constants.SQS_EXPIRY_LABEL in message.labels:
+            expiry = validate_expiry(message.labels[constants.SQS_EXPIRY_LABEL])
+            kwargs["message_attributes"] = {
+                constants.SQS_EXPIRY_LABEL: {"data_type": "Number", "string_value": str(expiry)},
+            }
         return kwargs
 
-    @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
-
-    @staticmethod
-    def _validate_message_group_id(group_id: Any) -> str:
-        if not isinstance(group_id, str) or not (1 <= len(group_id) <= constants.MAX_FIFO_ID_LENGTH):
-            raise InvalidMessageGroupIdError(group_id=group_id, max_length=constants.MAX_FIFO_ID_LENGTH)
-        return group_id
-
-    @staticmethod
-    def _validate_message_deduplication_id(deduplication_id: Any) -> str:
-        if not isinstance(deduplication_id, str) or not (1 <= len(deduplication_id) <= constants.MAX_FIFO_ID_LENGTH):
-            raise InvalidMessageDeduplicationIdError(
-                deduplication_id=deduplication_id,
-                max_length=constants.MAX_FIFO_ID_LENGTH,
-            )
-        return deduplication_id
-
     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))
@@ -223,6 +190,20 @@ async def ack() -> None:
 
         return ack
 
+    @staticmethod
+    def _is_expired(message: Mapping[str, Any]) -> bool:
+        expiry_attribute = message.get("message_attributes", {}).get(constants.SQS_EXPIRY_LABEL)
+        if expiry_attribute is None:
+            return False
+        string_value = expiry_attribute.get("string_value")
+        if string_value is None:
+            return False
+        try:
+            expiry = float(string_value)
+        except ValueError:
+            return False
+        return time.time() > expiry
+
     async def _poll_queue(self, queue: SQSQueue, incoming: "asyncio.Queue[_QueueItem]") -> None:
         """Continuously receive messages from a single queue and forward them to the shared incoming queue."""
         try:
@@ -233,17 +214,24 @@ async def _poll_queue(self, queue: SQSQueue, incoming: "asyncio.Queue[_QueueItem
                         queue_url=queue_url,
                         max_number_of_messages=queue.get("max_number_of_messages", 1),
                         wait_time_seconds=queue.get("wait_time_seconds", 0),
+                        message_attribute_names=[constants.SQS_EXPIRY_LABEL],
                     )
                 for message in results.get("messages", []):
                     body = message.get("body")
                     receipt_handle = message.get("receipt_handle")
-                    if body and receipt_handle:
-                        await incoming.put(
-                            AckableMessage(
-                                data=body.encode("utf-8"),
-                                ack=self._build_ack_function(queue["name"], queue_url, receipt_handle),
-                            ),
-                        )
+                    if not (body and receipt_handle):
+                        continue
+                    if self._is_expired(message):
+                        logger.info("Discarding expired message from queue '%s'", queue["name"])
+                        with self._handle_exceptions(queue["name"]):
+                            await self._sqs_client.delete_message(queue_url=queue_url, receipt_handle=receipt_handle)
+                        continue
+                    await incoming.put(
+                        AckableMessage(
+                            data=body.encode("utf-8"),
+                            ack=self._build_ack_function(queue["name"], queue_url, receipt_handle),
+                        ),
+                    )
         except asyncio.CancelledError:
             raise
         except Exception as exc:  # noqa: BLE001

src/taskiq_sqs/constants.py

@@ -15,3 +15,4 @@
 SQS_DELAY_SECONDS_LABEL: Final[str] = "delay"
 SQS_MESSAGE_GROUP_ID_LABEL: Final[str] = "group_id"
 SQS_MESSAGE_DEDUPLICATION_ID_LABEL: Final[str] = "deduplication_id"
+SQS_EXPIRY_LABEL: Final[str] = "expiry"

src/taskiq_sqs/exceptions.py

@@ -90,3 +90,10 @@ class InvalidMessageDeduplicationIdError(BaseTaskiqSQSError):
     )
     deduplication_id: object
     max_length: int
+
+
+class InvalidExpiryError(BaseTaskiqSQSError):
+    """Error if a message's expiry label isn't a valid unix timestamp."""
+
+    __template__ = "expiry must be a non-negative unix timestamp (int or float), got {expiry!r}"
+    expiry: object

src/taskiq_sqs/types/message.py

@@ -0,0 +1,55 @@
+from typing import Any
+
+from taskiq_sqs import constants
+from taskiq_sqs.exceptions import (
+    InvalidDelaySecondsError,
+    InvalidExpiryError,
+    InvalidMessageDeduplicationIdError,
+    InvalidMessageGroupIdError,
+)
+
+
+def validate_delay_seconds(delay_seconds: Any) -> int:
+    """Validate a message's delay label."""
+    if isinstance(delay_seconds, bool):
+        raise InvalidDelaySecondsError(delay_seconds=delay_seconds, max_delay_seconds=constants.MAX_DELAY_SECONDS)
+    try:
+        numeric = float(delay_seconds)
+    except (TypeError, ValueError):
+        raise InvalidDelaySecondsError(
+            delay_seconds=delay_seconds,
+            max_delay_seconds=constants.MAX_DELAY_SECONDS,
+        ) from None
+    if not numeric.is_integer() or numeric < 0 or numeric > constants.MAX_DELAY_SECONDS:
+        raise InvalidDelaySecondsError(delay_seconds=delay_seconds, max_delay_seconds=constants.MAX_DELAY_SECONDS)
+    return int(numeric)
+
+
+def validate_message_group_id(group_id: Any) -> str:
+    """Validate a message's group_id label."""
+    if not isinstance(group_id, str) or not (1 <= len(group_id) <= constants.MAX_FIFO_ID_LENGTH):
+        raise InvalidMessageGroupIdError(group_id=group_id, max_length=constants.MAX_FIFO_ID_LENGTH)
+    return group_id
+
+
+def validate_message_deduplication_id(deduplication_id: Any) -> str:
+    """Validate a message's deduplication_id label."""
+    if not isinstance(deduplication_id, str) or not (1 <= len(deduplication_id) <= constants.MAX_FIFO_ID_LENGTH):
+        raise InvalidMessageDeduplicationIdError(
+            deduplication_id=deduplication_id,
+            max_length=constants.MAX_FIFO_ID_LENGTH,
+        )
+    return deduplication_id
+
+
+def validate_expiry(expiry: Any) -> float:
+    """Validate a message's expiry label."""
+    if isinstance(expiry, bool):
+        raise InvalidExpiryError(expiry=expiry)
+    try:
+        value = float(expiry)
+    except (TypeError, ValueError):
+        raise InvalidExpiryError(expiry=expiry) from None
+    if value < 0:
+        raise InvalidExpiryError(expiry=expiry)
+    return value

src/taskiq_sqs/types/queue.py

@@ -1,5 +1,8 @@
 from typing import NotRequired, TypedDict
 
+from taskiq_sqs import constants
+from taskiq_sqs.exceptions import BrokerInitError
+
 
 class SQSQueue(TypedDict):
     """
@@ -16,3 +19,23 @@ class SQSQueue(TypedDict):
     max_number_of_messages: NotRequired[int]
     wait_time_seconds: NotRequired[int]
     is_fifo: NotRequired[bool]
+
+
+def validate_queue(queue: SQSQueue) -> None:
+    """Validate a single queue's own fields against SQS's constraints."""
+    max_number_of_messages = queue.get("max_number_of_messages", 1)
+    if max_number_of_messages > constants.MAX_NUMBER_OF_MESSAGES or max_number_of_messages < 1:
+        raise BrokerInitError(
+            details=f"MaxNumberOfMessages for queue '{queue['name']}' can be no greater than 10 or less than 1",
+        )
+    wait_time_seconds = queue.get("wait_time_seconds", 0)
+    if wait_time_seconds > constants.MAX_WAIT_TIME_SECONDS or wait_time_seconds < 0:
+        raise BrokerInitError(
+            details=f"WaitTimeSeconds for queue '{queue['name']}' can be no greater than 20 or less than 0",
+        )
+    ends_with_fifo_suffix = queue["name"].endswith(".fifo")
+    if "is_fifo" in queue and queue["is_fifo"] != ends_with_fifo_suffix:
+        raise BrokerInitError(
+            details=f"Queue '{queue['name']}' has is_fifo={queue['is_fifo']}, but SQS requires FIFO queue "
+            "names to end in '.fifo' and standard queue names not to",
+        )

tests/test_broker_kick.py

@@ -9,6 +9,7 @@
 from taskiq_sqs import SQSBroker
 from taskiq_sqs.constants import (
     SQS_DELAY_SECONDS_LABEL,
+    SQS_EXPIRY_LABEL,
     SQS_MESSAGE_DEDUPLICATION_ID_LABEL,
     SQS_MESSAGE_GROUP_ID_LABEL,
     SQS_QUEUE_LABEL,
@@ -17,6 +18,7 @@
     BrokerInitError,
     FifoDelayNotSupportedError,
     InvalidDelaySecondsError,
+    InvalidExpiryError,
     InvalidMessageDeduplicationIdError,
     InvalidMessageGroupIdError,
     UnknownQueueError,
@@ -101,7 +103,7 @@ async def test_when_kick_called_with_delay_label__then_message_is_delayed(
     assert len(delayed.get("messages", [])) == 1
 
 
-@pytest.mark.parametrize("delay_seconds", [-1, 901, "10", 10.5, True])
+@pytest.mark.parametrize("delay_seconds", [-1, 901, "not-a-number", 10.5, "10.5", True])
 async def test_when_kick_called_with_invalid_delay_label__then_should_raise_an_error(
     sqs_broker: SQSBroker,
     broker_message: BrokerMessage,
@@ -113,6 +115,46 @@ async def test_when_kick_called_with_invalid_delay_label__then_should_raise_an_e
         await sqs_broker.kick(broker_message)
 
 
+async def test_when_kick_called_with_stringified_delay_label__then_it_is_accepted(
+    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
+
+
+async def test_when_task_kicked_through_kicker_with_delay_label__then_it_is_delayed(
+    sqs_broker: SQSBroker,
+    sqs_client: capo_sqs.AsyncSQSClient,
+    sqs_queue: str,
+) -> None:
+    """End-to-end regression test for the real `@broker.task()` / `.kiq()` path, not a hand-built BrokerMessage."""
+
+    @sqs_broker.task()
+    async def sample_task() -> None: ...
+
+    await sample_task.kicker().with_labels(**{SQS_DELAY_SECONDS_LABEL: 1}).kiq()
+
+    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
+
+
 async def test_when_kick_called_on_standard_queue__then_no_fifo_attributes_are_sent(
     sqs_broker: SQSBroker,
     sqs_client: capo_sqs.AsyncSQSClient,
@@ -239,3 +281,33 @@ async def test_when_multiple_messages_kicked_to_same_group__then_order_is_preser
     response = await sqs_client.receive_message(queue_url=fifo_sqs_queue, max_number_of_messages=3)
     bodies = [message.get("body") for message in response.get("messages", [])]
     assert bodies == ["message-0", "message-1", "message-2"]
+
+
+@pytest.mark.parametrize("expiry", [-1, "soon", True])
+async def test_when_kick_called_with_invalid_expiry_label__then_should_raise_an_error(
+    sqs_broker: SQSBroker,
+    broker_message: BrokerMessage,
+    expiry: object,
+) -> None:
+    broker_message.labels[SQS_EXPIRY_LABEL] = expiry
+
+    with pytest.raises(InvalidExpiryError):
+        await sqs_broker.kick(broker_message)
+
+
+async def test_when_kick_called_with_stringified_expiry_label__then_it_is_accepted(
+    sqs_broker: SQSBroker,
+    sqs_client: capo_sqs.AsyncSQSClient,
+    sqs_queue: str,
+    broker_message: BrokerMessage,
+) -> None:
+    # same reasoning as test_when_kick_called_with_stringified_delay_label__then_it_is_accepted
+    broker_message.labels[SQS_EXPIRY_LABEL] = "1789505020.5"
+
+    await sqs_broker.kick(broker_message)
+
+    response = await sqs_client.receive_message(queue_url=sqs_queue, message_attribute_names=[SQS_EXPIRY_LABEL])
+    messages = response.get("messages", [])
+    assert len(messages) == 1
+    attribute = messages[0].get("message_attributes", {}).get(SQS_EXPIRY_LABEL, {})
+    assert attribute.get("string_value") == "1789505020.5"