feat: add sqs failure processors

Author: gmcrocetti

aws_lambda_powertools/utilities/batch/base.py

@@ -27,7 +27,7 @@ def _clean(self):
         raise NotImplementedError()
 
     @abstractmethod
-    def _process_record(self):
+    def _process_record(self, record):
         raise NotImplementedError()
 
     def process(self) -> List[Tuple]:

aws_lambda_powertools/utilities/batch/sqs.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+
+"""
+Batch SQS utilities
+"""
+
+from typing import List
+
+import boto3
+
+from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+
+from .base import BasePartialProcessor
+
+
+class BasePartialSQSProcessor(BasePartialProcessor):
+    def __init__(self):
+        self._client = boto3.client("sqs")
+        self.success_messages: List = []
+        self.fail_messages: List = []
+
+        super().__init__()
+
+    def get_queue_url(self):
+        *_, 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):
+        return [{"Id": msg["messageId"], "ReceiptHandle": msg["receiptHandle"]} for msg in self.success_messages]
+
+    def _process_record(self, record):
+        try:
+            result = self.handler(record)
+            return self.success_handler(record, result)
+        except Exception as exc:
+            return self.failure_handler(record, exc)
+
+    def _prepare(self):
+        """
+        """
+        self.success_messages.clear()
+        self.fail_messages.clear()
+
+    def _clean(self):
+        """
+        """
+        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)
+
+
+class DefaultPartialSQSProcessor(BasePartialSQSProcessor):
+    pass
+
+
+@lambda_handler_decorator
+def partial_sqs_processor(handler, event, context, record_handler, processor=None):
+    records = event["Records"]
+    processor = processor or DefaultPartialSQSProcessor()
+
+    with processor(records, record_handler) as ctx:
+        ctx.process()