feat: add message batcing support

README.md

@@ -22,7 +22,7 @@ from taskiq_sqs import S3ResultBackend, SQSBroker
 from taskiq_sqs.types import S3Bucket, SQSQueue
 
 broker = SQSBroker(
-    queues=SQSQueue(name="my-queue"),  # specify an existing queue
+    queues=SQSQueue(name="my-queue"),  # by default the broker creates the queue for you if it doesn't exist
     endpoint_url="http://localhost:4566",
     aws_region_name="us-east-1",
 ).with_result_backend(
@@ -72,6 +72,23 @@ async def urgent_task() -> None:
 
 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`.
 
+## Declaring queues
+
+By default the broker creates a queue on startup if it doesn't exist yet, the same way `S3Bucket` does for buckets. Set `is_declare=False` to require the queue to already exist instead (raises `QueueNotFoundError` if it doesn't). `options` are queue attributes (e.g. `VisibilityTimeout`, `MessageRetentionPeriod`) passed to `CreateQueue`, in AWS's own PascalCase naming, when the queue is declared — they have no effect on a queue that already exists:
+
+```python
+from taskiq_sqs import SQSBroker
+from taskiq_sqs.types import SQSQueue
+
+broker = SQSBroker(
+    queues=SQSQueue(name="my-queue", options={"VisibilityTimeout": "60", "MessageRetentionPeriod": "86400"}),
+)
+```
+
+FIFO queues get their `FifoQueue` attribute set automatically when declared — no need to include it in `options`.
+
+`S3Bucket` has the same `options` field, for parameters `CreateBucket` accepts beyond `name` (e.g. `acl`), passed through whenever `S3ResultBackend`/`S3OffloadMiddleware` create the bucket.
+
 ## Delayed tasks
 
 Set the `delay` label to delay delivery of a task by that many seconds (0-900, SQS's own limit)
@@ -128,6 +145,19 @@ await process_event.kicker().with_labels(expiry=time.time() + 300).kiq()  # disc
 
 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.
 
+## Message batching
+
+Set `is_batching_enabled` on a queue to buffer kicked messages in memory and flush them together via `SendMessageBatch` (up to `batch_size` messages, or after `batch_timeout` seconds, whichever comes first) instead of sending each one immediately:
+
+```python
+from taskiq_sqs import SQSBroker
+from taskiq_sqs.types import SQSQueue
+
+broker = SQSBroker(
+    queues=SQSQueue(name="my-queue", is_batching_enabled=True, batch_size=10, batch_timeout=1.0),
+)
+```
+
 ## 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

@@ -11,8 +11,14 @@
 from taskiq.message import BrokerMessage
 
 from taskiq_sqs import constants
-from taskiq_sqs.exceptions import BrokerInitError, FifoDelayNotSupportedError, UnknownQueueError
+from taskiq_sqs.exceptions import (
+    BrokerInitError,
+    FifoDelayNotSupportedError,
+    QueueNotFoundError,
+    UnknownQueueError,
+)
 from taskiq_sqs.types.message import (
+    is_label_true,
     validate_delay_seconds,
     validate_expiry,
     validate_message_deduplication_id,
@@ -56,6 +62,8 @@ def __init__(
         self._default_queue_name = self._queues[0]["name"]
         self._queues_by_name = {queue["name"]: queue for queue in self._queues}
         self._queue_urls: dict[str, str] = {}
+        self._batch_queues: dict[str, asyncio.Queue[dict[str, Any]]] = {}
+        self._batch_worker_tasks: dict[str, asyncio.Task[None]] = {}
 
     @staticmethod
     def _normalize_queues(queues: SQSQueue | Sequence[SQSQueue]) -> list[SQSQueue]:
@@ -104,8 +112,11 @@ async def startup(self) -> None:
         await self._sqs_client.__aenter__()
         try:
             for queue in self._queues:
-                queue_url = await self._get_queue_url(queue["name"])
+                queue_url = await self._get_queue_url(queue)
                 logger.info("Resolved queue '%s' URL: %s", queue["name"], queue_url)
+                if queue.get("is_batching_enabled", False):
+                    self._batch_queues[queue["name"]] = asyncio.Queue()
+                    self._batch_worker_tasks[queue["name"]] = asyncio.create_task(self._batch_worker(queue))
         except Exception:
             await self._sqs_client.__aexit__(None, None, None)
             raise
@@ -114,15 +125,39 @@ async def startup(self) -> None:
 
     async def shutdown(self) -> None:
         """Shuts down the SQS broker."""
+        for task in self._batch_worker_tasks.values():
+            task.cancel()
+        await asyncio.gather(*self._batch_worker_tasks.values(), return_exceptions=True)
+        for queue_name, batch_queue in self._batch_queues.items():
+            remaining = self._drain_batch_queue(batch_queue)
+            if remaining:
+                await self._send_batch(self._queues_by_name[queue_name], remaining)
         await self._sqs_client.__aexit__(None, None, None)
         await super().shutdown()
 
-    async def _get_queue_url(self, queue_name: str) -> str:
-        if queue_name not in self._queue_urls:
-            with self._handle_exceptions(queue_name):
-                result = await self._sqs_client.get_queue_url(queue_name=queue_name)
-            self._queue_urls[queue_name] = result["queue_url"]
-        return self._queue_urls[queue_name]
+    async def _get_queue_url(self, queue: SQSQueue) -> str:
+        name = queue["name"]
+        if name not in self._queue_urls:
+            result: Any
+            try:
+                result = await self._sqs_client.get_queue_url(queue_name=name)
+            except capo_sqs.errors.QueueDoesNotExist as exc:
+                if not queue.get("is_declare", True):
+                    raise QueueNotFoundError(queue_name=name) from exc
+                result = await self._create_queue(queue)
+            except capo_sqs.errors.ServiceError as exc:
+                raise BrokerInitError(details=exc.code or "") from exc
+            self._queue_urls[name] = result["queue_url"]
+        return self._queue_urls[name]
+
+    async def _create_queue(self, queue: SQSQueue) -> Any:
+        attributes: dict[Any, Any] = dict(queue.get("options", {}))
+        if queue.get("is_fifo", queue["name"].endswith(".fifo")):
+            attributes.setdefault("FifoQueue", "true")
+        try:
+            return await self._sqs_client.create_queue(queue_name=queue["name"], attributes=attributes or None)
+        except capo_sqs.errors.ServiceError as exc:
+            raise BrokerInitError(details=exc.code or "") from exc
 
     async def _build_kick_kwargs(
         self,
@@ -159,14 +194,86 @@ async def _build_kick_kwargs(
             }
         return kwargs
 
+    def _should_batch(self, queue: SQSQueue, message: BrokerMessage) -> bool:
+        if not queue.get("is_batching_enabled", False):
+            return False
+        if is_label_true(message.labels.get(constants.SQS_SKIP_BATCHING_LABEL, False)):
+            return False
+        return constants.SQS_DELAY_SECONDS_LABEL not in message.labels
+
     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"])
+        queue_url = await self._get_queue_url(queue)
         kwargs = await self._build_kick_kwargs(message, queue, queue_url)
+        if self._should_batch(queue, message):
+            await self._batch_queues[queue["name"]].put(kwargs)
+            return
         with self._handle_exceptions(queue["name"]):
             await self._sqs_client.send_message(**kwargs)
 
+    @staticmethod
+    def _drain_batch_queue(batch_queue: "asyncio.Queue[dict[str, Any]]") -> list[dict[str, Any]]:
+        drained = []
+        while not batch_queue.empty():
+            try:
+                drained.append(batch_queue.get_nowait())
+            except asyncio.QueueEmpty:
+                break
+        return drained
+
+    async def _send_batch_to_sqs(self, queue: SQSQueue, queue_url: str, batch: list[dict[str, Any]]) -> None:
+        entries: list[Any] = [
+            {"id": str(index), **{key: value for key, value in kwargs.items() if key != "queue_url"}}
+            for index, kwargs in enumerate(batch)
+        ]
+        with self._handle_exceptions(queue["name"]):
+            response = await self._sqs_client.send_message_batch(queue_url=queue_url, entries=entries)
+        for failure in response.get("failed", []):
+            logger.error(
+                "Failed to send batched message to queue '%s': %s (%s)",
+                queue["name"],
+                failure.get("message"),
+                failure.get("code"),
+            )
+
+    async def _send_batch(self, queue: SQSQueue, batch: list[dict[str, Any]]) -> None:
+        queue_url = await self._get_queue_url(queue)
+        if not queue.get("is_fifo", queue["name"].endswith(".fifo")):
+            await self._send_batch_to_sqs(queue, queue_url, batch)
+            return
+        # Keep each FIFO group's messages together in their own batch call, to preserve their relative order.
+        groups: dict[str, list[dict[str, Any]]] = {}
+        for kwargs in batch:
+            groups.setdefault(kwargs.get("message_group_id", ""), []).append(kwargs)
+        for group in groups.values():
+            await self._send_batch_to_sqs(queue, queue_url, group)
+
+    async def _batch_worker(self, queue: SQSQueue) -> None:
+        """Buffer kicked messages for queue and flush them together via batch send."""
+        batch_queue = self._batch_queues[queue["name"]]
+        batch_size = queue.get("batch_size", constants.DEFAULT_BATCH_SIZE)
+        batch_timeout = queue.get("batch_timeout", constants.DEFAULT_BATCH_TIMEOUT)
+        batch: list[dict[str, Any]] = []
+        try:
+            while True:
+                batch = [await batch_queue.get()]
+                deadline = time.monotonic() + batch_timeout
+                while len(batch) < batch_size:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        break
+                    try:
+                        batch.append(await asyncio.wait_for(batch_queue.get(), timeout=remaining))
+                    except TimeoutError:
+                        break
+                await self._send_batch(queue, batch)
+                batch = []
+        except asyncio.CancelledError:
+            for kwargs in batch:
+                batch_queue.put_nowait(kwargs)
+            raise
+
     def _build_ack_function(
         self,
         queue_name: str,
@@ -207,7 +314,7 @@ def _is_expired(message: Mapping[str, Any]) -> bool:
     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:
-            queue_url = await self._get_queue_url(queue["name"])
+            queue_url = await self._get_queue_url(queue)
             while True:
                 with self._handle_exceptions(queue["name"]):
                     results = await self._sqs_client.receive_message(

src/taskiq_sqs/constants.py

@@ -7,6 +7,9 @@
 MAX_NUMBER_OF_MESSAGES: Final[int] = 10
 MAX_DELAY_SECONDS: Final[int] = 900
 MAX_FIFO_ID_LENGTH: Final[int] = 128
+MAX_BATCH_SIZE: Final[int] = 10
+DEFAULT_BATCH_SIZE: Final[int] = 10
+DEFAULT_BATCH_TIMEOUT: Final[float] = 1.0
 
 SQS_MAX_MESSAGE_SIZE_BYTES: Final[int] = 262_144
 DEFAULT_S3_OFFLOAD_THRESHOLD_BYTES: Final[int] = 200_000
@@ -16,3 +19,4 @@
 SQS_MESSAGE_GROUP_ID_LABEL: Final[str] = "group_id"
 SQS_MESSAGE_DEDUPLICATION_ID_LABEL: Final[str] = "deduplication_id"
 SQS_EXPIRY_LABEL: Final[str] = "expiry"
+SQS_SKIP_BATCHING_LABEL: Final[str] = "skip_batching"

src/taskiq_sqs/exceptions.py

@@ -29,7 +29,7 @@ class ResultBackendError(BaseTaskiqSQSError):
 class BucketNotFoundError(BaseTaskiqSQSError):
     """Error if bucket not found."""
 
-    __template__ = "Bucket '{bucket_name}' not found during initialization and declare=False"
+    __template__ = "Bucket '{bucket_name}' not found during initialization and is_declare=False"
     bucket_name: str
 
 
@@ -56,6 +56,13 @@ class UnknownQueueError(BaseTaskiqSQSError):
     queue_name: str
 
 
+class QueueNotFoundError(BaseTaskiqSQSError):
+    """Error if a queue doesn't exist and is_declare=False."""
+
+    __template__ = "Queue '{queue_name}' not found during initialization and is_declare=False"
+    queue_name: str
+
+
 class InvalidDelaySecondsError(BaseTaskiqSQSError):
     """Error if a message's delay label is outside SQS's allowed range."""
 

src/taskiq_sqs/middleware.py

@@ -92,12 +92,12 @@ async def _ensure_bucket_exists(self) -> None:
         try:
             await self._s3_client.head_bucket(bucket=self._bucket["name"])
         except capo_s3.errors.NotFound:
-            if not self._bucket.get("declare", True):
+            if not self._bucket.get("is_declare", True):
                 raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from None
             await self._create_bucket()
 
     async def _create_bucket(self) -> None:
-        create_kwargs: dict[str, Any] = {}
+        create_kwargs: dict[str, Any] = dict(self._bucket.get("options", {}))
         if self._aws_region and self._aws_region != constants.AWS_DEFAULT_REGION:
             create_kwargs["create_bucket_configuration"] = {"location_constraint": self._aws_region}
         with contextlib.suppress(capo_s3.errors.BucketAlreadyOwnedByYou):

src/taskiq_sqs/result_backend.py

@@ -73,14 +73,14 @@ async def _ensure_bucket_exists(self) -> None:
         try:
             await self._s3_client.head_bucket(bucket=self._bucket["name"])
         except capo_s3.errors.NotFound:
-            if not self._bucket.get("declare", True):
+            if not self._bucket.get("is_declare", True):
                 raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from None
             await self._create_bucket()
         except capo_s3.errors.ServiceError as exc:
             raise exceptions.ResultBackendError(code=exc.code) from exc
 
     async def _create_bucket(self) -> None:
-        create_kwargs: dict[str, Any] = {}
+        create_kwargs: dict[str, Any] = dict(self._bucket.get("options", {}))
         if self._aws_region and self._aws_region != constants.AWS_DEFAULT_REGION:
             create_kwargs["create_bucket_configuration"] = {"location_constraint": self._aws_region}
         with contextlib.suppress(capo_s3.errors.BucketAlreadyOwnedByYou):

src/taskiq_sqs/types/bucket.py

@@ -1,4 +1,5 @@
-from typing import NotRequired, TypedDict
+from collections.abc import Mapping
+from typing import Any, NotRequired, TypedDict
 
 
 class S3Bucket(TypedDict):
@@ -7,8 +8,10 @@ class S3Bucket(TypedDict):
 
     Attributes:
         name: The name of the bucket.
-        declare: Whether to create the bucket on startup if it not exists yet. Defaults to True.
+        is_declare: Whether to create the bucket on startup if it not exists yet. Defaults to True.
+        options: Extra keyword arguments merged into the create bucket call when the bucket is declared.
     """
 
     name: str
-    declare: NotRequired[bool]
+    is_declare: NotRequired[bool]
+    options: NotRequired[Mapping[str, Any]]

src/taskiq_sqs/types/message.py

@@ -53,3 +53,10 @@ def validate_expiry(expiry: Any) -> float:
     if value < 0:
         raise InvalidExpiryError(expiry=expiry)
     return value
+
+
+def is_label_true(value: Any) -> bool:
+    """Interpret a boolean-ish message label the way taskiq's own label round-trip does."""
+    if isinstance(value, str):
+        return value.strip().lower() == "true"
+    return bool(value)

src/taskiq_sqs/types/queue.py

@@ -1,4 +1,5 @@
-from typing import NotRequired, TypedDict
+from collections.abc import Mapping
+from typing import Any, NotRequired, TypedDict
 
 from taskiq_sqs import constants
 from taskiq_sqs.exceptions import BrokerInitError
@@ -13,12 +14,22 @@ class SQSQueue(TypedDict):
         max_number_of_messages: Maximum messages to retrieve per poll (1-10). Defaults to 1.
         wait_time_seconds: Long polling wait time in seconds (0-20). Defaults to 0.
         is_fifo: Whether this is a FIFO queue.
+        is_batching_enabled: Whether to buffer kicked messages in memory and flush them via batch send.
+        batch_size: Maximum messages per batch (1-10). Defaults to 10.
+        batch_timeout: Maximum seconds to wait for a batch to fill up before flushing it anyway. Defaults to 1.0.
+        is_declare: Whether to create the queue on startup if it doesn't exist yet. Defaults to True.
+        options: Queue attributes  passed during queue creation when the queue declaration is enabled.
     """
 
     name: str
     max_number_of_messages: NotRequired[int]
     wait_time_seconds: NotRequired[int]
     is_fifo: NotRequired[bool]
+    is_batching_enabled: NotRequired[bool]
+    batch_size: NotRequired[int]
+    batch_timeout: NotRequired[float]
+    is_declare: NotRequired[bool]
+    options: NotRequired[Mapping[str, Any]]
 
 
 def validate_queue(queue: SQSQueue) -> None:
@@ -39,3 +50,11 @@ def validate_queue(queue: SQSQueue) -> None:
             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",
         )
+    batch_size = queue.get("batch_size", constants.DEFAULT_BATCH_SIZE)
+    if batch_size > constants.MAX_BATCH_SIZE or batch_size < 1:
+        raise BrokerInitError(
+            details=f"BatchSize for queue '{queue['name']}' can be no greater than 10 or less than 1",
+        )
+    batch_timeout = queue.get("batch_timeout", constants.DEFAULT_BATCH_TIMEOUT)
+    if batch_timeout <= 0:
+        raise BrokerInitError(details=f"BatchTimeout for queue '{queue['name']}' must be greater than 0")

tests/benchmarks/test_broker.py

@@ -11,7 +11,7 @@
 @pytest.mark.benchmark
 async def test_build_kick_kwargs(bench_broker: SQSBroker, broker_message: BrokerMessage) -> None:
     queue = bench_broker._resolve_queue(None)
-    queue_url = await bench_broker._get_queue_url(queue["name"])
+    queue_url = await bench_broker._get_queue_url(queue)
     await bench_broker._build_kick_kwargs(broker_message, queue, queue_url)