diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/sqs/ChangeSQSMessageVisibilityTimeout.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/sqs/ChangeSQSMessageVisibilityTimeout.java new file mode 100644 index 000000000000..b558cb53c225 --- /dev/null +++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/sqs/ChangeSQSMessageVisibilityTimeout.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.aws.sqs; + +import org.apache.nifi.annotation.behavior.InputRequirement; +import org.apache.nifi.annotation.behavior.InputRequirement.Requirement; +import org.apache.nifi.annotation.behavior.SupportsBatching; +import org.apache.nifi.annotation.documentation.CapabilityDescription; +import org.apache.nifi.annotation.documentation.SeeAlso; +import org.apache.nifi.annotation.documentation.Tags; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.flowfile.attributes.CoreAttributes; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.util.StandardValidators; +import org.apache.nifi.processors.aws.AbstractAwsSyncProcessor; +import software.amazon.awssdk.services.sqs.SqsClient; +import software.amazon.awssdk.services.sqs.SqsClientBuilder; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequestEntry; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse; + +import java.util.List; + +import static org.apache.nifi.processors.aws.region.RegionUtil.CUSTOM_REGION; +import static org.apache.nifi.processors.aws.region.RegionUtil.REGION; + +@SupportsBatching +@SeeAlso({GetSQS.class, PutSQS.class}) +@InputRequirement(Requirement.INPUT_REQUIRED) +@Tags({"Amazon", "AWS", "SQS", "Queue", "Update"}) +@CapabilityDescription("Updates the visibility timeout for a message from an Amazon Simple Queuing Service Queue") +public class ChangeSQSMessageVisibilityTimeout extends AbstractAwsSyncProcessor { + + public static final PropertyDescriptor QUEUE_URL = new PropertyDescriptor.Builder() + .name("Queue URL") + .description("The URL of the queue delete from") + .addValidator(StandardValidators.URL_VALIDATOR) + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .required(true) + .build(); + + public static final PropertyDescriptor RECEIPT_HANDLE = new PropertyDescriptor.Builder() + .name("Receipt Handle") + .description("The identifier that specifies the receipt of the message") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .required(true) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .defaultValue("${sqs.receipt.handle}") + .build(); + + public static final PropertyDescriptor VISIBILITY_TIMEOUT = new PropertyDescriptor.Builder() + .name("Visibility Timeout") + .description("The amount of time after a message is received but not deleted that the message is hidden from other consumers") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .required(true) + .defaultValue("15 mins") + .addValidator(StandardValidators.TIME_PERIOD_VALIDATOR) + .build(); + + public static final List PROPERTY_DESCRIPTORS = List.of( + QUEUE_URL, + VISIBILITY_TIMEOUT, + REGION, + CUSTOM_REGION, + AWS_CREDENTIALS_PROVIDER_SERVICE, + SSL_CONTEXT_SERVICE, + RECEIPT_HANDLE, + TIMEOUT, + ENDPOINT_OVERRIDE, + PROXY_CONFIGURATION_SERVICE + ); + + @Override + protected List getSupportedPropertyDescriptors() { + return PROPERTY_DESCRIPTORS; + } + + @Override + public void onTrigger(final ProcessContext context, final ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + + final String queueUrl = context.getProperty(QUEUE_URL).evaluateAttributeExpressions(flowFile).getValue(); + + final SqsClient client = getClient(context); + + final String receiptHandle = context.getProperty(RECEIPT_HANDLE).evaluateAttributeExpressions(flowFile).getValue(); + final int visibilityTimeout = Math.toIntExact(context.getProperty(VISIBILITY_TIMEOUT).evaluateAttributeExpressions(flowFile).asDuration().toSeconds()); + final String entryId = flowFile.getAttribute(CoreAttributes.UUID.key()); + final ChangeMessageVisibilityBatchRequestEntry entry = ChangeMessageVisibilityBatchRequestEntry.builder() + .receiptHandle(receiptHandle) + .id(entryId) + .visibilityTimeout(visibilityTimeout) + .build(); + + final ChangeMessageVisibilityBatchRequest request = ChangeMessageVisibilityBatchRequest.builder() + .queueUrl(queueUrl) + .entries(entry) + .build(); + + try { + ChangeMessageVisibilityBatchResponse response = client.changeMessageVisibilityBatch(request); + + // check for errors + if (!response.failed().isEmpty()) { + throw new ProcessException(response.failed().getFirst().toString()); + } + + getLogger().info("Successfully updated visibility timeout for SQS message for {}", flowFile); + session.transfer(flowFile, REL_SUCCESS); + } catch (final Exception e) { + getLogger().error("Failed to update visiibility timeout for SQS message", e); + flowFile = session.penalize(flowFile); + session.transfer(flowFile, REL_FAILURE); + } + } + + @Override + protected SqsClientBuilder createClientBuilder(final ProcessContext context) { + return SqsClient.builder(); + } +} diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor index b8885827c0ce..eb230a1e547e 100644 --- a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor +++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor @@ -22,6 +22,7 @@ org.apache.nifi.processors.aws.s3.DeleteS3Object org.apache.nifi.processors.aws.s3.TagS3Object org.apache.nifi.processors.aws.s3.ListS3 org.apache.nifi.processors.aws.sns.PutSNS +org.apache.nifi.processors.aws.sqs.ChangeSQSMessageVisibilityTimeout org.apache.nifi.processors.aws.sqs.GetSQS org.apache.nifi.processors.aws.sqs.PutSQS org.apache.nifi.processors.aws.sqs.DeleteSQS diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/ITChangeSQSMessageVisibilityTimeout.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/ITChangeSQSMessageVisibilityTimeout.java new file mode 100644 index 000000000000..8f6e2f2a6fab --- /dev/null +++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/ITChangeSQSMessageVisibilityTimeout.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.aws.sqs; + +import org.apache.nifi.util.TestRunner; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.services.sqs.model.Message; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse; +import software.amazon.awssdk.services.sqs.model.SendMessageRequest; +import software.amazon.awssdk.services.sqs.model.SendMessageResponse; + +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class ITChangeSQSMessageVisibilityTimeout extends AbstractSQSIT { + + @Test + public void testSimpleUpdate() { + final SendMessageRequest request = SendMessageRequest.builder() + .queueUrl(getQueueUrl()) + .messageBody("Hello World") + .build(); + final SendMessageResponse response = getClient().sendMessage(request); + assertTrue(response.sdkHttpResponse().isSuccessful()); + + // Setup - receive message to get receipt handle + final ReceiveMessageRequest receiveMessageRequest = ReceiveMessageRequest.builder() + .queueUrl(getQueueUrl()) + .build(); + final ReceiveMessageResponse receiveMessageResult = getClient().receiveMessage(receiveMessageRequest); + assertEquals(200, receiveMessageResult.sdkHttpResponse().statusCode()); + final Message updateMessage = receiveMessageResult.messages().getFirst(); + final String receiptHandle = updateMessage.receiptHandle(); + + // Test - update message with ChangeSQSMessageVisibilityTimeout + final TestRunner runner = initRunner(ChangeSQSMessageVisibilityTimeout.class); + final Map ffAttributes = new HashMap<>(); + ffAttributes.put("filename", "1.txt"); + ffAttributes.put("sqs.receipt.handle", receiptHandle); + runner.enqueue("TestMessageBody", ffAttributes); + + runner.run(1); + + runner.assertAllFlowFilesTransferred(ChangeSQSMessageVisibilityTimeout.REL_SUCCESS, 1); + } +} diff --git a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/TestChangeSQSMessageVisibilityTimeout.java b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/TestChangeSQSMessageVisibilityTimeout.java new file mode 100644 index 000000000000..c9baffee016d --- /dev/null +++ b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/sqs/TestChangeSQSMessageVisibilityTimeout.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.processors.aws.sqs; + +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.testutil.AuthUtils; +import org.apache.nifi.util.TestRunner; +import org.apache.nifi.util.TestRunners; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import software.amazon.awssdk.services.sqs.SqsClient; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class TestChangeSQSMessageVisibilityTimeout { + + private TestRunner runner = null; + private SqsClient mockSQSClient = null; + + @BeforeEach + public void setUp() { + mockSQSClient = Mockito.mock(SqsClient.class); + ChangeMessageVisibilityBatchResponse mockResponse = ChangeMessageVisibilityBatchResponse.builder() + .failed(Collections.emptyList()) + .build(); + Mockito.when(mockSQSClient.changeMessageVisibilityBatch(Mockito.any(ChangeMessageVisibilityBatchRequest.class))).thenReturn(mockResponse); + ChangeSQSMessageVisibilityTimeout mockChangeSQSMessageVisibilityTimeout = new ChangeSQSMessageVisibilityTimeout() { + + @Override + protected SqsClient getClient(ProcessContext context) { + return mockSQSClient; + } + }; + runner = TestRunners.newTestRunner(mockChangeSQSMessageVisibilityTimeout); + AuthUtils.enableAccessKey(runner, "accessKeyId", "secretKey"); + } + + @Test + public void testChangeVisibilityTimeoutSingleMessage() { + runner.setProperty(ChangeSQSMessageVisibilityTimeout.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000"); + final Map 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 captureChangeMessageVisibilityTimeoutRequest = ArgumentCaptor.forClass(ChangeMessageVisibilityBatchRequest.class); + Mockito.verify(mockSQSClient, Mockito.times(1)).changeMessageVisibilityBatch(captureChangeMessageVisibilityTimeoutRequest.capture()); + ChangeMessageVisibilityBatchRequest changeMessageVisibilityTimeoutRequest = captureChangeMessageVisibilityTimeoutRequest.getValue(); + assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", changeMessageVisibilityTimeoutRequest.queueUrl()); + assertEquals("test-receipt-handle-1", changeMessageVisibilityTimeoutRequest.entries().getFirst().receiptHandle()); + assertEquals(15 * 60, changeMessageVisibilityTimeoutRequest.entries().getFirst().visibilityTimeout()); + + runner.assertAllFlowFilesTransferred(ChangeSQSMessageVisibilityTimeout.REL_SUCCESS, 1); + } + + @Test + public void testChangeVisibilityTimeoutWithCustomReceiptHandle() { + runner.setProperty(ChangeSQSMessageVisibilityTimeout.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000"); + runner.setProperty(ChangeSQSMessageVisibilityTimeout.RECEIPT_HANDLE, "${custom.receipt.handle}"); + runner.setProperty(ChangeSQSMessageVisibilityTimeout.VISIBILITY_TIMEOUT, "${custom.timeout.value}"); + final Map ffAttributes = new HashMap<>(); + ffAttributes.put("filename", "1.txt"); + ffAttributes.put("custom.receipt.handle", "test-receipt-handle-1"); + ffAttributes.put("custom.timeout.value", "1 min"); + runner.enqueue("TestMessageBody", ffAttributes); + + runner.assertValid(); + runner.run(1); + + ArgumentCaptor captureChangeMessageVisibilityRequest = ArgumentCaptor.forClass(ChangeMessageVisibilityBatchRequest.class); + Mockito.verify(mockSQSClient, Mockito.times(1)).changeMessageVisibilityBatch(captureChangeMessageVisibilityRequest.capture()); + ChangeMessageVisibilityBatchRequest changeMessageVisibilityRequest = captureChangeMessageVisibilityRequest.getValue(); + assertEquals(60, changeMessageVisibilityRequest.entries().getFirst().visibilityTimeout()); + + runner.assertAllFlowFilesTransferred(ChangeSQSMessageVisibilityTimeout.REL_SUCCESS, 1); + } + + @Test + public void testChangeMessageVisibilityTimeoutException() { + runner.setProperty(ChangeSQSMessageVisibilityTimeout.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000"); + final Map ff1Attributes = new HashMap<>(); + ff1Attributes.put("filename", "1.txt"); + ff1Attributes.put("sqs.receipt.handle", "test-receipt-handle-1"); + runner.enqueue("TestMessageBody1", ff1Attributes); + Mockito.when(mockSQSClient.changeMessageVisibilityBatch(Mockito.any(ChangeMessageVisibilityBatchRequest.class))) + .thenThrow(new RuntimeException()); + + runner.assertValid(); + runner.run(1); + + ArgumentCaptor captureChangeVisibilityTimeoutRequest = ArgumentCaptor.forClass(ChangeMessageVisibilityBatchRequest.class); + Mockito.verify(mockSQSClient, Mockito.times(1)).changeMessageVisibilityBatch(captureChangeVisibilityTimeoutRequest.capture()); + + runner.assertAllFlowFilesTransferred(ChangeSQSMessageVisibilityTimeout.REL_FAILURE, 1); + } + +}