Search in sources :

Example 1 with DeleteMessageBatchRequest

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");
}
Also used : MockFlowFile(org.apache.nifi.util.MockFlowFile) DeleteMessageBatchRequest(com.amazonaws.services.sqs.model.DeleteMessageBatchRequest) Message(com.amazonaws.services.sqs.model.Message) ReceiveMessageRequest(com.amazonaws.services.sqs.model.ReceiveMessageRequest) ReceiveMessageResult(com.amazonaws.services.sqs.model.ReceiveMessageResult) Test(org.junit.Test)

Example 2 with DeleteMessageBatchRequest

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 });
        }
    }
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) DeleteMessageBatchRequest(com.amazonaws.services.sqs.model.DeleteMessageBatchRequest) ReceiveMessageRequest(com.amazonaws.services.sqs.model.ReceiveMessageRequest) Message(com.amazonaws.services.sqs.model.Message) HashMap(java.util.HashMap) OutputStream(java.io.OutputStream) ArrayList(java.util.ArrayList) Charset(java.nio.charset.Charset) AmazonSQSClient(com.amazonaws.services.sqs.AmazonSQSClient) IOException(java.io.IOException) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) DeleteMessageBatchRequestEntry(com.amazonaws.services.sqs.model.DeleteMessageBatchRequestEntry) HashMap(java.util.HashMap) Map(java.util.Map) ReceiveMessageResult(com.amazonaws.services.sqs.model.ReceiveMessageResult) MessageAttributeValue(com.amazonaws.services.sqs.model.MessageAttributeValue)

Example 3 with DeleteMessageBatchRequest

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;
}
Also used : DeleteMessageBatchRequest(com.amazonaws.services.sqs.model.DeleteMessageBatchRequest) ReceiveMessageRequest(com.amazonaws.services.sqs.model.ReceiveMessageRequest) Message(com.amazonaws.services.sqs.model.Message) NetworkMessage(co.krypt.krypton.protocol.NetworkMessage) ArrayList(java.util.ArrayList) AmazonSQSClient(com.amazonaws.services.sqs.AmazonSQSClient) DeleteMessageBatchRequestEntry(com.amazonaws.services.sqs.model.DeleteMessageBatchRequestEntry) TransportException(co.krypt.krypton.exception.TransportException) ReceiveMessageResult(com.amazonaws.services.sqs.model.ReceiveMessageResult)

Example 4 with DeleteMessageBatchRequest

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);
}
Also used : DeleteMessageBatchRequest(com.amazonaws.services.sqs.model.DeleteMessageBatchRequest) HashMap(java.util.HashMap) AmazonSQSException(com.amazonaws.services.sqs.model.AmazonSQSException) Test(org.junit.Test)

Example 5 with DeleteMessageBatchRequest

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);
}
Also used : DeleteMessageBatchRequest(com.amazonaws.services.sqs.model.DeleteMessageBatchRequest) HashMap(java.util.HashMap) Test(org.junit.Test)

Aggregations

DeleteMessageBatchRequest (com.amazonaws.services.sqs.model.DeleteMessageBatchRequest)6 AmazonSQSClient (com.amazonaws.services.sqs.AmazonSQSClient)3 DeleteMessageBatchRequestEntry (com.amazonaws.services.sqs.model.DeleteMessageBatchRequestEntry)3 Message (com.amazonaws.services.sqs.model.Message)3 ReceiveMessageRequest (com.amazonaws.services.sqs.model.ReceiveMessageRequest)3 ReceiveMessageResult (com.amazonaws.services.sqs.model.ReceiveMessageResult)3 ArrayList (java.util.ArrayList)3 HashMap (java.util.HashMap)3 Test (org.junit.Test)3 FlowFile (org.apache.nifi.flowfile.FlowFile)2 TransportException (co.krypt.krypton.exception.TransportException)1 NetworkMessage (co.krypt.krypton.protocol.NetworkMessage)1 AmazonSQSException (com.amazonaws.services.sqs.model.AmazonSQSException)1 MessageAttributeValue (com.amazonaws.services.sqs.model.MessageAttributeValue)1 IOException (java.io.IOException)1 OutputStream (java.io.OutputStream)1 Charset (java.nio.charset.Charset)1 List (java.util.List)1 Map (java.util.Map)1 OutputStreamCallback (org.apache.nifi.processor.io.OutputStreamCallback)1