add support to set a prefix for the S3 key #7

Author: Evangilo Morais

src/main/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedClient.java

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -965,7 +966,7 @@ private SendMessageBatchRequestEntry storeMessageInS3(SendMessageBatchRequestEnt
             updateMessageAttributePayloadSize(batchEntry.messageAttributes(), messageContentSize));
 
         // Store the message content in S3.
-        String largeMessagePointer = payloadStore.storeOriginalPayload(messageContentStr);
+        String largeMessagePointer = storeOriginalPayload(messageContentStr);
         batchEntryBuilder.messageBody(largeMessagePointer);
 
         return batchEntryBuilder.build();
@@ -984,12 +985,20 @@ private SendMessageRequest storeMessageInS3(SendMessageRequest sendMessageReques
             updateMessageAttributePayloadSize(sendMessageRequest.messageAttributes(), messageContentSize));
 
         // Store the message content in S3.
-        String largeMessagePointer = payloadStore.storeOriginalPayload(messageContentStr);
+        String largeMessagePointer = storeOriginalPayload(messageContentStr);
         sendMessageRequestBuilder.messageBody(largeMessagePointer);
 
         return sendMessageRequestBuilder.build();
     }
 
+    private String storeOriginalPayload(String messageContentStr) {
+        String s3KeyPrefix = clientConfiguration.getS3KeyPrefix();
+        if (StringUtils.isBlank(s3KeyPrefix)) {
+            return payloadStore.storeOriginalPayload(messageContentStr);
+        }
+        return payloadStore.storeOriginalPayload(messageContentStr, s3KeyPrefix + UUID.randomUUID());
+    }
+
     private Map<String, MessageAttributeValue> updateMessageAttributePayloadSize(
         Map<String, MessageAttributeValue> messageAttributes, Long messageContentSize) {
         Map<String, MessageAttributeValue> updatedMessageAttributes = new HashMap<>(messageAttributes);

src/main/java/com/amazon/sqs/javamessaging/ExtendedClientConfiguration.java

@@ -18,6 +18,7 @@
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
+import software.amazon.awssdk.utils.StringUtils;
 import software.amazon.payloadoffloading.PayloadStorageConfiguration;
 import software.amazon.payloadoffloading.ServerSideEncryptionStrategy;
 
@@ -32,6 +33,7 @@ public class ExtendedClientConfiguration extends PayloadStorageConfiguration {
     private boolean cleanupS3Payload = true;
     private boolean useLegacyReservedAttributeName = true;
     private boolean ignorePayloadNotFound = false;
+    private String s3KeyPrefix;
 
     public ExtendedClientConfiguration() {
         super();
@@ -43,6 +45,7 @@ public ExtendedClientConfiguration(ExtendedClientConfiguration other) {
         this.cleanupS3Payload = other.doesCleanupS3Payload();
         this.useLegacyReservedAttributeName = other.usesLegacyReservedAttributeName();
         this.ignorePayloadNotFound = other.ignoresPayloadNotFound();
+        this.s3KeyPrefix = other.s3KeyPrefix;
     }
 
     /**
@@ -128,6 +131,37 @@ public ExtendedClientConfiguration withIgnorePayloadNotFound(boolean ignorePaylo
         return this;
     }
 
+    /**
+     * Sets a string that will be used as prefix of the S3 Key.
+     *
+     * @param s3KeyPrefix
+     *             A S3 key prefix value
+     */
+    public void setS3KeyPrefix(String s3KeyPrefix) {
+        this.s3KeyPrefix = StringUtils.trim(s3KeyPrefix);
+    }
+
+    /**
+     * Sets a string that will be used as prefix of the S3 Key.
+     *
+     * @param s3KeyPrefix
+     *             A S3 key prefix value
+     *
+     * @return the updated ExtendedClientConfiguration object.
+     */
+    public ExtendedClientConfiguration withS3KeyPrefix(String s3KeyPrefix) {
+        setS3KeyPrefix(s3KeyPrefix);
+        return this;
+    }
+
+    /**
+     * Gets the S3 key prefix
+     * @return the prefix value which is being used for compose the S3 key.
+     */
+    public String getS3KeyPrefix() {
+        return this.s3KeyPrefix;
+    }
+
     /**
      * Checks whether or not clean up large objects in S3 is enabled.
      *

src/test/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedClientTest.java

@@ -15,9 +15,11 @@
 
 package com.amazon.sqs.javamessaging;
 
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.ResponseInputStream;
@@ -62,9 +64,11 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.isA;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -82,11 +86,16 @@ public class AmazonSQSExtendedClientTest {
     private SqsClient extendedSqsWithDefaultKMS;
     private SqsClient extendedSqsWithGenericReservedAttributeName;
     private SqsClient extendedSqsWithDeprecatedMethods;
+    private SqsClient extendedSqsWithS3KeyPrefix;
     private SqsClient mockSqsBackend;
     private S3Client mockS3;
+
+    private MockedStatic<UUID> uuidMockStatic;
     private static final String S3_BUCKET_NAME = "test-bucket-name";
     private static final String SQS_QUEUE_URL = "test-queue-url";
     private static final String S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID = "test-customer-managed-kms-key-id";
+    private static final String S3_KEY_PREFIX = "test-s3-key-prefix";
+    private static final String S3_KEY_UUID = "test-s3-key-uuid";
 
     private static final int LESS_THAN_SQS_SIZE_LIMIT = 3;
     private static final int SQS_SIZE_LIMIT = 262144;
@@ -101,6 +110,7 @@ public class AmazonSQSExtendedClientTest {
 
     @BeforeEach
     public void setupClients() {
+        uuidMockStatic = mockStatic(UUID.class);
         mockS3 = mock(S3Client.class);
         mockSqsBackend = mock(SqsClient.class);
         when(mockS3.putObject(isA(PutObjectRequest.class), isA(RequestBody.class))).thenReturn(null);
@@ -121,11 +131,25 @@ public void setupClients() {
 
         ExtendedClientConfiguration extendedClientConfigurationDeprecated = new ExtendedClientConfiguration().withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
 
+        ExtendedClientConfiguration extendedClientConfigurationWithS3KeyPrefix = new ExtendedClientConfiguration()
+                .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
+                .withS3KeyPrefix(S3_KEY_PREFIX);
+
+        UUID uuidMock = mock(UUID.class);
+        when(uuidMock.toString()).thenReturn(S3_KEY_UUID);
+        uuidMockStatic.when(UUID::randomUUID).thenReturn(uuidMock);
+
         extendedSqsWithDefaultConfig = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
         extendedSqsWithCustomKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithCustomKMS));
         extendedSqsWithDefaultKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithDefaultKMS));
         extendedSqsWithGenericReservedAttributeName = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithGenericReservedAttributeName));
         extendedSqsWithDeprecatedMethods = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationDeprecated));
+        extendedSqsWithS3KeyPrefix = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithS3KeyPrefix));
+    }
+
+    @AfterEach
+    public void tearDown() {
+        uuidMockStatic.close();
     }
 
     @Test
@@ -617,6 +641,32 @@ public void testWhenSendMessageWIthCannedAccessControlListDefined() {
         assertEquals(expected, captor.getValue().acl());
     }
 
+    @Test
+    public void testWhenSendLargeMessageWithS3PrefixKeyDefined() {
+        String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
+
+        SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
+
+        extendedSqsWithS3KeyPrefix.sendMessage(messageRequest);
+
+        verify(mockS3, times(1)).putObject(
+            argThat((PutObjectRequest obj) -> obj.key().equals(S3_KEY_PREFIX + S3_KEY_UUID)),
+            isA(RequestBody.class));
+    }
+
+    @Test
+    public void testWhenSendLargeMessageWithUndefinedS3PrefixKey() {
+        String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
+
+        SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
+
+        extendedSqsWithDefaultConfig.sendMessage(messageRequest);
+
+        verify(mockS3, times(1)).putObject(
+            argThat((PutObjectRequest obj) -> obj.key().equals(S3_KEY_UUID)),
+            isA(RequestBody.class));
+    }
+
     private void testReceiveMessage_when_MessageIsLarge(String reservedAttributeName) {
         String pointer = new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson();
         Message message = Message.builder()

src/test/java/com/amazon/sqs/javamessaging/ExtendedClientConfigurationTest.java

@@ -37,6 +37,7 @@ public class ExtendedClientConfigurationTest {
     private static final String s3BucketName = "test-bucket-name";
     private static final String s3ServerSideEncryptionKMSKeyId = "test-customer-managed-kms-key-id";
     private static final ServerSideEncryptionStrategy serverSideEncryptionStrategy = ServerSideEncryptionFactory.customerKey(s3ServerSideEncryptionKMSKeyId);
+    private static final String s3KeyPrefix = "test-s3-key-prefix";
 
     @Test
     public void testCopyConstructor() {
@@ -164,4 +165,22 @@ public void testMessageSizeThreshold() {
         assertEquals(messageLength, extendedClientConfiguration.getPayloadSizeThreshold());
 
     }
+
+    @Test
+    public void testS3keyPrefix() {
+        ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
+
+        extendedClientConfiguration.withS3KeyPrefix(s3KeyPrefix);
+
+        assertEquals(s3KeyPrefix, extendedClientConfiguration.getS3KeyPrefix());
+    }
+
+    @Test
+    public void testTrimS3keyPrefix() {
+        ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
+
+        extendedClientConfiguration.withS3KeyPrefix(String.format("   %s  ", s3KeyPrefix));
+
+        assertEquals(s3KeyPrefix, extendedClientConfiguration.getS3KeyPrefix());
+    }
 }