Merge branch 'main' into sdkv2

Author: agummi-amazon

pom.xml

@@ -65,7 +65,7 @@
         <dependency>
             <groupId>junit</groupId>
             <artifactId>junit</artifactId>
-            <version>4.12</version>
+            <version>4.13.1</version>
             <scope>test</scope>
         </dependency>
         <dependency>
@@ -186,4 +186,4 @@
             </build>
         </profile>
     </profiles>
-</project>
+</project>

src/main/java/software/amazon/sns/AmazonSNSExtendedClient.java

@@ -47,6 +47,28 @@ public AmazonSNSExtendedClient(SnsClient snsClient, SNSExtendedClientConfigurati
         this.payloadStore = new S3BackedPayloadStore(s3Dao, this.snsExtendedClientConfiguration.getS3BucketName());
     }
 
+    /**
+     * Constructs a new Amazon SNS extended client to invoke service methods on
+     * Amazon SNS with extended functionality using the specified Amazon SNS
+     * client object and Payload Store object.
+     * <p>
+     * <p>
+     * All service calls made using this new client object are blocking, and
+     * will not return until the service call completes.
+     *
+     * @param snsClient                   The Amazon SNS client to use to connect to Amazon SNS.
+     * @param snsExtendedClientConfiguration The sns extended client configuration options controlling the
+     *                                    functionality of this client.
+     * @param payloadStore                The Payload Store that handles logic for saving to the desired
+     *                                    extended storage.
+     */
+    public AmazonSNSExtendedClient(AmazonSNS snsClient, SNSExtendedClientConfiguration snsExtendedClientConfiguration, PayloadStore payloadStore) {
+        super(snsClient);
+
+        this.snsExtendedClientConfiguration = snsExtendedClientConfiguration;
+        this.payloadStore = payloadStore;
+    }
+
     /**
      * <p>
      * Sends a message to an Amazon SNS topic, a text message (SMS message) directly to a phone number, or a message to
@@ -212,9 +234,18 @@ private long getMessageAttributeSize(String MessageAttributeKey, MessageAttribut
         return messageAttributeSize;
     }
 
+    private static String getS3keyAttribute(Map<String, MessageAttributeValue> messageAttributes) {
+        if (messageAttributes != null && messageAttributes.containsKey(S3_KEY)) {
+            MessageAttributeValue attributeS3KeyValue = messageAttributes.get(S3_KEY);
+            return (attributeS3KeyValue == null) ? null : attributeS3KeyValue.getStringValue();
+        }
+        return null;
+    }
+
     private PublishRequest storeMessageInExtendedStore(PublishRequest publishRequest, long messageAttributeSize) {
         String messageContentStr = publishRequest.message();
         Long messageContentSize = Util.getStringSizeInBytes(messageContentStr);
+        String s3Key = getS3keyAttribute(publishRequest.getMessageAttributes()) ;
 
         PublishRequest.Builder publishRequestBuilder = publishRequest.toBuilder();
         String largeMessagePointer = payloadStore.storeOriginalPayload(messageContentStr);
@@ -225,6 +256,7 @@ private PublishRequest storeMessageInExtendedStore(PublishRequest publishRequest
         messageAttributeValueBuilder.stringValue(messageContentSize.toString());
         MessageAttributeValue messageAttributeValue = messageAttributeValueBuilder.build();
 
+
         Map<String, MessageAttributeValue> attributes = new HashMap<>(publishRequest.messageAttributes());
         attributes.put(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, messageAttributeValue);
         publishRequestBuilder.messageAttributes(attributes);

src/test/java/software/amazon/sns/AmazonSNSExtendedClientTest.java

@@ -69,6 +69,26 @@ public void testPublishLargeMessageS3IsUsed() {
         Assert.assertEquals(messageBody.length(), (int) Integer.valueOf(attributes.get(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME).stringValue()));
     }
 
+    @Test
+    public void testPublishLargeMessageS3IsUsedWithS3Key() {
+        String messageBody = generateStringWithLength(MORE_THAN_SNS_SIZE_LIMIT);
+
+        PublishRequest publishRequest = new PublishRequest(SNS_TOPIC_ARN, messageBody);
+        HashMap<String, MessageAttributeValue> attrs = new HashMap<>();
+        attrs.put("S3Key", new MessageAttributeValue().withStringValue("value"));
+        publishRequest.setMessageAttributes(attrs);
+        extendedSnsWithDefaultConfig.publish(publishRequest);
+
+        verify(mockS3, times(1)).putObject(any(PutObjectRequest.class));
+        ArgumentCaptor<PublishRequest> publishRequestCaptor = ArgumentCaptor.forClass(PublishRequest.class);
+        verify(mockSnsBackend, times(1)).publish(publishRequestCaptor.capture());
+
+        Map<String, MessageAttributeValue> attributes = publishRequestCaptor.getValue().getMessageAttributes();
+        Assert.assertEquals("Number", attributes.get(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME).getDataType());
+
+        Assert.assertEquals(messageBody.length(), (int) Integer.valueOf(attributes.get(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME).getStringValue()));
+    }
+
     @Test
     public void testPublishSmallMessageS3IsNotUsed() {
         String messageBody = generateStringWithLength(SNSExtendedClientConfiguration.SNS_DEFAULT_MESSAGE_SIZE);