diff --git a/README.md b/README.md index f3e3bc2..ae8feca 100644 --- a/README.md +++ b/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. diff --git a/src/taskiq_sqs/broker.py b/src/taskiq_sqs/broker.py index fa22c74..466314e 100644 --- a/src/taskiq_sqs/broker.py +++ b/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 diff --git a/src/taskiq_sqs/constants.py b/src/taskiq_sqs/constants.py index ddf55a0..9655b27 100644 --- a/src/taskiq_sqs/constants.py +++ b/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" diff --git a/src/taskiq_sqs/exceptions.py b/src/taskiq_sqs/exceptions.py index 3f92829..b9af2ae 100644 --- a/src/taskiq_sqs/exceptions.py +++ b/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 diff --git a/src/taskiq_sqs/types/message.py b/src/taskiq_sqs/types/message.py new file mode 100644 index 0000000..609e86f --- /dev/null +++ b/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 diff --git a/src/taskiq_sqs/types/queue.py b/src/taskiq_sqs/types/queue.py index 2e5dc30..067cdf0 100644 --- a/src/taskiq_sqs/types/queue.py +++ b/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", + ) diff --git a/tests/test_broker_kick.py b/tests/test_broker_kick.py index afc5b1c..38d1471 100644 --- a/tests/test_broker_kick.py +++ b/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" diff --git a/tests/test_broker_listen.py b/tests/test_broker_listen.py index ed761b4..a0bdb3b 100644 --- a/tests/test_broker_listen.py +++ b/tests/test_broker_listen.py @@ -1,6 +1,12 @@ +import asyncio +import time + import capo_sqs +import pytest +from taskiq import BrokerMessage from taskiq_sqs import SQSBroker +from taskiq_sqs.constants import SQS_EXPIRY_LABEL async def test_when_listen__than_we_should_delete_message_from_queue( @@ -48,3 +54,48 @@ async def test_when_listen_with_multiple_queues__then_messages_from_both_are_rec await generator.aclose() assert {message.data for message in messages} == {b"from_first_queue", b"from_second_queue"} + + +async def test_when_message_expired__then_it_is_discarded_without_being_yielded( + sqs_broker: SQSBroker, + sqs_client: capo_sqs.AsyncSQSClient, + sqs_queue: str, +) -> None: + expired_message = BrokerMessage( + task_id="expired_task", + task_name="expired_task", + message=b"expired_message", + labels={SQS_EXPIRY_LABEL: time.time() - 10}, + ) + await sqs_broker.kick(expired_message) + + generator = sqs_broker.listen() + try: + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(generator.__anext__(), timeout=2) + finally: + await generator.aclose() + + response = await sqs_client.receive_message(queue_url=sqs_queue) + assert not response.get("messages") + + +async def test_when_message_not_yet_expired__then_it_is_yielded_normally( + sqs_broker: SQSBroker, +) -> None: + live_message = BrokerMessage( + task_id="live_task", + task_name="live_task", + message=b"live_message", + labels={SQS_EXPIRY_LABEL: time.time() + 60}, + ) + await sqs_broker.kick(live_message) + + generator = sqs_broker.listen() + try: + received = await asyncio.wait_for(generator.__anext__(), timeout=2) + await received.ack() + finally: + await generator.aclose() + + assert received.data == b"live_message"