Ability to control payload deletion from S3

Author: Pankaj Agrawal

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

@@ -14,7 +14,6 @@
 package software.amazon.lambda.powertools.sqs;
 
 import java.util.List;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -39,13 +38,29 @@ public class PowertoolsSqs {
     /**
      * This is a utility method when you want to avoid using {@code LargeMessageHandler} annotation.
      * Gives you access to enriched messages from S3 in the SQS event produced via extended client lib.
+     * If all the large S3 payload are successfully retrieved, it will delete them from S3 post success.
      *
      * @param sqsEvent        Event received from SQS Extended client library
      * @param messageFunction Function to execute you business logic which provides access to enriched messages from S3 when needed.
-     * @return                Return value from the function.
+     * @return Return value from the function.
      */
     public static <R> R enrichedMessageFromS3(final SQSEvent sqsEvent,
                                               final Function<List<SQSMessage>, R> messageFunction) {
+        return enrichedMessageFromS3(sqsEvent, true, messageFunction);
+    }
+
+    /**
+     * This is a utility method when you want to avoid using {@code LargeMessageHandler} annotation.
+     * Gives you access to enriched messages from S3 in the SQS event produced via extended client lib.
+     * if all the large S3 payload are successfully retrieved, Control if it will delete payload from S3 post success.
+     *
+     * @param sqsEvent        Event received from SQS Extended client library
+     * @param messageFunction Function to execute you business logic which provides access to enriched messages from S3 when needed.
+     * @return Return value from the function.
+     */
+    public static <R> R enrichedMessageFromS3(final SQSEvent sqsEvent,
+                                              final boolean deleteS3Payload,
+                                              final Function<List<SQSMessage>, R> messageFunction) {
 
         List<SQSMessage> sqsMessages = sqsEvent.getRecords().stream()
                 .map(PowertoolsSqs::clonedMessage)
@@ -55,7 +70,9 @@ public static <R> R enrichedMessageFromS3(final SQSEvent sqsEvent,
 
         R returnValue = messageFunction.apply(sqsMessages);
 
-        s3Pointers.forEach(SqsMessageAspect::deleteMessage);
+        if (deleteS3Payload) {
+            s3Pointers.forEach(SqsMessageAspect::deleteMessage);
+        }
 
         return returnValue;
     }

powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/PowertoolsSqsTest.java

@@ -19,6 +19,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mock;
 import software.amazon.lambda.powertools.sqs.internal.SqsMessageAspect;
 
@@ -68,6 +69,31 @@ public void testLargeMessage() {
         verify(amazonS3).deleteObject(BUCKET_NAME, BUCKET_KEY);
     }
 
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    public void testLargeMessageDeleteFromS3Toggle(boolean deleteS3Payload) {
+        S3Object s3Response = new S3Object();
+        s3Response.setObjectContent(new ByteArrayInputStream("A big message".getBytes()));
+
+        when(amazonS3.getObject(BUCKET_NAME, BUCKET_KEY)).thenReturn(s3Response);
+        SQSEvent sqsEvent = messageWithBody("[\"software.amazon.payloadoffloading.PayloadS3Pointer\",{\"s3BucketName\":\"" + BUCKET_NAME + "\",\"s3Key\":\"" + BUCKET_KEY + "\"}]");
+
+        Map<String, String> sqsMessage = PowertoolsSqs.enrichedMessageFromS3(sqsEvent, deleteS3Payload, sqsMessages -> {
+            Map<String, String> someBusinessLogic = new HashMap<>();
+            someBusinessLogic.put("Message", sqsMessages.get(0).getBody());
+            return someBusinessLogic;
+        });
+
+        assertThat(sqsMessage)
+                .hasSize(1)
+                .containsEntry("Message", "A big message");
+        if (deleteS3Payload) {
+            verify(amazonS3).deleteObject(BUCKET_NAME, BUCKET_KEY);
+        } else {
+            verify(amazonS3, never()).deleteObject(BUCKET_NAME, BUCKET_KEY);
+        }
+    }
+
     @Test
     public void shouldNotProcessSmallMessageBody() {
         S3Object s3Response = new S3Object();