Changes to make queryPath optional

Author: dup05

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

@@ -92,7 +92,7 @@ public interface DLPTextToBigQueryStreamingV2PipelineOptions
   void setTableRef(String tableRef);
 
   @Description("read method default, direct, export")
-  @Default.Enum("EXPORT")
+  @Default.Enum("DEFAULT")
   Method getReadMethod();
 
   void setReadMethod(Method method);

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

@@ -15,6 +15,7 @@
  */
 package com.google.swarm.tokenization.common;
 
+import autovalue.shaded.org.checkerframework.checker.nullness.qual.Nullable;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.auto.value.AutoValue;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
@@ -50,6 +51,7 @@ public abstract static class Builder {
 
     public abstract Builder setKeyRange(Integer keyRange);
 
+    @Nullable
     public abstract Builder setQuery(String query);
 
     public abstract BigQueryReadTransform build();
@@ -64,23 +66,28 @@ public PCollection<KV<String, TableRow>> expand(PBegin input) {
 
     switch (readMethod()) {
       case DEFAULT:
-        return input
-            .apply(
-                "ReadFromBigQuery",
-                BigQueryIO.readTableRows()
-                    .fromQuery(query())
-                    .usingStandardSql()
-                    .withMethod(Method.DEFAULT))
-            .apply("AddTableNameAsKey", WithKeys.of(tableRef()));
+        if(query()!=null) {
+          return input
+                  .apply(
+                          "ReadFromBigQuery",
+                          BigQueryIO.readTableRows()
+                                  .fromQuery(query())
+                                  .usingStandardSql()
+                                  .withMethod(Method.DEFAULT))
+                  .apply("AddTableNameAsKey", WithKeys.of(tableRef()));
+        }
+        else {
+          return input
+                  .apply(
+                          "ReadFromBigQuery",
+                          BigQueryIO.readTableRows()
+                                  .from(tableRef())
+                                  .withMethod(Method.DEFAULT))
+                  .apply("AddTableNameAsKey", WithKeys.of(tableRef()));
+        }
+
       case EXPORT:
-        return input
-            .apply(
-                "ReadFromBigQuery",
-                BigQueryIO.readTableRows()
-                    .fromQuery(query())
-                    .usingStandardSql()
-                    .withMethod(Method.DEFAULT))
-            .apply("AddTableNameAsKey", WithKeys.of(tableRef()));
+        throw new IllegalArgumentException("Export method not supported");
 
       case DIRECT_READ:
         throw new IllegalArgumentException("Direct read not supported");

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

@@ -439,6 +439,10 @@ private static void validateStorageWriteApiAtLeastOnce(BigQueryOptions options)
   }
 
   public static String getQueryFromGcs(String gcsPath) {
+    if(gcsPath == null){
+      LOG.debug("Query path not provided, entire table will be read");
+      return null;
+    }
     GcsPath path = GcsPath.fromUri(URI.create(gcsPath));
     Storage storage = StorageOptions.getDefaultInstance().getService();
     BlobId blobId = BlobId.of(path.getBucket(), path.getObject());