Inspect classify

src/main/java/com/google/swarm/tokenization/DLPTextToBigQueryStreamingV2PipelineOptions.java

@@ -15,54 +15,28 @@
  */
 package com.google.swarm.tokenization;
 
-import com.google.privacy.dlp.v2.LocationName;
 import com.google.swarm.tokenization.common.Util.DLPMethod;
 import com.google.swarm.tokenization.common.Util.DataSinkType;
 import com.google.swarm.tokenization.common.Util.FileType;
-import com.google.swarm.tokenization.common.Util.InputLocation;
+import com.google.swarm.tokenization.options.CommonPipelineOptions;
 import java.util.List;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
-import org.apache.beam.sdk.io.aws.options.S3Options;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
 import org.apache.beam.sdk.options.*;
 
-public interface DLPTextToBigQueryStreamingV2PipelineOptions
-    extends DataflowPipelineOptions, S3Options {
+public interface DLPTextToBigQueryStreamingV2PipelineOptions extends CommonPipelineOptions {
 
   @Validation.Required
   String getFilePattern();
 
   void setFilePattern(String csvFilePattern);
 
-  @Description("DLP Inspect Template Name")
-  String getInspectTemplateName();
-
-  void setInspectTemplateName(String value);
-
-  @Description("DLP DeIdentify Template Name")
-  String getDeidentifyTemplateName();
-
-  void setDeidentifyTemplateName(String value);
-
   @Description("DLP method deid,inspect,reid")
   @Default.Enum("INSPECT")
   DLPMethod getDLPMethod();
 
   void setDLPMethod(DLPMethod value);
 
-  @Description("Batch Size (max 524kb)")
-  @Default.Integer(500000)
-  Integer getBatchSize();
-
-  void setBatchSize(Integer value);
-
-  @Description("key range")
-  @Default.Integer(100)
-  Integer getKeyRange();
-
-  void setKeyRange(Integer value);
-
   @Description("BQ Dataset")
   String getDataset();
 
@@ -107,11 +81,6 @@ public interface DLPTextToBigQueryStreamingV2PipelineOptions
 
   void setTopic(String topic);
 
-  @Default.Integer(900 * 1000)
-  Integer getSplitSize();
-
-  void setSplitSize(Integer value);
-
   @Description("DLP Table Headers- Required for Json type")
   List<String> getHeaders();
 
@@ -122,31 +91,6 @@ public interface DLPTextToBigQueryStreamingV2PipelineOptions
 
   void setSchema(String schema);
 
-  @Description(
-      "Number of shards for DLP request batches. "
-          + "Can be used to controls parallelism of DLP requests.")
-  @Default.Integer(100)
-  int getNumShardsPerDLPRequestBatching();
-
-  void setNumShardsPerDLPRequestBatching(int value);
-
-  @Description("Number of retries in case of transient errors in DLP API")
-  @Default.Integer(10)
-  int getDlpApiRetryCount();
-
-  void setDlpApiRetryCount(int value);
-
-  @Description("Initial backoff (in seconds) for retries with exponential backoff")
-  @Default.Integer(5)
-  int getInitialBackoff();
-
-  /**
-   * Initial backoff (in seconds) for retries with exponential backoff. See {@link
-   * org.apache.beam.sdk.util.FluentBackoff.BackoffImpl#nextBackOffMillis()} for details on how the
-   * exponential backoff is implemented.
-   */
-  void setInitialBackoff(int value);
-
   @Description("Output bucket to write de-identified files in Cloud Storage buckets")
   String getOutputBucket();
 
@@ -212,45 +156,4 @@ public FileType create(PipelineOptions options) {
   FileType getFileType();
 
   void setFileType(FileType fileType);
-
-  class DLPConfigProjectFactory implements DefaultValueFactory<String> {
-    @Override
-    public String create(PipelineOptions options) {
-      return LocationName.of(
-              ((DLPTextToBigQueryStreamingV2PipelineOptions) options).getProject(), "global")
-          .toString();
-    }
-  }
-
-  @Default.InstanceFactory(DLPConfigProjectFactory.class)
-  String getDLPParent();
-
-  void setDLPParent(String parent);
-
-  class InputPollingFactory implements DefaultValueFactory<InputLocation> {
-    @Override
-    public InputLocation create(PipelineOptions options) {
-      if (((DLPTextToBigQueryStreamingV2PipelineOptions) options)
-          .getFilePattern()
-          .startsWith("gs://")) return InputLocation.GCS;
-      else return InputLocation.NOT_GCS;
-    }
-  }
-
-  @Validation.Required
-  @Default.InstanceFactory(InputPollingFactory.class)
-  InputLocation getInputProviderType();
-
-  void setInputProviderType(InputLocation input);
-
-  @Description("Topic to use for GCS Pub/Sub")
-  String getGcsNotificationTopic();
-
-  void setGcsNotificationTopic(String topic);
-
-  @Description("Flag to process existing files")
-  @Default.Boolean(true)
-  Boolean getProcessExistingFiles();
-
-  void setProcessExistingFiles(Boolean value);
 }

src/main/java/com/google/swarm/tokenization/beam/ShardRows.java

@@ -16,15 +16,13 @@
 package com.google.swarm.tokenization.beam;
 
 import com.google.privacy.dlp.v2.Table.Row;
-import com.google.swarm.tokenization.DLPTextToBigQueryStreamingV2PipelineOptions;
+import com.google.swarm.tokenization.options.CommonPipelineOptions;
 import java.util.Random;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.ShardedKeyCoder;
 import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.Element;
-import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.KV;
@@ -58,8 +56,8 @@ public PCollection<KV<ShardedKey<String>, Row>> expand(PCollection<KV<String, Ro
     }
 
     if (numberOfShards == UNDEFINED_NUMBER_OF_SHARDS) {
-      DLPTextToBigQueryStreamingV2PipelineOptions options =
-          input.getPipeline().getOptions().as(DLPTextToBigQueryStreamingV2PipelineOptions.class);
+      CommonPipelineOptions options =
+          input.getPipeline().getOptions().as(CommonPipelineOptions.class);
       numberOfShards = options.getNumShardsPerDLPRequestBatching();
     }
 

src/main/java/com/google/swarm/tokenization/classification/ClassifyFiles.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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 com.google.swarm.tokenization.classification;
+
+import com.google.auto.value.AutoValue;
+import com.google.privacy.dlp.v2.InspectContentResponse;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
+import org.apache.beam.sdk.transforms.*;
+import org.apache.beam.sdk.values.*;
+
+@AutoValue
+public abstract class ClassifyFiles
+    extends PTransform<PCollection<KV<String, InspectContentResponse>>, PDone> {
+
+  public abstract String outputPubSubTopic();
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+    public abstract Builder setOutputPubSubTopic(String value);
+
+    public abstract ClassifyFiles build();
+  }
+
+  public static ClassifyFiles.Builder newBuilder() {
+    return new AutoValue_ClassifyFiles.Builder();
+  }
+
+  @Override
+  public PDone expand(PCollection<KV<String, InspectContentResponse>> input) {
+
+    PCollectionTuple findings =
+        input.apply(
+            "ProcessFindings",
+            ParDo.of(new ProcessInspectFindingsDoFn())
+                .withOutputTags(
+                    ProcessInspectFindingsDoFn.getInspectSuccessTag(),
+                    TupleTagList.of(ProcessInspectFindingsDoFn.getInspectFailureTag())));
+
+    PCollection<PubsubMessage> inspectSuccess =
+        findings
+            .get(ProcessInspectFindingsDoFn.getInspectSuccessTag())
+            .apply("FindTotalFindings", Combine.<String, Long, Long>perKey(Sum.ofLongs()))
+            .apply("ApplySensitiveTag", ParDo.of(new CreatePubSubMessage("success")));
+
+    PCollection<PubsubMessage> inspectFailure =
+        findings
+            .get(ProcessInspectFindingsDoFn.getInspectFailureTag())
+            .apply("FindTotalErrors", Combine.<String, Long, Long>perKey(Sum.ofLongs()))
+            .apply("CreateMessage", ParDo.of(new CreatePubSubMessage("error")));
+
+    return PCollectionList.of(inspectSuccess)
+        .and(inspectFailure)
+        .apply("Flatten", Flatten.pCollections())
+        .apply("PublishToPubSub", PubsubIO.writeMessages().to(outputPubSubTopic()));
+  }
+}

src/main/java/com/google/swarm/tokenization/classification/CreatePubSubMessage.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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 com.google.swarm.tokenization.classification;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CreatePubSubMessage extends DoFn<KV<String, Long>, PubsubMessage> {
+
+  public String messageType;
+  public static final Logger LOG = LoggerFactory.getLogger(CreatePubSubMessage.class);
+
+  public CreatePubSubMessage(String messageType) {
+    this.messageType = messageType;
+  }
+
+  @ProcessElement
+  public void processElement(ProcessContext c) {
+
+    String filename = c.element().getKey();
+    Map<String, String> attributeMap = new HashMap<String, String>();
+    attributeMap.put("filename", filename);
+
+    if (messageType.equals("success")) {
+      Long findings = c.element().getValue();
+      attributeMap.put("findings", findings.toString());
+      if (findings > 0) attributeMap.put("DLPSensitiveTag", "SENSITIVE");
+      else attributeMap.put("DLPSensitiveTag", "NOT_SENSITIVE");
+      LOG.info("Inspect findings {} in file {}", findings.toString(), filename);
+      c.output(new PubsubMessage(filename.getBytes(), attributeMap));
+
+    } else {
+      attributeMap.put("findings", "ERROR");
+      LOG.info("Inspect errors in file {}", c.element().getValue().toString(), filename);
+      c.output(new PubsubMessage("Error in DLP Inspect".getBytes(), attributeMap));
+    }
+  }
+}

src/main/java/com/google/swarm/tokenization/classification/ProcessInspectFindingsDoFn.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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 com.google.swarm.tokenization.classification;
+
+import com.google.privacy.dlp.v2.InspectContentResponse;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ProcessInspectFindingsDoFn
+    extends DoFn<KV<String, InspectContentResponse>, KV<String, Long>> {
+
+  public static final TupleTag<KV<String, Long>> inspectSuccessTag =
+      new TupleTag<KV<String, Long>>() {};
+  public static final TupleTag<KV<String, Long>> inspectFailureTag =
+      new TupleTag<KV<String, Long>>() {};
+
+  public static final Logger LOG = LoggerFactory.getLogger(ProcessInspectFindingsDoFn.class);
+
+  public static TupleTag<KV<String, Long>> getInspectSuccessTag() {
+    return inspectSuccessTag;
+  }
+
+  public static TupleTag<KV<String, Long>> getInspectFailureTag() {
+    return inspectFailureTag;
+  }
+
+  @ProcessElement
+  public void processElement(
+      @Element KV<String, InspectContentResponse> element, MultiOutputReceiver out) {
+
+    String filename = element.getKey();
+    Long findings = Long.valueOf(element.getValue().getResult().getFindingsList().size());
+
+    out.get(inspectSuccessTag).output(KV.of(filename, findings));
+
+    for (String error : element.getValue().findInitializationErrors()) {
+      LOG.info("Found Error in file {}, error: {}", filename, error.toString());
+      out.get(inspectFailureTag).output(KV.of(filename, Long.valueOf(1)));
+    }
+  }
+}

src/main/java/com/google/swarm/tokenization/common/DLPTransform.java

@@ -263,7 +263,7 @@ public void processElement(
     }
   }
 
-  static class ConvertInspectResponse
+  public static class ConvertInspectResponse
       extends DoFn<KV<String, InspectContentResponse>, KV<String, TableRow>> {
 
     private String jobName;

src/main/java/com/google/swarm/tokenization/common/Util.java

@@ -73,7 +73,8 @@ public class Util {
   public enum DLPMethod {
     INSPECT,
     DEID,
-    REID
+    REID,
+    INSPECT_FILE
   }
 
   public enum FileType {