fix: throw exception by default if messages processing fails

Author: Tom McCarthy

aws_lambda_powertools/utilities/batch/base.py

@@ -16,6 +16,7 @@ class BasePartialProcessor(ABC):
     def __init__(self):
         self.success_messages: List = []
         self.fail_messages: List = []
+        self.exceptions: List = []
 
     @abstractmethod
     def _prepare(self):
@@ -89,5 +90,6 @@ def failure_handler(self, record: Any, exception: Exception):
             "fail", exceptions args, original record
         """
         entry = ("fail", exception.args, record)
+        self.exceptions.append(exception)
         self.fail_messages.append(record)
         return entry

aws_lambda_powertools/utilities/batch/exceptions.py

@@ -0,0 +1,7 @@
+"""
+Batch processing exceptions
+"""
+
+
+class SQSBatchProcessingError(Exception):
+    """When at least one message within a batch could not be processed"""

aws_lambda_powertools/utilities/batch/middlewares.py

@@ -60,10 +60,15 @@ def batch_processor(
 
 @lambda_handler_decorator
 def sqs_batch_processor(
-    handler: Callable, event: Dict, context: Dict, record_handler: Callable, config: Optional[Config] = None
+    handler: Callable,
+    event: Dict,
+    context: Dict,
+    record_handler: Callable,
+    config: Optional[Config] = None,
+    suppress_exception: bool = False,
 ):
     """
-    Middleware to handle batch event processing
+    Middleware to handle SQS batch event processing
 
     Parameters
     ----------
@@ -77,6 +82,8 @@ def sqs_batch_processor(
         Callable to process each record from the batch
     config: Config
             botocore config object
+    suppress_exception: bool, optional
+        Supress exception raised if any messages fail processing, by default False
 
     Examples
     --------
@@ -96,7 +103,7 @@ def sqs_batch_processor(
 
     """
     config = config or Config()
-    processor = PartialSQSProcessor(config=config)
+    processor = PartialSQSProcessor(config=config, suppress_exception=suppress_exception)
 
     records = event["Records"]
 

aws_lambda_powertools/utilities/batch/sqs.py

@@ -9,20 +9,24 @@
 from botocore.config import Config
 
 from .base import BasePartialProcessor
+from .exceptions import SQSBatchProcessingError
 
 
 class PartialSQSProcessor(BasePartialProcessor):
     """
     Amazon SQS batch processor to delete successes from the Queue.
 
-    Only the **special** case of partial failure is handled, thus a batch in
-    which all records failed is **not** going to be removed from the queue, and
-    the same is valid for a full success.
+    The whole batch will be processed, even if failures occur. After all records are processed,
+    SQSBatchProcessingError will be raised if there were any failures, causing messages to
+    be returned to the SQS queue. This behaviour can be disabled by passing suppress_exception.
 
     Parameters
     ----------
     config: Config
         botocore config object
+    suppress_exception: bool, optional
+        Supress exception raised if any messages fail processing, by default False
+
 
     Example
     -------
@@ -46,12 +50,13 @@ class PartialSQSProcessor(BasePartialProcessor):
         >>>     return result
     """
 
-    def __init__(self, config: Optional[Config] = None):
+    def __init__(self, config: Optional[Config] = None, suppress_exception: bool = False):
         """
         Initializes sqs client.
         """
         config = config or Config()
         self.client = boto3.client("sqs", config=config)
+        self.suppress_exception = suppress_exception
 
         super().__init__()
 
@@ -97,10 +102,17 @@ def _clean(self):
         """
         Delete messages from Queue in case of partial failure.
         """
-        if not (self.fail_messages and self.success_messages):
+        # If all messages were successful, fall back to the default SQS -
+        # Lambda behaviour which deletes messages if Lambda responds successfully
+        if not self.fail_messages:
             return
 
         queue_url = self._get_queue_url()
         entries_to_remove = self._get_entries_to_clean()
 
-        return self.client.delete_message_batch(QueueUrl=queue_url, Entries=entries_to_remove)
+        delete_message_response = self.client.delete_message_batch(QueueUrl=queue_url, Entries=entries_to_remove)
+
+        if self.fail_messages and not self.suppress_exception:
+            raise SQSBatchProcessingError(list(self.exceptions))
+
+        return delete_message_response

tests/functional/test_utilities_batch.py

@@ -6,6 +6,7 @@
 from botocore.stub import Stubber
 
 from aws_lambda_powertools.utilities.batch import PartialSQSProcessor, batch_processor, sqs_batch_processor
+from aws_lambda_powertools.utilities.batch.exceptions import SQSBatchProcessingError
 
 
 @pytest.fixture(scope="module")
@@ -47,13 +48,25 @@ def partial_processor(config) -> PartialSQSProcessor:
     return PartialSQSProcessor(config=config)
 
 
+@pytest.fixture(scope="function")
+def partial_processor_suppressed(config) -> PartialSQSProcessor:
+    return PartialSQSProcessor(config=config, suppress_exception=True)
+
+
 @pytest.fixture(scope="function")
 def stubbed_partial_processor(config) -> PartialSQSProcessor:
     processor = PartialSQSProcessor(config=config)
     with Stubber(processor.client) as stubber:
         yield stubber, processor
 
 
+@pytest.fixture(scope="function")
+def stubbed_partial_processor_suppressed(config) -> PartialSQSProcessor:
+    processor = PartialSQSProcessor(config=config, suppress_exception=True)
+    with Stubber(processor.client) as stubber:
+        yield stubber, processor
+
+
 def test_partial_sqs_processor_context_with_failure(sqs_event_factory, record_handler, partial_processor):
     """
     Test processor with one failing record
@@ -68,16 +81,13 @@ def test_partial_sqs_processor_context_with_failure(sqs_event_factory, record_ha
     with Stubber(partial_processor.client) as stubber:
         stubber.add_response("delete_message_batch", response)
 
-        with partial_processor(records, record_handler) as ctx:
-            result = ctx.process()
+        with pytest.raises(SQSBatchProcessingError) as error:
+            with partial_processor(records, record_handler) as ctx:
+                ctx.process()
 
+        assert len(error.value.args[0]) == 1
         stubber.assert_no_pending_responses()
 
-    assert result == [
-        ("fail", ("Failed to process record.",), fail_record),
-        ("success", success_record["body"], success_record),
-    ]
-
 
 def test_partial_sqs_processor_context_only_success(sqs_event_factory, record_handler, partial_processor):
     """
@@ -126,18 +136,17 @@ def lambda_handler(event, context):
 
     fail_record = sqs_event_factory("fail")
 
-    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("success")]}
+    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("fail"), sqs_event_factory("success")]}
     response = {"Successful": [{"Id": fail_record["messageId"]}], "Failed": []}
 
     with Stubber(partial_processor.client) as stubber:
         stubber.add_response("delete_message_batch", response)
+        with pytest.raises(SQSBatchProcessingError) as error:
+            lambda_handler(event, {})
 
-        result = lambda_handler(event, {})
-
+        assert len(error.value.args[0]) == 2
         stubber.assert_no_pending_responses()
 
-    assert result is True
-
 
 @patch("aws_lambda_powertools.utilities.batch.middlewares.PartialSQSProcessor")
 def test_sqs_batch_processor_middleware(
@@ -159,10 +168,11 @@ def lambda_handler(event, context):
     event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("success")]}
     response = {"Successful": [{"Id": fail_record["messageId"]}], "Failed": []}
     stubber.add_response("delete_message_batch", response)
-    result = lambda_handler(event, {})
-    stubber.assert_no_pending_responses()
+    with pytest.raises(SQSBatchProcessingError) as error:
+        lambda_handler(event, {})
 
-    assert result is True
+    assert len(error.value.args[0]) == 1
+    stubber.assert_no_pending_responses()
 
 
 def test_batch_processor_middleware_with_custom_processor(capsys, sqs_event_factory, record_handler, config):
@@ -188,10 +198,80 @@ def lambda_handler(event, context):
 
     with Stubber(processor.client) as stubber:
         stubber.add_response("delete_message_batch", response)
+        with pytest.raises(SQSBatchProcessingError) as error:
+            lambda_handler(event, {})
+
+        stubber.assert_no_pending_responses()
+
+    assert len(error.value.args[0]) == 1
+    assert capsys.readouterr().out == "Oh no ! It's a failure.\n"
+
+
+def test_batch_processor_middleware_suppressed_exceptions(
+    sqs_event_factory, record_handler, partial_processor_suppressed
+):
+    """
+    Test middleware's integration with PartialSQSProcessor
+    """
+
+    @batch_processor(record_handler=record_handler, processor=partial_processor_suppressed)
+    def lambda_handler(event, context):
+        return True
+
+    fail_record = sqs_event_factory("fail")
+
+    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("fail"), sqs_event_factory("success")]}
+    response = {"Successful": [{"Id": fail_record["messageId"]}], "Failed": []}
 
+    with Stubber(partial_processor_suppressed.client) as stubber:
+        stubber.add_response("delete_message_batch", response)
         result = lambda_handler(event, {})
 
         stubber.assert_no_pending_responses()
+        assert result is True
+
+
+def test_partial_sqs_processor_suppressed_exceptions(sqs_event_factory, record_handler, partial_processor_suppressed):
+    """
+    Test processor without failure
+    """
 
+    first_record = sqs_event_factory("success")
+    second_record = sqs_event_factory("fail")
+    records = [first_record, second_record]
+
+    fail_record = sqs_event_factory("fail")
+    response = {"Successful": [{"Id": fail_record["messageId"]}], "Failed": []}
+
+    with Stubber(partial_processor_suppressed.client) as stubber:
+        stubber.add_response("delete_message_batch", response)
+        with partial_processor_suppressed(records, record_handler) as ctx:
+            ctx.process()
+
+    assert partial_processor_suppressed.success_messages == [first_record]
+
+
+@patch("aws_lambda_powertools.utilities.batch.middlewares.PartialSQSProcessor")
+def test_sqs_batch_processor_middleware_suppressed_exception(
+    patched_sqs_processor, sqs_event_factory, record_handler, stubbed_partial_processor_suppressed
+):
+    """
+    Test middleware's integration with PartialSQSProcessor
+    """
+
+    @sqs_batch_processor(record_handler=record_handler)
+    def lambda_handler(event, context):
+        return True
+
+    stubber, processor = stubbed_partial_processor_suppressed
+    patched_sqs_processor.return_value = processor
+
+    fail_record = sqs_event_factory("fail")
+
+    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("success")]}
+    response = {"Successful": [{"Id": fail_record["messageId"]}], "Failed": []}
+    stubber.add_response("delete_message_batch", response)
+    result = lambda_handler(event, {})
+
+    stubber.assert_no_pending_responses()
     assert result is True
-    assert capsys.readouterr().out == "Oh no ! It's a failure.\n"

tests/unit/test_utilities_batch.py

@@ -2,6 +2,7 @@
 from botocore.config import Config
 
 from aws_lambda_powertools.utilities.batch import PartialSQSProcessor
+from aws_lambda_powertools.utilities.batch.exceptions import SQSBatchProcessingError
 
 
 @pytest.fixture(scope="function")
@@ -126,8 +127,8 @@ def test_partial_sqs_clean(monkeypatch, mocker, partial_sqs_processor):
     entries_to_clean_mock.return_value = mocker.sentinel.entries_to_clean
 
     client_mock = mocker.patch.object(partial_sqs_processor, "client", autospec=True)
-
-    partial_sqs_processor._clean()
+    with pytest.raises(SQSBatchProcessingError):
+        partial_sqs_processor._clean()
 
     client_mock.delete_message_batch.assert_called_once_with(
         QueueUrl=mocker.sentinel.queue_url, Entries=mocker.sentinel.entries_to_clean