Integrate large payload support for SQS

kombu/asynchronous/aws/ext.py

@@ -22,6 +22,11 @@ class BotoCoreError(Exception):
     get_response = _void()
     get_cert_path = _void()
 
+try:
+    import sqs_extended_client
+except ImportError:
+    sqs_extended_client = None
+
 
 __all__ = (
     'exceptions', 'AWSRequest', 'get_response', 'get_cert_path',

kombu/asynchronous/aws/sqs/ext.py

@@ -7,3 +7,8 @@
     import boto3
 except ImportError:
     boto3 = None
+
+try:
+    import sqs_extended_client
+except ImportError:
+    sqs_extended_client = None

kombu/transport/SQS.py

@@ -133,6 +133,7 @@
 from __future__ import annotations
 
 import base64
+import json
 import socket
 import string
 import uuid
@@ -144,7 +145,7 @@
 from vine import ensure_promise, promise, transform
 
 from kombu.asynchronous import get_event_loop
-from kombu.asynchronous.aws.ext import boto3, exceptions
+from kombu.asynchronous.aws.ext import boto3, exceptions, sqs_extended_client
 from kombu.asynchronous.aws.sqs.connection import AsyncSQSConnection
 from kombu.asynchronous.aws.sqs.message import AsyncMessage
 from kombu.log import get_logger
@@ -498,6 +499,24 @@
                 message['ReceiptHandle'],
             )
         else:
+
+            if (
+                sqs_extended_client and
+                isinstance(payload, list)
+                and payload[0] == sqs_extended_client.client.MESSAGE_POINTER_CLASS
+            ):
+                # Used the sqs_extended_client, so we need to fetch the file from S3 and use that as the payload
+                s3_details = payload[1]
+                s3_bucket_name, s3_key = s3_details["s3BucketName"], s3_details["s3Key"]
+
+                s3_client = self.s3()
+                response = s3_client.get_object(Bucket=s3_bucket_name, Key=s3_key)
+
+                # The message body is under a wrapper class called StreamingBody
+                streaming_body = response["Body"]
+                body = self._optional_b64_decode(streaming_body.read())
+                payload = json.loads(body)
+
             try:
                 properties = payload['properties']
                 delivery_info = payload['properties']['delivery_info']
@@ -713,6 +732,32 @@
         #         if "can't set attribute" not in str(exc):
         #             raise
 
+    def new_s3_client(
+        self, region, access_key_id, secret_access_key, session_token=None
+    ):
+        session = boto3.session.Session(
+            region_name=region,
+            aws_access_key_id=access_key_id,
+            aws_secret_access_key=secret_access_key,
+            aws_session_token=session_token,
+        )
+        is_secure = self.is_secure if self.is_secure is not None else True
+        client_kwargs = {"use_ssl": is_secure}
+
+        if self.endpoint_url is not None:
+            client_kwargs["endpoint_url"] = self.endpoint_url
+
+        client = session.client("s3", **client_kwargs)
+
+        return client
+
+    def s3(self):
+        return self.new_s3_client(
+            region=self.region,
+            access_key_id=self.conninfo.userid,
+            secret_access_key=self.conninfo.password,
+        )
+
     def new_sqs_client(self, region, access_key_id,
                        secret_access_key, session_token=None):
         session = boto3.session.Session(
@@ -729,7 +774,13 @@
             client_kwargs['endpoint_url'] = self.endpoint_url
         client_config = self.transport_options.get('client-config') or {}
         config = Config(**client_config)
-        return session.client('sqs', config=config, **client_kwargs)
+        client = session.client('sqs', config=config, **client_kwargs)
+
+        if self.transport_options.get('large_payload_bucket') and sqs_extended_client:
+            client.large_payload_support = self.transport_options.get('large_payload_bucket')
+            client.use_legacy_attribute = False
+
+        return client
 
     def sqs(self, queue=None):
         if queue is not None and self.predefined_queues:

requirements/extras/sqs.txt

@@ -1,2 +1,3 @@
 boto3>=1.26.143
 urllib3>=1.26.16
+amazon-sqs-extended-client>=1.0.1

t/unit/transport/test_SQS.py

@@ -7,10 +7,12 @@
 from __future__ import annotations
 
 import base64
+import json
 import os
 import random
 import string
 from datetime import datetime, timedelta
+from io import BytesIO
 from queue import Empty
 from unittest.mock import Mock, patch
 
@@ -19,6 +21,7 @@
 from kombu import Connection, Exchange, Queue, messaging
 
 boto3 = pytest.importorskip('boto3')
+sqs_extended_client = pytest.importorskip('sqs_extended_client')
 
 from botocore.exceptions import ClientError  # noqa
 
@@ -364,6 +367,51 @@ def test_optional_b64_decode(self):
         assert self.channel._optional_b64_decode(raw) == raw
         assert self.channel._optional_b64_decode(b"test123") == b"test123"
 
+    @patch('boto3.session.Session')
+    def test_new_s3_client_with_is_secure_false(self, mock_session):
+        self.channel.is_secure = False
+        self.channel.endpoint_url = None
+
+        self.channel.new_s3_client(
+            region='us-west-2',
+            access_key_id='test_access_key',
+            secret_access_key='test_secret_key'
+        )
+
+        # assert isinstance(client, boto3.client('s3').__class__)
+        mock_session.assert_called_once_with(
+            region_name='us-west-2',
+            aws_access_key_id='test_access_key',
+            aws_secret_access_key='test_secret_key',
+            aws_session_token=None
+        )
+        mock_session().client.assert_called_once_with(
+            's3', use_ssl=False
+        )
+
+    @patch('boto3.session.Session')
+    def test_new_s3_client_with_custom_endpoint(self, mock_session):
+        mock_client = Mock()
+        mock_session.return_value.client.return_value = mock_client
+
+        self.channel.is_secure = True
+        self.channel.endpoint_url = 'https://custom-endpoint.com'
+
+        result = self.channel.new_s3_client('us-west-2', 'access_key', 'secret_key')
+
+        mock_session.assert_called_once_with(
+            region_name='us-west-2',
+            aws_access_key_id='access_key',
+            aws_secret_access_key='secret_key',
+            aws_session_token=None
+        )
+        mock_session.return_value.client.assert_called_once_with(
+            's3',
+            use_ssl=True,
+            endpoint_url='https://custom-endpoint.com'
+        )
+        assert result == mock_client
+
     def test_messages_to_python(self):
         from kombu.asynchronous.aws.sqs.message import Message
 
@@ -1035,3 +1083,30 @@ def test_message_attribute(self):
         assert message == output_message.payload
         # It's not propagated to the properties
         assert 'message_attributes' not in output_message.properties
+
+    def test_message_to_python_with_sqs_extended_client(self):
+        message = [
+            sqs_extended_client.client.MESSAGE_POINTER_CLASS,
+            {'s3BucketName': 's3://large-payload-bucket', 's3Key': 'payload.json'}
+        ]
+
+        # Get the messages now
+        with patch('kombu.transport.SQS.Channel.s3') as s3_mock:
+            s3_client = Mock(
+                get_object=Mock(
+                    return_value={'Body': BytesIO(json.dumps({"my_key": "Hello, World!"}).encode()), })
+            )
+            s3_mock.return_value = s3_client
+
+            result = self.channel._message_to_python(
+                {'Body': json.dumps(message), 'ReceiptHandle': 'handle'}, self.queue_name,
+                'test',
+            )
+
+        assert s3_client.get_object.called
+
+        # Make sure they're payload-style objects
+        assert 'properties' in result
+
+        # Data from s3 is loaded into the return payload
+        assert 'my_key' in result