feat: add async_process_partial_response

Author: heitorlessa

aws_lambda_powertools/utilities/batch/__init__.py

@@ -15,6 +15,7 @@
 )
 from aws_lambda_powertools.utilities.batch.decorators import (
     async_batch_processor,
+    async_process_partial_response,
     batch_processor,
     process_partial_response,
 )
@@ -26,6 +27,7 @@
 
 __all__ = (
     "async_batch_processor",
+    "async_process_partial_response",
     "batch_processor",
     "process_partial_response",
     "BatchProcessor",

aws_lambda_powertools/utilities/batch/decorators.py

@@ -54,7 +54,7 @@ def async_batch_processor(
         >>>     return payload
         >>>
         >>> @async_batch_processor(record_handler=async_record_handler, processor=processor)
-        >>> async def lambda_handler(event, context: LambdaContext):
+        >>> def lambda_handler(event, context):
         >>>     return processor.response()
 
     Limitations
@@ -183,3 +183,67 @@ def handler(event, context):
         processor.process()
 
     return processor.response()
+
+
+def async_process_partial_response(
+    event: Dict,
+    record_handler: Callable,
+    processor: AsyncBatchProcessor,
+    context: LambdaContext | None = None,
+) -> PartialItemFailureResponse:
+    """
+    Higher level function to handle batch event processing asynchronously.
+
+    Parameters
+    ----------
+    event: Dict
+        Lambda's original event
+    record_handler: Callable
+        Callable to process each record from the batch
+    processor: AsyncBatchProcessor
+        Batch Processor to handle partial failure cases
+    context: LambdaContext
+        Lambda's context, used to optionally inject in record handler
+
+    Returns
+    -------
+    result: PartialItemFailureResponse
+        Lambda Partial Batch Response
+
+    Examples
+    --------
+    **Processes Lambda's SQS event**
+
+    ```python
+    from aws_lambda_powertools.utilities.batch import AsyncBatchProcessor, EventType, process_partial_response
+    from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
+
+    processor = BatchProcessor(EventType.SQS)
+
+    async def record_handler(record: SQSRecord):
+        return record.body
+
+    def handler(event, context):
+        return async_process_partial_response(
+            event=event, record_handler=record_handler, processor=processor, context=context
+        )
+    ```
+
+    Limitations
+    -----------
+    * Sync batch processors. Use `process_partial_response` instead.
+    """
+    try:
+        records: List[Dict] = event.get("Records", [])
+    except AttributeError:
+        event_types = ", ".join(list(EventType.__members__))
+        docs = "https://awslabs.github.io/aws-lambda-powertools-python/latest/utilities/batch/#processing-messages-from-sqs"  # noqa: E501 # long-line
+        raise ValueError(
+            f"Invalid event format. Please ensure batch event is a valid {processor.event_type.value} event. \n"
+            f"See sample events in our documentation for either {event_types}: \n {docs}"
+        )
+
+    with processor(records, record_handler, context):
+        processor.async_process()
+
+    return processor.response()

tests/functional/test_utilities_batch.py

@@ -12,6 +12,7 @@
     EventType,
     SqsFifoPartialProcessor,
     async_batch_processor,
+    async_process_partial_response,
     batch_processor,
     process_partial_response,
 )
@@ -808,3 +809,35 @@ def test_process_partial_response_invalid_input(record_handler: Callable, batch:
     # WHEN/THEN
     with pytest.raises(ValueError):
         process_partial_response(batch, record_handler, processor)
+
+
+def test_async_process_partial_response(sqs_event_factory, async_record_handler):
+    # GIVEN
+    records = [sqs_event_factory("success"), sqs_event_factory("success")]
+    batch = {"Records": records}
+    processor = AsyncBatchProcessor(event_type=EventType.SQS)
+
+    # WHEN
+    ret = async_process_partial_response(batch, async_record_handler, processor)
+
+    # THEN
+    assert ret == {"batchItemFailures": []}
+
+
+@pytest.mark.parametrize(
+    "batch",
+    [
+        pytest.param(123456789, id="num"),
+        pytest.param([], id="list"),
+        pytest.param(False, id="bool"),
+        pytest.param(object, id="object"),
+        pytest.param(lambda x: x, id="callable"),
+    ],
+)
+def test_async_process_partial_response_invalid_input(async_record_handler: Callable, batch: Any):
+    # GIVEN
+    processor = AsyncBatchProcessor(event_type=EventType.SQS)
+
+    # WHEN/THEN
+    with pytest.raises(ValueError):
+        async_process_partial_response(batch, record_handler, processor)