NIFI-16113 Adds new processor for changing visibility timeouts of mes…

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<SqsClient, SqsClientBuilder> {
+
+    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<PropertyDescriptor> 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<PropertyDescriptor> 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();
+    }
+}

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

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<String, String> 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);
+    }
+}

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<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<ChangeMessageVisibilityBatchRequest> 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<String, String> 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<ChangeMessageVisibilityBatchRequest> 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<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.changeMessageVisibilityBatch(Mockito.any(ChangeMessageVisibilityBatchRequest.class)))
+                .thenThrow(new RuntimeException());
+
+        runner.assertValid();
+        runner.run(1);
+
+        ArgumentCaptor<ChangeMessageVisibilityBatchRequest> captureChangeVisibilityTimeoutRequest = ArgumentCaptor.forClass(ChangeMessageVisibilityBatchRequest.class);
+        Mockito.verify(mockSQSClient, Mockito.times(1)).changeMessageVisibilityBatch(captureChangeVisibilityTimeoutRequest.capture());
+
+        runner.assertAllFlowFilesTransferred(ChangeSQSMessageVisibilityTimeout.REL_FAILURE, 1);
+    }
+
+}