entry : msgAttributes.entrySet()) {
150 | totalMsgAttributesSize += Util.getStringSizeInBytes(entry.getKey());
151 |
152 | MessageAttributeValue entryVal = entry.getValue();
153 | if (entryVal.dataType() != null) {
154 | totalMsgAttributesSize += Util.getStringSizeInBytes(entryVal.dataType());
155 | }
156 |
157 | String stringVal = entryVal.stringValue();
158 | if (stringVal != null) {
159 | totalMsgAttributesSize += Util.getStringSizeInBytes(entryVal.stringValue());
160 | }
161 |
162 | SdkBytes binaryVal = entryVal.binaryValue();
163 | if (binaryVal != null) {
164 | totalMsgAttributesSize += binaryVal.asByteArray().length;
165 | }
166 | }
167 | return totalMsgAttributesSize;
168 | }
169 |
170 | public static String trimAndValidateS3KeyPrefix(String s3KeyPrefix) {
171 | String trimmedPrefix = StringUtils.trimToEmpty(s3KeyPrefix);
172 |
173 | if (trimmedPrefix.length() > SQSExtendedClientConstants.MAX_S3_KEY_PREFIX_LENGTH) {
174 | String errorMessage = "The S3 key prefix length must not be greater than "
175 | + SQSExtendedClientConstants.MAX_S3_KEY_PREFIX_LENGTH;
176 | LOG.error(errorMessage);
177 | throw SdkClientException.create(errorMessage);
178 | }
179 |
180 | if (trimmedPrefix.startsWith(".") || trimmedPrefix.startsWith("/")) {
181 | String errorMessage = "The S3 key prefix must not starts with '.' or '/'";
182 | LOG.error(errorMessage);
183 | throw SdkClientException.create(errorMessage);
184 | }
185 |
186 | if (trimmedPrefix.contains("..")) {
187 | String errorMessage = "The S3 key prefix must not contains the string '..'";
188 | LOG.error(errorMessage);
189 | throw SdkClientException.create(errorMessage);
190 | }
191 |
192 | if (SQSExtendedClientConstants.INVALID_S3_PREFIX_KEY_CHARACTERS_PATTERN.matcher(trimmedPrefix).find()) {
193 | String errorMessage = "The S3 key prefix contain invalid characters. The allowed characters are: letters, digits, '/', '_', '-', and '.'";
194 | LOG.error(errorMessage);
195 | throw SdkClientException.create(errorMessage);
196 | }
197 |
198 | return trimmedPrefix;
199 | }
200 | }
201 |
--------------------------------------------------------------------------------
/src/main/java/com/amazon/sqs/javamessaging/ExtendedClientConfiguration.java:
--------------------------------------------------------------------------------
1 | /*
2 | * Copyright 2010-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 | *
4 | * Licensed under the Apache License, Version 2.0 (the "License").
5 | * You may not use this file except in compliance with the License.
6 | * A copy of the License is located at
7 | *
8 | * http://aws.amazon.com/apache2.0
9 | *
10 | * or in the "license" file accompanying this file. This file is distributed
11 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12 | * express or implied. See the License for the specific language governing
13 | * permissions and limitations under the License.
14 | */
15 |
16 | package com.amazon.sqs.javamessaging;
17 |
18 | import org.apache.commons.logging.Log;
19 | import org.apache.commons.logging.LogFactory;
20 | import software.amazon.awssdk.annotations.NotThreadSafe;
21 | import software.amazon.awssdk.services.s3.S3Client;
22 | import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
23 | import software.amazon.payloadoffloading.PayloadStorageConfiguration;
24 | import software.amazon.payloadoffloading.ServerSideEncryptionStrategy;
25 |
26 |
27 | /**
28 | * Amazon SQS extended client configuration options such as Amazon S3 client,
29 | * bucket name, and message size threshold for large-payload messages.
30 | */
31 | @NotThreadSafe
32 | public class ExtendedClientConfiguration extends PayloadStorageConfiguration {
33 | private static final Log LOG = LogFactory.getLog(ExtendedClientConfiguration.class);
34 |
35 | private boolean cleanupS3Payload = true;
36 | private boolean useLegacyReservedAttributeName = true;
37 | private boolean ignorePayloadNotFound = false;
38 | private String s3KeyPrefix = "";
39 |
40 | public ExtendedClientConfiguration() {
41 | super();
42 | this.setPayloadSizeThreshold(SQSExtendedClientConstants.DEFAULT_MESSAGE_SIZE_THRESHOLD);
43 | }
44 |
45 | public ExtendedClientConfiguration(ExtendedClientConfiguration other) {
46 | super(other);
47 | this.cleanupS3Payload = other.doesCleanupS3Payload();
48 | this.useLegacyReservedAttributeName = other.usesLegacyReservedAttributeName();
49 | this.ignorePayloadNotFound = other.ignoresPayloadNotFound();
50 | this.s3KeyPrefix = other.s3KeyPrefix;
51 | }
52 |
53 | /**
54 | * Enables support for payload messages.
55 | * @param s3
56 | * Amazon S3 client which is going to be used for storing
57 | * payload messages.
58 | * @param s3BucketName
59 | * Name of the bucket which is going to be used for storing
60 | * payload messages. The bucket must be already created and
61 | * configured in s3.
62 | * @param cleanupS3Payload
63 | * If set to true, would handle deleting the S3 object as part
64 | * of deleting the message from SQS queue. Otherwise, would not
65 | * attempt to delete the object from S3. If opted to not delete S3
66 | * objects its the responsibility to the message producer to handle
67 | * the clean up appropriately.
68 | */
69 | public void setPayloadSupportEnabled(S3Client s3, String s3BucketName, boolean cleanupS3Payload) {
70 | setPayloadSupportEnabled(s3, s3BucketName);
71 | this.cleanupS3Payload = cleanupS3Payload;
72 | }
73 |
74 | /**
75 | * Enables support for payload messages.
76 | * @param s3
77 | * Amazon S3 client which is going to be used for storing
78 | * payload messages.
79 | * @param s3BucketName
80 | * Name of the bucket which is going to be used for storing
81 | * payload messages. The bucket must be already created and
82 | * configured in s3.
83 | * @param cleanupS3Payload
84 | * If set to true, would handle deleting the S3 object as part
85 | * of deleting the message from SQS queue. Otherwise, would not
86 | * attempt to delete the object from S3. If opted to not delete S3
87 | * objects its the responsibility to the message producer to handle
88 | * the clean up appropriately.
89 | */
90 | public ExtendedClientConfiguration withPayloadSupportEnabled(S3Client s3, String s3BucketName, boolean cleanupS3Payload) {
91 | setPayloadSupportEnabled(s3, s3BucketName, cleanupS3Payload);
92 | return this;
93 | }
94 |
95 | /**
96 | * Disables the utilization legacy payload attribute name when sending messages.
97 | */
98 | public void setLegacyReservedAttributeNameDisabled() {
99 | this.useLegacyReservedAttributeName = false;
100 | }
101 |
102 | /**
103 | * Disables the utilization legacy payload attribute name when sending messages.
104 | */
105 | public ExtendedClientConfiguration withLegacyReservedAttributeNameDisabled() {
106 | setLegacyReservedAttributeNameDisabled();
107 | return this;
108 | }
109 |
110 | /**
111 | * Sets whether or not messages should be removed from Amazon SQS
112 | * when payloads are not found in Amazon S3.
113 | *
114 | * @param ignorePayloadNotFound
115 | * Whether or not messages should be removed from Amazon SQS
116 | * when payloads are not found in Amazon S3. Default: false
117 | */
118 | public void setIgnorePayloadNotFound(boolean ignorePayloadNotFound) {
119 | this.ignorePayloadNotFound = ignorePayloadNotFound;
120 | }
121 |
122 | /**
123 | * Sets whether or not messages should be removed from Amazon SQS
124 | * when payloads are not found in Amazon S3.
125 | *
126 | * @param ignorePayloadNotFound
127 | * Whether or not messages should be removed from Amazon SQS
128 | * when payloads are not found in Amazon S3. Default: false
129 | * @return the updated ExtendedClientConfiguration object.
130 | */
131 | public ExtendedClientConfiguration withIgnorePayloadNotFound(boolean ignorePayloadNotFound) {
132 | setIgnorePayloadNotFound(ignorePayloadNotFound);
133 | return this;
134 | }
135 |
136 | /**
137 | * Sets a string that will be used as prefix of the S3 Key.
138 | *
139 | * @param s3KeyPrefix
140 | * A S3 key prefix value
141 | */
142 | public void setS3KeyPrefix(String s3KeyPrefix) {
143 | this.s3KeyPrefix = AmazonSQSExtendedClientUtil.trimAndValidateS3KeyPrefix(s3KeyPrefix);
144 | }
145 |
146 | /**
147 | * Sets a string that will be used as prefix of the S3 Key.
148 | *
149 | * @param s3KeyPrefix
150 | * A S3 key prefix value
151 | *
152 | * @return the updated ExtendedClientConfiguration object.
153 | */
154 | public ExtendedClientConfiguration withS3KeyPrefix(String s3KeyPrefix) {
155 | setS3KeyPrefix(s3KeyPrefix);
156 | return this;
157 | }
158 |
159 | /**
160 | * Gets the S3 key prefix
161 | * @return the prefix value which is being used for compose the S3 key.
162 | */
163 | public String getS3KeyPrefix() {
164 | return this.s3KeyPrefix;
165 | }
166 |
167 | /**
168 | * Checks whether or not clean up large objects in S3 is enabled.
169 | *
170 | * @return True if clean up is enabled when deleting the concerning SQS message.
171 | * Default: true
172 | */
173 | public boolean doesCleanupS3Payload() {
174 | return cleanupS3Payload;
175 | }
176 |
177 | /**
178 | * Checks whether or not the configuration uses the legacy reserved attribute name.
179 | *
180 | * @return True if legacy reserved attribute name is used.
181 | * Default: true
182 | */
183 |
184 | public boolean usesLegacyReservedAttributeName() {
185 | return useLegacyReservedAttributeName;
186 | }
187 |
188 | /**
189 | * Checks whether or not messages should be removed from Amazon SQS
190 | * when payloads are not found in Amazon S3.
191 | *
192 | * @return True if messages should be removed from Amazon SQS
193 | * when payloads are not found in Amazon S3. Default: false
194 | */
195 | public boolean ignoresPayloadNotFound() {
196 | return ignorePayloadNotFound;
197 | }
198 |
199 | @Override
200 | public ExtendedClientConfiguration withAlwaysThroughS3(boolean alwaysThroughS3) {
201 | setAlwaysThroughS3(alwaysThroughS3);
202 | return this;
203 | }
204 |
205 | @Override
206 | public ExtendedClientConfiguration withPayloadSupportEnabled(S3Client s3, String s3BucketName) {
207 | this.setPayloadSupportEnabled(s3, s3BucketName);
208 | return this;
209 | }
210 |
211 | @Override
212 | public ExtendedClientConfiguration withObjectCannedACL(ObjectCannedACL objectCannedACL) {
213 | this.setObjectCannedACL(objectCannedACL);
214 | return this;
215 | }
216 |
217 | @Override
218 | public ExtendedClientConfiguration withPayloadSizeThreshold(int payloadSizeThreshold) {
219 | this.setPayloadSizeThreshold(payloadSizeThreshold);
220 | return this;
221 | }
222 |
223 | @Override
224 | public ExtendedClientConfiguration withPayloadSupportDisabled() {
225 | this.setPayloadSupportDisabled();
226 | return this;
227 | }
228 |
229 | @Override
230 | public ExtendedClientConfiguration withServerSideEncryption(ServerSideEncryptionStrategy serverSideEncryption) {
231 | this.setServerSideEncryptionStrategy(serverSideEncryption);
232 | return this;
233 | }
234 |
235 | /**
236 | * Enables support for large-payload messages.
237 | *
238 | * @param s3
239 | * Amazon S3 client which is going to be used for storing
240 | * large-payload messages.
241 | * @param s3BucketName
242 | * Name of the bucket which is going to be used for storing
243 | * large-payload messages. The bucket must be already created and
244 | * configured in s3.
245 | *
246 | * @deprecated Instead use {@link #setPayloadSupportEnabled(S3Client, String, boolean)}
247 | */
248 | @Deprecated
249 | public void setLargePayloadSupportEnabled(S3Client s3, String s3BucketName) {
250 | this.setPayloadSupportEnabled(s3, s3BucketName);
251 | }
252 |
253 | /**
254 | * Enables support for large-payload messages.
255 | *
256 | * @param s3
257 | * Amazon S3 client which is going to be used for storing
258 | * large-payload messages.
259 | * @param s3BucketName
260 | * Name of the bucket which is going to be used for storing
261 | * large-payload messages. The bucket must be already created and
262 | * configured in s3.
263 | * @return the updated ExtendedClientConfiguration object.
264 | *
265 | * @deprecated Instead use {@link #withPayloadSupportEnabled(S3Client, String)}
266 | */
267 | @Deprecated
268 | public ExtendedClientConfiguration withLargePayloadSupportEnabled(S3Client s3, String s3BucketName) {
269 | setLargePayloadSupportEnabled(s3, s3BucketName);
270 | return this;
271 | }
272 |
273 | /**
274 | * Disables support for large-payload messages.
275 | *
276 | * @deprecated Instead use {@link #setPayloadSupportDisabled()}
277 | */
278 | @Deprecated
279 | public void setLargePayloadSupportDisabled() {
280 | this.setPayloadSupportDisabled();
281 | }
282 |
283 | /**
284 | * Disables support for large-payload messages.
285 | * @return the updated ExtendedClientConfiguration object.
286 | *
287 | * @deprecated Instead use {@link #withPayloadSupportDisabled()}
288 | */
289 | @Deprecated
290 | public ExtendedClientConfiguration withLargePayloadSupportDisabled() {
291 | setLargePayloadSupportDisabled();
292 | return this;
293 | }
294 |
295 | /**
296 | * Check if the support for large-payload message if enabled.
297 | * @return true if support for large-payload messages is enabled.
298 | *
299 | * @deprecated Instead use {@link #isPayloadSupportEnabled()}
300 | */
301 | @Deprecated
302 | public boolean isLargePayloadSupportEnabled() {
303 | return isPayloadSupportEnabled();
304 | }
305 |
306 | /**
307 | * Sets the message size threshold for storing message payloads in Amazon
308 | * S3.
309 | *
310 | * @param messageSizeThreshold
311 | * Message size threshold to be used for storing in Amazon S3.
312 | * Default: 256KB.
313 | *
314 | * @deprecated Instead use {@link #setPayloadSizeThreshold(int)}
315 | */
316 | @Deprecated
317 | public void setMessageSizeThreshold(int messageSizeThreshold) {
318 | this.setPayloadSizeThreshold(messageSizeThreshold);
319 | }
320 |
321 | /**
322 | * Sets the message size threshold for storing message payloads in Amazon
323 | * S3.
324 | *
325 | * @param messageSizeThreshold
326 | * Message size threshold to be used for storing in Amazon S3.
327 | * Default: 256KB.
328 | * @return the updated ExtendedClientConfiguration object.
329 | *
330 | * @deprecated Instead use {@link #withPayloadSizeThreshold(int)}
331 | */
332 | @Deprecated
333 | public ExtendedClientConfiguration withMessageSizeThreshold(int messageSizeThreshold) {
334 | setMessageSizeThreshold(messageSizeThreshold);
335 | return this;
336 | }
337 |
338 | /**
339 | * Gets the message size threshold for storing message payloads in Amazon
340 | * S3.
341 | *
342 | * @return Message size threshold which is being used for storing in Amazon
343 | * S3. Default: 256KB.
344 | *
345 | * @deprecated Instead use {@link #getPayloadSizeThreshold()}
346 | */
347 | @Deprecated
348 | public int getMessageSizeThreshold() {
349 | return getPayloadSizeThreshold();
350 | }
351 | }
--------------------------------------------------------------------------------
/src/main/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedAsyncClient.java:
--------------------------------------------------------------------------------
1 | package com.amazon.sqs.javamessaging;
2 |
3 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.checkMessageAttributes;
4 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.embedS3PointerInReceiptHandle;
5 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.getMessagePointerFromModifiedReceiptHandle;
6 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.getOrigReceiptHandle;
7 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.getReservedAttributeNameIfPresent;
8 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.isLarge;
9 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.isS3ReceiptHandle;
10 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClientUtil.updateMessageAttributePayloadSize;
11 |
12 | import java.util.ArrayList;
13 | import java.util.HashMap;
14 | import java.util.List;
15 | import java.util.Map;
16 | import java.util.Objects;
17 | import java.util.Optional;
18 | import java.util.UUID;
19 | import java.util.concurrent.CompletableFuture;
20 | import java.util.concurrent.CompletionException;
21 | import java.util.stream.Collectors;
22 | import org.apache.commons.logging.Log;
23 | import org.apache.commons.logging.LogFactory;
24 | import software.amazon.awssdk.awscore.AwsRequest;
25 | import software.amazon.awssdk.core.exception.SdkClientException;
26 | import software.amazon.awssdk.core.util.VersionInfo;
27 | import software.amazon.awssdk.services.sqs.SqsAsyncClient;
28 | import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest;
29 | import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry;
30 | import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse;
31 | import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityRequest;
32 | import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityResponse;
33 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
34 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
35 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
36 | import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest;
37 | import software.amazon.awssdk.services.sqs.model.DeleteMessageResponse;
38 | import software.amazon.awssdk.services.sqs.model.Message;
39 | import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
40 | import software.amazon.awssdk.services.sqs.model.PurgeQueueRequest;
41 | import software.amazon.awssdk.services.sqs.model.PurgeQueueResponse;
42 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
43 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
44 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest;
45 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
46 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchResponse;
47 | import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
48 | import software.amazon.awssdk.services.sqs.model.SendMessageResponse;
49 | import software.amazon.awssdk.utils.StringUtils;
50 | import software.amazon.payloadoffloading.PayloadStoreAsync;
51 | import software.amazon.payloadoffloading.S3AsyncDao;
52 | import software.amazon.payloadoffloading.S3BackedPayloadStoreAsync;
53 | import software.amazon.payloadoffloading.Util;
54 |
55 | /**
56 | * Amazon SQS Extended Async Client extends the functionality of Amazon Async SQS
57 | * client.
58 | *
59 | *
60 | * All service calls made using this client are asynchronous, and will return
61 | * immediately with a {@link CompletableFuture} that completes when the operation
62 | * completes or when an exception is thrown. Argument validation exceptions are thrown
63 | * immediately, and not through the future.
64 | *
65 | *
66 | *
67 | * The Amazon SQS extended client enables sending and receiving large messages
68 | * via Amazon S3. You can use this library to:
69 | *
70 | *
71 | *
72 | * - Specify whether messages are always stored in Amazon S3 or only when a
73 | * message size exceeds 256 KB.
74 | * - Send a message that references a single message object stored in an
75 | * Amazon S3 bucket.
76 | * - Get the corresponding message object from an Amazon S3 bucket.
77 | * - Delete the corresponding message object from an Amazon S3 bucket.
78 | *
79 | */
80 | public class AmazonSQSExtendedAsyncClient extends AmazonSQSExtendedAsyncClientBase implements SqsAsyncClient {
81 | static final String USER_AGENT_NAME = AmazonSQSExtendedAsyncClient.class.getSimpleName();
82 | static final String USER_AGENT_VERSION = VersionInfo.SDK_VERSION;
83 |
84 | private static final Log LOG = LogFactory.getLog(AmazonSQSExtendedAsyncClient.class);
85 | private ExtendedAsyncClientConfiguration clientConfiguration;
86 | private PayloadStoreAsync payloadStore;
87 |
88 | /**
89 | * Constructs a new Amazon SQS extended async client to invoke service methods on
90 | * Amazon SQS with extended functionality using the specified Amazon SQS
91 | * client object.
92 | *
93 | *
94 | * All service calls made using this client are asynchronous, and will return
95 | * immediately with a {@link CompletableFuture} that completes when the operation
96 | * completes or when an exception is thrown. Argument validation exceptions are thrown
97 | * immediately, and not through the future.
98 | *
99 | *
100 | * @param sqsClient
101 | * The Amazon SQS async client to use to connect to Amazon SQS.
102 | */
103 | public AmazonSQSExtendedAsyncClient(SqsAsyncClient sqsClient) {
104 | this(sqsClient, new ExtendedAsyncClientConfiguration());
105 | }
106 |
107 | /**
108 | * Constructs a new Amazon SQS extended client to invoke service methods on
109 | * Amazon SQS with extended functionality using the specified Amazon SQS
110 | * client object.
111 | *
112 | *
113 | * All service calls made using this client are asynchronous, and will return
114 | * immediately with a {@link CompletableFuture} that completes when the operation
115 | * completes or when an exception is thrown. Argument validation exceptions are thrown
116 | * immediately, and not through the future.
117 | *
118 | *
119 | * @param sqsClient
120 | * The Amazon SQS async client to use to connect to Amazon SQS.
121 | * @param extendedClientConfig
122 | * The extended client configuration options controlling the
123 | * functionality of this client.
124 | */
125 | public AmazonSQSExtendedAsyncClient(SqsAsyncClient sqsClient,
126 | ExtendedAsyncClientConfiguration extendedClientConfig) {
127 | super(sqsClient);
128 | this.clientConfiguration = new ExtendedAsyncClientConfiguration(extendedClientConfig);
129 | S3AsyncDao s3Dao = new S3AsyncDao(clientConfiguration.getS3AsyncClient(),
130 | clientConfiguration.getServerSideEncryptionStrategy(),
131 | clientConfiguration.getObjectCannedACL());
132 | this.payloadStore = new S3BackedPayloadStoreAsync(s3Dao, clientConfiguration.getS3BucketName());
133 | }
134 |
135 | /**
136 | * {@inheritDoc}
137 | */
138 | @Override
139 | public CompletableFuture sendMessage(SendMessageRequest sendMessageRequest) {
140 | // TODO: Clone request since it's modified in this method and will cause issues if the client reuses request
141 | // object.
142 | if (sendMessageRequest == null) {
143 | String errorMessage = "sendMessageRequest cannot be null.";
144 | LOG.error(errorMessage);
145 | throw SdkClientException.create(errorMessage);
146 | }
147 |
148 | SendMessageRequest.Builder sendMessageRequestBuilder = sendMessageRequest.toBuilder();
149 | sendMessageRequest = appendUserAgent(sendMessageRequestBuilder).build();
150 |
151 | if (!clientConfiguration.isPayloadSupportEnabled()) {
152 | return super.sendMessage(sendMessageRequest);
153 | }
154 |
155 | if (StringUtils.isEmpty(sendMessageRequest.messageBody())) {
156 | String errorMessage = "messageBody cannot be null or empty.";
157 | LOG.error(errorMessage);
158 | throw SdkClientException.create(errorMessage);
159 | }
160 |
161 | //Check message attributes for ExtendedClient related constraints
162 | checkMessageAttributes(clientConfiguration.getPayloadSizeThreshold(), sendMessageRequest.messageAttributes());
163 |
164 | if (clientConfiguration.isAlwaysThroughS3()
165 | || isLarge(clientConfiguration.getPayloadSizeThreshold(), sendMessageRequest)) {
166 | return storeMessageInS3(sendMessageRequest)
167 | .thenCompose(modifiedRequest -> super.sendMessage(modifiedRequest));
168 | }
169 |
170 | return super.sendMessage(sendMessageRequest);
171 | }
172 |
173 | /**
174 | * {@inheritDoc}
175 | */
176 | @Override
177 | public CompletableFuture receiveMessage(ReceiveMessageRequest receiveMessageRequest) {
178 | // TODO: Clone request since it's modified in this method and will cause issues if the client reuses request
179 | // object.
180 | if (receiveMessageRequest == null) {
181 | String errorMessage = "receiveMessageRequest cannot be null.";
182 | LOG.error(errorMessage);
183 | throw SdkClientException.create(errorMessage);
184 | }
185 |
186 | ReceiveMessageRequest.Builder receiveMessageRequestBuilder = receiveMessageRequest.toBuilder();
187 | appendUserAgent(receiveMessageRequestBuilder);
188 |
189 | if (!clientConfiguration.isPayloadSupportEnabled()) {
190 | return super.receiveMessage(receiveMessageRequestBuilder.build());
191 | }
192 |
193 | // Remove before adding to avoid any duplicates
194 | List messageAttributeNames = new ArrayList<>(receiveMessageRequest.messageAttributeNames());
195 | messageAttributeNames.removeAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES);
196 | messageAttributeNames.addAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES);
197 | receiveMessageRequestBuilder.messageAttributeNames(messageAttributeNames);
198 | String queueUrl = receiveMessageRequest.queueUrl();
199 | receiveMessageRequest = receiveMessageRequestBuilder.build();
200 |
201 | return super.receiveMessage(receiveMessageRequest)
202 | .thenCompose(receiveMessageResponse -> {
203 | List messages = receiveMessageResponse.messages();
204 |
205 | // Check for no messages. If so, no need to process further.
206 | if (messages.isEmpty()) {
207 | return CompletableFuture.completedFuture(messages);
208 | }
209 |
210 | List> modifiedMessageFutures = new ArrayList<>(messages.size());
211 | for (Message message : messages) {
212 | Message.Builder messageBuilder = message.toBuilder();
213 |
214 | // For each received message check if they are stored in S3.
215 | Optional largePayloadAttributeName = getReservedAttributeNameIfPresent(
216 | message.messageAttributes());
217 | if (!largePayloadAttributeName.isPresent()) {
218 | // Not S3
219 | modifiedMessageFutures.add(CompletableFuture.completedFuture(messageBuilder.build()));
220 | } else {
221 | // In S3
222 | final String largeMessagePointer = message.body()
223 | .replace("com.amazon.sqs.javamessaging.MessageS3Pointer",
224 | "software.amazon.payloadoffloading.PayloadS3Pointer");
225 |
226 | // Retrieve original payload
227 | modifiedMessageFutures.add(payloadStore.getOriginalPayload(largeMessagePointer)
228 | .handle((originalPayload,throwable) -> {
229 |
230 | if(throwable != null)
231 | {
232 | if(clientConfiguration.ignoresPayloadNotFound())
233 | {
234 | DeleteMessageRequest deleteMessageRequest = DeleteMessageRequest
235 | .builder()
236 | .queueUrl(queueUrl)
237 | .receiptHandle(message.receiptHandle())
238 | .build();
239 |
240 | deleteMessage(deleteMessageRequest).join();
241 | LOG.warn("Message deleted from SQS since payload with pointer could not be found in S3.");
242 | return null;
243 | }
244 | else
245 | {
246 | throw new CompletionException(throwable);
247 | }
248 | }
249 |
250 | // Set original payload
251 | messageBuilder.body(originalPayload);
252 |
253 | // Remove the additional attribute before returning the message
254 | // to user.
255 | Map messageAttributes = new HashMap<>(
256 | message.messageAttributes());
257 | messageAttributes.keySet().removeAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES);
258 | messageBuilder.messageAttributes(messageAttributes);
259 |
260 | // Embed s3 object pointer in the receipt handle.
261 | String modifiedReceiptHandle = embedS3PointerInReceiptHandle(
262 | message.receiptHandle(),
263 | largeMessagePointer);
264 | messageBuilder.receiptHandle(modifiedReceiptHandle);
265 |
266 | return messageBuilder.build();
267 | }));
268 | }
269 | }
270 |
271 | // Convert list of message futures to a future list of messages.
272 | return CompletableFuture.allOf(
273 | modifiedMessageFutures.toArray(new CompletableFuture[modifiedMessageFutures.size()]))
274 | .thenApply(v -> modifiedMessageFutures.stream()
275 | .map(CompletableFuture::join)
276 | .filter(Objects::nonNull)
277 | .collect(Collectors.toList()));
278 | })
279 | .thenApply(modifiedMessages -> {
280 | // Build response with modified message list.
281 | ReceiveMessageResponse.Builder receiveMessageResponseBuilder = ReceiveMessageResponse.builder();
282 | receiveMessageResponseBuilder.messages(modifiedMessages);
283 | return receiveMessageResponseBuilder.build();
284 | });
285 | }
286 |
287 | /**
288 | * {@inheritDoc}
289 | */
290 | @Override
291 | public CompletableFuture deleteMessage(DeleteMessageRequest deleteMessageRequest) {
292 | if (deleteMessageRequest == null) {
293 | String errorMessage = "deleteMessageRequest cannot be null.";
294 | LOG.error(errorMessage);
295 | throw SdkClientException.create(errorMessage);
296 | }
297 |
298 | DeleteMessageRequest.Builder deleteMessageRequestBuilder = deleteMessageRequest.toBuilder();
299 | appendUserAgent(deleteMessageRequestBuilder);
300 |
301 | String receiptHandle = deleteMessageRequest.receiptHandle();
302 | String origReceiptHandle = receiptHandle;
303 | String messagePointer = null;
304 |
305 | // Update original receipt handle if needed.
306 | if (clientConfiguration.isPayloadSupportEnabled() && isS3ReceiptHandle(receiptHandle)) {
307 | origReceiptHandle = getOrigReceiptHandle(receiptHandle);
308 |
309 | // Delete pay load from S3 if needed
310 | if (clientConfiguration.doesCleanupS3Payload()) {
311 | messagePointer = getMessagePointerFromModifiedReceiptHandle(receiptHandle);
312 | }
313 | }
314 |
315 | // The actual message to delete from SQS.
316 | deleteMessageRequestBuilder.receiptHandle(origReceiptHandle);
317 |
318 | // Check if message is in S3 or only in SQS.
319 | if (messagePointer == null) {
320 | // Delete only from SQS
321 | return super.deleteMessage(deleteMessageRequestBuilder.build());
322 | }
323 |
324 | // Delete from SQS first, then S3.
325 | final String messageToDeletePointer = messagePointer;
326 | return super.deleteMessage(deleteMessageRequestBuilder.build())
327 | .thenCompose(deleteMessageResponse ->
328 | payloadStore.deleteOriginalPayload(messageToDeletePointer)
329 | .thenApply(v -> deleteMessageResponse));
330 | }
331 |
332 | /**
333 | * {@inheritDoc}
334 | */
335 | @Override
336 | public CompletableFuture changeMessageVisibility(
337 | ChangeMessageVisibilityRequest changeMessageVisibilityRequest) {
338 |
339 | ChangeMessageVisibilityRequest.Builder changeMessageVisibilityRequestBuilder =
340 | changeMessageVisibilityRequest.toBuilder();
341 | if (isS3ReceiptHandle(changeMessageVisibilityRequest.receiptHandle())) {
342 | changeMessageVisibilityRequestBuilder.receiptHandle(
343 | getOrigReceiptHandle(changeMessageVisibilityRequest.receiptHandle()));
344 | }
345 | return amazonSqsToBeExtended.changeMessageVisibility(changeMessageVisibilityRequestBuilder.build());
346 | }
347 |
348 | /**
349 | * {@inheritDoc}
350 | */
351 | @Override
352 | public CompletableFuture sendMessageBatch(
353 | SendMessageBatchRequest sendMessageBatchRequestIn) {
354 |
355 | if (sendMessageBatchRequestIn == null) {
356 | String errorMessage = "sendMessageBatchRequest cannot be null.";
357 | LOG.error(errorMessage);
358 | throw SdkClientException.create(errorMessage);
359 | }
360 |
361 | SendMessageBatchRequest.Builder sendMessageBatchRequestBuilder = sendMessageBatchRequestIn.toBuilder();
362 | appendUserAgent(sendMessageBatchRequestBuilder);
363 | SendMessageBatchRequest sendMessageBatchRequest = sendMessageBatchRequestBuilder.build();
364 |
365 | if (!clientConfiguration.isPayloadSupportEnabled()) {
366 | return super.sendMessageBatch(sendMessageBatchRequest);
367 | }
368 |
369 | List> batchEntryFutures = new ArrayList<>(
370 | sendMessageBatchRequest.entries().size());
371 | boolean hasS3Entries = false;
372 | for (SendMessageBatchRequestEntry entry : sendMessageBatchRequest.entries()) {
373 | //Check message attributes for ExtendedClient related constraints
374 | checkMessageAttributes(clientConfiguration.getPayloadSizeThreshold(), entry.messageAttributes());
375 |
376 | if (clientConfiguration.isAlwaysThroughS3()
377 | || isLarge(clientConfiguration.getPayloadSizeThreshold(), entry)) {
378 | batchEntryFutures.add(storeMessageInS3(entry));
379 | hasS3Entries = true;
380 | } else {
381 | batchEntryFutures.add(CompletableFuture.completedFuture(entry));
382 | }
383 | }
384 |
385 | if (!hasS3Entries) {
386 | return super.sendMessageBatch(sendMessageBatchRequest);
387 | }
388 |
389 | // Convert list of entry futures to a future list of entries.
390 | return CompletableFuture.allOf(
391 | batchEntryFutures.toArray(new CompletableFuture[batchEntryFutures.size()]))
392 | .thenApply(v -> batchEntryFutures.stream()
393 | .map(CompletableFuture::join)
394 | .collect(Collectors.toList()))
395 | .thenCompose(batchEntries -> {
396 | SendMessageBatchRequest modifiedBatchRequest =
397 | sendMessageBatchRequest.toBuilder().entries(batchEntries).build();
398 | return super.sendMessageBatch(modifiedBatchRequest);
399 | });
400 | }
401 |
402 | /**
403 | * {@inheritDoc}
404 | */
405 | @Override
406 | public CompletableFuture deleteMessageBatch(
407 | DeleteMessageBatchRequest deleteMessageBatchRequest) {
408 |
409 | if (deleteMessageBatchRequest == null) {
410 | String errorMessage = "deleteMessageBatchRequest cannot be null.";
411 | LOG.error(errorMessage);
412 | throw SdkClientException.create(errorMessage);
413 | }
414 |
415 | DeleteMessageBatchRequest.Builder deleteMessageBatchRequestBuilder = deleteMessageBatchRequest.toBuilder();
416 | appendUserAgent(deleteMessageBatchRequestBuilder);
417 |
418 | if (!clientConfiguration.isPayloadSupportEnabled()) {
419 | return super.deleteMessageBatch(deleteMessageBatchRequest);
420 | }
421 |
422 | List entries = new ArrayList<>(deleteMessageBatchRequest.entries().size());
423 | for (DeleteMessageBatchRequestEntry entry : deleteMessageBatchRequest.entries()) {
424 | DeleteMessageBatchRequestEntry.Builder entryBuilder = entry.toBuilder();
425 | String receiptHandle = entry.receiptHandle();
426 | String origReceiptHandle = receiptHandle;
427 |
428 | // Update original receipt handle if needed
429 | if (isS3ReceiptHandle(receiptHandle)) {
430 | origReceiptHandle = getOrigReceiptHandle(receiptHandle);
431 | // Delete s3 payload if needed
432 | if (clientConfiguration.doesCleanupS3Payload()) {
433 | String messagePointer = getMessagePointerFromModifiedReceiptHandle(receiptHandle);
434 | payloadStore.deleteOriginalPayload(messagePointer);
435 | }
436 | }
437 |
438 | entryBuilder.receiptHandle(origReceiptHandle);
439 | entries.add(entryBuilder.build());
440 | }
441 |
442 | deleteMessageBatchRequestBuilder.entries(entries);
443 | return super.deleteMessageBatch(deleteMessageBatchRequestBuilder.build());
444 | }
445 |
446 | /**
447 | * {@inheritDoc}
448 | */
449 | @Override
450 | public CompletableFuture changeMessageVisibilityBatch(
451 | ChangeMessageVisibilityBatchRequest changeMessageVisibilityBatchRequest) {
452 | List entries = new ArrayList<>(
453 | changeMessageVisibilityBatchRequest.entries().size());
454 | for (ChangeMessageVisibilityBatchRequestEntry entry : changeMessageVisibilityBatchRequest.entries()) {
455 | ChangeMessageVisibilityBatchRequestEntry.Builder entryBuilder = entry.toBuilder();
456 | if (isS3ReceiptHandle(entry.receiptHandle())) {
457 | entryBuilder.receiptHandle(getOrigReceiptHandle(entry.receiptHandle()));
458 | }
459 | entries.add(entryBuilder.build());
460 | }
461 |
462 | return amazonSqsToBeExtended.changeMessageVisibilityBatch(
463 | changeMessageVisibilityBatchRequest.toBuilder().entries(entries).build());
464 | }
465 |
466 | /**
467 | * {@inheritDoc}
468 | */
469 | @Override
470 | public CompletableFuture purgeQueue(PurgeQueueRequest purgeQueueRequest) {
471 | LOG.warn("Calling purgeQueue deletes SQS messages without deleting their payload from S3.");
472 |
473 | if (purgeQueueRequest == null) {
474 | String errorMessage = "purgeQueueRequest cannot be null.";
475 | LOG.error(errorMessage);
476 | throw SdkClientException.create(errorMessage);
477 | }
478 |
479 | PurgeQueueRequest.Builder purgeQueueRequestBuilder = purgeQueueRequest.toBuilder();
480 | appendUserAgent(purgeQueueRequestBuilder);
481 |
482 | return super.purgeQueue(purgeQueueRequestBuilder.build());
483 | }
484 |
485 | private CompletableFuture storeMessageInS3(SendMessageBatchRequestEntry batchEntry) {
486 | // Read the content of the message from message body
487 | String messageContentStr = batchEntry.messageBody();
488 |
489 | Long messageContentSize = Util.getStringSizeInBytes(messageContentStr);
490 |
491 | SendMessageBatchRequestEntry.Builder batchEntryBuilder = batchEntry.toBuilder();
492 |
493 | batchEntryBuilder.messageAttributes(
494 | updateMessageAttributePayloadSize(batchEntry.messageAttributes(), messageContentSize,
495 | clientConfiguration.usesLegacyReservedAttributeName()));
496 |
497 | // Store the message content in S3.
498 | return storeOriginalPayload(messageContentStr)
499 | .thenApply(largeMessagePointer -> {
500 | batchEntryBuilder.messageBody(largeMessagePointer);
501 | return batchEntryBuilder.build();
502 | });
503 | }
504 |
505 | private CompletableFuture storeMessageInS3(SendMessageRequest sendMessageRequest) {
506 | // Read the content of the message from message body
507 | String messageContentStr = sendMessageRequest.messageBody();
508 |
509 | Long messageContentSize = Util.getStringSizeInBytes(messageContentStr);
510 |
511 | SendMessageRequest.Builder sendMessageRequestBuilder = sendMessageRequest.toBuilder();
512 |
513 | sendMessageRequestBuilder.messageAttributes(
514 | updateMessageAttributePayloadSize(sendMessageRequest.messageAttributes(), messageContentSize,
515 | clientConfiguration.usesLegacyReservedAttributeName()));
516 |
517 | // Store the message content in S3.
518 | return storeOriginalPayload(messageContentStr)
519 | .thenApply(largeMessagePointer -> {
520 | sendMessageRequestBuilder.messageBody(largeMessagePointer);
521 | return sendMessageRequestBuilder.build();
522 | });
523 | }
524 |
525 | private CompletableFuture storeOriginalPayload(String messageContentStr) {
526 | String s3KeyPrefix = clientConfiguration.getS3KeyPrefix();
527 | if (StringUtils.isBlank(s3KeyPrefix)) {
528 | return payloadStore.storeOriginalPayload(messageContentStr);
529 | }
530 | return payloadStore.storeOriginalPayload(messageContentStr, s3KeyPrefix + UUID.randomUUID());
531 | }
532 |
533 | private static T appendUserAgent(final T builder) {
534 | return AmazonSQSExtendedClientUtil.appendUserAgent(builder, USER_AGENT_NAME, USER_AGENT_VERSION);
535 | }
536 | }
537 |
--------------------------------------------------------------------------------
/src/test/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedAsyncClientTest.java:
--------------------------------------------------------------------------------
1 | package com.amazon.sqs.javamessaging;
2 |
3 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedAsyncClient.USER_AGENT_NAME;
4 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedAsyncClient.USER_AGENT_VERSION;
5 | import static org.junit.jupiter.api.Assertions.assertEquals;
6 | import static org.junit.jupiter.api.Assertions.assertFalse;
7 | import static org.junit.jupiter.api.Assertions.assertNotEquals;
8 | import static org.junit.jupiter.api.Assertions.assertNull;
9 | import static org.junit.jupiter.api.Assertions.assertTrue;
10 | import static org.junit.jupiter.api.Assertions.fail;
11 | import static org.mockito.ArgumentMatchers.any;
12 | import static org.mockito.ArgumentMatchers.eq;
13 | import static org.mockito.ArgumentMatchers.isA;
14 | import static org.mockito.Mockito.doThrow;
15 | import static org.mockito.Mockito.mock;
16 | import static org.mockito.Mockito.never;
17 | import static org.mockito.Mockito.spy;
18 | import static org.mockito.Mockito.times;
19 | import static org.mockito.Mockito.verify;
20 | import static org.mockito.Mockito.verifyNoInteractions;
21 | import static org.mockito.Mockito.when;
22 |
23 | import java.nio.charset.StandardCharsets;
24 | import java.util.ArrayList;
25 | import java.util.Arrays;
26 | import java.util.List;
27 | import java.util.Map;
28 | import java.util.UUID;
29 | import java.util.concurrent.CompletableFuture;
30 | import java.util.concurrent.CompletionException;
31 | import java.util.stream.Collectors;
32 | import java.util.stream.IntStream;
33 | import org.junit.jupiter.api.BeforeEach;
34 | import org.junit.jupiter.api.Test;
35 | import org.mockito.ArgumentCaptor;
36 | import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
37 | import software.amazon.awssdk.core.ApiName;
38 | import software.amazon.awssdk.core.ResponseBytes;
39 | import software.amazon.awssdk.core.async.AsyncRequestBody;
40 | import software.amazon.awssdk.core.async.AsyncResponseTransformer;
41 | import software.amazon.awssdk.core.exception.SdkException;
42 | import software.amazon.awssdk.services.s3.S3AsyncClient;
43 | import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
44 | import software.amazon.awssdk.services.s3.model.DeleteObjectResponse;
45 | import software.amazon.awssdk.services.s3.model.GetObjectRequest;
46 | import software.amazon.awssdk.services.s3.model.GetObjectResponse;
47 | import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
48 | import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
49 | import software.amazon.awssdk.services.s3.model.PutObjectRequest;
50 | import software.amazon.awssdk.services.sqs.SqsAsyncClient;
51 | import software.amazon.awssdk.services.sqs.SqsClient;
52 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
53 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
54 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse;
55 | import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest;
56 | import software.amazon.awssdk.services.sqs.model.DeleteMessageResponse;
57 | import software.amazon.awssdk.services.sqs.model.Message;
58 | import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
59 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
60 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
61 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest;
62 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
63 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchResponse;
64 | import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
65 | import software.amazon.awssdk.services.sqs.model.SendMessageResponse;
66 | import software.amazon.awssdk.utils.ImmutableMap;
67 | import software.amazon.payloadoffloading.PayloadS3Pointer;
68 | import software.amazon.payloadoffloading.ServerSideEncryptionFactory;
69 | import software.amazon.payloadoffloading.ServerSideEncryptionStrategy;
70 |
71 | public class AmazonSQSExtendedAsyncClientTest {
72 |
73 | private SqsAsyncClient extendedSqsWithDefaultConfig;
74 | private SqsAsyncClient extendedSqsWithCustomKMS;
75 | private SqsAsyncClient extendedSqsWithDefaultKMS;
76 | private SqsAsyncClient extendedSqsWithGenericReservedAttributeName;
77 | private SqsAsyncClient mockSqsBackend;
78 | private S3AsyncClient mockS3;
79 | private static final String S3_BUCKET_NAME = "test-bucket-name";
80 | private static final String SQS_QUEUE_URL = "test-queue-url";
81 | private static final String S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID = "test-customer-managed-kms-key-id";
82 |
83 | private static final int LESS_THAN_SQS_SIZE_LIMIT = 3;
84 | private static final int SQS_SIZE_LIMIT = 262144;
85 | private static final int MORE_THAN_SQS_SIZE_LIMIT = SQS_SIZE_LIMIT + 1;
86 | private static final ServerSideEncryptionStrategy SERVER_SIDE_ENCRYPTION_CUSTOM_STRATEGY = ServerSideEncryptionFactory.customerKey(S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID);
87 | private static final ServerSideEncryptionStrategy SERVER_SIDE_ENCRYPTION_DEFAULT_STRATEGY = ServerSideEncryptionFactory.awsManagedCmk();
88 |
89 | // should be > 1 and << SQS_SIZE_LIMIT
90 | private static final int ARBITRARY_SMALLER_THRESHOLD = 500;
91 |
92 | @BeforeEach
93 | public void setupClients() {
94 | mockS3 = mock(S3AsyncClient.class);
95 | mockSqsBackend = mock(SqsAsyncClient.class);
96 | when(mockS3.putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class))).thenReturn(
97 | CompletableFuture.completedFuture(null));
98 | when(mockS3.deleteObject(isA(DeleteObjectRequest.class))).thenReturn(
99 | CompletableFuture.completedFuture(DeleteObjectResponse.builder().build()));
100 | when(mockSqsBackend.sendMessage(isA(SendMessageRequest.class))).thenReturn(
101 | CompletableFuture.completedFuture(SendMessageResponse.builder().build()));
102 | when(mockSqsBackend.sendMessageBatch(isA(SendMessageBatchRequest.class))).thenReturn(
103 | CompletableFuture.completedFuture(SendMessageBatchResponse.builder().build()));
104 | when(mockSqsBackend.deleteMessage(isA(DeleteMessageRequest.class))).thenReturn(
105 | CompletableFuture.completedFuture(DeleteMessageResponse.builder().build()));
106 | when(mockSqsBackend.deleteMessageBatch(isA(DeleteMessageBatchRequest.class))).thenReturn(
107 | CompletableFuture.completedFuture(DeleteMessageBatchResponse.builder().build()));
108 |
109 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
110 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
111 |
112 | ExtendedAsyncClientConfiguration extendedClientConfigurationWithCustomKMS = new ExtendedAsyncClientConfiguration()
113 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
114 | .withServerSideEncryption(SERVER_SIDE_ENCRYPTION_CUSTOM_STRATEGY);
115 |
116 | ExtendedAsyncClientConfiguration extendedClientConfigurationWithDefaultKMS = new ExtendedAsyncClientConfiguration()
117 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
118 | .withServerSideEncryption(SERVER_SIDE_ENCRYPTION_DEFAULT_STRATEGY);
119 |
120 | ExtendedAsyncClientConfiguration extendedClientConfigurationWithGenericReservedAttributeName = new ExtendedAsyncClientConfiguration()
121 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withLegacyReservedAttributeNameDisabled();
122 |
123 | extendedSqsWithDefaultConfig = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
124 | extendedSqsWithCustomKMS = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfigurationWithCustomKMS));
125 | extendedSqsWithDefaultKMS = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfigurationWithDefaultKMS));
126 | extendedSqsWithGenericReservedAttributeName = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfigurationWithGenericReservedAttributeName));
127 | }
128 |
129 | @Test
130 | public void testWhenSendMessageWithLargePayloadSupportDisabledThenS3IsNotUsedAndSqsBackendIsResponsibleToFailItWithDeprecatedMethod() {
131 | int messageLength = MORE_THAN_SQS_SIZE_LIMIT;
132 | String messageBody = generateStringWithLength(messageLength);
133 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
134 | .withPayloadSupportDisabled();
135 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
136 |
137 | SendMessageRequest messageRequest = SendMessageRequest.builder()
138 | .queueUrl(SQS_QUEUE_URL)
139 | .messageBody(messageBody)
140 | .overrideConfiguration(
141 | AwsRequestOverrideConfiguration.builder()
142 | .addApiName(ApiName.builder().name(USER_AGENT_NAME).version(USER_AGENT_VERSION).build())
143 | .build())
144 | .build();
145 | sqsExtended.sendMessage(messageRequest);
146 |
147 | ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
148 |
149 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
150 | verify(mockSqsBackend).sendMessage(argumentCaptor.capture());
151 | assertEquals(messageRequest.queueUrl(), argumentCaptor.getValue().queueUrl());
152 | assertEquals(messageRequest.messageBody(), argumentCaptor.getValue().messageBody());
153 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).name(), argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).name());
154 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).version(), argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).version());
155 | }
156 |
157 | @Test
158 | public void testWhenSendMessageWithAlwaysThroughS3AndMessageIsSmallThenItIsStillStoredInS3WithDeprecatedMethod() {
159 | int messageLength = LESS_THAN_SQS_SIZE_LIMIT;
160 | String messageBody = generateStringWithLength(messageLength);
161 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
162 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
163 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
164 |
165 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
166 | sqsExtended.sendMessage(messageRequest).join();
167 |
168 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
169 | }
170 |
171 | @Test
172 | public void testWhenSendMessageWithSetMessageSizeThresholdThenThresholdIsHonoredWithDeprecatedMethod() {
173 | int messageLength = ARBITRARY_SMALLER_THRESHOLD * 2;
174 | String messageBody = generateStringWithLength(messageLength);
175 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
176 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withPayloadSizeThreshold(ARBITRARY_SMALLER_THRESHOLD);
177 |
178 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
179 |
180 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
181 | sqsExtended.sendMessage(messageRequest).join();
182 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
183 | }
184 |
185 | @Test
186 | public void testReceiveMessageMultipleTimesDoesNotAdditionallyAlterReceiveMessageRequestWithDeprecatedMethod() {
187 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
188 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
189 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
190 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
191 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().build()));
192 |
193 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
194 | ReceiveMessageRequest expectedRequest = ReceiveMessageRequest.builder().build();
195 |
196 | sqsExtended.receiveMessage(messageRequest).join();
197 | assertEquals(expectedRequest, messageRequest);
198 |
199 | sqsExtended.receiveMessage(messageRequest).join();
200 | assertEquals(expectedRequest, messageRequest);
201 | }
202 |
203 | @Test
204 | public void testWhenSendLargeMessageThenPayloadIsStoredInS3() {
205 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
206 |
207 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
208 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
209 |
210 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
211 | }
212 |
213 | @Test
214 | public void testWhenSendLargeMessage_WithoutKMS_ThenPayloadIsStoredInS3AndKMSKeyIdIsNotUsed() {
215 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
216 |
217 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
218 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
219 |
220 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
221 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(AsyncRequestBody.class);
222 | verify(mockS3, times(1)).putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
223 |
224 | assertNull(putObjectRequestArgumentCaptor.getValue().serverSideEncryption());
225 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
226 | }
227 |
228 | @Test
229 | public void testWhenSendLargeMessage_WithCustomKMS_ThenPayloadIsStoredInS3AndCorrectKMSKeyIdIsNotUsed() {
230 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
231 |
232 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
233 | extendedSqsWithCustomKMS.sendMessage(messageRequest).join();
234 |
235 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
236 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(AsyncRequestBody.class);
237 | verify(mockS3, times(1)).putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
238 |
239 | assertEquals(putObjectRequestArgumentCaptor.getValue().ssekmsKeyId(), S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID);
240 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
241 | }
242 |
243 | @Test
244 | public void testWhenSendLargeMessage_WithDefaultKMS_ThenPayloadIsStoredInS3AndCorrectKMSKeyIdIsNotUsed() {
245 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
246 |
247 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
248 | extendedSqsWithDefaultKMS.sendMessage(messageRequest).join();
249 |
250 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
251 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(AsyncRequestBody.class);
252 | verify(mockS3, times(1)).putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
253 |
254 | assertTrue(putObjectRequestArgumentCaptor.getValue().serverSideEncryption() != null &&
255 | putObjectRequestArgumentCaptor.getValue().ssekmsKeyId() == null);
256 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
257 | }
258 |
259 | @Test
260 | public void testSendLargeMessageWithDefaultConfigThenLegacyReservedAttributeNameIsUsed(){
261 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
262 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
263 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
264 |
265 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
266 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
267 |
268 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
269 | assertTrue(attributes.containsKey(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME));
270 | assertFalse(attributes.containsKey(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME));
271 |
272 | }
273 |
274 | @Test
275 | public void testSendLargeMessageWithGenericReservedAttributeNameConfigThenGenericReservedAttributeNameIsUsed(){
276 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
277 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
278 | extendedSqsWithGenericReservedAttributeName.sendMessage(messageRequest).join();
279 |
280 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
281 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
282 |
283 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
284 | assertTrue(attributes.containsKey(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME));
285 | assertFalse(attributes.containsKey(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME));
286 | }
287 |
288 | @Test
289 | public void testWhenSendSmallMessageThenS3IsNotUsed() {
290 | String messageBody = generateStringWithLength(SQS_SIZE_LIMIT);
291 |
292 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
293 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
294 |
295 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
296 | }
297 |
298 | @Test
299 | public void testWhenSendMessageWithLargePayloadSupportDisabledThenS3IsNotUsedAndSqsBackendIsResponsibleToFailIt() {
300 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
301 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
302 | .withPayloadSupportDisabled();
303 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
304 |
305 | SendMessageRequest messageRequest = SendMessageRequest.builder()
306 | .queueUrl(SQS_QUEUE_URL)
307 | .messageBody(messageBody)
308 | .overrideConfiguration(
309 | AwsRequestOverrideConfiguration.builder()
310 | .addApiName(ApiName.builder().name(USER_AGENT_NAME).version(USER_AGENT_VERSION).build())
311 | .build())
312 | .build();
313 | sqsExtended.sendMessage(messageRequest).join();
314 |
315 | ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
316 |
317 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
318 | verify(mockSqsBackend).sendMessage(argumentCaptor.capture());
319 | assertEquals(messageRequest.queueUrl(), argumentCaptor.getValue().queueUrl());
320 | assertEquals(messageRequest.messageBody(), argumentCaptor.getValue().messageBody());
321 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).name(), argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).name());
322 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).version(), argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).version());
323 | }
324 |
325 | @Test
326 | public void testWhenSendMessageWithAlwaysThroughS3AndMessageIsSmallThenItIsStillStoredInS3() {
327 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
328 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
329 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
330 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
331 |
332 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
333 | sqsExtended.sendMessage(messageRequest).join();
334 |
335 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
336 | }
337 |
338 | @Test
339 | public void testWhenSendMessageWithSetMessageSizeThresholdThenThresholdIsHonored() {
340 | int messageLength = ARBITRARY_SMALLER_THRESHOLD * 2;
341 | String messageBody = generateStringWithLength(messageLength);
342 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
343 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withPayloadSizeThreshold(ARBITRARY_SMALLER_THRESHOLD);
344 |
345 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
346 |
347 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
348 | sqsExtended.sendMessage(messageRequest).join();
349 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
350 | }
351 |
352 | @Test
353 | public void testReceiveMessageMultipleTimesDoesNotAdditionallyAlterReceiveMessageRequest() {
354 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
355 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().build()));
356 |
357 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
358 |
359 | ReceiveMessageRequest expectedRequest = ReceiveMessageRequest.builder().build();
360 |
361 | extendedSqsWithDefaultConfig.receiveMessage(messageRequest).join();
362 | assertEquals(expectedRequest, messageRequest);
363 |
364 | extendedSqsWithDefaultConfig.receiveMessage(messageRequest).join();
365 | assertEquals(expectedRequest, messageRequest);
366 | }
367 |
368 | @Test
369 | public void testReceiveMessage_when_MessageIsLarge_legacyReservedAttributeUsed() throws Exception {
370 | testReceiveMessage_when_MessageIsLarge(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME);
371 | }
372 |
373 | @Test
374 | public void testReceiveMessage_when_MessageIsLarge_ReservedAttributeUsed() throws Exception {
375 | testReceiveMessage_when_MessageIsLarge(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME);
376 | }
377 |
378 | @Test
379 | public void testReceiveMessage_when_MessageIsSmall() throws Exception {
380 | String expectedMessageAttributeName = "AnyMessageAttribute";
381 | String expectedMessage = "SmallMessage";
382 | Message message = Message.builder()
383 | .messageAttributes(ImmutableMap.of(expectedMessageAttributeName, MessageAttributeValue.builder().build()))
384 | .body(expectedMessage)
385 | .build();
386 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
387 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().messages(message).build()));
388 |
389 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
390 | ReceiveMessageResponse actualReceiveMessageResponse = extendedSqsWithDefaultConfig.receiveMessage(messageRequest).join();
391 | Message actualMessage = actualReceiveMessageResponse.messages().get(0);
392 |
393 | assertEquals(expectedMessage, actualMessage.body());
394 | assertTrue(actualMessage.messageAttributes().containsKey(expectedMessageAttributeName));
395 | assertFalse(actualMessage.messageAttributes().keySet().containsAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES));
396 | verifyNoInteractions(mockS3);
397 | }
398 |
399 | @Test
400 | public void testWhenMessageBatchIsSentThenOnlyMessagesLargerThanThresholdAreStoredInS3() {
401 | // This creates 10 messages, out of which only two are below the threshold (100K and 200K),
402 | // and the other 8 are above the threshold
403 |
404 | int[] messageLengthForCounter = new int[] {
405 | 100_000,
406 | 300_000,
407 | 400_000,
408 | 500_000,
409 | 600_000,
410 | 700_000,
411 | 800_000,
412 | 900_000,
413 | 200_000,
414 | 1000_000
415 | };
416 |
417 | List batchEntries = new ArrayList();
418 | for (int i = 0; i < 10; i++) {
419 | int messageLength = messageLengthForCounter[i];
420 | String messageBody = generateStringWithLength(messageLength);
421 | SendMessageBatchRequestEntry entry = SendMessageBatchRequestEntry.builder()
422 | .id("entry_" + i)
423 | .messageBody(messageBody)
424 | .build();
425 | batchEntries.add(entry);
426 | }
427 |
428 | SendMessageBatchRequest
429 | batchRequest = SendMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(batchEntries).build();
430 | extendedSqsWithDefaultConfig.sendMessageBatch(batchRequest).join();
431 |
432 | // There should be 8 puts for the 8 messages above the threshold
433 | verify(mockS3, times(8)).putObject(isA(PutObjectRequest.class), isA(AsyncRequestBody.class));
434 | }
435 |
436 | @Test
437 | public void testWhenMessageBatchIsLargeS3PointerIsCorrectlySentToSQSAndNotOriginalMessage() {
438 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
439 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
440 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
441 |
442 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
443 |
444 | List batchEntries = new ArrayList();
445 | for (int i = 0; i < 10; i++) {
446 | SendMessageBatchRequestEntry entry = SendMessageBatchRequestEntry.builder()
447 | .id("entry_" + i)
448 | .messageBody(messageBody)
449 | .build();
450 | batchEntries.add(entry);
451 | }
452 | SendMessageBatchRequest batchRequest = SendMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(batchEntries).build();
453 |
454 | sqsExtended.sendMessageBatch(batchRequest).join();
455 |
456 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
457 | verify(mockSqsBackend).sendMessageBatch(sendMessageRequestCaptor.capture());
458 |
459 | for (SendMessageBatchRequestEntry entry : sendMessageRequestCaptor.getValue().entries()) {
460 | assertNotEquals(messageBody, entry.messageBody());
461 | }
462 | }
463 |
464 | @Test
465 | public void testWhenSmallMessageIsSentThenNoAttributeIsAdded() {
466 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
467 |
468 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
469 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
470 |
471 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
472 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
473 |
474 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
475 | assertTrue(attributes.isEmpty());
476 | }
477 |
478 | @Test
479 | public void testWhenLargeMessageIsSentThenAttributeWithPayloadSizeIsAdded() {
480 | int messageLength = MORE_THAN_SQS_SIZE_LIMIT;
481 | String messageBody = generateStringWithLength(messageLength);
482 |
483 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
484 | extendedSqsWithDefaultConfig.sendMessage(messageRequest).join();
485 |
486 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
487 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
488 |
489 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
490 | assertEquals("Number", attributes.get(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME).dataType());
491 | assertEquals(messageLength, (int) Integer.parseInt(attributes.get(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME).stringValue()));
492 | }
493 |
494 | @Test
495 | public void testDefaultExtendedClientDeletesSmallMessage() {
496 | // given
497 | String receiptHandle = UUID.randomUUID().toString();
498 | DeleteMessageRequest
499 | deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(receiptHandle).build();
500 |
501 | // when
502 | extendedSqsWithDefaultConfig.deleteMessage(deleteRequest).join();
503 |
504 | // then
505 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
506 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
507 | assertEquals(receiptHandle, deleteRequestCaptor.getValue().receiptHandle());
508 | verifyNoInteractions(mockS3);
509 | }
510 |
511 | @Test
512 | public void testDefaultExtendedClientDeletesObjectS3UponMessageDelete() {
513 | // given
514 | String randomS3Key = UUID.randomUUID().toString();
515 | String originalReceiptHandle = UUID.randomUUID().toString();
516 | String largeMessageReceiptHandle = getLargeReceiptHandle(randomS3Key, originalReceiptHandle);
517 | DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(largeMessageReceiptHandle).build();
518 |
519 | // when
520 | extendedSqsWithDefaultConfig.deleteMessage(deleteRequest).join();
521 |
522 | // then
523 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
524 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
525 | assertEquals(originalReceiptHandle, deleteRequestCaptor.getValue().receiptHandle());
526 | DeleteObjectRequest deleteObjectRequest = DeleteObjectRequest.builder().bucket(S3_BUCKET_NAME).key(randomS3Key).build();
527 | verify(mockS3).deleteObject(eq(deleteObjectRequest));
528 | }
529 |
530 | @Test
531 | public void testExtendedClientConfiguredDoesNotDeleteObjectFromS3UponDelete() {
532 | // given
533 | String randomS3Key = UUID.randomUUID().toString();
534 | String originalReceiptHandle = UUID.randomUUID().toString();
535 | String largeMessageReceiptHandle = getLargeReceiptHandle(randomS3Key, originalReceiptHandle);
536 | DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(largeMessageReceiptHandle).build();
537 |
538 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
539 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME, false);
540 |
541 | SqsAsyncClient extendedSqs = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
542 |
543 | // when
544 | extendedSqs.deleteMessage(deleteRequest).join();
545 |
546 | // then
547 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
548 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
549 | assertEquals(originalReceiptHandle, deleteRequestCaptor.getValue().receiptHandle());
550 | verifyNoInteractions(mockS3);
551 | }
552 |
553 | @Test
554 | public void testExtendedClientConfiguredDoesNotDeletesObjectsFromS3UponDeleteBatch() {
555 | // given
556 | int batchSize = 10;
557 | List originalReceiptHandles = IntStream.range(0, batchSize)
558 | .mapToObj(i -> UUID.randomUUID().toString())
559 | .collect(Collectors.toList());
560 | DeleteMessageBatchRequest deleteBatchRequest = generateLargeDeleteBatchRequest(originalReceiptHandles);
561 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
562 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME, false);
563 | SqsAsyncClient extendedSqs = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
564 |
565 | // when
566 | extendedSqs.deleteMessageBatch(deleteBatchRequest).join();
567 |
568 | // then
569 | ArgumentCaptor deleteBatchRequestCaptor = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
570 | verify(mockSqsBackend, times(1)).deleteMessageBatch(deleteBatchRequestCaptor.capture());
571 | DeleteMessageBatchRequest request = deleteBatchRequestCaptor.getValue();
572 | assertEquals(originalReceiptHandles.size(), request.entries().size());
573 | IntStream.range(0, originalReceiptHandles.size()).forEach(i -> assertEquals(
574 | originalReceiptHandles.get(i),
575 | request.entries().get(i).receiptHandle()));
576 | verifyNoInteractions(mockS3);
577 | }
578 |
579 | @Test
580 | public void testDefaultExtendedClientDeletesObjectsFromS3UponDeleteBatch() {
581 | // given
582 | int batchSize = 10;
583 | List originalReceiptHandles = IntStream.range(0, batchSize)
584 | .mapToObj(i -> UUID.randomUUID().toString())
585 | .collect(Collectors.toList());
586 | DeleteMessageBatchRequest deleteBatchRequest = generateLargeDeleteBatchRequest(originalReceiptHandles);
587 |
588 | // when
589 | extendedSqsWithDefaultConfig.deleteMessageBatch(deleteBatchRequest).join();
590 |
591 | // then
592 | ArgumentCaptor deleteBatchRequestCaptor = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
593 | verify(mockSqsBackend, times(1)).deleteMessageBatch(deleteBatchRequestCaptor.capture());
594 | DeleteMessageBatchRequest request = deleteBatchRequestCaptor.getValue();
595 | assertEquals(originalReceiptHandles.size(), request.entries().size());
596 | IntStream.range(0, originalReceiptHandles.size()).forEach(i -> assertEquals(
597 | originalReceiptHandles.get(i),
598 | request.entries().get(i).receiptHandle()));
599 | verify(mockS3, times(batchSize)).deleteObject(any(DeleteObjectRequest.class));
600 | }
601 |
602 | @Test
603 | public void testWhenSendMessageWIthCannedAccessControlListDefined() {
604 | ObjectCannedACL expected = ObjectCannedACL.BUCKET_OWNER_FULL_CONTROL;
605 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
606 | ExtendedAsyncClientConfiguration extendedClientConfiguration = new ExtendedAsyncClientConfiguration()
607 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withObjectCannedACL(expected);
608 | SqsAsyncClient sqsExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedClientConfiguration));
609 |
610 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
611 | sqsExtended.sendMessage(messageRequest).join();
612 |
613 | ArgumentCaptor captor = ArgumentCaptor.forClass(PutObjectRequest.class);
614 |
615 | verify(mockS3).putObject(captor.capture(), any(AsyncRequestBody.class));
616 |
617 | assertEquals(expected, captor.getValue().acl());
618 | }
619 |
620 | private void testReceiveMessage_when_MessageIsLarge(String reservedAttributeName) throws Exception {
621 | String pointer = new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson();
622 | Message message = Message.builder()
623 | .messageAttributes(ImmutableMap.of(reservedAttributeName, MessageAttributeValue.builder().build()))
624 | .body(pointer)
625 | .build();
626 | String expectedMessage = "LargeMessage";
627 | GetObjectRequest getObjectRequest = GetObjectRequest.builder()
628 | .bucket(S3_BUCKET_NAME)
629 | .key("S3Key")
630 | .build();
631 |
632 | ResponseBytes s3Object = ResponseBytes.fromByteArray(
633 | GetObjectResponse.builder().build(),
634 | expectedMessage.getBytes(StandardCharsets.UTF_8));
635 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
636 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().messages(message).build()));
637 | when(mockS3.getObject(isA(GetObjectRequest.class), isA(AsyncResponseTransformer.class))).thenReturn(
638 | CompletableFuture.completedFuture(s3Object));
639 |
640 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
641 | ReceiveMessageResponse actualReceiveMessageResponse = extendedSqsWithDefaultConfig.receiveMessage(messageRequest).join();
642 | Message actualMessage = actualReceiveMessageResponse.messages().get(0);
643 |
644 | assertEquals(expectedMessage, actualMessage.body());
645 | assertFalse(actualMessage.messageAttributes().keySet().containsAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES));
646 | verify(mockS3, times(1)).getObject(isA(GetObjectRequest.class), isA(AsyncResponseTransformer.class));
647 | }
648 |
649 | @Test
650 | public void testReceiveMessage_when_ignorePayloadNotFound_then_messageWithPayloadNotFoundIsDeletedFromSQS() {
651 | ExtendedAsyncClientConfiguration extendedAsyncClientConfiguration = new ExtendedAsyncClientConfiguration()
652 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
653 | .withIgnorePayloadNotFound(true);
654 | SqsAsyncClient sqsAsyncExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedAsyncClientConfiguration));
655 |
656 | String receiptHandle = "receipt-handle";
657 | Message message = Message.builder()
658 | .messageAttributes(ImmutableMap.of(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, MessageAttributeValue.builder().build()))
659 | .body(new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson())
660 | .receiptHandle(receiptHandle)
661 | .build();
662 |
663 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
664 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().messages(message).build()));
665 | doThrow(NoSuchKeyException.class).when(mockS3).getObject((GetObjectRequest) any(), any(AsyncResponseTransformer.class));
666 |
667 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().queueUrl(SQS_QUEUE_URL).build();
668 | ReceiveMessageResponse receiveMessageResponse = sqsAsyncExtended.receiveMessage(messageRequest).join();
669 |
670 | assertTrue(receiveMessageResponse.messages().isEmpty());
671 |
672 | ArgumentCaptor deleteMessageRequestArgumentCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
673 | verify(mockSqsBackend).deleteMessage(deleteMessageRequestArgumentCaptor.capture());
674 | assertEquals(SQS_QUEUE_URL, deleteMessageRequestArgumentCaptor.getValue().queueUrl());
675 | assertEquals(receiptHandle, deleteMessageRequestArgumentCaptor.getValue().receiptHandle());
676 | }
677 |
678 | @Test
679 | public void testReceiveMessage_when_ignorePayloadNotFoundIsFalse_then_messageWithPayloadNotFoundThrowsException() {
680 | ExtendedAsyncClientConfiguration extendedAsyncClientConfiguration = new ExtendedAsyncClientConfiguration()
681 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
682 | .withIgnorePayloadNotFound(false);
683 | SqsAsyncClient sqsAsyncExtended = spy(new AmazonSQSExtendedAsyncClient(mockSqsBackend, extendedAsyncClientConfiguration));
684 |
685 | String receiptHandle = "receipt-handle";
686 | Message message = Message.builder()
687 | .messageAttributes(ImmutableMap.of(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, MessageAttributeValue.builder().build()))
688 | .body(new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson())
689 | .receiptHandle(receiptHandle)
690 | .build();
691 |
692 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
693 | CompletableFuture.completedFuture(ReceiveMessageResponse.builder().messages(message).build()));
694 | doThrow(NoSuchKeyException.class).when(mockS3).getObject((GetObjectRequest) any(), any(AsyncResponseTransformer.class));
695 |
696 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
697 | try {
698 | sqsAsyncExtended.receiveMessage(messageRequest).join();
699 | fail("Expected exception after receiving NoSuchKeyException from S3 was not thrown.");
700 | } catch (CompletionException e) {
701 | assertEquals(NoSuchKeyException.class.getName(), e.getCause().getClass().getName());
702 | verify(mockSqsBackend, never()).deleteMessage(any(DeleteMessageRequest.class));
703 | }
704 | }
705 |
706 | private DeleteMessageBatchRequest generateLargeDeleteBatchRequest(List originalReceiptHandles) {
707 | List deleteEntries = IntStream.range(0, originalReceiptHandles.size())
708 | .mapToObj(i -> DeleteMessageBatchRequestEntry.builder()
709 | .id(Integer.toString(i))
710 | .receiptHandle(getSampleLargeReceiptHandle(originalReceiptHandles.get(i)))
711 | .build())
712 | .collect(Collectors.toList());
713 |
714 | return DeleteMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(deleteEntries).build();
715 | }
716 |
717 | private String getLargeReceiptHandle(String s3Key, String originalReceiptHandle) {
718 | return SQSExtendedClientConstants.S3_BUCKET_NAME_MARKER + S3_BUCKET_NAME
719 | + SQSExtendedClientConstants.S3_BUCKET_NAME_MARKER + SQSExtendedClientConstants.S3_KEY_MARKER
720 | + s3Key + SQSExtendedClientConstants.S3_KEY_MARKER + originalReceiptHandle;
721 | }
722 |
723 | private String getSampleLargeReceiptHandle(String originalReceiptHandle) {
724 | return getLargeReceiptHandle(UUID.randomUUID().toString(), originalReceiptHandle);
725 | }
726 |
727 | private String generateStringWithLength(int messageLength) {
728 | char[] charArray = new char[messageLength];
729 | Arrays.fill(charArray, 'x');
730 | return new String(charArray);
731 | }
732 | }
733 |
--------------------------------------------------------------------------------
/src/test/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedClientTest.java:
--------------------------------------------------------------------------------
1 | /*
2 | * Copyright 2010-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 | *
4 | * Licensed under the Apache License, Version 2.0 (the "License").
5 | * You may not use this file except in compliance with the License.
6 | * A copy of the License is located at
7 | *
8 | * http://aws.amazon.com/apache2.0
9 | *
10 | * or in the "license" file accompanying this file. This file is distributed
11 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12 | * express or implied. See the License for the specific language governing
13 | * permissions and limitations under the License.
14 | */
15 |
16 | package com.amazon.sqs.javamessaging;
17 |
18 | import static com.amazon.sqs.javamessaging.StringTestUtil.generateStringWithLength;
19 |
20 | import org.junit.jupiter.api.AfterEach;
21 | import org.junit.jupiter.api.BeforeEach;
22 | import org.junit.jupiter.api.Test;
23 | import org.mockito.ArgumentCaptor;
24 | import org.mockito.MockedStatic;
25 | import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
26 | import software.amazon.awssdk.core.ApiName;
27 | import software.amazon.awssdk.core.ResponseInputStream;
28 | import software.amazon.awssdk.core.exception.SdkException;
29 | import software.amazon.awssdk.core.sync.RequestBody;
30 | import software.amazon.awssdk.http.AbortableInputStream;
31 | import software.amazon.awssdk.services.s3.S3Client;
32 | import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
33 | import software.amazon.awssdk.services.s3.model.GetObjectRequest;
34 | import software.amazon.awssdk.services.s3.model.GetObjectResponse;
35 | import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
36 | import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
37 | import software.amazon.awssdk.services.s3.model.PutObjectRequest;
38 | import software.amazon.awssdk.services.sqs.SqsClient;
39 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest;
40 | import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry;
41 | import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest;
42 | import software.amazon.awssdk.services.sqs.model.Message;
43 | import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
44 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
45 | import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
46 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest;
47 | import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
48 | import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
49 | import software.amazon.awssdk.utils.ImmutableMap;
50 | import software.amazon.awssdk.utils.StringInputStream;
51 | import software.amazon.payloadoffloading.PayloadS3Pointer;
52 | import software.amazon.payloadoffloading.ServerSideEncryptionFactory;
53 | import software.amazon.payloadoffloading.ServerSideEncryptionStrategy;
54 |
55 | import java.util.ArrayList;
56 | import java.util.List;
57 | import java.util.Map;
58 | import java.util.UUID;
59 | import java.util.stream.Collectors;
60 | import java.util.stream.IntStream;
61 |
62 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClient.USER_AGENT_NAME;
63 | import static com.amazon.sqs.javamessaging.AmazonSQSExtendedClient.USER_AGENT_VERSION;
64 |
65 | import static org.junit.jupiter.api.Assertions.assertEquals;
66 | import static org.junit.jupiter.api.Assertions.assertFalse;
67 | import static org.junit.jupiter.api.Assertions.assertNotEquals;
68 | import static org.junit.jupiter.api.Assertions.assertNull;
69 | import static org.junit.jupiter.api.Assertions.assertTrue;
70 | import static org.junit.jupiter.api.Assertions.fail;
71 | import static org.mockito.ArgumentMatchers.any;
72 | import static org.mockito.ArgumentMatchers.argThat;
73 | import static org.mockito.ArgumentMatchers.eq;
74 | import static org.mockito.Mockito.doThrow;
75 | import static org.mockito.Mockito.isA;
76 | import static org.mockito.Mockito.mock;
77 | import static org.mockito.Mockito.mockStatic;
78 | import static org.mockito.Mockito.never;
79 | import static org.mockito.Mockito.spy;
80 | import static org.mockito.Mockito.times;
81 | import static org.mockito.Mockito.verify;
82 | import static org.mockito.Mockito.verifyNoInteractions;
83 | import static org.mockito.Mockito.when;
84 |
85 | /**
86 | * Tests the AmazonSQSExtendedClient class.
87 | */
88 | public class AmazonSQSExtendedClientTest {
89 |
90 | private SqsClient extendedSqsWithDefaultConfig;
91 | private SqsClient extendedSqsWithCustomKMS;
92 | private SqsClient extendedSqsWithDefaultKMS;
93 | private SqsClient extendedSqsWithGenericReservedAttributeName;
94 | private SqsClient extendedSqsWithDeprecatedMethods;
95 | private SqsClient extendedSqsWithS3KeyPrefix;
96 | private SqsClient mockSqsBackend;
97 | private S3Client mockS3;
98 |
99 | private MockedStatic uuidMockStatic;
100 | private static final String S3_BUCKET_NAME = "test-bucket-name";
101 | private static final String SQS_QUEUE_URL = "test-queue-url";
102 | private static final String S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID = "test-customer-managed-kms-key-id";
103 | private static final String S3_KEY_PREFIX = "test-s3-key-prefix";
104 | private static final String S3_KEY_UUID = "test-s3-key-uuid";
105 |
106 | private static final int LESS_THAN_SQS_SIZE_LIMIT = 3;
107 | private static final int SQS_SIZE_LIMIT = 262144;
108 | private static final int MORE_THAN_SQS_SIZE_LIMIT = SQS_SIZE_LIMIT + 1;
109 | private static final ServerSideEncryptionStrategy SERVER_SIDE_ENCRYPTION_CUSTOM_STRATEGY =
110 | ServerSideEncryptionFactory.customerKey(S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID);
111 | private static final ServerSideEncryptionStrategy SERVER_SIDE_ENCRYPTION_DEFAULT_STRATEGY =
112 | ServerSideEncryptionFactory.awsManagedCmk();
113 |
114 | // should be > 1 and << SQS_SIZE_LIMIT
115 | private static final int ARBITRARY_SMALLER_THRESHOLD = 500;
116 |
117 | @BeforeEach
118 | public void setupClients() {
119 | uuidMockStatic = mockStatic(UUID.class);
120 | mockS3 = mock(S3Client.class);
121 | mockSqsBackend = mock(SqsClient.class);
122 | when(mockS3.putObject(isA(PutObjectRequest.class), isA(RequestBody.class))).thenReturn(null);
123 |
124 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
125 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
126 |
127 | ExtendedClientConfiguration extendedClientConfigurationWithCustomKMS = new ExtendedClientConfiguration()
128 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
129 | .withServerSideEncryption(SERVER_SIDE_ENCRYPTION_CUSTOM_STRATEGY);
130 |
131 | ExtendedClientConfiguration extendedClientConfigurationWithDefaultKMS = new ExtendedClientConfiguration()
132 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
133 | .withServerSideEncryption(SERVER_SIDE_ENCRYPTION_DEFAULT_STRATEGY);
134 |
135 | ExtendedClientConfiguration extendedClientConfigurationWithGenericReservedAttributeName = new ExtendedClientConfiguration()
136 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withLegacyReservedAttributeNameDisabled();
137 |
138 | ExtendedClientConfiguration extendedClientConfigurationDeprecated = new ExtendedClientConfiguration().withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
139 |
140 | ExtendedClientConfiguration extendedClientConfigurationWithS3KeyPrefix = new ExtendedClientConfiguration()
141 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
142 | .withS3KeyPrefix(S3_KEY_PREFIX);
143 |
144 | UUID uuidMock = mock(UUID.class);
145 | when(uuidMock.toString()).thenReturn(S3_KEY_UUID);
146 | uuidMockStatic.when(UUID::randomUUID).thenReturn(uuidMock);
147 |
148 | extendedSqsWithDefaultConfig = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
149 | extendedSqsWithCustomKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithCustomKMS));
150 | extendedSqsWithDefaultKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithDefaultKMS));
151 | extendedSqsWithGenericReservedAttributeName = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithGenericReservedAttributeName));
152 | extendedSqsWithDeprecatedMethods = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationDeprecated));
153 | extendedSqsWithS3KeyPrefix = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithS3KeyPrefix));
154 | }
155 |
156 | @AfterEach
157 | public void tearDown() {
158 | uuidMockStatic.close();
159 | }
160 |
161 | @Test
162 | public void testWhenSendMessageWithLargePayloadSupportDisabledThenS3IsNotUsedAndSqsBackendIsResponsibleToFailItWithDeprecatedMethod() {
163 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
164 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration().withPayloadSupportDisabled();
165 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
166 |
167 | SendMessageRequest messageRequest = SendMessageRequest.builder()
168 | .queueUrl(SQS_QUEUE_URL)
169 | .messageBody(messageBody)
170 | .overrideConfiguration(
171 | AwsRequestOverrideConfiguration.builder()
172 | .addApiName(ApiName.builder().name(USER_AGENT_NAME).version(USER_AGENT_VERSION).build())
173 | .build())
174 | .build();
175 | sqsExtended.sendMessage(messageRequest);
176 |
177 | ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
178 |
179 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
180 | verify(mockSqsBackend).sendMessage(argumentCaptor.capture());
181 | assertEquals(messageRequest.queueUrl(), argumentCaptor.getValue().queueUrl());
182 | assertEquals(messageRequest.messageBody(), argumentCaptor.getValue().messageBody());
183 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).name(),
184 | argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).name());
185 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).version(),
186 | argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).version());
187 | }
188 |
189 | @Test
190 | public void testWhenSendMessageWithAlwaysThroughS3AndMessageIsSmallThenItIsStillStoredInS3WithDeprecatedMethod() {
191 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
192 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
193 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
194 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mock(SqsClient.class), extendedClientConfiguration));
195 |
196 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
197 | sqsExtended.sendMessage(messageRequest);
198 |
199 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
200 | }
201 |
202 | @Test
203 | public void testWhenSendMessageWithSetMessageSizeThresholdThenThresholdIsHonoredWithDeprecatedMethod() {
204 | int messageLength = ARBITRARY_SMALLER_THRESHOLD * 2;
205 | String messageBody = generateStringWithLength(messageLength);
206 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
207 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withPayloadSizeThreshold(ARBITRARY_SMALLER_THRESHOLD);
208 |
209 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mock(SqsClient.class), extendedClientConfiguration));
210 |
211 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
212 | sqsExtended.sendMessage(messageRequest);
213 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
214 | }
215 |
216 | @Test
217 | public void testReceiveMessageMultipleTimesDoesNotAdditionallyAlterReceiveMessageRequestWithDeprecatedMethod() {
218 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
219 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
220 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
221 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(ReceiveMessageResponse.builder().build());
222 |
223 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
224 | ReceiveMessageRequest expectedRequest = ReceiveMessageRequest.builder().build();
225 |
226 | sqsExtended.receiveMessage(messageRequest);
227 | assertEquals(expectedRequest, messageRequest);
228 |
229 | sqsExtended.receiveMessage(messageRequest);
230 | assertEquals(expectedRequest, messageRequest);
231 | }
232 |
233 | @Test
234 | public void testWhenSendLargeMessageThenPayloadIsStoredInS3() {
235 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
236 |
237 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
238 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
239 |
240 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
241 | }
242 |
243 | @Test
244 | public void testWhenSendLargeMessage_WithoutKMS_ThenPayloadIsStoredInS3AndKMSKeyIdIsNotUsed() {
245 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
246 |
247 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
248 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
249 |
250 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
251 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(RequestBody.class);
252 | verify(mockS3, times(1))
253 | .putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
254 |
255 | assertNull(putObjectRequestArgumentCaptor.getValue().serverSideEncryption());
256 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
257 | }
258 |
259 | @Test
260 | public void testWhenSendLargeMessage_WithCustomKMS_ThenPayloadIsStoredInS3AndCorrectKMSKeyIdIsNotUsed() {
261 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
262 |
263 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
264 | extendedSqsWithCustomKMS.sendMessage(messageRequest);
265 |
266 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
267 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(RequestBody.class);
268 | verify(mockS3, times(1))
269 | .putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
270 |
271 | assertEquals(putObjectRequestArgumentCaptor.getValue().ssekmsKeyId(), S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID);
272 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
273 | }
274 |
275 | @Test
276 | public void testWhenSendLargeMessage_WithDefaultKMS_ThenPayloadIsStoredInS3AndCorrectKMSKeyIdIsNotUsed() {
277 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
278 |
279 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
280 | extendedSqsWithDefaultKMS.sendMessage(messageRequest);
281 |
282 | ArgumentCaptor putObjectRequestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
283 | ArgumentCaptor requestBodyArgumentCaptor = ArgumentCaptor.forClass(RequestBody.class);
284 | verify(mockS3, times(1))
285 | .putObject(putObjectRequestArgumentCaptor.capture(), requestBodyArgumentCaptor.capture());
286 |
287 | assertTrue(putObjectRequestArgumentCaptor.getValue().serverSideEncryption() != null &&
288 | putObjectRequestArgumentCaptor.getValue().ssekmsKeyId() == null);
289 | assertEquals(putObjectRequestArgumentCaptor.getValue().bucket(), S3_BUCKET_NAME);
290 | }
291 |
292 | @Test
293 | public void testSendLargeMessageWithDefaultConfigThenLegacyReservedAttributeNameIsUsed(){
294 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
295 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
296 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
297 |
298 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
299 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
300 |
301 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
302 | assertTrue(attributes.containsKey(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME));
303 | assertFalse(attributes.containsKey(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME));
304 | }
305 |
306 | @Test
307 | public void testSendLargeMessageWithGenericReservedAttributeNameConfigThenGenericReservedAttributeNameIsUsed(){
308 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
309 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
310 | extendedSqsWithGenericReservedAttributeName.sendMessage(messageRequest);
311 |
312 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
313 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
314 |
315 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
316 | assertTrue(attributes.containsKey(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME));
317 | assertFalse(attributes.containsKey(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME));
318 | }
319 |
320 | @Test
321 | public void testWhenSendSmallMessageThenS3IsNotUsed() {
322 | String messageBody = generateStringWithLength(SQS_SIZE_LIMIT);
323 |
324 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
325 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
326 |
327 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
328 | }
329 |
330 | @Test
331 | public void testWhenSendMessageWithLargePayloadSupportDisabledThenS3IsNotUsedAndSqsBackendIsResponsibleToFailIt() {
332 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
333 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
334 | .withPayloadSupportDisabled();
335 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
336 |
337 | SendMessageRequest messageRequest = SendMessageRequest.builder()
338 | .queueUrl(SQS_QUEUE_URL)
339 | .messageBody(messageBody)
340 | .overrideConfiguration(
341 | AwsRequestOverrideConfiguration.builder()
342 | .addApiName(ApiName.builder().name(USER_AGENT_NAME).version(USER_AGENT_VERSION).build())
343 | .build())
344 | .build();
345 | sqsExtended.sendMessage(messageRequest);
346 |
347 | ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
348 |
349 | verify(mockS3, never()).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
350 | verify(mockSqsBackend).sendMessage(argumentCaptor.capture());
351 | assertEquals(messageRequest.queueUrl(), argumentCaptor.getValue().queueUrl());
352 | assertEquals(messageRequest.messageBody(), argumentCaptor.getValue().messageBody());
353 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).name(),
354 | argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).name());
355 | assertEquals(messageRequest.overrideConfiguration().get().apiNames().get(0).version(),
356 | argumentCaptor.getValue().overrideConfiguration().get().apiNames().get(0).version());
357 | }
358 |
359 | @Test
360 | public void testWhenSendMessageWithAlwaysThroughS3AndMessageIsSmallThenItIsStillStoredInS3() {
361 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
362 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
363 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
364 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mock(SqsClient.class), extendedClientConfiguration));
365 |
366 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
367 | sqsExtended.sendMessage(messageRequest);
368 |
369 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
370 | }
371 |
372 | @Test
373 | public void testWhenSendMessageWithSetMessageSizeThresholdThenThresholdIsHonored() {
374 | int messageLength = ARBITRARY_SMALLER_THRESHOLD * 2;
375 | String messageBody = generateStringWithLength(messageLength);
376 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
377 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withPayloadSizeThreshold(ARBITRARY_SMALLER_THRESHOLD);
378 |
379 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mock(SqsClient.class), extendedClientConfiguration));
380 |
381 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
382 | sqsExtended.sendMessage(messageRequest);
383 | verify(mockS3, times(1)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
384 | }
385 |
386 | @Test
387 | public void testReceiveMessageMultipleTimesDoesNotAdditionallyAlterReceiveMessageRequest() {
388 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(ReceiveMessageResponse.builder().build());
389 |
390 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
391 |
392 | ReceiveMessageRequest expectedRequest = ReceiveMessageRequest.builder().build();
393 |
394 | extendedSqsWithDefaultConfig.receiveMessage(messageRequest);
395 | assertEquals(expectedRequest, messageRequest);
396 |
397 | extendedSqsWithDefaultConfig.receiveMessage(messageRequest);
398 | assertEquals(expectedRequest, messageRequest);
399 | }
400 |
401 | @Test
402 | public void testReceiveMessage_when_MessageIsLarge_legacyReservedAttributeUsed() {
403 | testReceiveMessage_when_MessageIsLarge(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME);
404 | }
405 |
406 | @Test
407 | public void testReceiveMessage_when_MessageIsLarge_ReservedAttributeUsed() {
408 | testReceiveMessage_when_MessageIsLarge(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME);
409 | }
410 |
411 | @Test
412 | public void testReceiveMessage_when_MessageIsSmall() {
413 | String expectedMessageAttributeName = "AnyMessageAttribute";
414 | String expectedMessage = "SmallMessage";
415 | Message message = Message.builder()
416 | .messageAttributes(ImmutableMap.of(expectedMessageAttributeName, MessageAttributeValue.builder().build()))
417 | .body(expectedMessage)
418 | .build();
419 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(ReceiveMessageResponse.builder().messages(message).build());
420 |
421 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
422 | ReceiveMessageResponse actualReceiveMessageResponse = extendedSqsWithDefaultConfig.receiveMessage(messageRequest);
423 | Message actualMessage = actualReceiveMessageResponse.messages().get(0);
424 |
425 | assertEquals(expectedMessage, actualMessage.body());
426 | assertTrue(actualMessage.messageAttributes().containsKey(expectedMessageAttributeName));
427 | assertFalse(actualMessage.messageAttributes().keySet().containsAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES));
428 | verifyNoInteractions(mockS3);
429 | }
430 |
431 | @Test
432 | public void testWhenMessageBatchIsSentThenOnlyMessagesLargerThanThresholdAreStoredInS3() {
433 | // This creates 10 messages, out of which only two are below the threshold (100K and 200K),
434 | // and the other 8 are above the threshold
435 |
436 | int[] messageLengthForCounter = new int[] {
437 | 100_000,
438 | 300_000,
439 | 400_000,
440 | 500_000,
441 | 600_000,
442 | 700_000,
443 | 800_000,
444 | 900_000,
445 | 200_000,
446 | 1000_000
447 | };
448 |
449 | List batchEntries = new ArrayList<>();
450 | for (int i = 0; i < 10; i++) {
451 | int messageLength = messageLengthForCounter[i];
452 | String messageBody = generateStringWithLength(messageLength);
453 | SendMessageBatchRequestEntry entry = SendMessageBatchRequestEntry.builder()
454 | .id("entry_" + i)
455 | .messageBody(messageBody)
456 | .build();
457 | batchEntries.add(entry);
458 | }
459 |
460 | SendMessageBatchRequest batchRequest = SendMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(batchEntries).build();
461 | extendedSqsWithDefaultConfig.sendMessageBatch(batchRequest);
462 |
463 | // There should be 8 puts for the 8 messages above the threshold
464 | verify(mockS3, times(8)).putObject(isA(PutObjectRequest.class), isA(RequestBody.class));
465 | }
466 |
467 | @Test
468 | public void testWhenMessageBatchIsLargeS3PointerIsCorrectlySentToSQSAndNotOriginalMessage() {
469 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
470 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
471 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withAlwaysThroughS3(true);
472 |
473 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
474 |
475 | List batchEntries = new ArrayList<>();
476 | for (int i = 0; i < 10; i++) {
477 | SendMessageBatchRequestEntry entry = SendMessageBatchRequestEntry.builder()
478 | .id("entry_" + i)
479 | .messageBody(messageBody)
480 | .build();
481 | batchEntries.add(entry);
482 | }
483 | SendMessageBatchRequest batchRequest = SendMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(batchEntries).build();
484 |
485 | sqsExtended.sendMessageBatch(batchRequest);
486 |
487 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
488 | verify(mockSqsBackend).sendMessageBatch(sendMessageRequestCaptor.capture());
489 |
490 | for (SendMessageBatchRequestEntry entry : sendMessageRequestCaptor.getValue().entries()) {
491 | assertNotEquals(messageBody, entry.messageBody());
492 | }
493 | }
494 |
495 | @Test
496 | public void testWhenSmallMessageIsSentThenNoAttributeIsAdded() {
497 | String messageBody = generateStringWithLength(LESS_THAN_SQS_SIZE_LIMIT);
498 |
499 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
500 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
501 |
502 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
503 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
504 |
505 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
506 | assertTrue(attributes.isEmpty());
507 | }
508 |
509 | @Test
510 | public void testWhenLargeMessageIsSentThenAttributeWithPayloadSizeIsAdded() {
511 | int messageLength = MORE_THAN_SQS_SIZE_LIMIT;
512 | String messageBody = generateStringWithLength(messageLength);
513 |
514 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
515 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
516 |
517 | ArgumentCaptor sendMessageRequestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
518 | verify(mockSqsBackend).sendMessage(sendMessageRequestCaptor.capture());
519 |
520 | Map attributes = sendMessageRequestCaptor.getValue().messageAttributes();
521 | assertEquals("Number", attributes.get(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME).dataType());
522 | assertEquals(messageLength, Integer.parseInt(attributes.get(AmazonSQSExtendedClientUtil.LEGACY_RESERVED_ATTRIBUTE_NAME).stringValue()));
523 | }
524 |
525 | @Test
526 | public void testDefaultExtendedClientDeletesSmallMessage() {
527 | // given
528 | String receiptHandle = UUID.randomUUID().toString();
529 | DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(receiptHandle).build();
530 |
531 | // when
532 | extendedSqsWithDefaultConfig.deleteMessage(deleteRequest);
533 |
534 | // then
535 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
536 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
537 | assertEquals(receiptHandle, deleteRequestCaptor.getValue().receiptHandle());
538 | verifyNoInteractions(mockS3);
539 | }
540 |
541 | @Test
542 | public void testDefaultExtendedClientDeletesObjectS3UponMessageDelete() {
543 | // given
544 | String randomS3Key = UUID.randomUUID().toString();
545 | String originalReceiptHandle = UUID.randomUUID().toString();
546 | String largeMessageReceiptHandle = getLargeReceiptHandle(randomS3Key, originalReceiptHandle);
547 | DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(largeMessageReceiptHandle).build();
548 |
549 | // when
550 | extendedSqsWithDefaultConfig.deleteMessage(deleteRequest);
551 |
552 | // then
553 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
554 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
555 | assertEquals(originalReceiptHandle, deleteRequestCaptor.getValue().receiptHandle());
556 | DeleteObjectRequest deleteObjectRequest = DeleteObjectRequest.builder().bucket(S3_BUCKET_NAME).key(randomS3Key).build();
557 | verify(mockS3).deleteObject(eq(deleteObjectRequest));
558 | }
559 |
560 | @Test
561 | public void testExtendedClientConfiguredDoesNotDeleteObjectFromS3UponDelete() {
562 | // given
563 | String randomS3Key = UUID.randomUUID().toString();
564 | String originalReceiptHandle = UUID.randomUUID().toString();
565 | String largeMessageReceiptHandle = getLargeReceiptHandle(randomS3Key, originalReceiptHandle);
566 | DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(SQS_QUEUE_URL).receiptHandle(largeMessageReceiptHandle).build();
567 |
568 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
569 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME, false);
570 |
571 | SqsClient extendedSqs = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
572 |
573 | // when
574 | extendedSqs.deleteMessage(deleteRequest);
575 |
576 | // then
577 | ArgumentCaptor deleteRequestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
578 | verify(mockSqsBackend).deleteMessage(deleteRequestCaptor.capture());
579 | assertEquals(originalReceiptHandle, deleteRequestCaptor.getValue().receiptHandle());
580 | verifyNoInteractions(mockS3);
581 | }
582 |
583 | @Test
584 | public void testExtendedClientConfiguredDoesNotDeletesObjectsFromS3UponDeleteBatch() {
585 | // given
586 | int batchSize = 10;
587 | List originalReceiptHandles = IntStream.range(0, batchSize)
588 | .mapToObj(i -> UUID.randomUUID().toString())
589 | .collect(Collectors.toList());
590 | DeleteMessageBatchRequest deleteBatchRequest = generateLargeDeleteBatchRequest(originalReceiptHandles);
591 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
592 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME, false);
593 | SqsClient extendedSqs = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
594 |
595 | // when
596 | extendedSqs.deleteMessageBatch(deleteBatchRequest);
597 |
598 | // then
599 | ArgumentCaptor deleteBatchRequestCaptor = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
600 | verify(mockSqsBackend, times(1)).deleteMessageBatch(deleteBatchRequestCaptor.capture());
601 | DeleteMessageBatchRequest request = deleteBatchRequestCaptor.getValue();
602 | assertEquals(originalReceiptHandles.size(), request.entries().size());
603 | IntStream.range(0, originalReceiptHandles.size()).forEach(i -> assertEquals(
604 | originalReceiptHandles.get(i),
605 | request.entries().get(i).receiptHandle()));
606 | verifyNoInteractions(mockS3);
607 | }
608 |
609 | @Test
610 | public void testDefaultExtendedClientDeletesObjectsFromS3UponDeleteBatch() {
611 | // given
612 | int batchSize = 10;
613 | List originalReceiptHandles = IntStream.range(0, batchSize)
614 | .mapToObj(i -> UUID.randomUUID().toString())
615 | .collect(Collectors.toList());
616 | DeleteMessageBatchRequest deleteBatchRequest = generateLargeDeleteBatchRequest(originalReceiptHandles);
617 |
618 | // when
619 | extendedSqsWithDefaultConfig.deleteMessageBatch(deleteBatchRequest);
620 |
621 | // then
622 | ArgumentCaptor deleteBatchRequestCaptor = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
623 | verify(mockSqsBackend, times(1)).deleteMessageBatch(deleteBatchRequestCaptor.capture());
624 | DeleteMessageBatchRequest request = deleteBatchRequestCaptor.getValue();
625 | assertEquals(originalReceiptHandles.size(), request.entries().size());
626 | IntStream.range(0, originalReceiptHandles.size()).forEach(i -> assertEquals(
627 | originalReceiptHandles.get(i),
628 | request.entries().get(i).receiptHandle()));
629 | verify(mockS3, times(batchSize)).deleteObject(any(DeleteObjectRequest.class));
630 | }
631 |
632 | @Test
633 | public void testWhenSendMessageWIthCannedAccessControlListDefined() {
634 | ObjectCannedACL expected = ObjectCannedACL.BUCKET_OWNER_FULL_CONTROL;
635 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
636 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
637 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME).withObjectCannedACL(expected);
638 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
639 |
640 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
641 | sqsExtended.sendMessage(messageRequest);
642 |
643 | ArgumentCaptor captor = ArgumentCaptor.forClass(PutObjectRequest.class);
644 |
645 | verify(mockS3).putObject(captor.capture(), any(RequestBody.class));
646 |
647 | assertEquals(expected, captor.getValue().acl());
648 | }
649 |
650 | @Test
651 | public void testWhenSendLargeMessageWithS3PrefixKeyDefined() {
652 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
653 |
654 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
655 |
656 | extendedSqsWithS3KeyPrefix.sendMessage(messageRequest);
657 |
658 | verify(mockS3, times(1)).putObject(
659 | argThat((PutObjectRequest obj) -> obj.key().equals(S3_KEY_PREFIX + S3_KEY_UUID)),
660 | isA(RequestBody.class));
661 | }
662 |
663 | @Test
664 | public void testWhenSendLargeMessageWithUndefinedS3PrefixKey() {
665 | String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
666 |
667 | SendMessageRequest messageRequest = SendMessageRequest.builder().queueUrl(SQS_QUEUE_URL).messageBody(messageBody).build();
668 |
669 | extendedSqsWithDefaultConfig.sendMessage(messageRequest);
670 |
671 | verify(mockS3, times(1)).putObject(
672 | argThat((PutObjectRequest obj) -> obj.key().equals(S3_KEY_UUID)),
673 | isA(RequestBody.class));
674 | }
675 |
676 | private void testReceiveMessage_when_MessageIsLarge(String reservedAttributeName) {
677 | String pointer = new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson();
678 | Message message = Message.builder()
679 | .messageAttributes(ImmutableMap.of(reservedAttributeName, MessageAttributeValue.builder().build()))
680 | .body(pointer)
681 | .build();
682 | String expectedMessage = "LargeMessage";
683 | GetObjectRequest getObjectRequest = GetObjectRequest.builder()
684 | .bucket(S3_BUCKET_NAME)
685 | .key("S3Key")
686 | .build();
687 |
688 | ResponseInputStream s3Object = new ResponseInputStream<>(GetObjectResponse.builder().build(), AbortableInputStream.create(new StringInputStream(expectedMessage)));
689 | // S3Object s3Object = S3Object.builder().build();
690 | // s3Object.setObjectContent(new StringInputStream(expectedMessage));
691 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(
692 | ReceiveMessageResponse.builder().messages(message).build());
693 | when(mockS3.getObject(isA(GetObjectRequest.class))).thenReturn(s3Object);
694 |
695 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
696 | ReceiveMessageResponse actualReceiveMessageResponse = extendedSqsWithDefaultConfig.receiveMessage(messageRequest);
697 | Message actualMessage = actualReceiveMessageResponse.messages().get(0);
698 |
699 | assertEquals(expectedMessage, actualMessage.body());
700 | assertFalse(actualMessage.messageAttributes().keySet().containsAll(AmazonSQSExtendedClientUtil.RESERVED_ATTRIBUTE_NAMES));
701 | verify(mockS3, times(1)).getObject(isA(GetObjectRequest.class));
702 | }
703 |
704 | @Test
705 | public void testReceiveMessage_when_ignorePayloadNotFound_then_messageWithPayloadNotFoundIsDeletedFromSQS() {
706 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
707 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
708 | .withIgnorePayloadNotFound(true);
709 |
710 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
711 |
712 | String receiptHandle = "receipt-handle";
713 | Message message = Message.builder()
714 | .messageAttributes(ImmutableMap.of(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, MessageAttributeValue.builder().build()))
715 | .body(new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson())
716 | .receiptHandle(receiptHandle)
717 | .build();
718 |
719 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(ReceiveMessageResponse.builder().messages(message).build());
720 |
721 | doThrow(NoSuchKeyException.class).when(mockS3).getObject(any(GetObjectRequest.class));
722 |
723 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().queueUrl(SQS_QUEUE_URL).build();
724 | ReceiveMessageResponse receiveMessageResponse = sqsExtended.receiveMessage(messageRequest);
725 |
726 | assertTrue(receiveMessageResponse.messages().isEmpty());
727 |
728 | ArgumentCaptor deleteMessageRequestArgumentCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
729 | verify(mockSqsBackend).deleteMessage(deleteMessageRequestArgumentCaptor.capture());
730 | assertEquals(SQS_QUEUE_URL, deleteMessageRequestArgumentCaptor.getValue().queueUrl());
731 | assertEquals(receiptHandle, deleteMessageRequestArgumentCaptor.getValue().receiptHandle());
732 | }
733 |
734 | @Test
735 | public void testReceiveMessage_when_ignorePayloadNotFoundIsFalse_then_messageWithPayloadNotFoundThrowsException() {
736 | ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
737 | .withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
738 | .withIgnorePayloadNotFound(false);
739 | SqsClient sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
740 |
741 | Message message = Message.builder()
742 | .messageAttributes(ImmutableMap.of(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, MessageAttributeValue.builder().build()))
743 | .body(new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson())
744 | .receiptHandle("receipt-handle")
745 | .build();
746 |
747 | when(mockSqsBackend.receiveMessage(isA(ReceiveMessageRequest.class))).thenReturn(ReceiveMessageResponse.builder().messages(message).build());
748 | doThrow(NoSuchKeyException.class).when(mockS3).getObject(any(GetObjectRequest.class));
749 |
750 | ReceiveMessageRequest messageRequest = ReceiveMessageRequest.builder().build();
751 | try {
752 | sqsExtended.receiveMessage(messageRequest);
753 | fail("Expected exception after receiving NoSuchKeyException from S3 was not thrown.");
754 | } catch (SdkException e) {
755 | assertEquals(NoSuchKeyException.class.getName(), e.getCause().getClass().getName());
756 | verify(mockSqsBackend, never()).deleteMessage(any(DeleteMessageRequest.class));
757 | }
758 | }
759 |
760 | private DeleteMessageBatchRequest generateLargeDeleteBatchRequest(List originalReceiptHandles) {
761 | List deleteEntries = IntStream.range(0, originalReceiptHandles.size())
762 | .mapToObj(i -> DeleteMessageBatchRequestEntry.builder()
763 | .id(Integer.toString(i))
764 | .receiptHandle(getSampleLargeReceiptHandle(originalReceiptHandles.get(i)))
765 | .build())
766 | .collect(Collectors.toList());
767 |
768 | return DeleteMessageBatchRequest.builder().queueUrl(SQS_QUEUE_URL).entries(deleteEntries).build();
769 | }
770 |
771 | private String getLargeReceiptHandle(String s3Key, String originalReceiptHandle) {
772 | return SQSExtendedClientConstants.S3_BUCKET_NAME_MARKER + S3_BUCKET_NAME
773 | + SQSExtendedClientConstants.S3_BUCKET_NAME_MARKER + SQSExtendedClientConstants.S3_KEY_MARKER
774 | + s3Key + SQSExtendedClientConstants.S3_KEY_MARKER + originalReceiptHandle;
775 | }
776 |
777 | private String getSampleLargeReceiptHandle(String originalReceiptHandle) {
778 | return getLargeReceiptHandle(UUID.randomUUID().toString(), originalReceiptHandle);
779 | }
780 | }
781 |
--------------------------------------------------------------------------------