Merge pull request #100 from gmcrocetti/partial-sqs-batch

feat: SQS Partial failure

Author: heitorlessa

aws_lambda_powertools/utilities/batch/__init__.py

@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+
+"""
+Batch processing utility
+"""
+
+from .base import BasePartialProcessor
+from .middlewares import batch_processor
+from .sqs import PartialSQSProcessor
+
+__all__ = (
+    "BasePartialProcessor",
+    "PartialSQSProcessor",
+    "batch_processor",
+)

aws_lambda_powertools/utilities/batch/base.py

@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+
+"""
+Batch processing utilities
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Iterable, List, Tuple
+
+
+class BasePartialProcessor(ABC):
+    """
+    Abstract class for batch processors.
+    """
+
+    def __init__(self):
+        self.success_messages: List = []
+        self.fail_messages: List = []
+
+    @abstractmethod
+    def _prepare(self):
+        """
+        Prepare context manager.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _clean(self):
+        """
+        Clear context manager.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _process_record(self, record: Any):
+        """
+        Process record with handler.
+        """
+        raise NotImplementedError()
+
+    def process(self) -> List[Tuple]:
+        """
+        Call instance's handler for each record.
+        """
+        return [self._process_record(record) for record in self.records]
+
+    def __enter__(self):
+        self._prepare()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self._clean()
+
+    def __call__(self, records: Iterable[Any], handler: Callable):
+        """
+        Set instance attributes before execution
+
+        Parameters
+        ----------
+        records: Iterable[Any]
+            Iterable with objects to be processed.
+        handler: Callable
+            Callable to process "records" entries.
+        """
+        self.records = records
+        self.handler = handler
+        return self
+
+    def success_handler(self, record: Any, result: Any):
+        """
+        Success callback
+
+        Returns
+        -------
+        tuple
+            "success", result, original record
+        """
+        entry = ("success", result, record)
+        self.success_messages.append(record)
+        return entry
+
+    def failure_handler(self, record: Any, exception: Exception):
+        """
+        Failure callback
+
+        Returns
+        -------
+        tuple
+            "fail", exceptions args, original record
+        """
+        entry = ("fail", exception.args, record)
+        self.fail_messages.append(record)
+        return entry

aws_lambda_powertools/utilities/batch/middlewares.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+
+"""
+Middlewares for batch utilities
+"""
+
+from typing import Callable, Dict
+
+from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+
+from .base import BasePartialProcessor
+
+
+@lambda_handler_decorator
+def batch_processor(
+    handler: Callable, event: Dict, context: Dict, record_handler: Callable, processor: BasePartialProcessor = None
+):
+    """
+    Middleware to handle batch event processing
+
+    Parameters
+    ----------
+    handler: Callable
+        Lambda's handler
+    event: Dict
+        Lambda's Event
+    context: Dict
+        Lambda's Context
+    record_handler: Callable
+        Callable to process each record from the batch
+    processor: PartialSQSProcessor
+        Batch Processor to handle partial failure cases
+
+    Examples
+    --------
+    **Processes Lambda's event with PartialSQSProcessor**
+        >>> from aws_lambda_powertools.utilities.batch import batch_processor
+        >>>
+        >>> def record_handler(record):
+        >>>     return record["body"]
+        >>>
+        >>> @batch_processor(record_handler=record_handler, processor=PartialSQSProcessor())
+        >>> def handler(event, context):
+        >>>     return {"StatusCode": 200}
+
+    Limitations
+    -----------
+    * Async batch processors
+
+    """
+    records = event["Records"]
+
+    with processor(records, record_handler):
+        processor.process()
+
+    return handler(event, context)

aws_lambda_powertools/utilities/batch/sqs.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+
+"""
+Batch SQS utilities
+"""
+from typing import List, Optional, Tuple
+
+import boto3
+from botocore.config import Config
+
+from .base import BasePartialProcessor
+
+
+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.
+
+    Parameters
+    ----------
+    config: Config
+        botocore config object
+
+    Example
+    -------
+    **Process batch triggered by SQS**
+
+        >>> from aws_lambda_powertools.utilities.batch import PartialSQSProcessor
+        >>>
+        >>> def record_handler(record):
+        >>>     return record["body"]
+        >>>
+        >>> def handler(event, context):
+        >>>     records = event["Records"]
+        >>>     processor = PartialSQSProcessor()
+        >>>
+        >>>     with processor(records=records, handler=record_handler):
+        >>>         result = processor.process()
+        >>>
+        >>>     # Case a partial failure occurred, all successful executions
+        >>>     # have been deleted from the queue after context's exit.
+        >>>
+        >>>     return result
+    """
+
+    def __init__(self, config: Optional[Config] = None):
+        """
+        Initializes sqs client.
+        """
+        config = config or Config()
+        self.client = boto3.client("sqs", config=config)
+
+        super().__init__()
+
+    def _get_queue_url(self) -> str:
+        """
+        Format QueueUrl from first records entry
+        """
+        if not getattr(self, "records", None):
+            return
+
+        *_, account_id, queue_name = self.records[0]["eventSourceARN"].split(":")
+        return f"{self.client._endpoint.host}/{account_id}/{queue_name}"
+
+    def _get_entries_to_clean(self) -> List:
+        """
+        Format messages to use in batch deletion
+        """
+        return [{"Id": msg["messageId"], "ReceiptHandle": msg["receiptHandle"]} for msg in self.success_messages]
+
+    def _process_record(self, record) -> Tuple:
+        """
+        Process a record with instance's handler
+
+        Parameters
+        ----------
+        record: Any
+            An object to be processed.
+        """
+        try:
+            result = self.handler(record)
+            return self.success_handler(record, result)
+        except Exception as exc:
+            return self.failure_handler(record, exc)
+
+    def _prepare(self):
+        """
+        Remove results from previous execution.
+        """
+        self.success_messages.clear()
+        self.fail_messages.clear()
+
+    def _clean(self):
+        """
+        Delete messages from Queue in case of partial failure.
+        """
+        if not (self.fail_messages and self.success_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)

aws_lambda_powertools/utilities/parameters/secrets.py

@@ -27,7 +27,7 @@ class SecretsProvider(BaseProvider):
         >>> from aws_lambda_powertools.utilities.parameters import SecretsProvider
         >>> secrets_provider = SecretsProvider()
         >>>
-        >>> value secrets_provider.get("my-parameter")
+        >>> value = secrets_provider.get("my-parameter")
         >>>
         >>> print(value)
         My parameter value

docs/content/index.mdx

@@ -24,6 +24,12 @@ Powertools is available in PyPi. You can use your favourite dependency managemen
 ```bash:title=hello_world.sh
 sam init --location https://github.com/aws-samples/cookiecutter-aws-sam-python
 ```
+* [Tracing](./core/tracer) - Decorators and utilities to trace Lambda function handlers, and both synchronous and asynchronous functions
+* [Logging](./core/logger) - Structured logging made easier, and decorator to enrich structured logging with key Lambda context details
+* [Metrics](./core/metrics) - Custom Metrics created asynchronously via CloudWatch Embedded Metric Format (EMF)
+* [Bring your own middleware](./utilities/middleware_factory) - Decorator factory to create your own middleware to run logic before, and after each Lambda invocation
+* [Parameters utility](./utilities/parameters) - Retrieve parameter values from AWS Systems Manager Parameter Store, AWS Secrets Manager, or Amazon DynamoDB, and cache them for a specific amount of time
+* [Batch utility](./utilities/batch) - Batch processing for AWS SQS, handles partial failure.
 
 ### Lambda Layer