Merge branch 'master' into inspect_classify

Author: dup05

.github/workflows/dlp-pipelines.yml

@@ -506,6 +506,9 @@ To run the pipeline for JSONL files, the list of comma-separated headers also ne
 ./gradlew run ... -Pargs="... --filePattern=gs://${PROJECT_ID}-demo-data/CCRecords_sample.jsonl --headers=<comma_separated_list_of_headers>"
 ```
 
+The original schema of the input file will not be preserved. This solution simplifies the data when it converts the data
+to a DLP API request object.
+
 #### 4. Avro
 
 The pipeline handles Avro files similarly to how it handles CSV files. No additional changes are required to run
@@ -515,6 +518,9 @@ the pipeline except updating the `--filePattern` parameter. For example:
 ./gradlew run ... -Pargs="... --filePattern=gs://${PROJECT_ID}-demo-data/*.avro"
 ```
 
+The original schema of the input file will not be preserved. This solution simplifies the data when it converts the data
+to a DLP API request object.
+
 #### 5. CSV files with custom delimiters
 
 The pipeline supports CSV files with a custom delimiter. The delimiter has to be passed in the pipeline option as `--columnDelimiter`.
@@ -535,6 +541,9 @@ No additional changes are required to run the pipeline except updating the `--fi
 ./gradlew run ... -Pargs="... --filePattern=gs://${PROJECT_ID}-demo-data/*.parquet"
 ```
 
+The original schema of the input file will not be preserved. This solution simplifies the data when it converts the data
+to a DLP API request object.
+
 #### 6. ORC
 
 ##### Inspection of ORC files
@@ -562,11 +571,12 @@ written in a BigQuery dataset as tables or an output Cloud Storage bucket as ORC
 --outputBucket=<output_storage_bucket> ..."
 ```
 
-In the above command, replace <i>output_storage_bucket</i> with the URI of the Cloud Storage bucket where you want to store the 
-de-identified ORC files.
+In the above command, replace <i>output_storage_bucket</i> with the URI of the Cloud Storage bucket where you want to 
+store the de-identified ORC files. The de-identification pipeline allows input files with varying schemas to be 
+processed in the same pipeline.
 
-The de-identification pipeline supports input files with varying schemas. Currently, it can process only primitive
-data types available in ORC format when the results are stored in an output Cloud Storage bucket.
+Currently, this solution can process only primitive data types available in ORC format when the results are stored in an
+output Cloud Storage bucket.
 
 For sample data in ORC file format, refer to [mock-data](.github/mock-data).
 

README.md

@@ -12,7 +12,7 @@
  # See the License for the specific language governing permissions and
  # limitations under the License.
 steps:
-- name: gradle:6-jdk11
+- name: gradle:8.4.0-jdk17
   entrypoint: gradle
   args: ['run', '-DmainClass=com.google.swarm.tokenization.tool.DLPTemplateHelper', '--x', 'test', '-Pargs=${_KEK_CONFIG_FILE}']
 - name: gcr.io/cloud-builders/gsutil
@@ -29,4 +29,4 @@ steps:
 - name: gcr.io/cloud-builders/gsutil
   args: ['cp', 'reid-template.json','${_GCS_BUCKET_NAME}/reid-template.json']
 - name: gcr.io/cloud-builders/gsutil
-  args: ['cp', 'inspect-template.json','${_GCS_BUCKET_NAME}/inspect-template.json']     
+  args: ['cp', 'inspect-template.json','${_GCS_BUCKET_NAME}/inspect-template.json']     

dlp-demo-part-2-dlp-template.yaml

@@ -23,7 +23,7 @@ export KEY_RING_NAME=demo-key-ring
 export KEY_NAME=demo-key
 export KEK_FILE_NAME=kek.json
 export TOKENIZING_ROLE_NAME="dlp_tokenizing_runner"
-export PROJECT_NUMBER=$(gcloud projects list --filter=${PROJECT_ID} --format="value(PROJECT_NUMBER)") 
+export PROJECT_NUMBER=$(gcloud projects list --filter="PROJECT_ID: ${PROJECT_ID}" --format="value(PROJECT_NUMBER)")
 export SERVICE_ACCOUNT_NAME=demo-service-account
 export REGION=us-central1
 export BQ_DATASET_NAME=demo_dataset

setup-data-tokeninzation-solution-v2.sh

@@ -243,8 +243,7 @@ private static void runInspectAndDeidPipeline(
 
       case ORC:
         records =
-            inputFiles.apply(
-                "ReadFromORCFilesAsOrcStruct", ParDo.of(new ORCReaderDoFn(options.getProject())));
+            inputFiles.apply("ReadORCFiles", ParDo.of(new ORCReaderDoFn(options.getProject())));
         break;
 
       default:

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

@@ -145,7 +145,7 @@ public interface DLPTextToBigQueryStreamingV2PipelineOptions extends CommonPipel
    */
   void setInitialBackoff(int value);
 
-  @Description("Output bucket to write DEID output as csv file")
+  @Description("Output bucket to write de-identified files in Cloud Storage buckets")
   String getOutputBucket();
 
   void setOutputBucket(String outputBucket);

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

@@ -42,8 +42,8 @@ public class AvroReaderSplittableDoFn
     extends DoFn<KV<String, ReadableFile>, KV<String, GenericRecord>> {
 
   public static final Logger LOG = LoggerFactory.getLogger(AvroReaderSplittableDoFn.class);
-  private final Counter numberOfAvroRecordsIngested =
-      Metrics.counter(AvroReaderSplittableDoFn.class, "numberOfAvroRecordsIngested");
+  private final Counter numberOfRowsRead =
+      Metrics.counter(AvroReaderSplittableDoFn.class, "numberOfRowsRead");
   private final Integer splitSize;
   private final Integer keyRange;
 
@@ -82,7 +82,7 @@ public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Lon
 
         // Output the Avro record
         c.outputWithTimestamp(KV.of(fileName, record), Instant.now());
-        numberOfAvroRecordsIngested.inc();
+        numberOfRowsRead.inc();
       }
     }
 

src/main/java/com/google/swarm/tokenization/avro/AvroReaderSplittableDoFn.java

@@ -245,9 +245,15 @@ static class DeidentifyText
       extends DoFn<KV<String, Iterable<Table.Row>>, KV<String, DeidentifyContentResponse>> {
 
     public static final Logger LOG = LoggerFactory.getLogger(DeidentifyText.class);
+
+    // Counter to track total number of times DLP Content API calls for DEID failed
     private final Counter numberOfDLPRowBagsFailedDeid =
         Metrics.counter(DeidentifyText.class, "numberOfDLPRowBagsFailedDeid");
 
+    // Counter to track total number of rows that failed to deidentify
+    private final Counter numberOfDLPRowsFailedDeid =
+        Metrics.counter(DeidentifyText.class, "numberOfDLPRowsFailedDeid");
+
     private final String projectId;
     private final String inspectTemplateName;
     private final String deidentifyTemplateName;
@@ -361,6 +367,7 @@ public void processElement(ProcessContext c) throws IOException, InterruptedExce
             LOG.warn("Error in DLP API, Retrying...");
           } else {
             numberOfDLPRowBagsFailedDeid.inc();
+            numberOfDLPRowsFailedDeid.inc(table.getRowsCount());
             LOG.error(
                 "Retried {} times unsuccessfully. Some records were not de-identified. Exception: {}",
                 this.dlpApiRetryCount,

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

@@ -231,6 +231,9 @@ static class InspectData
     private final Counter numberOfDLPRowBagsFailedInspection =
         Metrics.counter(DLPInspectText.InspectData.class, "numberOfDLPRowBagsFailedInspection");
 
+    private final Counter numberOfDLPRowsFailedInspection =
+        Metrics.counter(DLPInspectText.InspectData.class, "numberOfDLPRowsFailedInspection");
+
     /**
      * @param projectId ID of GCP project that should be used for data inspection.
      * @param inspectTemplateName Template name for inspection.
@@ -321,6 +324,7 @@ public void processElement(ProcessContext c) throws IOException, InterruptedExce
             LOG.warn("Error in DLP API, Retrying...");
           } else {
             numberOfDLPRowBagsFailedInspection.inc();
+            numberOfDLPRowsFailedInspection.inc(table.getRowsCount());
             LOG.error(
                 "Retried {} times unsuccessfully. Not able to inspect some records. Exception: {}",
                 this.dlpApiRetryCount,

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

@@ -261,6 +261,9 @@ static class ReidentifyText
     private final Counter numberOfDLPRowBagsFailedReid =
         Metrics.counter(DLPInspectText.InspectData.class, "numberOfDLPRowBagsFailedReid");
 
+    private final Counter numberOfDLPRowsFailedReid =
+        Metrics.counter(DLPInspectText.InspectData.class, "numberOfDLPRowsFailedReid");
+
     @Setup
     public void setup() throws IOException {
       requestBuilder = ReidentifyContentRequest.newBuilder().setParent(projectId);
@@ -363,6 +366,7 @@ public void processElement(ProcessContext context) throws IOException, Interrupt
             LOG.warn("Error in DLP API, Retrying...");
           } else {
             numberOfDLPRowBagsFailedReid.inc();
+            numberOfDLPRowsFailedReid.inc(table.getRowsCount());
             LOG.error(
                 "Retried {} times unsuccessfully. Not able to reidentify some records. Exception: {}",
                 this.dlpApiRetryCount,