Refactoring logic to skip execution when messages are part of a groupid with failed messages

Author: leandrodamascena

aws_lambda_powertools/utilities/batch/sqs_fifo_partial_processor.py

@@ -1,8 +1,11 @@
+import logging
 from typing import List, Optional, Tuple
 
 from aws_lambda_powertools.utilities.batch import BatchProcessor, EventType
 from aws_lambda_powertools.utilities.batch.types import BatchSqsTypeModel
 
+logger = logging.getLogger(__name__)
+
 
 class SQSFifoCircuitBreakerError(Exception):
     """
@@ -57,19 +60,20 @@ def lambda_handler(event, context: LambdaContext):
         None,
     )
 
-    def __init__(self, model: Optional["BatchSqsTypeModel"] = None, return_on_first_error: bool = True):
+    def __init__(self, model: Optional["BatchSqsTypeModel"] = None, skip_group_on_error: bool = False):
         """
         Initialize the SqsFifoProcessor.
 
         Parameters
         ----------
         model: Optional["BatchSqsTypeModel"]
             An optional model for batch processing.
-        return_on_first_error: bool
-            Flag to determine whether to return on the first error encountered. Default is True
+        skip_group_on_error: bool
+            # TODO: Alterar
+            Determine whether to return on the first error encountered. Default is True
 
         """
-        self.return_on_first_error = return_on_first_error
+        self._skip_group_on_error = skip_group_on_error
         super().__init__(EventType.SQS, model)
 
     def process(self) -> List[Tuple]:
@@ -78,25 +82,41 @@ def process(self) -> List[Tuple]:
         the process is short-circuited, and the remaining messages are reported as failed items.
         """
         result: List[Tuple] = []
+        skip_messages_group_id: List = []
 
         for i, record in enumerate(self.records):
             # 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:
+            if self.fail_messages and not self._skip_group_on_error:
+                logger.debug("Processing of failed messages stopped due to the 'skip_group_on_error' is set to False")
                 return self._short_circuit_processing(i, result)
 
-            # Process the current record
+            msg_id = record.get("messageId")
+
+            # skip_group_on_error is True:
+            # Skip processing the current message if its ID belongs to a group with failed messages
+            if msg_id in skip_messages_group_id:
+                logger.debug(
+                    f"Skipping message with ID '{msg_id}' as it is part of a group containing failed messages.",
+                )
+                continue
+
             processed_messages = self._process_record(record)
 
-            # If a processed message fail,
+            # If a processed message fail and skip_group_on_error is True,
             # mark subsequent messages from the same MessageGroupId as skipped
-            if processed_messages[0] == "fail":
+            if processed_messages[0] == "fail" and self._skip_group_on_error:
+                _attributes_record = record.get("attributes", {})
                 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
+                    _attributes = subsequent_record.get("attributes", {})
+                    if _attributes.get("MessageGroupId") == _attributes_record.get("MessageGroupId"):
+                        skip_messages_group_id.append(subsequent_record.get("messageId"))
+                        data = self._to_batch_type(
+                            record=subsequent_record,
+                            event_type=self.event_type,
+                            model=self.model,
+                        )
+                        result.append(self.failure_handler(record=data, exception=self.circuit_breaker_exc))
 
             # Append the processed message normally
             result.append(processed_messages)

docs/utilities/batch.md

@@ -141,7 +141,9 @@ Processing batches from SQS works in three stages:
 
 #### FIFO queues
 
-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.
+When working with [SQS FIFO queues](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/FIFO-queues.html){target="_blank"}, you should know that a batch may include messages from different group IDs.
+
+By default, we will stop processing at the first failure and mark unprocessed messages as failed to preserve ordering.
 
 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.
 

tests/functional/test_utilities_batch.py

@@ -28,6 +28,7 @@
 from aws_lambda_powertools.utilities.parser.models import (
     DynamoDBStreamChangedRecordModel,
     DynamoDBStreamRecordModel,
+    SqsRecordModel,
 )
 from aws_lambda_powertools.utilities.parser.types import Literal
 from tests.functional.batch.sample_models import (
@@ -728,13 +729,13 @@ 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):
+def test_sqs_fifo_batch_processor_middleware_with_skip_group_on_error(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"))
+    fourth_record = SQSRecord(sqs_event_fifo_factory("success", "2"))
+    fifth_record = SQSRecord(sqs_event_fifo_factory("fail", "3"))
     event = {
         "Records": [
             first_record.raw_event,
@@ -746,7 +747,7 @@ def test_sqs_fifo_batch_processor_middleware_without_first_failure(sqs_event_fif
     }
 
     # WHEN the FIFO processor is set to continue processing even after encountering errors in specific MessageGroupID
-    processor = SqsFifoPartialProcessor(return_on_first_error=False)
+    processor = SqsFifoPartialProcessor(skip_group_on_error=True)
 
     @batch_processor(record_handler=record_handler, processor=processor)
     def lambda_handler(event, context):
@@ -755,10 +756,53 @@ def lambda_handler(event, context):
     # WHEN
     result = lambda_handler(event, {})
 
-    # THEN only failed messages should originate from MessageGroupID 2
-    assert len(result["batchItemFailures"]) == 2
+    # THEN only failed messages should originate from MessageGroupID 3
+    assert len(result["batchItemFailures"]) == 3
+    assert result["batchItemFailures"][0]["itemIdentifier"] == third_record.message_id
+    assert result["batchItemFailures"][1]["itemIdentifier"] == fourth_record.message_id
+    assert result["batchItemFailures"][2]["itemIdentifier"] == fifth_record.message_id
+
+
+def test_sqs_fifo_batch_processor_middleware_with_skip_group_on_error_and_model(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("success", "2"))
+    fifth_record = SQSRecord(sqs_event_fifo_factory("fail", "3"))
+    event = {
+        "Records": [
+            first_record.raw_event,
+            second_record.raw_event,
+            third_record.raw_event,
+            fourth_record.raw_event,
+            fifth_record.raw_event,
+        ],
+    }
+
+    class OrderSqsRecord(SqsRecordModel):
+        receiptHandle: str
+
+    # WHEN the FIFO processor is set to continue processing even after encountering errors in specific MessageGroupID
+    # WHEN processor is using a Pydantic Model we must be able to access MessageGroupID property
+    processor = SqsFifoPartialProcessor(skip_group_on_error=True, model=OrderSqsRecord)
+
+    def record_handler(record: OrderSqsRecord):
+        if record.body == "fail":
+            raise ValueError("blah")
+
+    @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 3
+    assert len(result["batchItemFailures"]) == 3
     assert result["batchItemFailures"][0]["itemIdentifier"] == third_record.message_id
     assert result["batchItemFailures"][1]["itemIdentifier"] == fourth_record.message_id
+    assert result["batchItemFailures"][2]["itemIdentifier"] == fifth_record.message_id
 
 
 def test_async_batch_processor_middleware_success_only(sqs_event_factory, async_record_handler):