use of com.amazonaws.services.sqs.model.DeleteMessageBatchRequest in project nifi by apache.
the class TestGetSQS method testGetMessageAndAutoDelete.
@Test
public void testGetMessageAndAutoDelete() {
runner.setProperty(GetSQS.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000");
runner.setProperty(GetSQS.AUTO_DELETE, "true");
Message message1 = new Message();
message1.setBody("TestMessage1");
message1.setMessageId("test-message-id-1");
message1.setReceiptHandle("test-receipt-handle-1");
Message message2 = new Message();
message2.setBody("TestMessage2");
message2.setMessageId("test-message-id-2");
message2.setReceiptHandle("test-receipt-handle-2");
ReceiveMessageResult receiveMessageResult = new ReceiveMessageResult().withMessages(message1, message2);
Mockito.when(mockSQSClient.receiveMessage(Mockito.any(ReceiveMessageRequest.class))).thenReturn(receiveMessageResult);
runner.run(1);
ArgumentCaptor<ReceiveMessageRequest> captureReceiveRequest = ArgumentCaptor.forClass(ReceiveMessageRequest.class);
Mockito.verify(mockSQSClient, Mockito.times(1)).receiveMessage(captureReceiveRequest.capture());
ReceiveMessageRequest receiveRequest = captureReceiveRequest.getValue();
assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", receiveRequest.getQueueUrl());
ArgumentCaptor<DeleteMessageBatchRequest> captureDeleteRequest = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
Mockito.verify(mockSQSClient, Mockito.times(1)).deleteMessageBatch(captureDeleteRequest.capture());
DeleteMessageBatchRequest deleteRequest = captureDeleteRequest.getValue();
assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", deleteRequest.getQueueUrl());
assertEquals("test-message-id-1", deleteRequest.getEntries().get(0).getId());
assertEquals("test-message-id-2", deleteRequest.getEntries().get(1).getId());
runner.assertAllFlowFilesTransferred(GetSQS.REL_SUCCESS, 2);
List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(GetSQS.REL_SUCCESS);
MockFlowFile ff0 = flowFiles.get(0);
ff0.assertAttributeEquals("sqs.message.id", "test-message-id-1");
MockFlowFile ff1 = flowFiles.get(1);
ff1.assertAttributeEquals("sqs.message.id", "test-message-id-2");
}
use of com.amazonaws.services.sqs.model.DeleteMessageBatchRequest in project nifi by apache.
the class GetSQS method onTrigger.
@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
final String queueUrl = context.getProperty(DYNAMIC_QUEUE_URL).evaluateAttributeExpressions().getValue();
final AmazonSQSClient client = getClient();
final ReceiveMessageRequest request = new ReceiveMessageRequest();
request.setAttributeNames(Collections.singleton("All"));
request.setMessageAttributeNames(Collections.singleton("All"));
request.setMaxNumberOfMessages(context.getProperty(BATCH_SIZE).asInteger());
request.setVisibilityTimeout(context.getProperty(VISIBILITY_TIMEOUT).asTimePeriod(TimeUnit.SECONDS).intValue());
request.setQueueUrl(queueUrl);
request.setWaitTimeSeconds(context.getProperty(RECEIVE_MSG_WAIT_TIME).asTimePeriod(TimeUnit.SECONDS).intValue());
final Charset charset = Charset.forName(context.getProperty(CHARSET).getValue());
final ReceiveMessageResult result;
try {
result = client.receiveMessage(request);
} catch (final Exception e) {
getLogger().error("Failed to receive messages from Amazon SQS due to {}", new Object[] { e });
context.yield();
return;
}
final List<Message> messages = result.getMessages();
if (messages.isEmpty()) {
context.yield();
return;
}
final boolean autoDelete = context.getProperty(AUTO_DELETE).asBoolean();
for (final Message message : messages) {
FlowFile flowFile = session.create();
final Map<String, String> attributes = new HashMap<>();
for (final Map.Entry<String, String> entry : message.getAttributes().entrySet()) {
attributes.put("sqs." + entry.getKey(), entry.getValue());
}
for (final Map.Entry<String, MessageAttributeValue> entry : message.getMessageAttributes().entrySet()) {
attributes.put("sqs." + entry.getKey(), entry.getValue().getStringValue());
}
attributes.put("hash.value", message.getMD5OfBody());
attributes.put("hash.algorithm", "md5");
attributes.put("sqs.message.id", message.getMessageId());
attributes.put("sqs.receipt.handle", message.getReceiptHandle());
flowFile = session.putAllAttributes(flowFile, attributes);
flowFile = session.write(flowFile, new OutputStreamCallback() {
@Override
public void process(final OutputStream out) throws IOException {
out.write(message.getBody().getBytes(charset));
}
});
session.transfer(flowFile, REL_SUCCESS);
session.getProvenanceReporter().receive(flowFile, queueUrl);
getLogger().info("Successfully received {} from Amazon SQS", new Object[] { flowFile });
}
if (autoDelete) {
// If we want to auto-delete messages, we must fist commit the session to ensure that the data
// is persisted in NiFi's repositories.
session.commit();
final DeleteMessageBatchRequest deleteRequest = new DeleteMessageBatchRequest();
deleteRequest.setQueueUrl(queueUrl);
final List<DeleteMessageBatchRequestEntry> deleteRequestEntries = new ArrayList<>();
for (final Message message : messages) {
final DeleteMessageBatchRequestEntry entry = new DeleteMessageBatchRequestEntry();
entry.setId(message.getMessageId());
entry.setReceiptHandle(message.getReceiptHandle());
deleteRequestEntries.add(entry);
}
deleteRequest.setEntries(deleteRequestEntries);
try {
client.deleteMessageBatch(deleteRequest);
} catch (final Exception e) {
getLogger().error("Received {} messages from Amazon SQS but failed to delete the messages; these messages" + " may be duplicated. Reason for deletion failure: {}", new Object[] { messages.size(), e });
}
}
}
use of com.amazonaws.services.sqs.model.DeleteMessageBatchRequest in project krypton-android by kryptco.
the class SQSTransport method receiveMessages.
public static List<byte[]> receiveMessages(final Pairing pairing) throws TransportException {
final AmazonSQSClient client = getClient();
ReceiveMessageRequest request = new ReceiveMessageRequest(sendQueueURL(pairing));
request.setWaitTimeSeconds(10);
request.setMaxNumberOfMessages(10);
ReceiveMessageResult result = client.receiveMessage(request);
final List<DeleteMessageBatchRequestEntry> deleteEntries = new ArrayList<>();
ArrayList<byte[]> messages = new ArrayList<byte[]>();
for (Message m : result.getMessages()) {
deleteEntries.add(new DeleteMessageBatchRequestEntry(m.getMessageId(), m.getReceiptHandle()));
try {
messages.add(Base64.decode(m.getBody()));
} catch (Exception e) {
Log.e(TAG, "failed to decode message: " + e.getMessage());
}
}
if (!deleteEntries.isEmpty()) {
deleteThreadPool.submit(() -> {
try {
DeleteMessageBatchRequest deleteRequest = new DeleteMessageBatchRequest(sendQueueURL(pairing)).withEntries(deleteEntries);
client.deleteMessageBatch(deleteRequest);
} catch (Exception e) {
Log.e(TAG, "failed to delete messages: " + e.getMessage());
}
});
}
return messages;
}
use of com.amazonaws.services.sqs.model.DeleteMessageBatchRequest in project nifi by apache.
the class TestDeleteSQS method testDeleteException.
@Test
public void testDeleteException() {
runner.setProperty(DeleteSQS.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000");
final Map<String, String> ff1Attributes = new HashMap<>();
ff1Attributes.put("filename", "1.txt");
ff1Attributes.put("sqs.receipt.handle", "test-receipt-handle-1");
runner.enqueue("TestMessageBody1", ff1Attributes);
Mockito.when(mockSQSClient.deleteMessageBatch(Mockito.any(DeleteMessageBatchRequest.class))).thenThrow(new AmazonSQSException("TestFail"));
runner.assertValid();
runner.run(1);
ArgumentCaptor<DeleteMessageBatchRequest> captureDeleteRequest = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
Mockito.verify(mockSQSClient, Mockito.times(1)).deleteMessageBatch(captureDeleteRequest.capture());
runner.assertAllFlowFilesTransferred(DeleteSQS.REL_FAILURE, 1);
}
use of com.amazonaws.services.sqs.model.DeleteMessageBatchRequest in project nifi by apache.
the class TestDeleteSQS method testDeleteSingleMessage.
@Test
public void testDeleteSingleMessage() {
runner.setProperty(DeleteSQS.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000");
final Map<String, String> ffAttributes = new HashMap<>();
ffAttributes.put("filename", "1.txt");
ffAttributes.put("sqs.receipt.handle", "test-receipt-handle-1");
runner.enqueue("TestMessageBody", ffAttributes);
runner.assertValid();
runner.run(1);
ArgumentCaptor<DeleteMessageBatchRequest> captureDeleteRequest = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
Mockito.verify(mockSQSClient, Mockito.times(1)).deleteMessageBatch(captureDeleteRequest.capture());
DeleteMessageBatchRequest deleteRequest = captureDeleteRequest.getValue();
assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", deleteRequest.getQueueUrl());
assertEquals("test-receipt-handle-1", deleteRequest.getEntries().get(0).getReceiptHandle());
runner.assertAllFlowFilesTransferred(DeleteSQS.REL_SUCCESS, 1);
}
Aggregations