feat: add middleware to automatically offload large task payload

.github/workflows/code-check.yml

@@ -28,7 +28,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest] # eventually add `windows-latest` and `macos-latest`
-        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+        python-version: [ "3.11", "3.12", "3.13", "3.14" ]
     services:
       ministack:
         image: ministackorg/ministack:1.3.53

README.md

@@ -81,3 +81,33 @@ async def main() -> None:
 if __name__ == "__main__":
     asyncio.run(main())
 ```
+
+## 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.
+
+```python
+import asyncio
+from taskiq_sqs import S3Bucket, S3OffloadMiddleware, SQSBroker
+
+broker = SQSBroker("http://localhost:4566/000000000000/my-queue")
+broker.add_middlewares(
+    S3OffloadMiddleware(
+        bucket=S3Bucket(name="offload-bucket"),  # created automatically if it doesn't exist
+        max_message_size=200_000,  # payloads larger than this many bytes are offloaded to S3
+    ),
+)
+
+@broker.task
+async def process_document(content: str) -> int:
+    return len(content)
+
+
+async def main() -> None:
+    await broker.startup()
+    await process_document.kiq("x" * 1_000_000)  # too large for SQS, transparently offloaded to S3
+    await broker.shutdown()
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```

pyproject.toml

@@ -14,7 +14,6 @@ classifiers = [
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
@@ -28,10 +27,11 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 keywords = ["taskiq", "broker", "aws", "sqs"]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
-    "taskiq>=0.12.1",
+    "taskiq>=0.12.6",
     "aiobotocore>=2.13.3",
+    "capo-s3>=0.15.0",
 ]
 
 [project.urls]
@@ -44,29 +44,28 @@ dev = [
     {include-group = "lint"},
     {include-group = "test"},
     {include-group = "types"},
-    "prek>=0.5.0",
+    "prek>=0.5.3",
 ]
 test = [
-    "pytest>=9.0.3",
-    "pytest-asyncio>=0.23.8",
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
     "pytest-codspeed>=5.0.3",
 ]
 lint = [
-    "bandit>=1.9.4",
-    "ruff>=0.16.5",
-    "zizmor>=1.29.0",
+    "ruff>=0.16.7",
+    "zizmor>=1.30.1",
 ]
 types = [
     "mypy>=2.3.1",
     "types-aiobotocore[essential]>=3.7.0",
 ]
 examples = [
-    "python-dotenv>=1.2.2",
+    "python-dotenv>=1.2.3",
 ]
 
 
 [build-system]
-requires = ["uv_build>=0.11,<0.12"]
+requires = ["uv_build>=0.12,<0.13"]
 build-backend = "uv_build"
 
 [tool.uv.build-backend]
@@ -104,7 +103,7 @@ omit = [
 
 [tool.ruff]
 line-length = 120
-target-version = "py310"
+target-version = "py311"
 
 [tool.ruff.lint]
 select = ["ALL"]
@@ -146,6 +145,10 @@ ignore = [
     "D",
     "INP001",
 ]
+"src/taskiq_sqs/middleware.py" = [
+    "PLR0913",  # too many arguments
+    "PLR0917",  # too many positional arguments
+]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"

src/taskiq_sqs/__init__.py

@@ -1,10 +1,10 @@
 from taskiq_sqs.broker import SQSBroker
-from taskiq_sqs.bucket import S3Bucket
+from taskiq_sqs.middleware import S3OffloadMiddleware
 from taskiq_sqs.result_backend import S3ResultBackend
 
 
 __all__ = [
-    "S3Bucket",
+    "S3OffloadMiddleware",
     "S3ResultBackend",
     "SQSBroker",
 ]

src/taskiq_sqs/broker.py

@@ -11,7 +11,7 @@
 
 from taskiq_sqs import constants
 from taskiq_sqs.exceptions import BrokerInitError
-from taskiq_sqs.queue import SQSQueue
+from taskiq_sqs.types import SQSQueue
 
 
 if TYPE_CHECKING:

src/taskiq_sqs/constants.py

@@ -5,3 +5,6 @@
 
 MAX_WAIT_TIME_SECONDS: Final[int] = 20
 MAX_NUMBER_OF_MESSAGES: Final[int] = 10
+
+SQS_MAX_MESSAGE_SIZE_BYTES: Final[int] = 262_144
+DEFAULT_S3_OFFLOAD_THRESHOLD_BYTES: Final[int] = 200_000

src/taskiq_sqs/exceptions.py

@@ -38,3 +38,12 @@ class ResultIsMissingError(BaseTaskiqSQSError):
 
     __template__ = "Result for task {task_id} is missing in the result backend"
     task_id: str
+
+
+class OffloadedPayloadMissingError(BaseTaskiqSQSError):
+    """Error if a message references an S3-offloaded payload that can't be found."""
+
+    __template__ = "Offloaded payload for task {task_id} is missing in bucket '{bucket_name}' (key: {key})"
+    task_id: str
+    bucket_name: str
+    key: str

src/taskiq_sqs/middleware.py

@@ -0,0 +1,168 @@
+import contextlib
+import logging
+from typing import Any
+
+import capo_s3
+from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.message import TaskiqMessage
+from taskiq.serializers import JSONSerializer
+
+from taskiq_sqs import constants, exceptions
+from taskiq_sqs.types import S3Bucket
+
+
+logger = logging.getLogger(__name__)
+
+OFFLOAD_KEY_LABEL = "s3_offload_key"
+
+
+class S3OffloadMiddleware(TaskiqMiddleware):
+    """
+    Offloads large task payloads to S3 instead of sending them through SQS.
+
+    SQS messages are limited to 256 KiB, so payloads that would exceed a configured threshold are uploaded to S3 before
+    the message is sent, and the message itself only carries a reference (bucket key) to the uploaded payload.
+    On the worker side, the original arguments are downloaded back from S3 before task execution.
+    """
+
+    def __init__(
+        self,
+        bucket: S3Bucket,
+        max_message_size: int = constants.DEFAULT_S3_OFFLOAD_THRESHOLD_BYTES,
+        base_path: str = "",
+        endpoint_url: str | None = None,
+        aws_region_name: str = constants.AWS_DEFAULT_REGION,
+        aws_access_key_id: str | None = None,
+        aws_secret_access_key: str | None = None,
+        delete_after_execute: bool = True,
+        serializer: TaskiqSerializer | None = None,
+    ) -> None:
+        """
+        Constructs a new S3 offload middleware.
+
+        :param bucket: S3 bucket configuration.
+        :param max_message_size: payloads larger than this many bytes are offloaded to S3.
+        :param base_path: base path (prefix) for offloaded payloads.
+        :param endpoint_url: endpoint URL for S3.
+        :param aws_region_name: AWS region, default is 'us-east-1'.
+        :param aws_access_key_id: AWS access key ID.
+        :param aws_secret_access_key: AWS secret access key.
+        :param delete_after_execute: whether to delete the offloaded payload from S3 after the task has been executed
+            by the worker.
+        :param serializer: serializer used to encode/decode offloaded payloads.
+        """
+        super().__init__()
+        self._bucket = bucket
+        self._max_message_size = max_message_size
+        self._base_path = base_path
+        self._aws_endpoint_url = endpoint_url
+        self._aws_region = aws_region_name
+        self._aws_access_key_id = aws_access_key_id
+        self._aws_secret_access_key = aws_secret_access_key
+        self._delete_after_execute = delete_after_execute
+        self._serializer = serializer or JSONSerializer()
+
+    async def startup(self) -> None:
+        """Initialize the S3 client and ensure the bucket exists."""
+        credentials = None
+        if self._aws_access_key_id and self._aws_secret_access_key:
+            credentials = capo_s3.Credentials(
+                access_key=self._aws_access_key_id,
+                secret_key=self._aws_secret_access_key,
+            )
+        self._s3_client = capo_s3.AsyncS3Client(
+            region=self._aws_region,
+            endpoint=self._aws_endpoint_url,
+            credentials=credentials,
+            force_path_style=True,
+        )
+        await self._s3_client.__aenter__()
+        try:
+            await self._ensure_bucket_exists()
+        except Exception:
+            await self._s3_client.__aexit__(None, None, None)
+            raise
+
+    async def shutdown(self) -> None:
+        """Shut down the S3 client."""
+        await self._s3_client.__aexit__(None, None, None)
+
+    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):
+                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] = {}
+        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):
+            await self._s3_client.create_bucket(bucket=self._bucket["name"], **create_kwargs)
+
+    def _build_key(self, task_id: str) -> str:
+        key = f"{task_id}.json"
+        if self._base_path:
+            key = f"{self._base_path.rstrip('/')}/{key}"
+        return key
+
+    async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
+        """
+        Offload the message payload to S3 if it's too large to send through SQS.
+
+        :param message: message to send.
+        :return: message with args/kwargs replaced by an S3 reference, if offloaded.
+        """
+        payload = self._serializer.dumpb({"args": message.args, "kwargs": message.kwargs})
+        if len(payload) <= self._max_message_size:
+            return message
+
+        key = self._build_key(message.task_id)
+        await self._s3_client.put_object(bucket=self._bucket["name"], key=key, body=payload)
+        logger.debug("Offloaded payload of task '%s' to s3://%s/%s", message.task_id, self._bucket["name"], key)
+
+        message.args = []
+        message.kwargs = {}
+        message.labels[OFFLOAD_KEY_LABEL] = key
+        return message
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+        """
+        Restore the original message payload from S3, if it was offloaded.
+
+        :param message: incoming parsed taskiq message.
+        :return: message with the original args/kwargs restored.
+        """
+        key = message.labels.get(OFFLOAD_KEY_LABEL)
+        if key is None:
+            return message
+
+        try:
+            async with self._s3_client.get_object(bucket=self._bucket["name"], key=key) as output:
+                body = b"".join([chunk async for chunk in output["body"]])
+        except capo_s3.errors.NoSuchKey as exc:
+            raise exceptions.OffloadedPayloadMissingError(
+                task_id=message.task_id,
+                bucket_name=self._bucket["name"],
+                key=key,
+            ) from exc
+
+        payload = self._serializer.loadb(body)
+        message.args = payload["args"]
+        message.kwargs = payload["kwargs"]
+        return message
+
+    async def post_execute(self, message: TaskiqMessage, result: Any) -> None:  # noqa: ARG002
+        """
+        Delete the offloaded payload from S3 once the task has been executed.
+
+        :param message: processed message.
+        :param result: result of execution for current task, not used.
+        """
+        key = message.labels.pop(OFFLOAD_KEY_LABEL, None)
+        if key is None or not self._delete_after_execute:
+            return
+        await self._s3_client.delete_object(bucket=self._bucket["name"], key=key)

src/taskiq_sqs/result_backend.py

@@ -9,7 +9,7 @@
 from taskiq.serializers import JSONSerializer
 
 from taskiq_sqs import constants, exceptions
-from taskiq_sqs.bucket import S3Bucket
+from taskiq_sqs.types import S3Bucket
 
 
 if TYPE_CHECKING:
@@ -79,17 +79,17 @@ async def startup(self) -> None:
 
     async def _ensure_bucket_exists(self) -> None:
         try:
-            await self._s3_client.head_bucket(Bucket=self._bucket.name)
+            await self._s3_client.head_bucket(Bucket=self._bucket["name"])
         except ClientError as exc:
             code = exc.response.get("Error", {}).get("Code")
             if code not in ("404", "NoSuchBucket"):
                 raise exceptions.ResultBackendError(code=code) from exc
-            if not self._bucket.declare:
-                raise exceptions.BucketNotFoundError(bucket_name=self._bucket.name) from exc
+            if not self._bucket.get("declare", True):
+                raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from exc
             await self._create_bucket()
 
     async def _create_bucket(self) -> None:
-        create_kwargs: dict[str, Any] = {"Bucket": self._bucket.name}
+        create_kwargs: dict[str, Any] = {"Bucket": self._bucket["name"]}
         if self._aws_region and self._aws_region != constants.AWS_DEFAULT_REGION:
             create_kwargs["CreateBucketConfiguration"] = {"LocationConstraint": self._aws_region}
         try:
@@ -118,7 +118,7 @@ async def set_result(
             task_id = f"{self._base_path.rstrip('/')}/{task_id}"
 
         await self._s3_client.put_object(
-            Bucket=self._bucket.name,
+            Bucket=self._bucket["name"],
             Key=task_id,
             Body=self._serializer.dumpb(model_dump(result)),
         )
@@ -143,7 +143,7 @@ async def get_result(
             task_id = f"{self._base_path.rstrip('/')}/{task_id}"
         try:
             if response := await self._s3_client.get_object(
-                Bucket=self._bucket.name,
+                Bucket=self._bucket["name"],
                 Key=task_id,
             ):
                 async with response["Body"] as stream:
@@ -176,7 +176,7 @@ async def is_result_ready(self, task_id: str) -> bool:
         if self._base_path:
             task_id = f"{self._base_path.rstrip('/')}/{task_id}"
         try:
-            if await self._s3_client.head_object(Bucket=self._bucket.name, Key=task_id):
+            if await self._s3_client.head_object(Bucket=self._bucket["name"], Key=task_id):
                 return True
         except ClientError as exc:
             code = exc.response.get("Error", {}).get("Code")

src/taskiq_sqs/types/__init__.py

@@ -0,0 +1,8 @@
+from taskiq_sqs.types.bucket import S3Bucket
+from taskiq_sqs.types.queue import SQSQueue
+
+
+__all__ = [
+    "S3Bucket",
+    "SQSQueue",
+]