diff --git a/.github/workflows/dlp-pipelines.yml b/.github/workflows/dlp-pipelines.yml index 582373e9..2c600069 100644 --- a/.github/workflows/dlp-pipelines.yml +++ b/.github/workflows/dlp-pipelines.yml @@ -31,6 +31,44 @@ env: OUTPUT_GCS_BUCKET: "deidentified_files" jobs: + skip-workflow: + runs-on: ubuntu-latest + + outputs: + skip_inspect: ${{ steps.skip-inspect-workflow.outputs.skip }} + skip_deid: ${{ steps.skip-deid-workflow.outputs.skip }} + skip_reid: ${{ steps.skip-reid-workflow.outputs.skip }} + + steps: + - uses: actions/checkout@v2 + + - name: skip-inspect-workflow + id: skip-inspect-workflow # id used for referencing step + uses: saulmaldonado/skip-workflow@v1 + with: + phrase: '[skip-inspect]' + search: '["pull_request"]' + pr-message: 'body' + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: skip-deid-workflow + id: skip-deid-workflow # id used for referencing step + uses: saulmaldonado/skip-workflow@v1 + with: + phrase: '[skip-deid]' + search: '["pull_request"]' + pr-message: 'body' + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: skip-reid-workflow + id: skip-reid-workflow # id used for referencing step + uses: saulmaldonado/skip-workflow@v1 + with: + phrase: '[skip-reid]' + search: '["pull_request"]' + pr-message: 'body' + github-token: ${{ secrets.GITHUB_TOKEN }} + build-and-unit-test: permissions: @@ -58,7 +96,9 @@ jobs: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} generate-uuid: - needs: build-and-unit-test + needs: + - build-and-unit-test + - skip-workflow runs-on: ubuntu-latest timeout-minutes: 5 @@ -129,6 +169,9 @@ jobs: needs: - generate-uuid - create-dataset + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_inspect }} runs-on: - self-hosted @@ -214,6 +257,9 @@ jobs: needs: - generate-uuid - create-dataset + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_inspect }} runs-on: - self-hosted @@ -292,6 +338,9 @@ jobs: needs: - generate-uuid - create-dataset + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_deid }} runs-on: - self-hosted @@ -404,6 +453,9 @@ jobs: needs: - generate-uuid - create-dataset + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_deid }} runs-on: - self-hosted @@ -485,6 +537,9 @@ jobs: needs: - generate-uuid - create-dataset + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_deid }} runs-on: - self-hosted @@ -644,6 +699,9 @@ jobs: - generate-uuid - create-dataset - de-identification + - skip-workflow + + if: ${{ !needs.skip-workflow.outputs.skip_reid }} runs-on: - self-hosted diff --git a/README.md b/README.md index 4508c50d..a9d920af 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ * [De-identification](#de-identification) * [Validation](#validation) * [Re-identification](#re-identification-from-bigquery) + * [Inspect and Classify](#inspect-and-classify) * [Pipeline parameters](#pipeline-parameters) * [Supported file formats](#supported-file-formats) * [Amazon S3 scanner](#amazon-s3-scanner) @@ -447,6 +448,24 @@ to validate de-identified results: stored in `reid_query.sql`. The re-identified results can be found in the BigQuery dataset (`dataset` parameter) with the name of the input table as the suffix. +#### Inspect and Classify + +```commandline +./gradlew run -DmainClass=com.google.swarm.tokenization.InspectClassify \ +-Pargs=" --region=${REGION} \ +--project=${PROJECT_ID} \ +--tempLocation=gs://${DATA_STORAGE_BUCKET}/temp \ +--numWorkers=1 --maxNumWorkers=2 \ +--runner=DataflowRunner \ +--filePattern=gs://dlp-orc-support-398810-demo-data/inspect_classify/* \ +--dataset=inspect_classify \ +--inspectTemplateName=${INSPECT_TEMPLATE_NAME} \ +--fileTypes=fileType1,fileType2,fileTypeN \ +--batchSize=200000 \ +--DLPMethod=INSPECT \ +--serviceAccount=${SERVICE_ACCOUNT_EMAIL}" +``` + ### Pipeline Parameters diff --git a/src/main/java/com/google/swarm/tokenization/InspectClassify.java b/src/main/java/com/google/swarm/tokenization/InspectClassify.java new file mode 100644 index 00000000..630f62c6 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/InspectClassify.java @@ -0,0 +1,334 @@ +/* + * 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; + +import com.google.api.services.bigquery.model.TableRow; +import com.google.privacy.dlp.v2.FileType; +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.beam.ConvertCSVRecordToDLPRow; +import com.google.swarm.tokenization.classification.NewExtractColumnNamesTransform; +import com.google.swarm.tokenization.coders.DeterministicTableRowJsonCoder; +import com.google.swarm.tokenization.common.*; +import com.google.swarm.tokenization.json.JsonColumnNameDoFn; +import com.google.swarm.tokenization.orc.ORCColumnNameDoFn; +import com.google.swarm.tokenization.parquet.ParquetColumnNamesDoFn; +import com.google.swarm.tokenization.txt.TxtColumnNameDoFn; +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.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.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.*; +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); + + /** PubSub configuration for default batch size in number of messages */ + public static final Integer PUB_SUB_BATCH_SIZE = 1000; + + /** PubSub configuration for default batch size in bytes */ + public static final Integer PUB_SUB_BATCH_SIZE_BYTES = 10000; + + public static void main(String[] args) throws CannotProvideCoderException { + + InspectClassifyPipelineOptions options = + PipelineOptionsFactory.fromArgs(args).as(InspectClassifyPipelineOptions.class); + run(options); + } + + public static PCollectionView>> createHeaders( + PCollectionList> input, + Pipeline p, + List fileTypes, + InspectClassifyPipelineOptions options) { + + PCollection>> readHeader; + PCollectionList>> headerList = PCollectionList.empty(p); + + for (Util.FileType fileType : fileTypes) { + switch (fileType) { + case AVRO: + PCollection>> readHeaderAvro = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new AvroColumnNamesDoFn())); + headerList.and(readHeaderAvro); + break; + + case CSV: + PCollection>> readHeaderCsv = + input + .get(fileType.ordinal()) + .apply( + "ReadHeader", ParDo.of(new CSVColumnNamesDoFn(options.getColumnDelimiter()))); + headerList.and(readHeaderCsv); + break; + + case JSONL: + PCollection>> readHeaderJsonl = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new JsonColumnNameDoFn(options.getHeaders()))); + headerList.and(readHeaderJsonl); + break; + + case ORC: + PCollection>> readHeaderOrc = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new ORCColumnNameDoFn(options.getProject()))); + headerList.and(readHeaderOrc); + break; + + case PARQUET: + PCollection>> readHeaderParquet = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new ParquetColumnNamesDoFn())); + headerList.and(readHeaderParquet); + break; + + case TSV: + PCollection>> readHeaderTsv = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new CSVColumnNamesDoFn('\t'))); + headerList.and(readHeaderTsv); + break; + + case TXT: + PCollection>> readHeaderTxt = + input + .get(fileType.ordinal()) + .apply("ReadHeader", ParDo.of(new TxtColumnNameDoFn(options.getHeaders()))); + headerList.and(readHeaderTxt); + break; + + default: + throw new IllegalStateException("Unexpected value: " + fileType); + } + } + + readHeader = headerList.apply(Flatten.pCollections()) + .setCoder(KvCoder.of(StringUtf8Coder.of(), ListCoder.of(StringUtf8Coder.of()))); + + return readHeader.apply("ViewAsList", View.asMap()); + } + + public static PipelineResult run(InspectClassifyPipelineOptions options) throws CannotProvideCoderException { + Pipeline p = Pipeline.create(options); + p.getCoderRegistry().registerCoderForClass(TableRow.class, DeterministicTableRowJsonCoder.of()); + + PCollection> inputFiles = + p.apply( + "Read Existing Files", + ReadExistingFilesTransform.newBuilder() + .setFilePattern(options.getFilePattern()) + .setProcessExistingFiles(options.getProcessExistingFiles()) + .build()) + .apply("Fixed Window", Window.into(FixedWindows.of(WINDOW_INTERVAL))); + + /** Create a List of PCollections where list index represents the FileType enum's ordinal */ + PCollectionList> files = + inputFiles.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(); + } + })); + + /** + * Store Table.Row objects of every FileType in a Mapping from FileType to PCollection before + * DLPTransform + */ + PCollectionList> records = PCollectionList.empty(p); + + List fileExtensionList = + options.getFileTypes().stream() + .map(x -> Util.FileType.valueOf(x.toUpperCase())) + .collect(Collectors.toList()); + + /** + * Store header objects of every FileType in a single PCollectionView assuming no two input + * files can have same name. + */ + final PCollectionView>> headers = + files.apply( + "Extract Column Names", + NewExtractColumnNamesTransform.newBuilder() + .setHeaders(options.getHeaders()) + .setColumnDelimiter(options.getColumnDelimiter()) + .setPubSubGcs(false) + .setProjectId(options.getProject()) + .setFileTypes(fileExtensionList) + .build()); +// PCollectionView>> headers = +// createHeaders(files, p, fileExtensionList, options); + + for (Util.FileType fileType : fileExtensionList) { + switch (fileType) { + case AVRO: + PCollection> inputFilesavro = files.get(Util.FileType.AVRO.ordinal()); + PCollection> avroRecords = inputFilesavro.apply(ParDo.of(new AvroReaderSplittableDoFn(options.getKeyRange(), options.getSplitSize()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), GenericRecordCoder.of())) + .apply(ParDo.of(new ConvertAvroRecordToDlpRowDoFn())); + records = records.and(avroRecords); + break; + + case CSV: + PCollection> inputFilescsv = + files.get(Util.FileType.CSV.ordinal()); + PCollection> csvRecords = + files + .get(Util.FileType.CSV.ordinal()) + .apply( + "SplitCSVFile", + ParDo.of( + new CSVFileReaderSplitDoFn( + options.getRecordDelimiter(), options.getSplitSize()))) + .apply( + "ConvertToDLPRow", + ParDo.of(new ConvertCSVRecordToDLPRow(options.getColumnDelimiter(), headers)) + .withSideInputs(headers)); + + records = records.and(csvRecords); + break; + + // case JSONL: + // records = + // inputFiles + // .apply( + // "SplitJSONFile", + // ParDo.of( + // new JsonReaderSplitDoFn( + // options.getKeyRange(), + // + // options.getRecordDelimiter(), + // options.getSplitSize()))) + // .apply("ConvertToDLPRow", ParDo.of(new + // ConvertJsonRecordToDLPRow())); + // break; + // case TXT: + // PCollectionTuple recordTuple = + // inputFiles + // .apply( + // "SplitTextFile", + // ParDo.of( + // new TxtReaderSplitDoFn( + // options.getKeyRange(), + // + // options.getRecordDelimiter(), + // options.getSplitSize()))) + // .apply( + // "ParseTextFile", + // ParDo.of(new ParseTextLogDoFn()) + // .withOutputTags( + // Util.agentTranscriptTuple, + // + // TupleTagList.of(Util.customerTranscriptTuple))); + // + // records = + // + // PCollectionList.of(recordTuple.get(Util.agentTranscriptTuple)) + // .and(recordTuple.get(Util.customerTranscriptTuple)) + // .apply("Flatten", Flatten.pCollections()) + // .apply( + // "ConvertToDLPRow", + // ParDo.of(new + // ConvertTxtToDLPRow(options.getColumnDelimiter(), headers)) + // .withSideInputs(headers)); + // break; + // case PARQUET: + //// TODO: Remove KeyRange parameter, as it is unused + // records = inputFiles + // .apply(ParDo.of(new + // ParquetReaderSplittableDoFn(options.getKeyRange(), options.getSplitSize()))); + // break; + // + // case ORC: + // records = inputFiles + // .apply("ReadFromORCFilesAsOrcStruct", + // ParDo.of(new ORCReaderDoFn(options.getProject()))); + // break; + + case TEXT: + + default: + throw new IllegalArgumentException("Please validate FileType parameter"); + } + } + + PCollection> flattenedRecords = null; + try { + flattenedRecords = records.apply(Flatten.pCollections()) + .setCoder(KvCoder.of(StringUtf8Coder.of(), p.getCoderRegistry().getCoder(Table.Row.class))); + } catch (CannotProvideCoderException e) { + throw new RuntimeException(e); + } + ; + + flattenedRecords + .apply( + "DLPTransform", + DLPTransform.newBuilder() + .setBatchSize(options.getBatchSize()) + .setInspectTemplateName(options.getInspectTemplateName()) + .setDeidTemplateName(options.getDeidentifyTemplateName()) + .setDlpmethod(options.getDLPMethod()) + .setProjectId(options.getDLPParent()) + .setHeaders(headers) + .setColumnDelimiter(options.getColumnDelimiter()) + .setJobName(options.getJobName()) + .setDlpApiRetryCount(options.getDlpApiRetryCount()) + .setInitialBackoff(options.getInitialBackoff()) + .setDataSinkType(options.getDataSinkType()) + .build()) + .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/InspectClassifyPipelineOptions.java b/src/main/java/com/google/swarm/tokenization/InspectClassifyPipelineOptions.java new file mode 100644 index 00000000..16cdb778 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/InspectClassifyPipelineOptions.java @@ -0,0 +1,24 @@ +/* + * 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; + +import java.util.List; + +public interface InspectClassifyPipelineOptions extends DLPTextToBigQueryStreamingV2PipelineOptions { + + List getFileTypes(); + void setFileTypes(List fileTypes); +} 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..d6896118 100644 --- a/src/main/java/com/google/swarm/tokenization/beam/ShardRows.java +++ b/src/main/java/com/google/swarm/tokenization/beam/ShardRows.java @@ -16,7 +16,7 @@ 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.InspectClassifyPipelineOptions; import java.util.Random; import org.apache.beam.sdk.coders.CannotProvideCoderException; import org.apache.beam.sdk.coders.Coder; @@ -58,8 +58,8 @@ public PCollection, Row>> expand(PCollection>, + PCollectionView>>> { + + @Nullable + public abstract List headers(); + + public abstract Character columnDelimiter(); + + public abstract Boolean pubSubGcs(); + + public abstract String projectId(); + + public abstract List fileTypes(); + + @AutoValue.Builder + public abstract static class Builder { + + public abstract NewExtractColumnNamesTransform.Builder setHeaders(List headers); + + public abstract NewExtractColumnNamesTransform.Builder setColumnDelimiter( + Character columnDelimiter); + + public abstract NewExtractColumnNamesTransform.Builder setPubSubGcs(Boolean pubSubGcs); + + public abstract NewExtractColumnNamesTransform.Builder setProjectId(String projectId); + + public abstract NewExtractColumnNamesTransform.Builder setFileTypes(List fileTypes); + + public abstract NewExtractColumnNamesTransform build(); + } + + public static NewExtractColumnNamesTransform.Builder newBuilder() { + return new AutoValue_NewExtractColumnNamesTransform.Builder(); + } + + @Override + public PCollectionView>> expand( + PCollectionList> input) { + PCollectionList>> headerList = PCollectionList.empty(input.getPipeline()); + + for (FileType fileType : fileTypes()) { + switch (fileType) { + case AVRO: + PCollection>> readHeaderAvro = + input + .get(fileType.ordinal()) + .apply("ReadHeaderAVRO", ParDo.of(new AvroColumnNamesDoFn())); + headerList = headerList.and(readHeaderAvro); + break; + + case CSV: + PCollection>> readHeaderCsv = + input + .get(fileType.ordinal()) + .apply("ReadHeaderCSV", ParDo.of(new CSVColumnNamesDoFn(columnDelimiter()))); + headerList = headerList.and(readHeaderCsv); + break; + + case JSONL: + PCollection>> readHeaderJsonl = + input + .get(fileType.ordinal()) + .apply("ReadHeaderJSONL", ParDo.of(new JsonColumnNameDoFn(headers()))); + headerList = headerList.and(readHeaderJsonl); + break; + + case ORC: + PCollection>> readHeaderOrc = + input + .get(fileType.ordinal()) + .apply("ReadHeaderORC", ParDo.of(new ORCColumnNameDoFn(projectId()))); + headerList = headerList.and(readHeaderOrc); + break; + + case PARQUET: + PCollection>> readHeaderParquet = + input + .get(fileType.ordinal()) + .apply("ReadHeaderPARQUET", ParDo.of(new ParquetColumnNamesDoFn())); + headerList = headerList.and(readHeaderParquet); + break; + + case TSV: + PCollection>> readHeaderTsv = + input + .get(fileType.ordinal()) + .apply("ReadHeaderTSV", ParDo.of(new CSVColumnNamesDoFn('\t'))); + headerList = headerList.and(readHeaderTsv); + break; + + case TXT: + PCollection>> readHeaderTxt = + input + .get(fileType.ordinal()) + .apply("ReadHeaderTXT", ParDo.of(new TxtColumnNameDoFn(headers()))); + headerList = headerList.and(readHeaderTxt); + break; + + default: + throw new IllegalStateException("Unexpected value: " + fileType); + } + } + +// readHeader = headerList.apply("Merge headers",Flatten.pCollections()); +// +// if (!pubSubGcs()) { +// return readHeader.apply("ViewAsList", View.asMap()); +// } +// +// return readHeader +// .apply(Combine., List>perKey(new ResolveDuplicatesCombineFn())) +// .apply("ViewAsList", View.asMap()); + + return headerList.apply("Merge Headers",Flatten.pCollections()).setCoder(KvCoder.of(StringUtf8Coder.of(), ListCoder.of(StringUtf8Coder.of()))) + .apply("ViewAsList", View.asMap()); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/classification/SanitizeFileKeyDoFn.java b/src/main/java/com/google/swarm/tokenization/classification/SanitizeFileKeyDoFn.java new file mode 100644 index 00000000..e8597101 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/classification/SanitizeFileKeyDoFn.java @@ -0,0 +1,85 @@ +/* + * 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.classification; + +import com.google.common.io.Files; +import com.google.swarm.tokenization.common.Util; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.values.KV; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class SanitizeFileKeyDoFn + extends DoFn> { + public static final Logger LOG = LoggerFactory.getLogger(SanitizeFileKeyDoFn.class); + + // Regular expression that matches valid BQ table IDs + private static final String TABLE_REGEXP = "[-\\w$@]{1,1024}"; + private Util.InputLocation inputProviderType; + + public SanitizeFileKeyDoFn(Util.InputLocation inputType) { + this.inputProviderType = inputType; + } + + public static FileKeyObject sanitizeFileName(String file) { + String extension = Files.getFileExtension(file); + if (Util.ALLOWED_TEXT_FILE_EXTENSIONS.contains(extension)) { + extension = "text"; + } + if (!Util.ALLOWED_FILE_EXTENSIONS.contains(extension)) { + throw new RuntimeException( + "Invalid file name '" + + file + + "': must have one of these extensions: " + + Util.ALLOWED_FILE_EXTENSIONS); + } + + String sanitizedName = file.substring(0, file.length() - extension.length() - 1); + sanitizedName = sanitizedName.replace(".", "_"); + sanitizedName = sanitizedName.replace("-", "_"); + + if (!sanitizedName.matches(TABLE_REGEXP)) { + throw new RuntimeException( + "Invalid file name '" + + file + + "': base name must be a valid BigQuery table name -" + + " can contain only letters, numbers, or underscores"); + } + + // Return sanitized file name without extension + // return new FileKeyObject(sanitizedName,Util.getExtension(extension)); + return new FileKeyObject(sanitizedName, Util.FileType.valueOf(extension.toUpperCase())); + } + + @ProcessElement + public void processElement(ProcessContext c) { + FileIO.ReadableFile file = c.element(); + long lastModified = file.getMetadata().lastModifiedMillis(); + if (lastModified == 0L) { + lastModified = Instant.now().getMillis(); + } + FileKeyObject fileKey = sanitizeFileName(file.getMetadata().resourceId().getFilename()); + /* For files whose metadata is coming through pub/sub notification and being read, + * the last modified timestamp would be older than the current window time and hence + * these files are windowed with timestamp when the file is actually being processed. + */ + if (this.inputProviderType == Util.InputLocation.GCS) + c.outputWithTimestamp(KV.of(fileKey, file), Instant.ofEpochMilli(Instant.now().getMillis())); + else c.outputWithTimestamp(KV.of(fileKey, file), Instant.ofEpochMilli(lastModified)); + } +} diff --git a/src/main/java/com/google/swarm/tokenization/common/BigQueryDynamicWriteTransform.java b/src/main/java/com/google/swarm/tokenization/common/BigQueryDynamicWriteTransform.java index 09ceb974..610b15e6 100644 --- a/src/main/java/com/google/swarm/tokenization/common/BigQueryDynamicWriteTransform.java +++ b/src/main/java/com/google/swarm/tokenization/common/BigQueryDynamicWriteTransform.java @@ -71,7 +71,7 @@ public WriteResult expand(PCollection> input) { Write> transform = BigQueryIO.>write() .to(new BQDestination(datasetId(), projectId())) - .withFormatFunction(KV::getValue) + .withFormatFunction(x -> x.getValue()) .withWriteDisposition(WriteDisposition.WRITE_APPEND) .withoutValidation() .ignoreInsertIds() 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..9c854cd8 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 { @@ -83,7 +84,8 @@ public enum FileType { ORC, PARQUET, TSV, - TXT + TXT, + TEXT } public enum InputLocation { @@ -133,6 +135,9 @@ public enum DataSinkType { public static final String BQ_ERROR_TABLE_NAME = String.valueOf("error_log"); public static final String BQ_REID_TABLE_EXT = String.valueOf("re_id"); + public static final Set ALLOWED_TEXT_FILE_EXTENSIONS = + Arrays.asList("txt", "json", "xml", "yaml").stream().collect(Collectors.toUnmodifiableSet()); + public static final Set ALLOWED_FILE_EXTENSIONS = Arrays.asList("avro", "csv", "jsonl", "orc", "parquet", "tsv", "txt").stream() .collect(Collectors.toUnmodifiableSet()); @@ -520,4 +525,23 @@ public static String parseChatlog(String fileName, String chatLog) { return StringUtils.EMPTY; } + + public static List getExtension(List fileTypes) { + List extensionList = new ArrayList<>(); + for (String filetype : fileTypes) { + + if (filetype == "avro") extensionList.add(Util.FileType.AVRO); + else if (filetype == "csv") extensionList.add(Util.FileType.CSV); + else if (filetype == "jsonl") extensionList.add(Util.FileType.JSONL); + else if (filetype == "text") extensionList.add(Util.FileType.TXT); + } + return extensionList; + } + + public static Util.FileType getExtension(String fileType) { + if (fileType == "avro") return Util.FileType.AVRO; + else if (fileType == "csv") return Util.FileType.CSV; + else if (fileType == "jsonl") return Util.FileType.JSONL; + else return Util.FileType.TXT; + } } diff --git a/src/main/java/com/google/swarm/tokenization/text/TextFileReader.java b/src/main/java/com/google/swarm/tokenization/text/TextFileReader.java new file mode 100644 index 00000000..395114d0 --- /dev/null +++ b/src/main/java/com/google/swarm/tokenization/text/TextFileReader.java @@ -0,0 +1,18 @@ +/* + * 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.text; + +public class TextFileReader {}