Adding a flag to SqsFifoProcessor to allow message processing to continue

Author: leandrodamascena

aws_lambda_powertools/utilities/batch/sqs_fifo_partial_processor.py

@@ -57,7 +57,8 @@ def lambda_handler(event, context: LambdaContext):
         None,
     )
 
-    def __init__(self, model: Optional["BatchSqsTypeModel"] = None):
+    def __init__(self, model: Optional["BatchSqsTypeModel"] = None, return_on_first_error: bool = True):
+        self.return_on_first_error = return_on_first_error
         super().__init__(EventType.SQS, model)
 
     def process(self) -> List[Tuple]:
@@ -68,13 +69,26 @@ def process(self) -> List[Tuple]:
         result: List[Tuple] = []
 
         for i, record in enumerate(self.records):
-            # If we have failed messages, it means that the last message failed.
-            # We then short circuit the process, failing the remaining messages
-            if self.fail_messages:
+            # If we have failed messages and we are set to return on the first error,
+            # short circuit the process and return the remaining messages as failed items
+            if self.fail_messages and self.return_on_first_error:
                 return self._short_circuit_processing(i, result)
 
-            # Otherwise, process the message normally
-            result.append(self._process_record(record))
+            # Process the current record
+            processed_messages = self._process_record(record)
+
+            # If a processed message fail,
+            # mark subsequent messages from the same MessageGroupId as skipped
+            if processed_messages[0] == "fail":
+                for subsequent_record in self.records[i + 1 :]:
+                    if subsequent_record.get("attributes", {}).get("MessageGroupId") == record.get(
+                        "attributes",
+                        {},
+                    ).get("MessageGroupId"):
+                        continue  # Skip subsequent message from the same MessageGroupId
+
+            # Append the processed message normally
+            result.append(processed_messages)
 
         return result
 

docs/utilities/batch.md

@@ -141,8 +141,12 @@ Processing batches from SQS works in three stages:
 
 #### FIFO queues
 
-When using [SQS FIFO queues](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/FIFO-queues.html){target="_blank" rel="nofollow"}, we will stop processing messages after the first failure, and return all failed and unprocessed messages in `batchItemFailures`.
-This helps preserve the ordering of messages in your queue.
+When working with [SQS FIFO queues](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/FIFO-queues.html){target="_blank"}, it's important to know that a batch sent from SQS to Lambda can include multiple messages from different group IDs.
+
+By default, message processing halts after the initial failure, returning all failed and unprocessed messages in `batchItemFailures` to preserve the ordering of messages in your queue. However, customers can opt to continue processing messages and retrieve failed messages within a message group ID by setting `return_on_first_error` to False.
+
+???+ notice "Having problems with DLQ?"
+    `AsyncBatchProcessor` uses `asyncio.gather`. This might cause [side effects and reach trace limits at high concurrency](../core/tracer.md#concurrent-asynchronous-functions){target="_blank"}.
 
 === "Recommended"
 

tests/functional/test_utilities_batch.py

@@ -38,6 +38,32 @@
 from tests.functional.utils import b64_to_str, str_to_b64
 
 
+@pytest.fixture(scope="module")
+def sqs_event_fifo_factory() -> Callable:
+    def factory(body: str, message_group_id: str = ""):
+        return {
+            "messageId": f"{uuid.uuid4()}",
+            "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a",
+            "body": body,
+            "attributes": {
+                "ApproximateReceiveCount": "1",
+                "SentTimestamp": "1703675223472",
+                "SequenceNumber": "18882884930918384133",
+                "MessageGroupId": message_group_id,
+                "SenderId": "SenderId",
+                "MessageDeduplicationId": "1eea03c3f7e782c7bdc2f2a917f40389314733ff39f5ab16219580c0109ade98",
+                "ApproximateFirstReceiveTimestamp": "1703675223484",
+            },
+            "messageAttributes": {},
+            "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+            "eventSource": "aws:sqs",
+            "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+            "awsRegion": "us-east-1",
+        }
+
+    return factory
+
+
 @pytest.fixture(scope="module")
 def sqs_event_factory() -> Callable:
     def factory(body: str):
@@ -48,7 +74,7 @@ def factory(body: str):
             "attributes": {
                 "ApproximateReceiveCount": "1",
                 "SentTimestamp": "1545082649183",
-                "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+                "SenderId": "SenderId",
                 "ApproximateFirstReceiveTimestamp": "1545082649185",
             },
             "messageAttributes": {},
@@ -660,10 +686,10 @@ def lambda_handler(event, context):
     assert "All records failed processing. " in str(e.value)
 
 
-def test_sqs_fifo_batch_processor_middleware_success_only(sqs_event_factory, record_handler):
+def test_sqs_fifo_batch_processor_middleware_success_only(sqs_event_fifo_factory, record_handler):
     # GIVEN
-    first_record = SQSRecord(sqs_event_factory("success"))
-    second_record = SQSRecord(sqs_event_factory("success"))
+    first_record = SQSRecord(sqs_event_fifo_factory("success"))
+    second_record = SQSRecord(sqs_event_fifo_factory("success"))
     event = {"Records": [first_record.raw_event, second_record.raw_event]}
 
     processor = SqsFifoPartialProcessor()
@@ -679,12 +705,12 @@ def lambda_handler(event, context):
     assert result["batchItemFailures"] == []
 
 
-def test_sqs_fifo_batch_processor_middleware_with_failure(sqs_event_factory, record_handler):
+def test_sqs_fifo_batch_processor_middleware_with_failure(sqs_event_fifo_factory, record_handler):
     # GIVEN
-    first_record = SQSRecord(sqs_event_factory("success"))
-    second_record = SQSRecord(sqs_event_factory("fail"))
+    first_record = SQSRecord(sqs_event_fifo_factory("success"))
+    second_record = SQSRecord(sqs_event_fifo_factory("fail"))
     # this would normally succeed, but since it's a FIFO queue, it will be marked as failure
-    third_record = SQSRecord(sqs_event_factory("success"))
+    third_record = SQSRecord(sqs_event_fifo_factory("success"))
     event = {"Records": [first_record.raw_event, second_record.raw_event, third_record.raw_event]}
 
     processor = SqsFifoPartialProcessor()
@@ -702,6 +728,39 @@ def lambda_handler(event, context):
     assert result["batchItemFailures"][1]["itemIdentifier"] == third_record.message_id
 
 
+def test_sqs_fifo_batch_processor_middleware_without_first_failure(sqs_event_fifo_factory, record_handler):
+    # GIVEN a batch of 5 records with 3 different MessageGroupID
+    first_record = SQSRecord(sqs_event_fifo_factory("success", "1"))
+    second_record = SQSRecord(sqs_event_fifo_factory("success", "1"))
+    third_record = SQSRecord(sqs_event_fifo_factory("fail", "2"))
+    fourth_record = SQSRecord(sqs_event_fifo_factory("fail", "2"))
+    fifth_record = SQSRecord(sqs_event_fifo_factory("success", "3"))
+    event = {
+        "Records": [
+            first_record.raw_event,
+            second_record.raw_event,
+            third_record.raw_event,
+            fourth_record.raw_event,
+            fifth_record.raw_event,
+        ],
+    }
+
+    # WHEN the FIFO processor is set to continue processing even after encountering errors in specific MessageGroupID
+    processor = SqsFifoPartialProcessor(return_on_first_error=False)
+
+    @batch_processor(record_handler=record_handler, processor=processor)
+    def lambda_handler(event, context):
+        return processor.response()
+
+    # WHEN
+    result = lambda_handler(event, {})
+
+    # THEN only failed messages should originate from MessageGroupID 2
+    assert len(result["batchItemFailures"]) == 2
+    assert result["batchItemFailures"][0]["itemIdentifier"] == third_record.message_id
+    assert result["batchItemFailures"][1]["itemIdentifier"] == fourth_record.message_id
+
+
 def test_async_batch_processor_middleware_success_only(sqs_event_factory, async_record_handler):
     # GIVEN
     first_record = SQSRecord(sqs_event_factory("success"))