NIFI-10953: Implement GCP Vision AI processors

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/pom.xml

@@ -187,6 +187,28 @@
             <artifactId>jackson-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-vision</artifactId>
+            <exclusions>
+                <exclusion>
+                    <groupId>commons-logging</groupId>
+                    <artifactId>commons-logging</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.bouncycastle</groupId>
+                    <artifactId>bcprov-jdk15on</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>org.bouncycastle</groupId>
+                    <artifactId>bcpkix-jdk15on</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
     </dependencies>
 
     <build>

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/vision/AbstractGcpVisionProcessor.java

@@ -0,0 +1,96 @@
+/*
+ * 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.gcp.vision;
+
+import static org.apache.nifi.processors.gcp.util.GoogleUtils.GCP_CREDENTIALS_PROVIDER_SERVICE;
+
+import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.vision.v1.ImageAnnotatorClient;
+import com.google.cloud.vision.v1.ImageAnnotatorSettings;
+import java.io.IOException;
+import java.io.InputStream;
+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.annotation.lifecycle.OnScheduled;
+import org.apache.nifi.components.PropertyDescriptor;
+import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.gcp.credentials.service.GCPCredentialsService;
+import org.apache.nifi.processor.AbstractProcessor;
+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;
+
+public abstract class AbstractGcpVisionProcessor extends AbstractProcessor  {
+    public static final String GCP_OPERATION_KEY = "operationKey";
+
+    public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success")
+            .description("FlowFiles are routed to success relationship").build();
+    public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure")
+            .description("FlowFiles are routed to failure relationship").build();
+
+    protected static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
+            REL_SUCCESS,
+            REL_FAILURE
+    )));
+
+    private ImageAnnotatorClient vision;
+
+    @Override
+    public List<PropertyDescriptor> getSupportedPropertyDescriptors() {
+        return Collections.unmodifiableList(Arrays.asList(
+                GCP_CREDENTIALS_PROVIDER_SERVICE)
+        );
+    }
+
+    @Override
+    public Set<Relationship> getRelationships() {
+        return relationships;
+    }
+
+    @OnScheduled
+    public void onScheduled(ProcessContext context) {
+        final GCPCredentialsService gcpCredentialsService =
+                context.getProperty(GCP_CREDENTIALS_PROVIDER_SERVICE).asControllerService(GCPCredentialsService.class);
+        try {
+            GoogleCredentials credentials = gcpCredentialsService.getGoogleCredentials();
+            FixedCredentialsProvider credentialsProvider = FixedCredentialsProvider.create(credentials);
+            ImageAnnotatorSettings.Builder builder = ImageAnnotatorSettings.newBuilder().setCredentialsProvider(credentialsProvider);
+            vision = ImageAnnotatorClient.create(builder.build());
+        } catch (Exception e) {
+            getLogger().error("Failed to create vision client.", e);
+        }
+    }
+
+    protected String readFlowFile(final ProcessSession session, final FlowFile flowFile) {
+        try (InputStream inputStream = session.read(flowFile)) {
+            return new String(IOUtils.toByteArray(inputStream));
+        } catch (final IOException e) {
+            throw new ProcessException("Read FlowFile Failed", e);
+        }
+    }
+
+    protected ImageAnnotatorClient getVisionClient() {
+        return this.vision;
+    }
+}

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/vision/AbstractGetGcpVisionAnnotateOperationStatus.java

@@ -0,0 +1,104 @@
+/*
+ * 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.gcp.vision;
+
+import static org.apache.nifi.processors.gcp.util.GoogleUtils.GCP_CREDENTIALS_PROVIDER_SERVICE;
+
+import com.google.longrunning.Operation;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.GeneratedMessageV3;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import com.google.rpc.Status;
+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.Relationship;
+import org.apache.nifi.processor.exception.ProcessException;
+
+abstract public class AbstractGetGcpVisionAnnotateOperationStatus extends AbstractGcpVisionProcessor {
+
+    private static final String GCP_OPERATION_KEY = "operationKey";
+
+    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_ORIGINAL = new Relationship.Builder()
+            .name("original")
+            .description("Upon successful completion, the original FlowFile will be routed to this relationship.")
+            .autoTerminateDefault(true)
+            .build();
+
+    private static final Set<Relationship> relationships = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
+            REL_ORIGINAL,
+            REL_SUCCESS,
+            REL_FAILURE,
+            REL_RUNNING
+    )));
+
+    @Override
+    public List<PropertyDescriptor> getSupportedPropertyDescriptors() {
+        return Collections.unmodifiableList(Arrays.asList(
+                GCP_CREDENTIALS_PROVIDER_SERVICE)
+        );
+    }
+
+    @Override
+    public Set<Relationship> getRelationships() {
+        return relationships;
+    }
+
+    @Override
+    public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
+        FlowFile flowFile = session.get();
+        if (flowFile == null) {
+            return;
+        }
+        try {
+            String operationKey = flowFile.getAttribute(GCP_OPERATION_KEY);
+            Operation operation = getVisionClient().getOperationsClient().getOperation(operationKey);
+            getLogger().info(operation.toString());
+            if (operation.getDone() && !operation.hasError()) {
+                GeneratedMessageV3 response = deserializeResponse(operation.getResponse().getValue());
+                FlowFile childFlowFile = session.create(flowFile);
+                session.write(childFlowFile, out -> out.write(JsonFormat.printer().print(response).getBytes(StandardCharsets.UTF_8)));
+                session.transfer(flowFile, REL_ORIGINAL);
+                session.transfer(childFlowFile, REL_SUCCESS);
+            } else if (!operation.getDone()) {
+                session.transfer(flowFile, REL_RUNNING);
+            } else {
+                Status error = operation.getError();
+                getLogger().error("Failed to execute vision operation. Error code: {}, Error message: {}", error.getCode(), error.getMessage());
+                session.transfer(flowFile, REL_FAILURE);
+            }
+        } catch (Exception e) {
+            getLogger().error("Fail to get GCP Vision operation's status", e);
+            session.transfer(flowFile, REL_FAILURE);
+        }
+    }
+
+    abstract protected GeneratedMessageV3 deserializeResponse(ByteString responseValue) throws InvalidProtocolBufferException;
+}

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/vision/AbstractStartGcpVisionOperation.java

@@ -0,0 +1,53 @@
+/*
+ * 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.gcp.vision;
+
+import com.google.api.gax.longrunning.OperationFuture;
+import com.google.api.gax.longrunning.OperationSnapshot;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.Message;
+import java.io.IOException;
+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;
+
+public abstract class AbstractStartGcpVisionOperation extends AbstractGcpVisionProcessor  {
+
+
+    @Override
+    public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
+        FlowFile flowFile = session.get();
+        if (flowFile == null) {
+            return;
+        }
+        try  {
+            OperationFuture asyncResponse = startOperation(session, flowFile);
+            String operationName = ((OperationSnapshot) asyncResponse.getInitialFuture().get()).getName();
+            session.putAttribute(flowFile, GCP_OPERATION_KEY, operationName);
+            session.transfer(flowFile, REL_SUCCESS);
+        } catch (Exception e) {
+            getLogger().error("Fail to start GCP Vision operation", e);
+            session.transfer(flowFile, REL_FAILURE);
+        }
+    }
+
+    abstract protected Message fromJson(String json) throws IOException;
+
+    abstract protected OperationFuture startOperation(ProcessSession session, FlowFile flowFile) throws InvalidProtocolBufferException;
+}

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/vision/GetGcpVisionAnnotateFilesOperationStatus.java

@@ -0,0 +1,36 @@
+/*
+ * 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.gcp.vision;
+
+import com.google.cloud.vision.v1p2beta1.AsyncBatchAnnotateFilesResponse;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.GeneratedMessageV3;
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.nifi.annotation.documentation.CapabilityDescription;
+import org.apache.nifi.annotation.documentation.SeeAlso;
+import org.apache.nifi.annotation.documentation.Tags;
+
+@Tags({"Google", "Cloud", "Vision", "Machine Learning"})
+@CapabilityDescription("Retrieves the current status of an Google Vision operation.")
+@SeeAlso({StartGcpVisionAnnotateFilesOperation.class})
+public class GetGcpVisionAnnotateFilesOperationStatus extends AbstractGetGcpVisionAnnotateOperationStatus {
+    @Override
+    protected GeneratedMessageV3 deserializeResponse(ByteString responseValue) throws InvalidProtocolBufferException {
+        return AsyncBatchAnnotateFilesResponse.parseFrom(responseValue);
+    }
+}

nifi-nar-bundles/nifi-gcp-bundle/nifi-gcp-processors/src/main/java/org/apache/nifi/processors/gcp/vision/GetGcpVisionAnnotateImagesOperationStatus.java

@@ -0,0 +1,36 @@
+/*
+ * 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.gcp.vision;
+
+import com.google.cloud.vision.v1.AsyncBatchAnnotateImagesResponse;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.GeneratedMessageV3;
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.nifi.annotation.documentation.CapabilityDescription;
+import org.apache.nifi.annotation.documentation.SeeAlso;
+import org.apache.nifi.annotation.documentation.Tags;
+
+@Tags({"Google", "Cloud", "Vision", "Machine Learning"})
+@CapabilityDescription("Retrieves the current status of an Google Vision operation.")
+@SeeAlso({StartGcpVisionAnnotateImagesOperation.class})
+public class GetGcpVisionAnnotateImagesOperationStatus extends AbstractGetGcpVisionAnnotateOperationStatus {
+    @Override
+    protected GeneratedMessageV3 deserializeResponse(ByteString responseValue) throws InvalidProtocolBufferException {
+        return AsyncBatchAnnotateImagesResponse.parseFrom(responseValue);
+    }
+}