exceptionfactory commented on code in PR #6589: URL: https://github.com/apache/nifi/pull/6589#discussion_r1030613669
########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml: ########## @@ -117,6 +117,26 @@ <version>1.19.0-SNAPSHOT</version> <scope>provided</scope> </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-translate</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-polly</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-transcribe</artifactId> + <version>${aws.ml.version}</version> Review Comment: ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + RESPONSE response; Review Comment: Recommend adding `final` keyword: ```suggestion final RESPONSE response; ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + RESPONSE response; + FlowFile childFlowFile; + try { + response = sendRequest(buildRequest(session, context, flowFile), context); + childFlowFile = writeToFlowFile(session, flowFile, response); + postProcessFlowFile(context, session, childFlowFile, response); + } catch (Exception e) { + session.transfer(flowFile, REL_FAILURE); + getLogger().error("Sending AWS ML Request failed", e); + return; + } + session.transfer(flowFile, REL_ORIGINAL); + session.transfer(childFlowFile, REL_SUCCESS); + } + + protected void postProcessFlowFile(ProcessContext context, ProcessSession session, FlowFile flowFile, RESPONSE response) { + session.putAttribute(flowFile, AWS_TASK_ID_PROPERTY, getAwsTaskId(context, response)); + getLogger().debug("AWS ML task has been started with task id: {}", getAwsTaskId(context, response)); + } + + protected REQUEST buildRequest(ProcessSession session, ProcessContext context, FlowFile flowFile) throws JsonProcessingException { + return mapper.readValue(getPayload(session, context, flowFile), getAwsRequestClass(context)); + } + + private String getPayload(ProcessSession session, ProcessContext context, FlowFile flowFile) { + String payloadPropertyValue = context.getProperty(JSON_PAYLOAD).evaluateAttributeExpressions(flowFile).getValue(); + if (payloadPropertyValue == null) { + payloadPropertyValue = readFlowFile(session, flowFile); + } + return payloadPropertyValue; + } + + @Override + protected T createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + throw new UnsupportedOperationException("Tried to create client in a deprecated way."); + } + + abstract protected RESPONSE sendRequest(REQUEST request, ProcessContext context) throws JsonProcessingException; + + abstract protected Class<? extends REQUEST> getAwsRequestClass(ProcessContext context); + + abstract protected String getAwsTaskId(ProcessContext context, RESPONSE response); + + protected FlowFile writeToFlowFile(ProcessSession session, FlowFile flowFile, RESPONSE response) { + FlowFile childFlowFile = session.create(flowFile); + session.write(childFlowFile, out -> { + try (OutputStreamWriter outputStreamWriter = new OutputStreamWriter(out, StandardCharsets.UTF_8)) { + outputStreamWriter.write(mapper.writeValueAsString(response)); + outputStreamWriter.flush(); + } Review Comment: The ObjectMapper.writeValue() method can be used: ```suggestion mapper.writeValue(out, response)); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyJobStatus.java: ########## @@ -0,0 +1,114 @@ +/* + * 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.ml.polly; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.AmazonPollyClientBuilder; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.TaskStatus; +import com.amazonaws.services.textract.model.ThrottlingException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Polly"}) +@CapabilityDescription("Retrieves the current status of an AWS Polly job.") +@SeeAlso({StartAwsPollyJob.class}) +public class GetAwsPollyJobStatus extends AwsMachineLearningJobStatusGetter<AmazonPollyClient> { + private static final String BUCKET = "bucket"; + private static final String KEY = "key"; + private static final Pattern S3_PATH = Pattern.compile("https://s3.*amazonaws.com/(?<" + BUCKET + ">[^/]+)/(?<" + KEY + ">.*)"); + private static final String AWS_S3_BUCKET = "PollyS3OutputBucket"; + private static final String AWS_S3_KEY = "PollyS3OutputKey"; Review Comment: These values should be defined using the `WritesAttributes` documentation annotation on the Processor. ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyJobStatus.java: ########## @@ -0,0 +1,114 @@ +/* + * 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.ml.polly; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.AmazonPollyClientBuilder; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.TaskStatus; +import com.amazonaws.services.textract.model.ThrottlingException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Polly"}) +@CapabilityDescription("Retrieves the current status of an AWS Polly job.") +@SeeAlso({StartAwsPollyJob.class}) +public class GetAwsPollyJobStatus extends AwsMachineLearningJobStatusGetter<AmazonPollyClient> { + private static final String BUCKET = "bucket"; + private static final String KEY = "key"; + private static final Pattern S3_PATH = Pattern.compile("https://s3.*amazonaws.com/(?<" + BUCKET + ">[^/]+)/(?<" + KEY + ">.*)"); + private static final String AWS_S3_BUCKET = "PollyS3OutputBucket"; + private static final String AWS_S3_KEY = "PollyS3OutputKey"; + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonPollyClient) AmazonPollyClientBuilder.standard() + .withCredentials(credentialsProvider) + .withRegion(context.getProperty(REGION).getValue()) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetSpeechSynthesisTaskResult speechSynthesisTask; + try { + speechSynthesisTask = getSynthesisTask(flowFile); + } catch (ThrottlingException e) { + getLogger().info("Request Rate Limit exceeded", e); + session.transfer(flowFile, REL_THROTTLED); + return; + } catch (Exception e) { + getLogger().info("Failed to get Polly Job status", e); + session.transfer(flowFile, REL_FAILURE); + return; + } + + TaskStatus taskStatus = TaskStatus.fromValue(speechSynthesisTask.getSynthesisTask().getTaskStatus()); + + if (taskStatus == TaskStatus.InProgress || taskStatus == TaskStatus.Scheduled) { + session.penalize(flowFile); + session.transfer(flowFile, REL_RUNNING); + } + + if (taskStatus == TaskStatus.Completed) { + String outputUri = speechSynthesisTask.getSynthesisTask().getOutputUri(); + + Matcher matcher = S3_PATH.matcher(outputUri); + if (matcher.find()) { + session.putAttribute(flowFile, AWS_S3_BUCKET, matcher.group(BUCKET)); + session.putAttribute(flowFile, AWS_S3_KEY, matcher.group(KEY)); + } + FlowFile childFlowFile = session.create(flowFile); + writeToFlowFile(session, childFlowFile, speechSynthesisTask); + session.putAttribute(childFlowFile, AWS_TASK_OUTPUT_LOCATION, outputUri); + session.transfer(flowFile, REL_ORIGINAL); + session.transfer(childFlowFile, REL_SUCCESS); + getLogger().info("Amazon Polly reported that the task completed for {}", flowFile); + return; + } + + if (taskStatus == TaskStatus.Failed) { + final String failureReason = speechSynthesisTask.getSynthesisTask().getTaskStatusReason(); + session.putAttribute(flowFile, FAILURE_REASON_ATTRIBUTE, failureReason); Review Comment: ```suggestion flowFile = session.putAttribute(flowFile, FAILURE_REASON_ATTRIBUTE, failureReason); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { Review Comment: ```suggestion public void onTrigger(final ProcessContext context, final ProcessSession session) { ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml: ########## @@ -117,6 +117,26 @@ <version>1.19.0-SNAPSHOT</version> <scope>provided</scope> </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-translate</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-polly</artifactId> + <version>${aws.ml.version}</version> Review Comment: ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + +public class GetAwsPollyStatusTest { + private static final String TEST_TASK_ID = "testTaskId"; + private TestRunner runner = null; + private AmazonPollyClient mockPollyClient = null; + private AWSCredentialsProviderService mockAwsCredentialsProvider = null; + + @BeforeEach + public void setUp() throws InitializationException { + mockPollyClient = Mockito.mock(AmazonPollyClient.class); + mockAwsCredentialsProvider = Mockito.mock(AWSCredentialsProviderService.class); + when(mockAwsCredentialsProvider.getIdentifier()).thenReturn("awsCredetialProvider"); + final GetAwsPollyJobStatus mockGetAwsPollyStatus = new GetAwsPollyJobStatus() { + protected AmazonPollyClient getClient() { + return mockPollyClient; + } + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + return mockPollyClient; + } + }; + runner = TestRunners.newTestRunner(mockGetAwsPollyStatus); + runner.addControllerService("awsCredetialProvider", mockAwsCredentialsProvider); + runner.enableControllerService(mockAwsCredentialsProvider); + runner.setProperty(AWS_CREDENTIALS_PROVIDER_SERVICE, "awsCredetialProvider"); + } + + @Test + public void testPollyTaskInProgress() { + ArgumentCaptor<GetSpeechSynthesisTaskRequest> requestCaptor = ArgumentCaptor.forClass(GetSpeechSynthesisTaskRequest.class); Review Comment: The `ArgumentCaptor` can be defined with an `@Captor` annotation at the class level. ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + RESPONSE response; + FlowFile childFlowFile; + try { + response = sendRequest(buildRequest(session, context, flowFile), context); + childFlowFile = writeToFlowFile(session, flowFile, response); + postProcessFlowFile(context, session, childFlowFile, response); + } catch (Exception e) { + session.transfer(flowFile, REL_FAILURE); + getLogger().error("Sending AWS ML Request failed", e); + return; + } + session.transfer(flowFile, REL_ORIGINAL); + session.transfer(childFlowFile, REL_SUCCESS); + } + + protected void postProcessFlowFile(ProcessContext context, ProcessSession session, FlowFile flowFile, RESPONSE response) { + session.putAttribute(flowFile, AWS_TASK_ID_PROPERTY, getAwsTaskId(context, response)); + getLogger().debug("AWS ML task has been started with task id: {}", getAwsTaskId(context, response)); + } + + protected REQUEST buildRequest(ProcessSession session, ProcessContext context, FlowFile flowFile) throws JsonProcessingException { + return mapper.readValue(getPayload(session, context, flowFile), getAwsRequestClass(context)); + } + + private String getPayload(ProcessSession session, ProcessContext context, FlowFile flowFile) { + String payloadPropertyValue = context.getProperty(JSON_PAYLOAD).evaluateAttributeExpressions(flowFile).getValue(); + if (payloadPropertyValue == null) { + payloadPropertyValue = readFlowFile(session, flowFile); + } + return payloadPropertyValue; + } + + @Override + protected T createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + throw new UnsupportedOperationException("Tried to create client in a deprecated way."); Review Comment: Recommend adjusting the message: ```suggestion throw new UnsupportedOperationException("createClient(ProcessContext, AWSCredentials, ClientConfiguration) is not supported"); ``` ########## pom.xml: ########## @@ -147,6 +147,7 @@ <swagger.annotations.version>1.6.6</swagger.annotations.version> <h2.version>2.1.214</h2.version> <zookeeper.version>3.8.0</zookeeper.version> + <aws.ml.version>1.12.328</aws.ml.version> Review Comment: This property is not necessary since the BOM dependency sets the version. ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml: ########## @@ -117,6 +117,26 @@ <version>1.19.0-SNAPSHOT</version> <scope>provided</scope> </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-translate</artifactId> + <version>${aws.ml.version}</version> Review Comment: This version number can be removed since it is already managed by the `aws-java-sdk-bom` in the root Maven configuration. ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/pom.xml: ########## @@ -117,6 +117,26 @@ <version>1.19.0-SNAPSHOT</version> <scope>provided</scope> </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-translate</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-polly</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-transcribe</artifactId> + <version>${aws.ml.version}</version> + </dependency> + <dependency> + <groupId>com.amazonaws</groupId> + <artifactId>aws-java-sdk-textract</artifactId> + <version>${aws.ml.version}</version> Review Comment: ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStatusGetter.java: ########## @@ -0,0 +1,138 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.ResponseMetadata; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.http.SdkHttpMetadata; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.ProcessorInitializationContext; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStatusGetter<T extends AmazonWebServiceClient> Review Comment: ```suggestion public abstract class AwsMachineLearningJobStatusProcessor<T extends AmazonWebServiceClient> ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") Review Comment: Recommend indicating that when this property is not provided, the JSON request will be read from the FlowFile. The AWS API reference is not very helpful without a link, and the details are covered in additionalDetails, so recommend removing that sentence. ```suggestion .description("JSON request for AWS Machine Learning services. The Processor will use FlowFile content for the request when this property is not specified.") ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") Review Comment: Recommend simplifying the property name: ```suggestion .name("aws-region") ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStatusGetter.java: ########## @@ -0,0 +1,138 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.ResponseMetadata; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.http.SdkHttpMetadata; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.ProcessorInitializationContext; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStatusGetter<T extends AmazonWebServiceClient> + extends AbstractAWSCredentialsProviderProcessor<T> { + public static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final String AWS_TASK_OUTPUT_LOCATION = "outputLocation"; + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + public static final Relationship REL_RUNNING = new Relationship.Builder() + .name("running") + .description("The job is currently still being processed") + .build(); + public static final Relationship REL_SUCCESS = new Relationship.Builder() + .name("success") + .description("Job successfully finished. FlowFile will be routed to this relation.") + .build(); + public static final Relationship REL_THROTTLED = new Relationship.Builder() + .name("throttled") + .description("Retrieving results failed for some reason, but the issue is likely to resolve on its own, such as Provisioned Throughput Exceeded or a Throttling failure. " + + "It is generally expected to retry this relationship.") + .build(); + public static final Relationship REL_FAILURE = new Relationship.Builder() + .name("failure") + .description("The job failed, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + protected final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + protected static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE, + PROXY_CONFIGURATION_SERVICE)); + + public static final String FAILURE_REASON_ATTRIBUTE = "failure.reason"; + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_RUNNING, + REL_THROTTLED, + REL_FAILURE + ))); + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + protected T createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + throw new UnsupportedOperationException("Client creation not supported"); + } + + @Override + protected void init(ProcessorInitializationContext context) { + SimpleModule awsResponseModule = new SimpleModule(); + awsResponseModule.addDeserializer(ResponseMetadata.class, new AwsResponseMetadataDeserializer()); + SimpleModule sdkHttpModule = new SimpleModule(); + awsResponseModule.addDeserializer(SdkHttpMetadata.class, new SdkHttpMetadataDeserializer()); + mapper.registerModule(awsResponseModule); + mapper.registerModule(sdkHttpModule); + } + + + protected void writeToFlowFile(ProcessSession session, FlowFile flowFile, Object response) { + session.write(flowFile, out -> { + try (OutputStreamWriter outputStreamWriter = new OutputStreamWriter(out, StandardCharsets.UTF_8)) { + outputStreamWriter.write(mapper.writeValueAsString(response)); + outputStreamWriter.flush(); + } Review Comment: ```suggestion mapper.writeValue(out, response); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + RESPONSE response; + FlowFile childFlowFile; + try { + response = sendRequest(buildRequest(session, context, flowFile), context); + childFlowFile = writeToFlowFile(session, flowFile, response); + postProcessFlowFile(context, session, childFlowFile, response); + } catch (Exception e) { + session.transfer(flowFile, REL_FAILURE); + getLogger().error("Sending AWS ML Request failed", e); + return; + } + session.transfer(flowFile, REL_ORIGINAL); + session.transfer(childFlowFile, REL_SUCCESS); + } + + protected void postProcessFlowFile(ProcessContext context, ProcessSession session, FlowFile flowFile, RESPONSE response) { + session.putAttribute(flowFile, AWS_TASK_ID_PROPERTY, getAwsTaskId(context, response)); + getLogger().debug("AWS ML task has been started with task id: {}", getAwsTaskId(context, response)); + } + + protected REQUEST buildRequest(ProcessSession session, ProcessContext context, FlowFile flowFile) throws JsonProcessingException { + return mapper.readValue(getPayload(session, context, flowFile), getAwsRequestClass(context)); + } + + private String getPayload(ProcessSession session, ProcessContext context, FlowFile flowFile) { + String payloadPropertyValue = context.getProperty(JSON_PAYLOAD).evaluateAttributeExpressions(flowFile).getValue(); + if (payloadPropertyValue == null) { + payloadPropertyValue = readFlowFile(session, flowFile); + } + return payloadPropertyValue; + } + + @Override + protected T createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + throw new UnsupportedOperationException("Tried to create client in a deprecated way."); + } + + abstract protected RESPONSE sendRequest(REQUEST request, ProcessContext context) throws JsonProcessingException; + + abstract protected Class<? extends REQUEST> getAwsRequestClass(ProcessContext context); + + abstract protected String getAwsTaskId(ProcessContext context, RESPONSE response); + + protected FlowFile writeToFlowFile(ProcessSession session, FlowFile flowFile, RESPONSE response) { + FlowFile childFlowFile = session.create(flowFile); + session.write(childFlowFile, out -> { Review Comment: ```suggestion childFlowFile = session.write(childFlowFile, out -> { ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyJobStatus.java: ########## @@ -0,0 +1,114 @@ +/* + * 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.ml.polly; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.AmazonPollyClientBuilder; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.TaskStatus; +import com.amazonaws.services.textract.model.ThrottlingException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Polly"}) +@CapabilityDescription("Retrieves the current status of an AWS Polly job.") +@SeeAlso({StartAwsPollyJob.class}) +public class GetAwsPollyJobStatus extends AwsMachineLearningJobStatusGetter<AmazonPollyClient> { + private static final String BUCKET = "bucket"; + private static final String KEY = "key"; + private static final Pattern S3_PATH = Pattern.compile("https://s3.*amazonaws.com/(?<" + BUCKET + ">[^/]+)/(?<" + KEY + ">.*)"); + private static final String AWS_S3_BUCKET = "PollyS3OutputBucket"; + private static final String AWS_S3_KEY = "PollyS3OutputKey"; + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonPollyClient) AmazonPollyClientBuilder.standard() + .withCredentials(credentialsProvider) + .withRegion(context.getProperty(REGION).getValue()) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetSpeechSynthesisTaskResult speechSynthesisTask; + try { + speechSynthesisTask = getSynthesisTask(flowFile); + } catch (ThrottlingException e) { + getLogger().info("Request Rate Limit exceeded", e); + session.transfer(flowFile, REL_THROTTLED); + return; + } catch (Exception e) { + getLogger().info("Failed to get Polly Job status", e); Review Comment: ```suggestion getLogger().warn("Failed to get Polly Job status", e); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStarter.java: ########## @@ -0,0 +1,168 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.commons.io.IOUtils; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStarter<T extends AmazonWebServiceClient, REQUEST extends AmazonWebServiceRequest, RESPONSE extends AmazonWebServiceResult> + extends AbstractAWSCredentialsProviderProcessor<T> { + protected static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final PropertyDescriptor JSON_PAYLOAD = new PropertyDescriptor.Builder() + .name("json-payload") + .displayName("JSON Payload") + .description("JSON Payload that represent an AWS ML Request. See more details in AWS API documentation.") + .expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES) + .build(); + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") + .required(true) + .allowableValues(getAvailableRegions()) + .defaultValue(createAllowableValue(Regions.DEFAULT_REGION).getValue()) + .build(); + private static final List<PropertyDescriptor> PROPERTIES = Collections.unmodifiableList(Arrays.asList( + JSON_PAYLOAD, + MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE, + REGION, + TIMEOUT, + SSL_CONTEXT_SERVICE, + ENDPOINT_OVERRIDE)); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + private final ObjectMapper mapper = JsonMapper.builder() + .configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true) + .build(); + + private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + REL_ORIGINAL, + REL_SUCCESS, + REL_FAILURE + ))); + + @Override + public Set<Relationship> getRelationships() { + return relationships; + } + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return PROPERTIES; + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + RESPONSE response; + FlowFile childFlowFile; + try { + response = sendRequest(buildRequest(session, context, flowFile), context); + childFlowFile = writeToFlowFile(session, flowFile, response); + postProcessFlowFile(context, session, childFlowFile, response); + } catch (Exception e) { + session.transfer(flowFile, REL_FAILURE); + getLogger().error("Sending AWS ML Request failed", e); + return; + } + session.transfer(flowFile, REL_ORIGINAL); + session.transfer(childFlowFile, REL_SUCCESS); + } + + protected void postProcessFlowFile(ProcessContext context, ProcessSession session, FlowFile flowFile, RESPONSE response) { + session.putAttribute(flowFile, AWS_TASK_ID_PROPERTY, getAwsTaskId(context, response)); + getLogger().debug("AWS ML task has been started with task id: {}", getAwsTaskId(context, response)); + } + + protected REQUEST buildRequest(ProcessSession session, ProcessContext context, FlowFile flowFile) throws JsonProcessingException { + return mapper.readValue(getPayload(session, context, flowFile), getAwsRequestClass(context)); + } + + private String getPayload(ProcessSession session, ProcessContext context, FlowFile flowFile) { + String payloadPropertyValue = context.getProperty(JSON_PAYLOAD).evaluateAttributeExpressions(flowFile).getValue(); + if (payloadPropertyValue == null) { + payloadPropertyValue = readFlowFile(session, flowFile); + } + return payloadPropertyValue; + } + + @Override + protected T createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + throw new UnsupportedOperationException("Tried to create client in a deprecated way."); + } + + abstract protected RESPONSE sendRequest(REQUEST request, ProcessContext context) throws JsonProcessingException; + + abstract protected Class<? extends REQUEST> getAwsRequestClass(ProcessContext context); + + abstract protected String getAwsTaskId(ProcessContext context, RESPONSE response); Review Comment: Recommend moving these abstract methods to the end of the class. ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/AwsMachineLearningJobStatusGetter.java: ########## @@ -0,0 +1,138 @@ +/* + * 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.ml; + +import com.amazonaws.AmazonWebServiceClient; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.ResponseMetadata; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.http.SdkHttpMetadata; +import com.amazonaws.regions.Regions; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; +import java.io.OutputStreamWriter; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.ProcessorInitializationContext; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor; + +public abstract class AwsMachineLearningJobStatusGetter<T extends AmazonWebServiceClient> + extends AbstractAWSCredentialsProviderProcessor<T> { + public static final String AWS_TASK_ID_PROPERTY = "awsTaskId"; + public static final String AWS_TASK_OUTPUT_LOCATION = "outputLocation"; + public static final PropertyDescriptor MANDATORY_AWS_CREDENTIALS_PROVIDER_SERVICE = + new PropertyDescriptor.Builder().fromPropertyDescriptor(AWS_CREDENTIALS_PROVIDER_SERVICE) + .required(true) + .build(); + public static final Relationship REL_ORIGINAL = new Relationship.Builder() + .name("original") + .description("Upon successful completion, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + public static final Relationship REL_RUNNING = new Relationship.Builder() + .name("running") + .description("The job is currently still being processed") + .build(); + public static final Relationship REL_SUCCESS = new Relationship.Builder() + .name("success") + .description("Job successfully finished. FlowFile will be routed to this relation.") + .build(); + public static final Relationship REL_THROTTLED = new Relationship.Builder() + .name("throttled") + .description("Retrieving results failed for some reason, but the issue is likely to resolve on its own, such as Provisioned Throughput Exceeded or a Throttling failure. " + + "It is generally expected to retry this relationship.") + .build(); + public static final Relationship REL_FAILURE = new Relationship.Builder() + .name("failure") + .description("The job failed, the original FlowFile will be routed to this relationship.") + .autoTerminateDefault(true) + .build(); + public static final PropertyDescriptor REGION = new PropertyDescriptor.Builder() + .displayName("Region") + .name("aws-ml-region") Review Comment: ```suggestion .name("aws-region") ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyJobStatus.java: ########## @@ -0,0 +1,114 @@ +/* + * 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.ml.polly; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.AmazonPollyClientBuilder; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.TaskStatus; +import com.amazonaws.services.textract.model.ThrottlingException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Polly"}) +@CapabilityDescription("Retrieves the current status of an AWS Polly job.") +@SeeAlso({StartAwsPollyJob.class}) +public class GetAwsPollyJobStatus extends AwsMachineLearningJobStatusGetter<AmazonPollyClient> { + private static final String BUCKET = "bucket"; + private static final String KEY = "key"; + private static final Pattern S3_PATH = Pattern.compile("https://s3.*amazonaws.com/(?<" + BUCKET + ">[^/]+)/(?<" + KEY + ">.*)"); + private static final String AWS_S3_BUCKET = "PollyS3OutputBucket"; + private static final String AWS_S3_KEY = "PollyS3OutputKey"; + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonPollyClient) AmazonPollyClientBuilder.standard() + .withCredentials(credentialsProvider) + .withRegion(context.getProperty(REGION).getValue()) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetSpeechSynthesisTaskResult speechSynthesisTask; + try { + speechSynthesisTask = getSynthesisTask(flowFile); + } catch (ThrottlingException e) { + getLogger().info("Request Rate Limit exceeded", e); + session.transfer(flowFile, REL_THROTTLED); + return; + } catch (Exception e) { + getLogger().info("Failed to get Polly Job status", e); + session.transfer(flowFile, REL_FAILURE); + return; + } + + TaskStatus taskStatus = TaskStatus.fromValue(speechSynthesisTask.getSynthesisTask().getTaskStatus()); + + if (taskStatus == TaskStatus.InProgress || taskStatus == TaskStatus.Scheduled) { + session.penalize(flowFile); + session.transfer(flowFile, REL_RUNNING); + } + + if (taskStatus == TaskStatus.Completed) { + String outputUri = speechSynthesisTask.getSynthesisTask().getOutputUri(); + + Matcher matcher = S3_PATH.matcher(outputUri); + if (matcher.find()) { + session.putAttribute(flowFile, AWS_S3_BUCKET, matcher.group(BUCKET)); + session.putAttribute(flowFile, AWS_S3_KEY, matcher.group(KEY)); + } + FlowFile childFlowFile = session.create(flowFile); + writeToFlowFile(session, childFlowFile, speechSynthesisTask); + session.putAttribute(childFlowFile, AWS_TASK_OUTPUT_LOCATION, outputUri); Review Comment: ```suggestion childFlowFile = session.putAttribute(childFlowFile, AWS_TASK_OUTPUT_LOCATION, outputUri); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyJobStatus.java: ########## @@ -0,0 +1,114 @@ +/* + * 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.ml.polly; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.AmazonPollyClientBuilder; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.TaskStatus; +import com.amazonaws.services.textract.model.ThrottlingException; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Polly"}) +@CapabilityDescription("Retrieves the current status of an AWS Polly job.") +@SeeAlso({StartAwsPollyJob.class}) +public class GetAwsPollyJobStatus extends AwsMachineLearningJobStatusGetter<AmazonPollyClient> { + private static final String BUCKET = "bucket"; + private static final String KEY = "key"; + private static final Pattern S3_PATH = Pattern.compile("https://s3.*amazonaws.com/(?<" + BUCKET + ">[^/]+)/(?<" + KEY + ">.*)"); + private static final String AWS_S3_BUCKET = "PollyS3OutputBucket"; + private static final String AWS_S3_KEY = "PollyS3OutputKey"; + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonPollyClient) AmazonPollyClientBuilder.standard() + .withCredentials(credentialsProvider) + .withRegion(context.getProperty(REGION).getValue()) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetSpeechSynthesisTaskResult speechSynthesisTask; + try { + speechSynthesisTask = getSynthesisTask(flowFile); + } catch (ThrottlingException e) { + getLogger().info("Request Rate Limit exceeded", e); + session.transfer(flowFile, REL_THROTTLED); + return; + } catch (Exception e) { + getLogger().info("Failed to get Polly Job status", e); + session.transfer(flowFile, REL_FAILURE); + return; + } + + TaskStatus taskStatus = TaskStatus.fromValue(speechSynthesisTask.getSynthesisTask().getTaskStatus()); + + if (taskStatus == TaskStatus.InProgress || taskStatus == TaskStatus.Scheduled) { + session.penalize(flowFile); + session.transfer(flowFile, REL_RUNNING); + } + + if (taskStatus == TaskStatus.Completed) { Review Comment: It looks like this should be defined as an `else if` with with previous conditional. ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/textract/GetAwsTextractJobStatus.java: ########## @@ -0,0 +1,139 @@ +/* + * 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.ml.textract; + +import static org.apache.nifi.expression.ExpressionLanguageScope.FLOWFILE_ATTRIBUTES; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.textract.AmazonTextractClient; +import com.amazonaws.services.textract.model.GetDocumentAnalysisRequest; +import com.amazonaws.services.textract.model.GetDocumentTextDetectionRequest; +import com.amazonaws.services.textract.model.GetExpenseAnalysisRequest; +import com.amazonaws.services.textract.model.JobStatus; +import com.google.common.collect.ImmutableList; +import java.util.List; +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.flowfile.FlowFile; +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.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Textract"}) +@CapabilityDescription("Retrieves the current status of an AWS Textract job.") +@SeeAlso({StartAwsTextractJob.class}) +public class GetAwsTextractJobStatus extends AwsMachineLearningJobStatusGetter<AmazonTextractClient> { + public static final String DOCUMENT_ANALYSIS = "Document Analysis"; + public static final String DOCUMENT_TEXT_DETECTION = "Document Text Detection"; + public static final String EXPENSE_ANALYSIS = "Expense Analysis"; + public static final PropertyDescriptor TYPE = new PropertyDescriptor.Builder() + .name("type-of-service") + .displayName("Type of textract") + .expressionLanguageSupported(FLOWFILE_ATTRIBUTES) + .allowableValues(DOCUMENT_ANALYSIS, DOCUMENT_TEXT_DETECTION, EXPENSE_ANALYSIS) + .required(true) + .defaultValue("Document Analysis") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + private static final List<PropertyDescriptor> TEXTRACT_PROPERTIES = ImmutableList.<PropertyDescriptor>builder() + .addAll(PROPERTIES) + .add(TYPE) + .build();; + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return TEXTRACT_PROPERTIES; + } + + @Override + protected AmazonTextractClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonTextractClient) AmazonTextractClient.builder() + .withRegion(context.getProperty(REGION).getValue()) + .withCredentials(credentialsProvider) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + String typeOfTextract = context.getProperty(TYPE).evaluateAttributeExpressions().getValue(); + + String awsTaskId = flowFile.getAttribute(AWS_TASK_ID_PROPERTY); + JobStatus jobStatus = getTaskStatus(typeOfTextract, getClient(), awsTaskId); + if (JobStatus.SUCCEEDED == jobStatus) { + Object task = getTask(typeOfTextract, getClient(), awsTaskId); + writeToFlowFile(session, flowFile, task); + session.transfer(flowFile, REL_SUCCESS); + } + + if (JobStatus.IN_PROGRESS == jobStatus) { + session.transfer(flowFile, REL_RUNNING); + } + + if (JobStatus.PARTIAL_SUCCESS == jobStatus) { + session.transfer(flowFile, REL_THROTTLED); + } + + if (JobStatus.FAILED == jobStatus) { Review Comment: It looks like these conditionals shoudl be chained together with `else if` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/textract/GetAwsTextractJobStatus.java: ########## @@ -0,0 +1,139 @@ +/* + * 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.ml.textract; + +import static org.apache.nifi.expression.ExpressionLanguageScope.FLOWFILE_ATTRIBUTES; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.textract.AmazonTextractClient; +import com.amazonaws.services.textract.model.GetDocumentAnalysisRequest; +import com.amazonaws.services.textract.model.GetDocumentTextDetectionRequest; +import com.amazonaws.services.textract.model.GetExpenseAnalysisRequest; +import com.amazonaws.services.textract.model.JobStatus; +import com.google.common.collect.ImmutableList; Review Comment: Google Guava classes should not be used, this can be replaced with `Collections.unmodifiableList()` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/textract/StartAwsTextractJob.java: ########## @@ -0,0 +1,134 @@ +/* + * 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.ml.textract; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.textract.AmazonTextractClient; +import com.amazonaws.services.textract.model.StartDocumentAnalysisRequest; +import com.amazonaws.services.textract.model.StartDocumentAnalysisResult; +import com.amazonaws.services.textract.model.StartDocumentTextDetectionRequest; +import com.amazonaws.services.textract.model.StartDocumentTextDetectionResult; +import com.amazonaws.services.textract.model.StartExpenseAnalysisRequest; +import com.amazonaws.services.textract.model.StartExpenseAnalysisResult; +import com.google.common.collect.ImmutableList; Review Comment: Should be replaced with `Collections.unmodifiableList()` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/textract/StartAwsTextractJob.java: ########## @@ -0,0 +1,134 @@ +/* + * 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.ml.textract; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.AmazonWebServiceResult; +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.textract.AmazonTextractClient; +import com.amazonaws.services.textract.model.StartDocumentAnalysisRequest; +import com.amazonaws.services.textract.model.StartDocumentAnalysisResult; +import com.amazonaws.services.textract.model.StartDocumentTextDetectionRequest; +import com.amazonaws.services.textract.model.StartDocumentTextDetectionResult; +import com.amazonaws.services.textract.model.StartExpenseAnalysisRequest; +import com.amazonaws.services.textract.model.StartExpenseAnalysisResult; +import com.google.common.collect.ImmutableList; +import java.util.List; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.util.StandardValidators; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStarter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Textract"}) +@CapabilityDescription("Trigger a AWS Textract job. It should be followed by GetAwsTextractJobStatus processor in order to monitor job status.") +@SeeAlso({GetAwsTextractJobStatus.class}) +public class StartAwsTextractJob extends AwsMachineLearningJobStarter<AmazonTextractClient, AmazonWebServiceRequest, AmazonWebServiceResult> { + private static final String DOCUMENT_ANALYSIS = "Document Analysis"; + private static final String DOCUMENT_TEXT_DETECTION = "Document Text Detection"; + private static final String EXPENSE_ANALYSIS = "Expense Analysis"; + public static final PropertyDescriptor TYPE = new PropertyDescriptor.Builder() + .name("type-of-service") + .displayName("Type of textract") + .allowableValues(DOCUMENT_ANALYSIS, DOCUMENT_TEXT_DETECTION, EXPENSE_ANALYSIS) + .required(true) + .defaultValue("Document Analysis") + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + @Override + public List<PropertyDescriptor> getSupportedPropertyDescriptors() { + return new ImmutableList.Builder().add(TYPE).add(super.getSupportedPropertyDescriptors().toArray()).build(); + } + + @Override + protected void postProcessFlowFile(ProcessContext context, ProcessSession session, FlowFile flowFile, AmazonWebServiceResult response) { + super.postProcessFlowFile(context, session, flowFile, response); + session.putAttribute(flowFile, TYPE.getName(), context.getProperty(TYPE.getName()).getValue()); + } + + @Override + protected AmazonTextractClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonTextractClient) AmazonTextractClient.builder() + .withRegion(context.getProperty(REGION).getValue()) + .withCredentials(credentialsProvider) + .build(); + } + + @Override + protected AmazonWebServiceResult sendRequest(AmazonWebServiceRequest request, ProcessContext context) { + String typeOfTextract = context.getProperty(TYPE.getName()).getValue(); + AmazonWebServiceResult result; + switch (typeOfTextract) { + case DOCUMENT_ANALYSIS : + result = getClient().startDocumentAnalysis((StartDocumentAnalysisRequest) request); + break; + case DOCUMENT_TEXT_DETECTION: + result = getClient().startDocumentTextDetection((StartDocumentTextDetectionRequest) request); + break; + case EXPENSE_ANALYSIS : + result = getClient().startExpenseAnalysis((StartExpenseAnalysisRequest) request); + break; + default: throw new UnsupportedOperationException("Unsupported textract type."); + } + return result; + } + + @Override + protected Class<? extends AmazonWebServiceRequest> getAwsRequestClass(ProcessContext context) { + String typeOfTextract = context.getProperty(TYPE.getName()).getValue(); + Class<? extends AmazonWebServiceRequest> result; + switch (typeOfTextract) { + case DOCUMENT_ANALYSIS: + result = StartDocumentAnalysisRequest.class; + break; + case DOCUMENT_TEXT_DETECTION: + result = StartDocumentTextDetectionRequest.class; + break; + case EXPENSE_ANALYSIS: + result = StartExpenseAnalysisRequest.class; + break; + default: throw new UnsupportedOperationException("Unsupported textract type."); + } + return result; + } + + @Override + protected String getAwsTaskId(ProcessContext context, AmazonWebServiceResult amazonWebServiceResult) { + String typeOfTextract = context.getProperty(TYPE.getName()).getValue(); + String result; + switch (typeOfTextract) { + case DOCUMENT_ANALYSIS: + result = ((StartDocumentAnalysisResult) amazonWebServiceResult).getJobId(); + break; + case DOCUMENT_TEXT_DETECTION: + result = ((StartDocumentTextDetectionResult) amazonWebServiceResult).getJobId(); + break; + case EXPENSE_ANALYSIS: + result = ((StartExpenseAnalysisResult) amazonWebServiceResult).getJobId(); + break; + default: throw new UnsupportedOperationException("Unsupported textract type."); Review Comment: The `typeOfTextract` string should be included: ```suggestion default: throw new UnsupportedOperationException("Unsupported textract type: " + typeOfTextract); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + +public class GetAwsPollyStatusTest { + private static final String TEST_TASK_ID = "testTaskId"; + private TestRunner runner = null; + private AmazonPollyClient mockPollyClient = null; + private AWSCredentialsProviderService mockAwsCredentialsProvider = null; + + @BeforeEach + public void setUp() throws InitializationException { + mockPollyClient = Mockito.mock(AmazonPollyClient.class); + mockAwsCredentialsProvider = Mockito.mock(AWSCredentialsProviderService.class); Review Comment: These mocks can be configured using the `@Mock` annotation on the property ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/translate/GetAwsTranslateJobStatus.java: ########## @@ -0,0 +1,80 @@ +/* + * 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.ml.translate; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.translate.AmazonTranslateClient; +import com.amazonaws.services.translate.model.DescribeTextTranslationJobRequest; +import com.amazonaws.services.translate.model.DescribeTextTranslationJobResult; +import com.amazonaws.services.translate.model.JobStatus; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Translate"}) +@CapabilityDescription("Retrieves the current status of an AWS Translate job.") +@SeeAlso({StartAwsTranslateJob.class}) +public class GetAwsTranslateJobStatus extends AwsMachineLearningJobStatusGetter<AmazonTranslateClient> { + @Override + protected AmazonTranslateClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonTranslateClient) AmazonTranslateClient.builder() + .withRegion(context.getProperty(REGION).getValue()) + .withCredentials(credentialsProvider) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + String awsTaskId = flowFile.getAttribute(AWS_TASK_ID_PROPERTY); + DescribeTextTranslationJobResult describeTextTranslationJobResult = getStatusString(awsTaskId); + JobStatus status = JobStatus.fromValue(describeTextTranslationJobResult.getTextTranslationJobProperties().getJobStatus()); + + if (status == JobStatus.IN_PROGRESS || status == JobStatus.SUBMITTED) { + writeToFlowFile(session, flowFile, describeTextTranslationJobResult); + session.penalize(flowFile); + session.transfer(flowFile, REL_RUNNING); + } + + if (status == JobStatus.COMPLETED) { Review Comment: ```suggestion } else if (status == JobStatus.COMPLETED) { ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + +public class GetAwsPollyStatusTest { + private static final String TEST_TASK_ID = "testTaskId"; + private TestRunner runner = null; + private AmazonPollyClient mockPollyClient = null; Review Comment: ```suggestion @Mock private AmazonPollyClient mockPollyClient; ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/transcribe/GetAwsTranscribeJobStatus.java: ########## @@ -0,0 +1,80 @@ +/* + * 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.ml.transcribe; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.transcribe.AmazonTranscribeClient; +import com.amazonaws.services.transcribe.model.GetTranscriptionJobRequest; +import com.amazonaws.services.transcribe.model.GetTranscriptionJobResult; +import com.amazonaws.services.transcribe.model.TranscriptionJobStatus; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Transcribe"}) +@CapabilityDescription("Retrieves the current status of an AWS Transcribe job.") +@SeeAlso({StartAwsTranscribeJob.class}) +public class GetAwsTranscribeJobStatus extends AwsMachineLearningJobStatusGetter<AmazonTranscribeClient> { + @Override + protected AmazonTranscribeClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonTranscribeClient) AmazonTranscribeClient.builder() + .withRegion(context.getProperty(REGION).getValue()) + .withCredentials(credentialsProvider) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetTranscriptionJobResult job = getJob(flowFile); + TranscriptionJobStatus jobStatus = TranscriptionJobStatus.fromValue(job.getTranscriptionJob().getTranscriptionJobStatus()); + + if (TranscriptionJobStatus.COMPLETED == jobStatus) { + writeToFlowFile(session, flowFile, job); + session.putAttribute(flowFile, AWS_TASK_OUTPUT_LOCATION, job.getTranscriptionJob().getTranscript().getTranscriptFileUri()); + session.transfer(flowFile, REL_SUCCESS); + } + + if (TranscriptionJobStatus.IN_PROGRESS == jobStatus) { + session.transfer(flowFile, REL_RUNNING); + } + + if (TranscriptionJobStatus.FAILED == jobStatus) { + final String failureReason = job.getTranscriptionJob().getFailureReason(); + session.putAttribute(flowFile, FAILURE_REASON_ATTRIBUTE, failureReason); + session.transfer(flowFile, REL_FAILURE); + getLogger().error("Transcribe Task Failed for {}: {}", flowFile, failureReason); + return; Review Comment: This can be removed ```suggestion ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + Review Comment: Using the Mockito Extension enables annotated mocks: ```suggestion @ExtendWith(MockitoExtension.class) ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/main/java/org/apache/nifi/processors/aws/ml/transcribe/GetAwsTranscribeJobStatus.java: ########## @@ -0,0 +1,80 @@ +/* + * 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.ml.transcribe; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.services.transcribe.AmazonTranscribeClient; +import com.amazonaws.services.transcribe.model.GetTranscriptionJobRequest; +import com.amazonaws.services.transcribe.model.GetTranscriptionJobResult; +import com.amazonaws.services.transcribe.model.TranscriptionJobStatus; +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.flowfile.FlowFile; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter; + +@Tags({"Amazon", "AWS", "ML", "Machine Learning", "Transcribe"}) +@CapabilityDescription("Retrieves the current status of an AWS Transcribe job.") +@SeeAlso({StartAwsTranscribeJob.class}) +public class GetAwsTranscribeJobStatus extends AwsMachineLearningJobStatusGetter<AmazonTranscribeClient> { + @Override + protected AmazonTranscribeClient createClient(ProcessContext context, AWSCredentialsProvider credentialsProvider, ClientConfiguration config) { + return (AmazonTranscribeClient) AmazonTranscribeClient.builder() + .withRegion(context.getProperty(REGION).getValue()) + .withCredentials(credentialsProvider) + .build(); + } + + @Override + public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException { + FlowFile flowFile = session.get(); + if (flowFile == null) { + return; + } + GetTranscriptionJobResult job = getJob(flowFile); + TranscriptionJobStatus jobStatus = TranscriptionJobStatus.fromValue(job.getTranscriptionJob().getTranscriptionJobStatus()); + + if (TranscriptionJobStatus.COMPLETED == jobStatus) { + writeToFlowFile(session, flowFile, job); + session.putAttribute(flowFile, AWS_TASK_OUTPUT_LOCATION, job.getTranscriptionJob().getTranscript().getTranscriptFileUri()); + session.transfer(flowFile, REL_SUCCESS); + } + + if (TranscriptionJobStatus.IN_PROGRESS == jobStatus) { Review Comment: ```suggestion } else if (TranscriptionJobStatus.IN_PROGRESS == jobStatus) { ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + +public class GetAwsPollyStatusTest { + private static final String TEST_TASK_ID = "testTaskId"; + private TestRunner runner = null; + private AmazonPollyClient mockPollyClient = null; + private AWSCredentialsProviderService mockAwsCredentialsProvider = null; + + @BeforeEach + public void setUp() throws InitializationException { + mockPollyClient = Mockito.mock(AmazonPollyClient.class); + mockAwsCredentialsProvider = Mockito.mock(AWSCredentialsProviderService.class); + when(mockAwsCredentialsProvider.getIdentifier()).thenReturn("awsCredetialProvider"); + final GetAwsPollyJobStatus mockGetAwsPollyStatus = new GetAwsPollyJobStatus() { + protected AmazonPollyClient getClient() { + return mockPollyClient; + } + + @Override + protected AmazonPollyClient createClient(ProcessContext context, AWSCredentials credentials, ClientConfiguration config) { + return mockPollyClient; + } + }; + runner = TestRunners.newTestRunner(mockGetAwsPollyStatus); + runner.addControllerService("awsCredetialProvider", mockAwsCredentialsProvider); + runner.enableControllerService(mockAwsCredentialsProvider); + runner.setProperty(AWS_CREDENTIALS_PROVIDER_SERVICE, "awsCredetialProvider"); + } + + @Test + public void testPollyTaskInProgress() { + ArgumentCaptor<GetSpeechSynthesisTaskRequest> requestCaptor = ArgumentCaptor.forClass(GetSpeechSynthesisTaskRequest.class); + GetSpeechSynthesisTaskResult taskResult = new GetSpeechSynthesisTaskResult(); + SynthesisTask task = new SynthesisTask().withTaskId(TEST_TASK_ID) + .withTaskStatus(TaskStatus.InProgress); + taskResult.setSynthesisTask(task); + when(mockPollyClient.getSpeechSynthesisTask(requestCaptor.capture())).thenReturn(taskResult); + runner.enqueue("content", Collections.singletonMap(AWS_TASK_ID_PROPERTY, TEST_TASK_ID)); + runner.run(); + + runner.assertAllFlowFilesTransferred(REL_RUNNING); + assertEquals(requestCaptor.getValue().getTaskId(), TEST_TASK_ID); + } + + @Test + public void testPollyTaskCompleted() { + ArgumentCaptor<GetSpeechSynthesisTaskRequest> requestCaptor = ArgumentCaptor.forClass(GetSpeechSynthesisTaskRequest.class); + GetSpeechSynthesisTaskResult taskResult = new GetSpeechSynthesisTaskResult(); + SynthesisTask task = new SynthesisTask().withTaskId(TEST_TASK_ID) + .withTaskStatus(TaskStatus.Completed) + .withOutputUri("outputLocationPath"); + taskResult.setSynthesisTask(task); + when(mockPollyClient.getSpeechSynthesisTask(requestCaptor.capture())).thenReturn(taskResult); + runner.enqueue("content", Collections.singletonMap(AWS_TASK_ID_PROPERTY, TEST_TASK_ID)); + runner.run(); + + runner.assertTransferCount(REL_SUCCESS, 1); + runner.assertTransferCount(REL_ORIGINAL, 1); + runner.assertAllFlowFilesContainAttribute(REL_SUCCESS, AWS_TASK_OUTPUT_LOCATION); + assertEquals(requestCaptor.getValue().getTaskId(), TEST_TASK_ID); + } + + + @Test + public void testPollyTaskFailed() { + ArgumentCaptor<GetSpeechSynthesisTaskRequest> requestCaptor = ArgumentCaptor.forClass(GetSpeechSynthesisTaskRequest.class); + GetSpeechSynthesisTaskResult taskResult = new GetSpeechSynthesisTaskResult(); + SynthesisTask task = new SynthesisTask().withTaskId(TEST_TASK_ID) + .withTaskStatus(TaskStatus.Failed) + .withTaskStatusReason("reasonOfFailure"); + taskResult.setSynthesisTask(task); + when(mockPollyClient.getSpeechSynthesisTask(requestCaptor.capture())).thenReturn(taskResult); + runner.enqueue("content", Collections.singletonMap(AWS_TASK_ID_PROPERTY, TEST_TASK_ID)); Review Comment: Recommend defining and reusing a static variable the placeholder content string. ```suggestion runner.enqueue(PLACEHOLDER_CONTENT, Collections.singletonMap(AWS_TASK_ID_PROPERTY, TEST_TASK_ID)); ``` ########## nifi-nar-bundles/nifi-aws-bundle/nifi-aws-processors/src/test/java/org/apache/nifi/processors/aws/ml/polly/GetAwsPollyStatusTest.java: ########## @@ -0,0 +1,124 @@ +/* + * 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.ml.polly; + +import static org.apache.nifi.processors.aws.AbstractAWSCredentialsProviderProcessor.AWS_CREDENTIALS_PROVIDER_SERVICE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_ID_PROPERTY; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.AWS_TASK_OUTPUT_LOCATION; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_FAILURE; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_RUNNING; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_ORIGINAL; +import static org.apache.nifi.processors.aws.ml.AwsMachineLearningJobStatusGetter.REL_SUCCESS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.amazonaws.ClientConfiguration; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.services.polly.AmazonPollyClient; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskRequest; +import com.amazonaws.services.polly.model.GetSpeechSynthesisTaskResult; +import com.amazonaws.services.polly.model.SynthesisTask; +import com.amazonaws.services.polly.model.TaskStatus; +import java.util.Collections; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processors.aws.credentials.provider.service.AWSCredentialsProviderService; +import org.apache.nifi.reporting.InitializationException; +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; + +public class GetAwsPollyStatusTest { + private static final String TEST_TASK_ID = "testTaskId"; + private TestRunner runner = null; + private AmazonPollyClient mockPollyClient = null; + private AWSCredentialsProviderService mockAwsCredentialsProvider = null; Review Comment: ```suggestion @Mock private AWSCredentialsProviderService mockAwsCredentialsProvider; ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@nifi.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org