ISSUE-30: added ignorePayloadNotFound optional parameter in case you want to get rid of SQS messages that do not have their corresponding payload

Author: Damian Nardelli

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

@@ -365,7 +365,17 @@ public ReceiveMessageResult receiveMessage(ReceiveMessageRequest receiveMessageR
 				String s3MsgBucketName = s3Pointer.getS3BucketName();
 				String s3MsgKey = s3Pointer.getS3Key();
 
-				String origMsgBody = getTextFromS3(s3MsgBucketName, s3MsgKey);
+				final String origMsgBody;
+				try {
+					origMsgBody = getTextFromS3(s3MsgBucketName, s3MsgKey);
+				} catch (AmazonServiceException e) {
+					if (((AmazonServiceException) e.getCause()).getErrorCode().equals("NoSuchKey")) {
+						deleteMessage(receiveMessageRequest.getQueueUrl(), message.getReceiptHandle());
+						LOG.warn("SQS message deleted as it could not be found in S3");
+						continue;
+					}
+					throw e;
+				}
 				LOG.info("S3 object read, Bucket name: " + s3MsgBucketName + ", Object key: " + s3MsgKey + ".");
 
 				message.setBody(origMsgBody);

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

@@ -35,6 +35,7 @@ public class ExtendedClientConfiguration {
 	private String s3BucketName;
 	private boolean largePayloadSupport = false;
 	private boolean alwaysThroughS3 = false;
+	private boolean ignorePayloadNotFound = false;
 	private int messageSizeThreshold = SQSExtendedClientConstants.DEFAULT_MESSAGE_SIZE_THRESHOLD;
 
 	public ExtendedClientConfiguration() {
@@ -47,6 +48,7 @@ public ExtendedClientConfiguration(ExtendedClientConfiguration other) {
 		this.s3BucketName = other.s3BucketName;
 		this.largePayloadSupport = other.largePayloadSupport;
 		this.alwaysThroughS3 = other.alwaysThroughS3;
+		this.ignorePayloadNotFound = other.ignorePayloadNotFound;
 		this.messageSizeThreshold = other.messageSizeThreshold;
 	}
 
@@ -214,4 +216,43 @@ public ExtendedClientConfiguration withAlwaysThroughS3(boolean alwaysThroughS3)
 	public boolean isAlwaysThroughS3() {
 		return alwaysThroughS3;
 	}
+
+	/**
+	 * Sets whether or not messages should be removed from Amazon SQS
+	 * when payloads are not found in Amazon S3.
+	 *
+	 * @param ignorePayloadNotFound
+	 *            Whether or not messages should be removed from Amazon SQS
+	 *            when payloads are not found in Amazon S3. Default: false
+	 */
+	public void setIgnorePayloadNotFound(boolean ignorePayloadNotFound) {
+		this.ignorePayloadNotFound = ignorePayloadNotFound;
+	}
+
+	/**
+	 * Sets whether or not messages should be removed from Amazon SQS
+	 * when payloads are not found in Amazon S3.
+	 *
+	 * @param ignorePayloadNotFound
+	 *            Whether or not messages should be removed from Amazon SQS
+	 *            when payloads are not found in Amazon S3. Default: false
+	 * @return the updated ExtendedClientConfiguration object.
+	 */
+	public ExtendedClientConfiguration withIgnorePayloadNotFound(boolean ignorePayloadNotFound) {
+		setIgnorePayloadNotFound(ignorePayloadNotFound);
+		return this;
+	}
+
+	/**
+	 * Checks whether or not messages should be removed from Amazon SQS
+	 * when payloads are not found in Amazon S3.
+	 *
+	 * @return True if messages should be removed from Amazon SQS
+	 *         when payloads are not found in Amazon S3. Default: false
+	 */
+	public boolean isIgnorePayloadNotFound() {
+		return ignorePayloadNotFound;
+	}
+
+
 }

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

@@ -17,13 +17,19 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import com.amazonaws.AmazonServiceException;
 import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.s3.model.GetObjectRequest;
 import com.amazonaws.services.s3.model.PutObjectRequest;
 import com.amazonaws.services.sqs.AmazonSQS;
 import com.amazonaws.services.sqs.AmazonSQSClient;
+import com.amazonaws.services.sqs.model.DeleteMessageRequest;
+import com.amazonaws.services.sqs.model.Message;
 import com.amazonaws.services.sqs.model.MessageAttributeValue;
 import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
 import com.amazonaws.services.sqs.model.ReceiveMessageResult;
@@ -36,7 +42,9 @@
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 
+import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.isA;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -224,6 +232,39 @@ public void testWhenLargeMessgaeIsSentThenAttributeWithPayloadSizeIsAdded() {
         Assert.assertEquals(messageLength, (int)Integer.valueOf(attributes.get(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME).getStringValue()));
     }
 
+    @Test
+    public void testWhenIgnorePayloadNotFoundIsSentThenNotFoundKeysInS3AreDeletedInSQS() throws Exception {
+        ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration()
+            .withIgnorePayloadNotFound(true).withLargePayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
+
+        AmazonServiceException mockException = mock(AmazonServiceException.class);
+        when(mockException.getErrorCode()).thenReturn("NoSuchKey");
+
+        Message mockMessage = mock(Message.class);
+        MessageS3Pointer messageS3Pointer = new MessageS3Pointer();
+        when(mockMessage.getBody()).thenReturn(new JsonDataConverter().serializeToJson(messageS3Pointer));
+        when(mockMessage.getReceiptHandle()).thenReturn("receipt-handle");
+        Map<String, MessageAttributeValue> messageAttributes = new HashMap<>();
+        messageAttributes.put(SQSExtendedClientConstants.RESERVED_ATTRIBUTE_NAME, new MessageAttributeValue());
+        when(mockMessage.getMessageAttributes()).thenReturn(messageAttributes);
+
+        ReceiveMessageResult mockReceiveMessageResult = mock(ReceiveMessageResult.class);
+        when(mockSqsBackend.receiveMessage(any(ReceiveMessageRequest.class))).thenReturn(mockReceiveMessageResult);
+        when(mockReceiveMessageResult.getMessages()).thenReturn(Collections.singletonList(mockMessage));
+
+        doThrow(mockException).when(mockS3).getObject(any(GetObjectRequest.class));
+
+        AmazonSQS sqsExtended = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
+
+        sqsExtended.receiveMessage(SQS_QUEUE_URL);
+
+        ArgumentCaptor<DeleteMessageRequest> deleteMessageRequestArgumentCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class);
+        verify(mockSqsBackend).deleteMessage(deleteMessageRequestArgumentCaptor.capture());
+
+        Assert.assertEquals(SQS_QUEUE_URL, deleteMessageRequestArgumentCaptor.getValue().getQueueUrl());
+        Assert.assertEquals("receipt-handle", deleteMessageRequestArgumentCaptor.getValue().getReceiptHandle());
+    }
+
     private String generateStringWithLength(int messageLength) {
         char[] charArray = new char[messageLength];
         Arrays.fill(charArray, 'x');