Initial API skeleton for Partial SQS batch util

Author: Pankaj Agrawal

powertools-sqs/pom.xml

@@ -57,6 +57,10 @@
             <groupId>software.amazon.payloadoffloading</groupId>
             <artifactId>payloadoffloading-common</artifactId>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sqs</artifactId>
+        </dependency>
 
         <dependency>
             <groupId>org.aspectj</groupId>

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/PowertoolsSqs.java

@@ -13,13 +13,17 @@
  */
 package software.amazon.lambda.powertools.sqs;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import com.amazonaws.services.lambda.runtime.events.SQSEvent;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import software.amazon.awssdk.services.sqs.SqsClient;
+import software.amazon.lambda.powertools.sqs.internal.BatchContext;
 import software.amazon.lambda.powertools.sqs.internal.SqsMessageAspect;
 import software.amazon.payloadoffloading.PayloadS3Pointer;
 
@@ -34,6 +38,7 @@
 public final class PowertoolsSqs {
 
     private static final ObjectMapper objectMapper = new ObjectMapper();
+    private static SqsClient client = SqsClient.create();
 
     private PowertoolsSqs() {
     }
@@ -80,6 +85,47 @@ public static <R> R enrichedMessageFromS3(final SQSEvent sqsEvent,
         return returnValue;
     }
 
+    public static void defaultSqsClient(SqsClient client) {
+        PowertoolsSqs.client = client;
+    }
+
+    public static SqsClient defaultSqsClient() {
+        return client;
+    }
+
+    public static <R> List<R> partialBatchProcessor(final SQSEvent event,
+                                                    boolean suppressException,
+                                                    final Class<? extends SqsMessageHandler<R>> handler) {
+
+        try {
+            return partialBatchProcessor(event, suppressException, handler.newInstance());
+        } catch (IllegalAccessException | InstantiationException e) {
+            // LOG something
+            return Collections.emptyList();
+        }
+    }
+
+    public static <R> List<R> partialBatchProcessor(final SQSEvent event,
+                                                    final boolean suppressException,
+                                                    final SqsMessageHandler<R> handler) {
+        final List<R> handlerReturn = new ArrayList<>();
+
+        BatchContext batchContext = new BatchContext(defaultSqsClient());
+
+        for (SQSMessage message : event.getRecords()) {
+            try {
+                handlerReturn.add(handler.process(message));
+                batchContext.addSuccess(message);
+            } catch (Exception e) {
+                batchContext.addFailure(message, e);
+            }
+        }
+
+        batchContext.processSuccessAndReset(suppressException);
+
+        return handlerReturn;
+    }
+
     private static SQSMessage clonedMessage(final SQSMessage sqsMessage) {
         try {
             return objectMapper
@@ -88,4 +134,5 @@ private static SQSMessage clonedMessage(final SQSMessage sqsMessage) {
             throw new RuntimeException(e);
         }
     }
+
 }

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SQSBatchProcessingException.java

@@ -0,0 +1,22 @@
+package software.amazon.lambda.powertools.sqs;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static java.util.stream.Collectors.joining;
+
+public class SQSBatchProcessingException extends RuntimeException {
+
+    private final List<Exception> exceptions;
+
+    public SQSBatchProcessingException(List<Exception> exceptions) {
+        this.exceptions = new ArrayList<>(exceptions);
+    }
+
+    @Override
+    public String getMessage() {
+        return exceptions.stream()
+                .map(Throwable::getMessage)
+                .collect(joining("\n"));
+    }
+}

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SqsBatchProcessor.java

@@ -0,0 +1,16 @@
+package software.amazon.lambda.powertools.sqs;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface SqsBatchProcessor {
+
+    Class<? extends SqsMessageHandler<Object>> value();
+
+    boolean suppressException() default false;
+}

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SqsMessageHandler.java

@@ -0,0 +1,9 @@
+package software.amazon.lambda.powertools.sqs;
+
+import static com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
+
+@FunctionalInterface
+public interface SqsMessageHandler<R> {
+
+    R process(SQSMessage message);
+}

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/BatchContext.java

@@ -0,0 +1,77 @@
+package software.amazon.lambda.powertools.sqs.internal;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import software.amazon.awssdk.services.sqs.SqsClient;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
+import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
+import software.amazon.awssdk.services.sqs.model.GetQueueUrlRequest;
+import software.amazon.lambda.powertools.sqs.SQSBatchProcessingException;
+
+import static java.util.stream.Collectors.toList;
+
+public final class BatchContext {
+    private List<SQSEvent.SQSMessage> success = new ArrayList<>();
+    private List<SQSEvent.SQSMessage> failures = new ArrayList<>();
+    private List<Exception> exceptions = new ArrayList<>();
+    private SqsClient client;
+
+    public BatchContext(SqsClient client) {
+        this.client = client;
+    }
+
+    public void addSuccess(SQSEvent.SQSMessage event) {
+        success.add(event);
+    }
+
+    public void addFailure(SQSEvent.SQSMessage event, Exception e) {
+        failures.add(event);
+        exceptions.add(e);
+    }
+
+    private boolean hasFailures() {
+        return !failures.isEmpty();
+    }
+
+    private void cleanUpAndReset() {
+        DeleteMessageBatchRequest request = DeleteMessageBatchRequest.builder()
+                .queueUrl(url())
+                .entries(success.stream().map(m -> DeleteMessageBatchRequestEntry.builder()
+                        .id(m.getMessageId())
+                        .receiptHandle(m.getReceiptHandle())
+                        .build()).collect(toList()))
+                .build();
+
+        client.deleteMessageBatch(request);
+    }
+
+    private String url() {
+        String[] arnArray = success.get(0).getEventSourceArn().split(":");
+        return client.getQueueUrl(GetQueueUrlRequest.builder()
+                .queueOwnerAWSAccountId(arnArray[1])
+                .queueName(arnArray[2])
+                .build())
+                .queueUrl();
+    }
+
+    private void reset() {
+        success = new ArrayList<>();
+        failures = new ArrayList<>();
+        exceptions = new ArrayList<>();
+    }
+
+    public void processSuccessAndReset(final boolean suppressException) {
+        if (hasFailures() && !suppressException) {
+            SQSBatchProcessingException exception = new SQSBatchProcessingException(exceptions);
+            cleanUpAndReset();
+            throw exception;
+        } else if (hasFailures()) {
+            // LOG
+            cleanUpAndReset();
+        } else {
+            reset();
+        }
+    }
+}

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/SqsMessageBatchProcessorAspect.java

@@ -0,0 +1,40 @@
+package software.amazon.lambda.powertools.sqs.internal;
+
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
+import software.amazon.awssdk.services.sqs.SqsClient;
+import software.amazon.lambda.powertools.sqs.PowertoolsSqs;
+import software.amazon.lambda.powertools.sqs.SqsBatchProcessor;
+
+import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.isHandlerMethod;
+import static software.amazon.lambda.powertools.sqs.internal.SqsMessageAspect.placedOnSqsEventRequestHandler;
+
+@Aspect
+public class SqsMessageBatchProcessorAspect {
+    private static final SqsClient client = SqsClient.create();
+    private static BatchContext details = new BatchContext(PowertoolsSqs.defaultSqsClient());
+
+    @SuppressWarnings({"EmptyMethod"})
+    @Pointcut("@annotation(sqsBatchProcessor)")
+    public void callAt(SqsBatchProcessor sqsBatchProcessor) {
+    }
+
+    @Around(value = "callAt(sqsBatchProcessor) && execution(@SqsBatchProcessor * *.*(..))", argNames = "pjp,sqsBatchProcessor")
+    public Object around(ProceedingJoinPoint pjp,
+                         SqsBatchProcessor sqsBatchProcessor) throws Throwable {
+        Object[] proceedArgs = pjp.getArgs();
+
+        if (isHandlerMethod(pjp)
+                && placedOnSqsEventRequestHandler(pjp)) {
+
+            SQSEvent sqsEvent = (SQSEvent) proceedArgs[0];
+
+            PowertoolsSqs.partialBatchProcessor(sqsEvent, sqsBatchProcessor.suppressException(), sqsBatchProcessor.value().newInstance());
+        }
+
+        return pjp.proceed(proceedArgs);
+    }
+}

powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/handlers/PartialBatchHandler.java

@@ -0,0 +1,37 @@
+package software.amazon.lambda.powertools.sqs.handlers;
+
+import java.util.List;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import software.amazon.lambda.powertools.sqs.LargeMessageHandler;
+import software.amazon.lambda.powertools.sqs.PowertoolsSqs;
+import software.amazon.lambda.powertools.sqs.SqsBatchProcessor;
+import software.amazon.lambda.powertools.sqs.SqsMessageHandler;
+
+public class PartialBatchHandler implements RequestHandler<SQSEvent, List<Object>> {
+
+
+    @Override
+    @LargeMessageHandler
+    @SqsBatchProcessor(HandlerSqs.class)
+    public List<Object> handleRequest(SQSEvent sqsEvent, Context context) {
+
+        List<Object> returnValues =
+                PowertoolsSqs.partialBatchProcessor(sqsEvent, false, HandlerSqs.class);
+
+        // Do some processing on processed message
+
+        return returnValues;
+    }
+
+    private class HandlerSqs implements SqsMessageHandler<Object> {
+
+        @Override
+        public String process(SQSEvent.SQSMessage message) {
+            // This is where you process message
+            return null;
+        }
+    }
+}