feat(batch): new BatchProcessor for SQS, DynamoDB, Kinesis

.pre-commit-config.yaml

@@ -11,10 +11,6 @@ repos:
           - id: trailing-whitespace
           - id: end-of-file-fixer
           - id: check-toml
-    - repo: https://github.com/pre-commit/pygrep-hooks
-      rev: v1.5.1
-      hooks:
-          - id: python-use-type-annotations
     - repo: local
       hooks:
           - id: black

aws_lambda_powertools/utilities/batch/__init__.py

@@ -4,7 +4,14 @@
 Batch processing utility
 """
 
-from .base import BasePartialProcessor, batch_processor
-from .sqs import PartialSQSProcessor, sqs_batch_processor
+from aws_lambda_powertools.utilities.batch.base import BasePartialProcessor, BatchProcessor, EventType, batch_processor
+from aws_lambda_powertools.utilities.batch.sqs import PartialSQSProcessor, sqs_batch_processor
 
-__all__ = ("BasePartialProcessor", "PartialSQSProcessor", "batch_processor", "sqs_batch_processor")
+__all__ = (
+    "BatchProcessor",
+    "BasePartialProcessor",
+    "EventType",
+    "PartialSQSProcessor",
+    "batch_processor",
+    "sqs_batch_processor",
+)

aws_lambda_powertools/utilities/batch/base.py

@@ -3,24 +3,62 @@
 """
 Batch processing utilities
 """
-
 import logging
+import sys
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Tuple
+from enum import Enum
+from types import TracebackType
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, overload
 
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+from aws_lambda_powertools.utilities.data_classes.dynamo_db_stream_event import DynamoDBRecord
+from aws_lambda_powertools.utilities.data_classes.kinesis_stream_event import KinesisStreamRecord
+from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
 
 logger = logging.getLogger(__name__)
 
 
+class EventType(Enum):
+    SQS = "SQS"
+    KinesisDataStreams = "KinesisDataStreams"
+    DynamoDBStreams = "DynamoDBStreams"
+
+
+#
+# type specifics
+#
+has_pydantic = "pydantic" in sys.modules
+_ExcInfo = Tuple[Type[BaseException], BaseException, TracebackType]
+_OptExcInfo = Union[_ExcInfo, Tuple[None, None, None]]
+
+# For IntelliSense and Mypy to work, we need to account for possible SQS, Kinesis and DynamoDB subclasses
+# We need them as subclasses as we must access their message ID or sequence number metadata via dot notation
+if has_pydantic:
+    from aws_lambda_powertools.utilities.parser.models import DynamoDBStreamRecordModel
+    from aws_lambda_powertools.utilities.parser.models import KinesisDataStreamRecord as KinesisDataStreamRecordModel
+    from aws_lambda_powertools.utilities.parser.models import SqsRecordModel
+
+    BatchTypeModels = Optional[
+        Union[Type[SqsRecordModel], Type[DynamoDBStreamRecordModel], Type[KinesisDataStreamRecordModel]]
+    ]
+
+# When using processor with default arguments, records will carry EventSourceDataClassTypes
+# and depending on what EventType it's passed it'll correctly map to the right record
+# When using Pydantic Models, it'll accept any
+EventSourceDataClassTypes = Union[SQSRecord, KinesisStreamRecord, DynamoDBRecord]
+BatchEventTypes = Union[EventSourceDataClassTypes, "BatchTypeModels"]
+SuccessCallback = Tuple[str, Any, BatchEventTypes]
+FailureCallback = Tuple[str, str, BatchEventTypes]
+
+
 class BasePartialProcessor(ABC):
     """
     Abstract class for batch processors.
     """
 
     def __init__(self):
-        self.success_messages: List = []
-        self.fail_messages: List = []
+        self.success_messages: List[BatchEventTypes] = []
+        self.fail_messages: List[BatchEventTypes] = []
         self.exceptions: List = []
 
     @abstractmethod
@@ -38,7 +76,7 @@ def _clean(self):
         raise NotImplementedError()
 
     @abstractmethod
-    def _process_record(self, record: Any):
+    def _process_record(self, record: dict):
         """
         Process record with handler.
         """
@@ -57,13 +95,13 @@ def __enter__(self):
     def __exit__(self, exception_type, exception_value, traceback):
         self._clean()
 
-    def __call__(self, records: List[Any], handler: Callable):
+    def __call__(self, records: List[dict], handler: Callable):
         """
         Set instance attributes before execution
 
         Parameters
         ----------
-        records: List[Any]
+        records: List[dict]
             List with objects to be processed.
         handler: Callable
             Callable to process "records" entries.
@@ -72,7 +110,7 @@ def __call__(self, records: List[Any], handler: Callable):
         self.handler = handler
         return self
 
-    def success_handler(self, record: Any, result: Any):
+    def success_handler(self, record, result: Any) -> SuccessCallback:
         """
         Success callback
 
@@ -85,7 +123,7 @@ def success_handler(self, record: Any, result: Any):
         self.success_messages.append(record)
         return entry
 
-    def failure_handler(self, record: Any, exception: Tuple):
+    def failure_handler(self, record, exception: _OptExcInfo) -> FailureCallback:
         """
         Failure callback
 
@@ -146,3 +184,108 @@ def batch_processor(
         processor.process()
 
     return handler(event, context)
+
+
+class BatchProcessor(BasePartialProcessor):
+    DEFAULT_RESPONSE: Dict[str, List[Optional[dict]]] = {"batchItemFailures": []}
+
+    def __init__(self, event_type: EventType, model: Optional["BatchTypeModels"] = None):
+        """Process batch and partially report failed items
+
+        Parameters
+        ----------
+        event_type: EventType
+            Whether this is a SQS, DynamoDB Streams, or Kinesis Data Stream event
+        model: Optional["BatchTypeModels"]
+            Parser's data model using either SqsRecordModel, DynamoDBStreamRecordModel, KinesisDataStreamRecord
+        """
+        self.event_type = event_type
+        self.model = model
+        self.batch_response = self.DEFAULT_RESPONSE
+        self._COLLECTOR_MAPPING = {
+            EventType.SQS: self._collect_sqs_failures,
+            EventType.KinesisDataStreams: self._collect_kinesis_failures,
+            EventType.DynamoDBStreams: self._collect_dynamodb_failures,
+        }
+        self._DATA_CLASS_MAPPING = {
+            EventType.SQS: SQSRecord,
+            EventType.KinesisDataStreams: KinesisStreamRecord,
+            EventType.DynamoDBStreams: DynamoDBRecord,
+        }
+
+        super().__init__()
+
+    def response(self):
+        """Batch items that failed processing, if any"""
+        return self.batch_response
+
+    def _prepare(self):
+        """
+        Remove results from previous execution.
+        """
+        self.success_messages.clear()
+        self.fail_messages.clear()
+        self.batch_response = self.DEFAULT_RESPONSE
+
+    def _process_record(self, record: dict) -> Union[SuccessCallback, FailureCallback]:
+        """
+        Process a record with instance's handler
+
+        Parameters
+        ----------
+        record: dict
+            A batch record to be processed.
+        """
+        data = self._to_batch_type(record=record, event_type=self.event_type, model=self.model)
+        try:
+            result = self.handler(record=data)
+            return self.success_handler(record=record, result=result)
+        except Exception:
+            return self.failure_handler(record=data, exception=sys.exc_info())
+
+    def _clean(self):
+        """
+        Report messages to be deleted in case of partial failure.
+        """
+
+        if not self._has_messages_to_report():
+            return
+
+        messages = self._get_messages_to_report()
+        self.batch_response = {"batchItemFailures": [messages]}
+
+    def _has_messages_to_report(self) -> bool:
+        if self.fail_messages:
+            return True
+
+        logger.debug(f"All {len(self.success_messages)} records successfully processed")
+        return False
+
+    def _get_messages_to_report(self) -> Dict[str, str]:
+        """
+        Format messages to use in batch deletion
+        """
+        return self._COLLECTOR_MAPPING[self.event_type]()
+
+    def _collect_sqs_failures(self):
+        return {"itemIdentifier": msg.message_id for msg in self.fail_messages}
+
+    def _collect_kinesis_failures(self):
+        return {"itemIdentifier": msg.kinesis.sequence_number for msg in self.fail_messages}
+
+    def _collect_dynamodb_failures(self):
+        return {"itemIdentifier": msg.dynamodb.sequence_number for msg in self.fail_messages}
+
+    @overload
+    def _to_batch_type(self, record: dict, event_type: EventType, model: "BatchTypeModels") -> "BatchTypeModels":
+        ...
+
+    @overload
+    def _to_batch_type(self, record: dict, event_type: EventType) -> EventSourceDataClassTypes:
+        ...
+
+    def _to_batch_type(self, record: dict, event_type: EventType, model: Optional["BatchTypeModels"] = None):
+        if model is not None:
+            return model.parse_obj(record)
+        else:
+            return self._DATA_CLASS_MAPPING[event_type](record)