diff --git a/src/main/java/com/google/swarm/tokenization/DLPTextToBigQueryStreamingV2PipelineOptions.java b/src/main/java/com/google/swarm/tokenization/DLPTextToBigQueryStreamingV2PipelineOptions.java index 61736edc..c5b75d6a 100644 --- a/src/main/java/com/google/swarm/tokenization/DLPTextToBigQueryStreamingV2PipelineOptions.java +++ b/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 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 { - @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 { - @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); } diff --git a/src/main/java/com/google/swarm/tokenization/beam/ShardRows.java b/src/main/java/com/google/swarm/tokenization/beam/ShardRows.java index c19ccf3b..d4fa5724 100644 --- a/src/main/java/com/google/swarm/tokenization/beam/ShardRows.java +++ b/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, Row>> expand(PCollection>, 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> input) { + + PCollectionTuple findings = + input.apply( + "ProcessFindings", + ParDo.of(new ProcessInspectFindingsDoFn()) + .withOutputTags( + ProcessInspectFindingsDoFn.getInspectSuccessTag(), + TupleTagList.of(ProcessInspectFindingsDoFn.getInspectFailureTag()))); + + PCollection inspectSuccess = + findings + .get(ProcessInspectFindingsDoFn.getInspectSuccessTag()) + .apply("FindTotalFindings", Combine.perKey(Sum.ofLongs())) + .apply("ApplySensitiveTag", ParDo.of(new CreatePubSubMessage("success"))); + + PCollection inspectFailure = + findings + .get(ProcessInspectFindingsDoFn.getInspectFailureTag()) + .apply("FindTotalErrors", Combine.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())); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/classification/CreatePubSubMessage.java b/src/main/java/com/google/swarm/tokenization/classification/CreatePubSubMessage.java new file mode 100644 index 00000000..41fc7a09 --- /dev/null +++ b/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, 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 attributeMap = new HashMap(); + 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)); + } + } +} diff --git a/src/main/java/com/google/swarm/tokenization/classification/ProcessInspectFindingsDoFn.java b/src/main/java/com/google/swarm/tokenization/classification/ProcessInspectFindingsDoFn.java new file mode 100644 index 00000000..ab2e96d0 --- /dev/null +++ b/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> { + + public static final TupleTag> inspectSuccessTag = + new TupleTag>() {}; + public static final TupleTag> inspectFailureTag = + new TupleTag>() {}; + + public static final Logger LOG = LoggerFactory.getLogger(ProcessInspectFindingsDoFn.class); + + public static TupleTag> getInspectSuccessTag() { + return inspectSuccessTag; + } + + public static TupleTag> getInspectFailureTag() { + return inspectFailureTag; + } + + @ProcessElement + public void processElement( + @Element KV 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))); + } + } +} diff --git a/src/main/java/com/google/swarm/tokenization/common/DLPTransform.java b/src/main/java/com/google/swarm/tokenization/common/DLPTransform.java index f6978b66..a8844a0c 100644 --- a/src/main/java/com/google/swarm/tokenization/common/DLPTransform.java +++ b/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> { private String jobName; diff --git a/src/main/java/com/google/swarm/tokenization/common/Util.java b/src/main/java/com/google/swarm/tokenization/common/Util.java index 37ccc5ef..9b5bb6b8 100644 --- a/src/main/java/com/google/swarm/tokenization/common/Util.java +++ b/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 { diff --git a/src/main/java/com/google/swarm/tokenization/options/CommonPipelineOptions.java b/src/main/java/com/google/swarm/tokenization/options/CommonPipelineOptions.java new file mode 100644 index 00000000..67fe0f9e --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/options/CommonPipelineOptions.java @@ -0,0 +1,119 @@ +/* + * Copyright 2020 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.options; + +import com.google.privacy.dlp.v2.LocationName; +import com.google.swarm.tokenization.common.Util; +import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; +import org.apache.beam.sdk.io.aws.options.S3Options; +import org.apache.beam.sdk.options.*; + +public interface CommonPipelineOptions extends DataflowPipelineOptions, S3Options { + + @Description("DLP Inspect Template Name") + String getInspectTemplateName(); + + void setInspectTemplateName(String value); + + @Description("DLP DeIdentify Template Name") + String getDeidentifyTemplateName(); + + void setDeidentifyTemplateName(String 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); + + @Default.Integer(900 * 1000) + Integer getSplitSize(); + + void setSplitSize(Integer value); + + @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); + + class DLPConfigProjectFactory implements DefaultValueFactory { + @Override + public String create(PipelineOptions options) { + return LocationName.of( + ((com.google.swarm.tokenization.options.CommonPipelineOptions) options).getProject(), + "global") + .toString(); + } + } + + @Default.InstanceFactory(DLPConfigProjectFactory.class) + String getDLPParent(); + + void setDLPParent(String parent); + + class InputPollingFactory implements DefaultValueFactory { + @Override + public Util.InputLocation create(PipelineOptions options) { + if (((com.google.swarm.tokenization.DLPTextToBigQueryStreamingV2PipelineOptions) options) + .getFilePattern() + .startsWith("gs://")) return Util.InputLocation.GCS; + else return Util.InputLocation.NOT_GCS; + } + } + + @Validation.Required + @Default.InstanceFactory(InputPollingFactory.class) + Util.InputLocation getInputProviderType(); + + void setInputProviderType(Util.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); +} diff --git a/src/main/java/com/google/swarm/tokenization/options/InspectClassifyPipelineOptions.java b/src/main/java/com/google/swarm/tokenization/options/InspectClassifyPipelineOptions.java new file mode 100644 index 00000000..040a91d4 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/options/InspectClassifyPipelineOptions.java @@ -0,0 +1,64 @@ +/* + * Copyright 2023 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.options; + +import com.google.swarm.tokenization.common.Util; +import java.util.List; +import org.apache.beam.sdk.options.*; + +public interface InspectClassifyPipelineOptions extends CommonPipelineOptions { + @Validation.Required + String getFilePattern(); + + void setFilePattern(String csvFilePattern); + + @Description("DLP method deid,inspect,reid") + @Default.Enum("INSPECT") + Util.DLPMethod getDLPMethod(); + + void setDLPMethod(Util.DLPMethod value); + + @Description("BQ Dataset") + String getDataset(); + + void setDataset(String value); + + List getFileTypes(); + + void setFileTypes(List fileTypes); + + @Description("Topic to publish the inspection results") + String getTopic(); + + void setTopic(String topic); + + @Description("Column Delimiter for csv files") + @Default.Character(',') + Character getColumnDelimiterForCsvFiles(); + + void setColumnDelimiterForCsvFiles(Character value); + + @Description("Record delimiter") + @Default.String("\n") + String getRecordDelimiter(); + + void setRecordDelimiter(String value); + + @Description("DLP Table Headers- Required for Jsonl type") + List getHeaders(); + + void setHeaders(List topic); +} diff --git a/src/main/java/com/google/swarm/tokenization/pipelines/InspectClassify.java b/src/main/java/com/google/swarm/tokenization/pipelines/InspectClassify.java new file mode 100644 index 00000000..3f57aa91 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/pipelines/InspectClassify.java @@ -0,0 +1,126 @@ +/* + * Copyright 2023 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.pipelines; + +import com.google.api.services.bigquery.model.TableRow; +import com.google.privacy.dlp.v2.InspectContentResponse; +import com.google.swarm.tokenization.beam.DLPInspectText; +import com.google.swarm.tokenization.classification.ClassifyFiles; +import com.google.swarm.tokenization.coders.DeterministicTableRowJsonCoder; +import com.google.swarm.tokenization.common.BigQueryDynamicWriteTransform; +import com.google.swarm.tokenization.common.DLPTransform; +import com.google.swarm.tokenization.common.Util; +import com.google.swarm.tokenization.options.InspectClassifyPipelineOptions; +import com.google.swarm.tokenization.transforms.ProcessFiles; +import com.google.swarm.tokenization.transforms.ReadFilesWithGivenExtensions; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.CannotProvideCoderException; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.*; +import org.joda.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class InspectClassify { + + public static final Logger LOG = LoggerFactory.getLogger(InspectClassify.class); + private static final Duration DEFAULT_POLL_INTERVAL = Duration.standardSeconds(3); + private static final Duration WINDOW_INTERVAL = Duration.standardSeconds(3); + + public static void main(String[] args) throws CannotProvideCoderException { + + InspectClassifyPipelineOptions options = + PipelineOptionsFactory.fromArgs(args).as(InspectClassifyPipelineOptions.class); + run(options); + } + + public static PipelineResult run(InspectClassifyPipelineOptions options) + throws CannotProvideCoderException { + Pipeline p = Pipeline.create(options); + p.getCoderRegistry().registerCoderForClass(TableRow.class, DeterministicTableRowJsonCoder.of()); + + List fileExtensionList = + options.getFileTypes().stream() + .map(x -> Util.FileType.valueOf(x.toUpperCase())) + .collect(Collectors.toList()); + + /** Fetch List of input files */ + PCollection> inputFiles = + p.apply( + "ReadExistingFiles", + ReadFilesWithGivenExtensions.newBuilder() + .setFilePattern(options.getFilePattern()) + .setFileTypes(options.getFileTypes()) + .build()) + .apply("Fixed Window", Window.into(FixedWindows.of(WINDOW_INTERVAL))); + + /** Process Files */ + PCollectionTuple records = + inputFiles.apply( + "ProcessFiles", ProcessFiles.newBuilder().setFileTypes(fileExtensionList).build()); + + final PCollectionView>> headers = + records.get(ProcessFiles.headersMap).apply("ViewAsList", View.asMap()); + + PCollection> inspectResponse = + records + .get(ProcessFiles.tableRows) + .apply( + "DLPInspection", + DLPInspectText.newBuilder() + .setBatchSizeBytes(options.getBatchSize()) + .setColumnDelimiter(options.getColumnDelimiterForCsvFiles()) + .setHeaderColumns(headers) + .setInspectTemplateName(options.getInspectTemplateName()) + .setProjectId(options.getDLPParent()) + .setDlpApiRetryCount(options.getDlpApiRetryCount()) + .setInitialBackoff(options.getInitialBackoff()) + .build()); + + if (options.getTopic() != null) { + inspectResponse.apply( + "ClassifyFiles", + ClassifyFiles.newBuilder().setOutputPubSubTopic(options.getTopic()).build()); + } + + if (options.getDataset() != null) { + inspectResponse + .apply( + "ConvertInspectResponse", + ParDo.of(new DLPTransform.ConvertInspectResponse(options.getJobName())) + .withOutputTags( + Util.inspectOrDeidSuccess, TupleTagList.of(Util.inspectOrDeidFailure))) + .get(Util.inspectOrDeidSuccess) + .apply( + "InsertToBQ", + BigQueryDynamicWriteTransform.newBuilder() + .setDatasetId(options.getDataset()) + .setProjectId(options.getProject()) + .build()); + } + + return p.run(); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ProcessAvro.java b/src/main/java/com/google/swarm/tokenization/transforms/ProcessAvro.java new file mode 100644 index 00000000..8ce73138 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ProcessAvro.java @@ -0,0 +1,80 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.privacy.dlp.v2.Table; +import com.google.swarm.tokenization.avro.AvroColumnNamesDoFn; +import com.google.swarm.tokenization.avro.AvroReaderSplittableDoFn; +import com.google.swarm.tokenization.avro.ConvertAvroRecordToDlpRowDoFn; +import com.google.swarm.tokenization.avro.GenericRecordCoder; +import com.google.swarm.tokenization.common.ResolveDuplicatesCombineFn; +import com.google.swarm.tokenization.options.InspectClassifyPipelineOptions; +import java.util.List; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; + +@AutoValue +public abstract class ProcessAvro + extends PTransform>, PCollectionTuple> { + + public static final TupleTag> recordsTuple = + new TupleTag>() {}; + public static final TupleTag>> headersTuple = + new TupleTag>>() {}; + + @AutoValue.Builder + public abstract static class Builder { + + public abstract ProcessAvro build(); + } + + public static ProcessAvro.Builder newBuilder() { + return new AutoValue_ProcessAvro.Builder(); + } + + @Override + public PCollectionTuple expand(PCollection> input) { + + InspectClassifyPipelineOptions options = + input.getPipeline().getOptions().as(InspectClassifyPipelineOptions.class); + + PCollection>> headers = + input + .apply("ReadHeader", ParDo.of(new AvroColumnNamesDoFn())) + .apply( + Combine., List>perKey( + new ResolveDuplicatesCombineFn())); + + PCollection> records = + input + .apply( + ParDo.of( + new AvroReaderSplittableDoFn(options.getKeyRange(), options.getSplitSize()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), GenericRecordCoder.of())) + .apply(ParDo.of(new ConvertAvroRecordToDlpRowDoFn())); + + return PCollectionTuple.of(recordsTuple, records).and(headersTuple, headers); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ProcessCSV.java b/src/main/java/com/google/swarm/tokenization/transforms/ProcessCSV.java new file mode 100644 index 00000000..85ad9b64 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ProcessCSV.java @@ -0,0 +1,87 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.privacy.dlp.v2.Table; +import com.google.swarm.tokenization.beam.ConvertCSVRecordToDLPRow; +import com.google.swarm.tokenization.common.CSVColumnNamesDoFn; +import com.google.swarm.tokenization.common.CSVFileReaderSplitDoFn; +import com.google.swarm.tokenization.common.ResolveDuplicatesCombineFn; +import com.google.swarm.tokenization.options.InspectClassifyPipelineOptions; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.values.*; + +@AutoValue +public abstract class ProcessCSV + extends PTransform>, PCollectionTuple> { + + public static final TupleTag> recordsTuple = + new TupleTag>() {}; + public static final TupleTag>> headersTuple = + new TupleTag>>() {}; + + @AutoValue.Builder + public abstract static class Builder { + + public abstract ProcessCSV build(); + } + + public static ProcessCSV.Builder newBuilder() { + return new AutoValue_ProcessCSV.Builder(); + } + + @Override + public PCollectionTuple expand(PCollection> input) { + + InspectClassifyPipelineOptions options = + input.getPipeline().getOptions().as(InspectClassifyPipelineOptions.class); + + PCollection>> headers = + input + .apply( + "ReadHeader", + ParDo.of(new CSVColumnNamesDoFn(options.getColumnDelimiterForCsvFiles()))) + .apply( + Combine., List>perKey( + new ResolveDuplicatesCombineFn())); + + PCollectionView>> headersMap = + headers.apply("ViewAsList", View.asMap()); + + PCollection> records = + input + .apply( + "SplitCSVFile", + ParDo.of( + new CSVFileReaderSplitDoFn( + options.getRecordDelimiter(), options.getSplitSize()))) + .apply( + "ConvertToDLPRow", + ParDo.of( + new ConvertCSVRecordToDLPRow( + options.getColumnDelimiterForCsvFiles(), headersMap)) + .withSideInputs(headersMap)); + + return PCollectionTuple.of(recordsTuple, records).and(headersTuple, headers); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ProcessFiles.java b/src/main/java/com/google/swarm/tokenization/transforms/ProcessFiles.java new file mode 100644 index 00000000..54524240 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ProcessFiles.java @@ -0,0 +1,143 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.privacy.dlp.v2.Table; +import com.google.swarm.tokenization.common.Util; +import java.util.List; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.*; +import org.apache.beam.sdk.values.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@AutoValue +public abstract class ProcessFiles + extends PTransform>, PCollectionTuple> { + + public abstract List fileTypes(); + + public static final Logger LOG = LoggerFactory.getLogger(ProcessFiles.class); + + public static final TupleTag> tableRows = + new TupleTag>() {}; + public static final TupleTag>> headersMap = + new TupleTag>>() {}; + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setFileTypes(List value); + + public abstract ProcessFiles build(); + } + + public static ProcessFiles.Builder newBuilder() { + return new AutoValue_ProcessFiles.Builder(); + } + + @Override + public PCollectionTuple expand(PCollection> input) { + + PCollectionList> fileList = + input.apply( + Partition.of( + Util.FileType.values().length, + new Partition.PartitionFn>() { + @Override + public int partitionFor(KV elem, int numPartitions) { + String resourceId = String.valueOf(elem.getValue().getMetadata().resourceId()); + String extension = + String.valueOf(resourceId.substring(resourceId.lastIndexOf(".") + 1)); + LOG.info( + "Found extension: {} and FileType: {} ", + extension, + Util.FileType.valueOf(extension.toUpperCase()).toString()); + return Util.FileType.valueOf(extension.toUpperCase()).ordinal(); + } + })); + + PCollectionList> recordsList = PCollectionList.empty(input.getPipeline()); + PCollectionList>> headerList = + PCollectionList.empty(input.getPipeline()); + + for (Util.FileType fileType : fileTypes()) { + switch (fileType) { + case AVRO: + PCollection> inputFilesAvro = + fileList.get(Util.FileType.AVRO.ordinal()); + PCollectionTuple avroRecords = + inputFilesAvro.apply("ProcessAVRO", ProcessAvro.newBuilder().build()); + recordsList = recordsList.and(avroRecords.get(ProcessAvro.recordsTuple)); + headerList = headerList.and(avroRecords.get(ProcessAvro.headersTuple)); + break; + + case CSV: + PCollection> inputFilesCsv = + fileList.get(Util.FileType.CSV.ordinal()); + PCollectionTuple csvRecords = + inputFilesCsv.apply("ProcessCSV", ProcessCSV.newBuilder().build()); + recordsList = recordsList.and(csvRecords.get(ProcessCSV.recordsTuple)); + headerList = headerList.and(csvRecords.get(ProcessCSV.headersTuple)); + break; + + case JSONL: + PCollection> inputFilesJsonl = + fileList.get(Util.FileType.CSV.ordinal()); + PCollectionTuple jsonlRecords = + inputFilesJsonl.apply("ProcessJSONL", ProcessJsonl.newBuilder().build()); + recordsList = recordsList.and(jsonlRecords.get(ProcessJsonl.recordsTuple)); + headerList = headerList.and(jsonlRecords.get(ProcessJsonl.headersTuple)); + break; + + case PARQUET: + PCollection> inputFilesParquet = + fileList.get(Util.FileType.CSV.ordinal()); + PCollectionTuple parquetRecords = + inputFilesParquet.apply("ProcessPARQUET", ProcessParquet.newBuilder().build()); + recordsList = recordsList.and(parquetRecords.get(ProcessParquet.recordsTuple)); + headerList = headerList.and(parquetRecords.get(ProcessParquet.headersTuple)); + break; + + default: + throw new IllegalArgumentException("Please validate FileType parameter"); + } + } + + try { + PCollection> flattenedRecords = + recordsList + .apply(Flatten.pCollections()) + .setCoder( + KvCoder.of( + StringUtf8Coder.of(), + input.getPipeline().getCoderRegistry().getCoder(Table.Row.class))); + + PCollection>> flattenedHeaders = + headerList + .apply("Merge Headers", Flatten.pCollections()) + .setCoder(KvCoder.of(StringUtf8Coder.of(), ListCoder.of(StringUtf8Coder.of()))); + + return PCollectionTuple.of(tableRows, flattenedRecords).and(headersMap, flattenedHeaders); + } catch (Exception e) { + LOG.info("Found exception " + e.toString()); + throw new RuntimeException(e); + } + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ProcessJsonl.java b/src/main/java/com/google/swarm/tokenization/transforms/ProcessJsonl.java new file mode 100644 index 00000000..3220b884 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ProcessJsonl.java @@ -0,0 +1,82 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.privacy.dlp.v2.Table; +import com.google.swarm.tokenization.common.ResolveDuplicatesCombineFn; +import com.google.swarm.tokenization.json.ConvertJsonRecordToDLPRow; +import com.google.swarm.tokenization.json.JsonColumnNameDoFn; +import com.google.swarm.tokenization.json.JsonReaderSplitDoFn; +import com.google.swarm.tokenization.options.InspectClassifyPipelineOptions; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.values.*; + +@AutoValue +public abstract class ProcessJsonl + extends PTransform>, PCollectionTuple> { + + public static final TupleTag> recordsTuple = + new TupleTag>() {}; + public static final TupleTag>> headersTuple = + new TupleTag>>() {}; + + @AutoValue.Builder + public abstract static class Builder { + + public abstract ProcessJsonl build(); + } + + public static ProcessJsonl.Builder newBuilder() { + return new AutoValue_ProcessJsonl.Builder(); + } + + @Override + public PCollectionTuple expand(PCollection> input) { + + InspectClassifyPipelineOptions options = + input.getPipeline().getOptions().as(InspectClassifyPipelineOptions.class); + + PCollection>> headers = + input + .apply("ReadHeader", ParDo.of(new JsonColumnNameDoFn(options.getHeaders()))) + .apply( + Combine., List>perKey( + new ResolveDuplicatesCombineFn())); + + PCollectionView>> headersMap = + headers.apply("ViewAsList", View.asMap()); + + PCollection> records = + input + .apply( + "SplitJSONFile", + ParDo.of( + new JsonReaderSplitDoFn( + options.getKeyRange(), + options.getRecordDelimiter(), + options.getSplitSize()))) + .apply("ConvertToDLPRow", ParDo.of(new ConvertJsonRecordToDLPRow())); + + return PCollectionTuple.of(recordsTuple, records).and(headersTuple, headers); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ProcessParquet.java b/src/main/java/com/google/swarm/tokenization/transforms/ProcessParquet.java new file mode 100644 index 00000000..f315abec --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ProcessParquet.java @@ -0,0 +1,75 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.privacy.dlp.v2.Table; +import com.google.swarm.tokenization.common.ResolveDuplicatesCombineFn; +import com.google.swarm.tokenization.options.InspectClassifyPipelineOptions; +import com.google.swarm.tokenization.parquet.ParquetColumnNamesDoFn; +import com.google.swarm.tokenization.parquet.ParquetReaderSplittableDoFn; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.values.*; + +@AutoValue +public abstract class ProcessParquet + extends PTransform>, PCollectionTuple> { + + public static final TupleTag> recordsTuple = + new TupleTag>() {}; + public static final TupleTag>> headersTuple = + new TupleTag>>() {}; + + @AutoValue.Builder + public abstract static class Builder { + + public abstract ProcessParquet build(); + } + + public static ProcessParquet.Builder newBuilder() { + return new AutoValue_ProcessParquet.Builder(); + } + + @Override + public PCollectionTuple expand(PCollection> input) { + + InspectClassifyPipelineOptions options = + input.getPipeline().getOptions().as(InspectClassifyPipelineOptions.class); + + PCollection>> headers = + input + .apply("ReadHeader", ParDo.of(new ParquetColumnNamesDoFn())) + .apply( + Combine., List>perKey( + new ResolveDuplicatesCombineFn())); + + PCollectionView>> headersMap = + headers.apply("ViewAsList", View.asMap()); + + PCollection> records = + input.apply( + ParDo.of( + new ParquetReaderSplittableDoFn(options.getKeyRange(), options.getSplitSize()))); + + return PCollectionTuple.of(recordsTuple, records).and(headersTuple, headers); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/transforms/ReadFilesWithGivenExtensions.java b/src/main/java/com/google/swarm/tokenization/transforms/ReadFilesWithGivenExtensions.java new file mode 100644 index 00000000..50264f25 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/transforms/ReadFilesWithGivenExtensions.java @@ -0,0 +1,71 @@ +/* + * 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.transforms; + +import com.google.auto.value.AutoValue; +import com.google.swarm.tokenization.common.SanitizeFileNameDoFn; +import com.google.swarm.tokenization.common.Util; +import java.util.List; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.io.fs.MatchResult; +import org.apache.beam.sdk.transforms.*; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; + +@AutoValue +public abstract class ReadFilesWithGivenExtensions + extends PTransform>> { + + public abstract String filePattern(); + + public abstract List fileTypes(); + + @AutoValue.Builder + public abstract static class Builder { + + public abstract ReadFilesWithGivenExtensions.Builder setFilePattern(String filePattern); + + public abstract ReadFilesWithGivenExtensions.Builder setFileTypes(List value); + + public abstract ReadFilesWithGivenExtensions build(); + } + + public static ReadFilesWithGivenExtensions.Builder newBuilder() { + return new AutoValue_ReadFilesWithGivenExtensions.Builder(); + } + + @Override + public PCollection> expand(PBegin input) { + return input + .apply("MatchExistingFiles", FileIO.match().filepattern(filePattern())) + .apply( + "FilterFiles", + Filter.by( + new SerializableFunction() { + @Override + public Boolean apply(MatchResult.Metadata input) { + String filename = input.resourceId().getFilename(); + String fileExtension = filename.substring(filename.lastIndexOf(".") + 1); + return fileTypes().contains(fileExtension); + } + })) + .apply("ReadExistingFiles", FileIO.readMatches()) + .apply( + "SanitizeFileNameExistingFiles", + ParDo.of(new SanitizeFileNameDoFn(Util.InputLocation.GCS))); + } +}