• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

localstack / localstack / 788218a1-7c88-4b9b-95b7-0b9b914593b4

10 Feb 2025 11:28PM UTC coverage: 86.929% (+0.007%) from 86.922%
788218a1-7c88-4b9b-95b7-0b9b914593b4

push

circleci

web-flow
Step Functions: Add Supported Integrations For Glue (#12248)

2 of 3 new or added lines in 1 file covered. (66.67%)

124 existing lines in 10 files now uncovered.

61509 of 70758 relevant lines covered (86.93%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

98.33
/localstack-core/localstack/services/lambda_/event_source_mapping/pollers/sqs_poller.py
1
import json
1✔
2
import logging
1✔
3
from collections import defaultdict
1✔
4
from functools import cached_property
1✔
5

6
from botocore.client import BaseClient
1✔
7

8
from localstack.aws.api.pipes import PipeSourceSqsQueueParameters
1✔
9
from localstack.aws.api.sqs import MessageSystemAttributeName
1✔
10
from localstack.config import internal_service_url
1✔
11
from localstack.services.lambda_.event_source_mapping.event_processor import (
1✔
12
    EventProcessor,
13
    PartialBatchFailureError,
14
)
15
from localstack.services.lambda_.event_source_mapping.pollers.poller import (
1✔
16
    Poller,
17
    parse_batch_item_failures,
18
)
19
from localstack.services.sqs.constants import HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT
1✔
20
from localstack.utils.aws.arns import parse_arn
1✔
21
from localstack.utils.strings import first_char_to_lower
1✔
22

23
LOG = logging.getLogger(__name__)
1✔
24

25
DEFAULT_MAX_RECEIVE_COUNT = 10
1✔
26

27

28
class SqsPoller(Poller):
1✔
29
    queue_url: str
1✔
30

31
    batch_size: int
1✔
32
    maximum_batching_window: int
1✔
33

34
    def __init__(
1✔
35
        self,
36
        source_arn: str,
37
        source_parameters: dict | None = None,
38
        source_client: BaseClient | None = None,
39
        processor: EventProcessor | None = None,
40
    ):
41
        super().__init__(source_arn, source_parameters, source_client, processor)
1✔
42
        self.queue_url = get_queue_url(self.source_arn)
1✔
43

44
        self.batch_size = self.sqs_queue_parameters.get("BatchSize", DEFAULT_MAX_RECEIVE_COUNT)
1✔
45
        # HACK: When the MaximumBatchingWindowInSeconds is not set, just default to short-polling.
46
        # While set in ESM (via the config factory) setting this param as a default in Pipes causes
47
        # parity issues with a retrieved config since no default value is returned.
48
        self.maximum_batching_window = self.sqs_queue_parameters.get(
1✔
49
            "MaximumBatchingWindowInSeconds", 0
50
        )
51

52
        self._register_client_hooks()
1✔
53

54
    @property
1✔
55
    def sqs_queue_parameters(self) -> PipeSourceSqsQueueParameters:
1✔
56
        # TODO: De-couple Poller configuration params from ESM/Pipes specific config (i.e PipeSourceSqsQueueParameters)
57
        return self.source_parameters["SqsQueueParameters"]
1✔
58

59
    @cached_property
1✔
60
    def is_fifo_queue(self) -> bool:
1✔
61
        # Alternative heuristic: self.queue_url.endswith(".fifo"), but we need the call to get_queue_attributes for IAM
62
        return self.get_queue_attributes().get("FifoQueue", "false").lower() == "true"
1✔
63

64
    def _register_client_hooks(self):
1✔
65
        event_system = self.source_client.meta.events
1✔
66

67
        def handle_message_count_override(params, context, **kwargs):
1✔
68
            requested_count = params.pop("sqs_override_max_message_count", None)
1✔
69
            if not requested_count or requested_count <= DEFAULT_MAX_RECEIVE_COUNT:
1✔
70
                return
1✔
71

72
            context[HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT] = str(requested_count)
1✔
73

74
        def handle_inject_headers(params, context, **kwargs):
1✔
75
            if override := context.pop(HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT, None):
1✔
76
                params["headers"][HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT] = override
1✔
77

78
        event_system.register(
1✔
79
            "provide-client-params.sqs.ReceiveMessage", handle_message_count_override
80
        )
81
        # Since we delete SQS messages after processing, this allows us to remove up to 10K entries at a time.
82
        event_system.register(
1✔
83
            "provide-client-params.sqs.DeleteMessageBatch", handle_message_count_override
84
        )
85

86
        event_system.register("before-call.sqs.ReceiveMessage", handle_inject_headers)
1✔
87
        event_system.register("before-call.sqs.DeleteMessageBatch", handle_inject_headers)
1✔
88

89
    def get_queue_attributes(self) -> dict:
1✔
90
        """The API call to sqs:GetQueueAttributes is required for IAM policy streamsing."""
91
        get_queue_attributes_response = self.source_client.get_queue_attributes(
1✔
92
            QueueUrl=self.queue_url,
93
            AttributeNames=["FifoQueue"],
94
        )
95
        return get_queue_attributes_response.get("Attributes", {})
1✔
96

97
    def event_source(self) -> str:
1✔
98
        return "aws:sqs"
1✔
99

100
    def poll_events(self) -> None:
1✔
101
        # SQS pipe source: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-sqs.html
102
        # "The 9 Ways an SQS Message can be Deleted": https://lucvandonkersgoed.com/2022/01/20/the-9-ways-an-sqs-message-can-be-deleted/
103
        # TODO: implement batch window expires based on MaximumBatchingWindowInSeconds
104
        # TODO: implement invocation payload size quota
105
        # TODO: consider long-polling vs. short-polling trade-off. AWS uses long-polling:
106
        #  https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-sqs.html#pipes-sqs-scaling
107
        response = self.source_client.receive_message(
1✔
108
            QueueUrl=self.queue_url,
109
            MaxNumberOfMessages=min(self.batch_size, DEFAULT_MAX_RECEIVE_COUNT),
110
            MessageAttributeNames=["All"],
111
            MessageSystemAttributeNames=[MessageSystemAttributeName.All],
112
            # Override how many messages we can receive per call
113
            sqs_override_max_message_count=self.batch_size,
114
        )
115
        if messages := response.get("Messages"):
1✔
116
            LOG.debug("Polled %d events from %s", len(messages), self.source_arn)
1✔
117
            try:
1✔
118
                if self.is_fifo_queue:
1✔
119
                    # TODO: think about starvation behavior because once failing message could block other groups
120
                    fifo_groups = split_by_message_group_id(messages)
1✔
121
                    for fifo_group_messages in fifo_groups.values():
1✔
122
                        self.handle_messages(fifo_group_messages)
1✔
123
                else:
124
                    self.handle_messages(messages)
1✔
125

126
            # TODO: unify exception handling across pollers: should we catch and raise?
127
            except Exception as e:
1✔
128
                # TODO: improve error messages (produce same failure and design better error messages)
129
                LOG.warning(
1✔
130
                    "Polling or batch processing failed: %s",
131
                    e,
132
                    exc_info=LOG.isEnabledFor(logging.DEBUG),
133
                )
134

135
    def handle_messages(self, messages):
1✔
136
        polled_events = transform_into_events(messages)
1✔
137
        # Filtering: matching vs. discarded (i.e., not matching filter criteria)
138
        # TODO: implement format detection behavior (e.g., for JSON body):
139
        #  https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html#pipes-filter-sqs
140
        #  Check whether we need poller-specific filter-preprocessing here without modifying the actual event!
141
        # convert to json for filtering (HACK for fixing parity with v1 and getting regression tests passing)
142
        for event in polled_events:
1✔
143
            try:
1✔
144
                event["body"] = json.loads(event["body"])
1✔
145
            except json.JSONDecodeError:
1✔
146
                LOG.debug(
1✔
147
                    "Unable to convert event body '%s' to json... Event might be dropped.",
148
                    event["body"],
149
                )
150
        matching_events = self.filter_events(polled_events)
1✔
151
        # convert them back (HACK for fixing parity with v1 and getting regression tests passing)
152
        for event in matching_events:
1✔
153
            event["body"] = (
1✔
154
                json.dumps(event["body"]) if not isinstance(event["body"], str) else event["body"]
155
            )
156

157
        all_message_ids = {message["MessageId"] for message in messages}
1✔
158
        matching_message_ids = {event["messageId"] for event in matching_events}
1✔
159
        discarded_message_ids = all_message_ids.difference(matching_message_ids)
1✔
160
        # Delete discarded events immediately:
161
        # https://lucvandonkersgoed.com/2022/01/20/the-9-ways-an-sqs-message-can-be-deleted/#7-event-source-mappings-with-filters
162
        self.delete_messages(messages, discarded_message_ids)
1✔
163

164
        # Don't trigger upon empty events
165
        if len(matching_events) == 0:
1✔
UNCOV
166
            return
×
167
        # Enrich events with metadata after filtering
168
        enriched_events = self.add_source_metadata(matching_events)
1✔
169

170
        # Invoke the processor (e.g., Pipe, ESM) and handle partial batch failures
171
        try:
1✔
172
            self.processor.process_events_batch(enriched_events)
1✔
173
            successful_message_ids = all_message_ids
1✔
174
        except PartialBatchFailureError as e:
1✔
175
            failed_message_ids = parse_batch_item_failures(
1✔
176
                e.partial_failure_payload, matching_message_ids
177
            )
178
            successful_message_ids = matching_message_ids.difference(failed_message_ids)
1✔
179

180
        # Only delete messages that are processed successfully as described here:
181
        # https://docs.aws.amazon.com/en_gb/lambda/latest/dg/with-sqs.html
182
        # When Lambda reads a batch, the messages stay in the queue but are hidden for the length of the queue's
183
        # visibility timeout. If your function successfully processes the batch, Lambda deletes the messages
184
        # from the queue. By default, if your function encounters an error while processing a batch,
185
        # all messages in that batch become visible in the queue again. For this reason, your function code must
186
        # be able to process the same message multiple times without unintended side effects.
187
        # Troubleshooting: https://repost.aws/knowledge-center/lambda-sqs-report-batch-item-failures
188
        # For FIFO queues, AWS also deletes successfully sent messages. Therefore, the AWS docs recommends:
189
        # "If you're using this feature with a FIFO queue, your function should stop processing messages after the first
190
        # failure and return all failed and unprocessed messages in batchItemFailures. This helps preserve the ordering
191
        # of messages in your queue."
192
        # Following this recommendation could result in the unsolved side effect that valid messages are continuously
193
        # placed in the same batch as failing messages:
194
        # * https://stackoverflow.com/questions/78694079/how-to-stop-fifo-sqs-messages-from-being-placed-in-a-batch-with-failing-messages
195
        # * https://stackoverflow.com/questions/76912394/can-i-report-only-messages-from-failing-group-id-in-reportbatchitemfailures-resp
196

197
        # TODO: Test blocking failure behavior for FIFO queues to guarantee strict ordering
198
        #  -> might require some checkpointing or retry control on the poller side?!
199
        # The poller should only proceed processing FIFO queues after having retried failing messages:
200
        # "If your pipe returns an error, the pipe attempts all retries on the affected messages before EventBridge
201
        # receives additional messages from the same group."
202
        # https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-sqs.html
203
        self.delete_messages(messages, successful_message_ids)
1✔
204

205
    def delete_messages(self, messages: list[dict], message_ids_to_delete: set):
1✔
206
        """Delete SQS `messages` from the source queue that match a MessageId within `message_ids_to_delete`"""
207
        # TODO: unclear how (partial) failures for deleting are handled, retry or fail batch? Hard to test against AWS
208
        if len(message_ids_to_delete) > 0:
1✔
209
            entries = [
1✔
210
                {"Id": str(count), "ReceiptHandle": message["ReceiptHandle"]}
211
                for count, message in enumerate(messages)
212
                if message["MessageId"] in message_ids_to_delete
213
            ]
214

215
            self.source_client.delete_message_batch(
1✔
216
                QueueUrl=self.queue_url,
217
                Entries=entries,
218
                # Override how many messages can be deleted at once
219
                sqs_override_max_message_count=self.batch_size,
220
            )
221

222

223
def split_by_message_group_id(messages) -> defaultdict[str, list[dict]]:
1✔
224
    """Splitting SQS messages by MessageGroupId to ensure strict ordering for FIFO queues"""
225
    fifo_groups = defaultdict(list)
1✔
226
    for message in messages:
1✔
227
        message_group_id = message["Attributes"]["MessageGroupId"]
1✔
228
        fifo_groups[message_group_id].append(message)
1✔
229
    return fifo_groups
1✔
230

231

232
def transform_into_events(messages: list[dict]) -> list[dict]:
1✔
233
    events = []
1✔
234
    for message in messages:
1✔
235
        # TODO: consolidate with SQS event source listener:
236
        #  localstack.services.lambda_.event_source_listeners.sqs_event_source_listener.SQSEventSourceListener._send_event_to_lambda
237
        message_attrs = message_attributes_to_lower(message.get("MessageAttributes"))
1✔
238
        event = {
1✔
239
            # Original SQS message attributes
240
            "messageId": message["MessageId"],
241
            "receiptHandle": message["ReceiptHandle"],
242
            # TODO: test with empty body
243
            # TODO: implement heuristic based on content type: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html#pipes-filter-sqs
244
            "body": message.get("Body", "MessageBody"),
245
            "attributes": message.get("Attributes", {}),
246
            "messageAttributes": message_attrs,
247
            # TODO: test with empty body
248
            "md5OfBody": message.get("MD5OfBody") or message.get("MD5OfMessageBody"),
249
        }
250
        # TODO: test Pipe with message attributes (only covered by Lambda ESM SQS test so far)
251
        if md5_of_message_attributes := message.get("MD5OfMessageAttributes"):
1✔
252
            event["md5OfMessageAttributes"] = md5_of_message_attributes
1✔
253
        events.append(event)
1✔
254
    return events
1✔
255

256

257
def get_queue_url(queue_arn: str) -> str:
1✔
258
    # TODO: consolidate this method with localstack.services.sqs.models.SqsQueue.url
259
    # * Do we need to support different endpoint strategies?
260
    # * If so, how can we achieve this without having a request context
261
    host_url = internal_service_url()
1✔
262
    host = host_url.rstrip("/")
1✔
263
    parsed_arn = parse_arn(queue_arn)
1✔
264
    account_id = parsed_arn["account"]
1✔
265
    name = parsed_arn["resource"]
1✔
266
    return f"{host}/{account_id}/{name}"
1✔
267

268

269
def message_attributes_to_lower(message_attrs):
1✔
270
    """Convert message attribute details (first characters) to lower case (e.g., stringValue, dataType)."""
271
    message_attrs = message_attrs or {}
1✔
272
    for _, attr in message_attrs.items():
1✔
273
        if not isinstance(attr, dict):
1✔
UNCOV
274
            continue
×
275
        for key, value in dict(attr).items():
1✔
276
            attr[first_char_to_lower(key)] = attr.pop(key)
1✔
277
    return message_attrs
1✔
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE TRIAL · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc