|
| 1 | +import contextlib |
| 2 | +import logging |
| 3 | +from typing import Any |
| 4 | + |
| 5 | +import capo_s3 |
| 6 | +from taskiq.abc.middleware import TaskiqMiddleware |
| 7 | +from taskiq.abc.serializer import TaskiqSerializer |
| 8 | +from taskiq.message import TaskiqMessage |
| 9 | +from taskiq.serializers import JSONSerializer |
| 10 | + |
| 11 | +from taskiq_sqs import constants, exceptions |
| 12 | +from taskiq_sqs.types import S3Bucket |
| 13 | + |
| 14 | + |
| 15 | +logger = logging.getLogger(__name__) |
| 16 | + |
| 17 | +OFFLOAD_KEY_LABEL = "s3_offload_key" |
| 18 | + |
| 19 | + |
| 20 | +class S3OffloadMiddleware(TaskiqMiddleware): |
| 21 | + """ |
| 22 | + Offloads large task payloads to S3 instead of sending them through SQS. |
| 23 | +
|
| 24 | + SQS messages are limited to 256 KiB, so payloads that would exceed a configured threshold are uploaded to S3 before |
| 25 | + the message is sent, and the message itself only carries a reference (bucket key) to the uploaded payload. |
| 26 | + On the worker side, the original arguments are downloaded back from S3 before task execution. |
| 27 | + """ |
| 28 | + |
| 29 | + def __init__( |
| 30 | + self, |
| 31 | + bucket: S3Bucket, |
| 32 | + max_message_size: int = constants.DEFAULT_S3_OFFLOAD_THRESHOLD_BYTES, |
| 33 | + base_path: str = "", |
| 34 | + endpoint_url: str | None = None, |
| 35 | + aws_region_name: str = constants.AWS_DEFAULT_REGION, |
| 36 | + aws_access_key_id: str | None = None, |
| 37 | + aws_secret_access_key: str | None = None, |
| 38 | + delete_after_execute: bool = True, |
| 39 | + serializer: TaskiqSerializer | None = None, |
| 40 | + ) -> None: |
| 41 | + """ |
| 42 | + Constructs a new S3 offload middleware. |
| 43 | +
|
| 44 | + :param bucket: S3 bucket configuration. |
| 45 | + :param max_message_size: payloads larger than this many bytes are offloaded to S3. |
| 46 | + :param base_path: base path (prefix) for offloaded payloads. |
| 47 | + :param endpoint_url: endpoint URL for S3. |
| 48 | + :param aws_region_name: AWS region, default is 'us-east-1'. |
| 49 | + :param aws_access_key_id: AWS access key ID. |
| 50 | + :param aws_secret_access_key: AWS secret access key. |
| 51 | + :param delete_after_execute: whether to delete the offloaded payload from S3 after the task has been executed |
| 52 | + by the worker. |
| 53 | + :param serializer: serializer used to encode/decode offloaded payloads. |
| 54 | + """ |
| 55 | + super().__init__() |
| 56 | + self._bucket = bucket |
| 57 | + self._max_message_size = max_message_size |
| 58 | + self._base_path = base_path |
| 59 | + self._aws_endpoint_url = endpoint_url |
| 60 | + self._aws_region = aws_region_name |
| 61 | + self._aws_access_key_id = aws_access_key_id |
| 62 | + self._aws_secret_access_key = aws_secret_access_key |
| 63 | + self._delete_after_execute = delete_after_execute |
| 64 | + self._serializer = serializer or JSONSerializer() |
| 65 | + |
| 66 | + async def startup(self) -> None: |
| 67 | + """Initialize the S3 client and ensure the bucket exists.""" |
| 68 | + credentials = None |
| 69 | + if self._aws_access_key_id and self._aws_secret_access_key: |
| 70 | + credentials = capo_s3.Credentials( |
| 71 | + access_key=self._aws_access_key_id, |
| 72 | + secret_key=self._aws_secret_access_key, |
| 73 | + ) |
| 74 | + self._s3_client = capo_s3.AsyncS3Client( |
| 75 | + region=self._aws_region, |
| 76 | + endpoint=self._aws_endpoint_url, |
| 77 | + credentials=credentials, |
| 78 | + force_path_style=True, |
| 79 | + ) |
| 80 | + await self._s3_client.__aenter__() |
| 81 | + try: |
| 82 | + await self._ensure_bucket_exists() |
| 83 | + except Exception: |
| 84 | + await self._s3_client.__aexit__(None, None, None) |
| 85 | + raise |
| 86 | + |
| 87 | + async def shutdown(self) -> None: |
| 88 | + """Shut down the S3 client.""" |
| 89 | + await self._s3_client.__aexit__(None, None, None) |
| 90 | + |
| 91 | + async def _ensure_bucket_exists(self) -> None: |
| 92 | + try: |
| 93 | + await self._s3_client.head_bucket(bucket=self._bucket["name"]) |
| 94 | + except capo_s3.errors.NotFound: |
| 95 | + if not self._bucket.get("declare", True): |
| 96 | + raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from None |
| 97 | + await self._create_bucket() |
| 98 | + |
| 99 | + async def _create_bucket(self) -> None: |
| 100 | + create_kwargs: dict[str, Any] = {} |
| 101 | + if self._aws_region and self._aws_region != constants.AWS_DEFAULT_REGION: |
| 102 | + create_kwargs["create_bucket_configuration"] = {"location_constraint": self._aws_region} |
| 103 | + with contextlib.suppress(capo_s3.errors.BucketAlreadyOwnedByYou): |
| 104 | + await self._s3_client.create_bucket(bucket=self._bucket["name"], **create_kwargs) |
| 105 | + |
| 106 | + def _build_key(self, task_id: str) -> str: |
| 107 | + key = f"{task_id}.json" |
| 108 | + if self._base_path: |
| 109 | + key = f"{self._base_path.rstrip('/')}/{key}" |
| 110 | + return key |
| 111 | + |
| 112 | + async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: |
| 113 | + """ |
| 114 | + Offload the message payload to S3 if it's too large to send through SQS. |
| 115 | +
|
| 116 | + :param message: message to send. |
| 117 | + :return: message with args/kwargs replaced by an S3 reference, if offloaded. |
| 118 | + """ |
| 119 | + payload = self._serializer.dumpb({"args": message.args, "kwargs": message.kwargs}) |
| 120 | + if len(payload) <= self._max_message_size: |
| 121 | + return message |
| 122 | + |
| 123 | + key = self._build_key(message.task_id) |
| 124 | + await self._s3_client.put_object(bucket=self._bucket["name"], key=key, body=payload) |
| 125 | + logger.debug("Offloaded payload of task '%s' to s3://%s/%s", message.task_id, self._bucket["name"], key) |
| 126 | + |
| 127 | + message.args = [] |
| 128 | + message.kwargs = {} |
| 129 | + message.labels[OFFLOAD_KEY_LABEL] = key |
| 130 | + return message |
| 131 | + |
| 132 | + async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: |
| 133 | + """ |
| 134 | + Restore the original message payload from S3, if it was offloaded. |
| 135 | +
|
| 136 | + :param message: incoming parsed taskiq message. |
| 137 | + :return: message with the original args/kwargs restored. |
| 138 | + """ |
| 139 | + key = message.labels.get(OFFLOAD_KEY_LABEL) |
| 140 | + if key is None: |
| 141 | + return message |
| 142 | + |
| 143 | + try: |
| 144 | + async with self._s3_client.get_object(bucket=self._bucket["name"], key=key) as output: |
| 145 | + body = b"".join([chunk async for chunk in output["body"]]) |
| 146 | + except capo_s3.errors.NoSuchKey as exc: |
| 147 | + raise exceptions.OffloadedPayloadMissingError( |
| 148 | + task_id=message.task_id, |
| 149 | + bucket_name=self._bucket["name"], |
| 150 | + key=key, |
| 151 | + ) from exc |
| 152 | + |
| 153 | + payload = self._serializer.loadb(body) |
| 154 | + message.args = payload["args"] |
| 155 | + message.kwargs = payload["kwargs"] |
| 156 | + return message |
| 157 | + |
| 158 | + async def post_execute(self, message: TaskiqMessage, result: Any) -> None: # noqa: ARG002 |
| 159 | + """ |
| 160 | + Delete the offloaded payload from S3 once the task has been executed. |
| 161 | +
|
| 162 | + :param message: processed message. |
| 163 | + :param result: result of execution for current task, not used. |
| 164 | + """ |
| 165 | + key = message.labels.pop(OFFLOAD_KEY_LABEL, None) |
| 166 | + if key is None or not self._delete_after_execute: |
| 167 | + return |
| 168 | + await self._s3_client.delete_object(bucket=self._bucket["name"], key=key) |
0 commit comments