Add support for WRITE_TRUNCATE_DATA in BQ Sink

Author: itsankit-google

docs/BigQueryMultiTable-batchsink.md

@@ -59,6 +59,9 @@ If the bucket was created automatically, it will be deleted after the run finish
 **Truncate Table:** Whether or not to truncate the table before writing to it.
 Should only be used with the Insert operation.
 
+**Write Disposition**: Describes whether a job should truncate table but preserve metadata or not.
+For more details, see [here](https://docs.cloud.google.com/bigquery/docs/reference/auditlogs/rest/Shared.Types/BigQueryAuditMetadata.WriteDisposition).
+
 **Location:** The location where the big query datasets will get created. This value is ignored
 if the dataset or temporary bucket already exist.
 

docs/BigQueryTable-batchsink.md

@@ -76,6 +76,9 @@ will be dropped.
 **Truncate Table**: Whether or not to truncate the table before writing to it.
 Should only be used with the Insert operation.
 
+**Write Disposition**: Describes whether a job should truncate table but preserve metadata or not.
+For more details, see [here](https://docs.cloud.google.com/bigquery/docs/reference/auditlogs/rest/Shared.Types/BigQueryAuditMetadata.WriteDisposition).
+
 **Table Key**: List of fields that determines relation between tables during Update and Upsert operations.
 
 **Dedupe By**: Column names and sort order used to choose which input record to update/upsert when there are

pom.xml

@@ -20,7 +20,7 @@
 
   <groupId>io.cdap.plugin</groupId>
   <artifactId>google-cloud</artifactId>
-  <version>0.23.5</version>
+  <version>0.23.6</version>
   <name>Google Cloud Plugins</name>
   <packaging>jar</packaging>
   <description>Plugins for Google Big Query</description>

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AbstractBigQuerySink.java

@@ -250,7 +250,7 @@ private Configuration getBaseConfiguration(@Nullable CryptoKeyName cmekKeyName)
     baseConfiguration.setBoolean(BigQueryConstants.CONFIG_ALLOW_SCHEMA_RELAXATION,
                                  config.isAllowSchemaRelaxation());
     baseConfiguration.setStrings(BigQueryConfiguration.OUTPUT_TABLE_WRITE_DISPOSITION.getKey(),
-                                 config.getWriteDisposition().name());
+                                 config.getWriteDisposition());
     baseConfiguration.setStrings(BigQueryConstants.CONFIG_JSON_STRING_FIELDS, config.getJsonStringFields());
     // this setting is needed because gcs has default chunk size of 64MB. This is large default chunk size which can
     // cause OOM issue if there are many tables being written. See this - CDAP-16670

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AbstractBigQuerySinkConfig.java

@@ -48,6 +48,7 @@ public abstract class AbstractBigQuerySinkConfig extends BigQueryBaseConfig {
                     Schema.Type.BOOLEAN, Schema.Type.BYTES, Schema.Type.ARRAY, Schema.Type.RECORD);
 
   public static final String NAME_TRUNCATE_TABLE = "truncateTable";
+  public static final String NAME_WRITE_DISPOSITION = "writeDisposition";
   public static final String NAME_LOCATION = "location";
   private static final String NAME_GCS_CHUNK_SIZE = "gcsChunkSize";
   public static final String NAME_BQ_JOB_LABELS = "jobLabels";
@@ -77,9 +78,19 @@ public abstract class AbstractBigQuerySinkConfig extends BigQueryBaseConfig {
   @Macro
   @Nullable
   @Description("Whether or not to truncate the table before writing to it. "
-    + "Should only be used with the Insert operation. This could overwrite the table schema")
+    + "Should only be used with the Insert operation. This could overwrite the table schema based "
+    + "on write disposition value.")
   protected Boolean truncateTable;
 
+  @Name(NAME_WRITE_DISPOSITION)
+  @Macro
+  @Nullable
+  @Description("WRITE_TRUNCATE_DATA preserves the table metadata where as WRITE_TRUNCATE does not. "
+      + "For more details, see "
+      + "https://docs.cloud.google.com/bigquery/docs/reference/auditlogs/rest/Shared.Types/"
+      + "BigQueryAuditMetadata.WriteDisposition.")
+  protected String writeDisposition;
+
   @Name(NAME_LOCATION)
   @Macro
   @Nullable
@@ -143,9 +154,16 @@ public boolean isAllowSchemaRelaxation() {
     return allowSchemaRelaxation == null ? false : allowSchemaRelaxation;
   }
 
-  public JobInfo.WriteDisposition getWriteDisposition() {
-    return isTruncateTableSet() ? JobInfo.WriteDisposition.WRITE_TRUNCATE
-      : JobInfo.WriteDisposition.WRITE_APPEND;
+  private String getTruncateTableWriteDisposition() {
+    if (writeDisposition == null) {
+      return JobInfo.WriteDisposition.WRITE_TRUNCATE.name();
+    }
+    return writeDisposition;
+  }
+
+  public String getWriteDisposition() {
+    return isTruncateTableSet() ? getTruncateTableWriteDisposition()
+      : JobInfo.WriteDisposition.WRITE_APPEND.name();
   }
 
   public boolean isTruncateTableSet() {

src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQueryOutputFormat.java

@@ -351,7 +351,7 @@ private void importFromGcs(String projectId, TableReference tableRef, @Nullable
 
       Map<String, String> fieldDescriptions = new HashMap<>();
       if (JobInfo.WriteDisposition.WRITE_TRUNCATE
-        .equals(JobInfo.WriteDisposition.valueOf(writeDisposition)) && tableExists) {
+        .equals(writeDisposition.toUpperCase()) && tableExists) {
           List<TableFieldSchema> tableFieldSchemas = Optional.ofNullable(bigQueryHelper.getTable(tableRef))
             .map(it -> it.getSchema())
             .map(it -> it.getFields())
@@ -399,8 +399,8 @@ private void importFromGcs(String projectId, TableReference tableRef, @Nullable
         // Schema update options should only be specified with WRITE_APPEND disposition,
         // or with WRITE_TRUNCATE disposition on a table partition - The logic below should change when we support
         // insertion into single partition
-        if (allowSchemaRelaxation && !JobInfo.WriteDisposition.WRITE_TRUNCATE
-          .equals(JobInfo.WriteDisposition.valueOf(writeDisposition))) {
+        if (allowSchemaRelaxation && !JobInfo.WriteDisposition.WRITE_TRUNCATE.name()
+          .equals(writeDisposition.toUpperCase())) {
           loadConfig.setSchemaUpdateOptions(Arrays.asList(
             JobInfo.SchemaUpdateOption.ALLOW_FIELD_ADDITION.name(),
             JobInfo.SchemaUpdateOption.ALLOW_FIELD_RELAXATION.name()));
@@ -428,7 +428,7 @@ private void importFromGcs(String projectId, TableReference tableRef, @Nullable
       if (operation.equals(Operation.INSERT) && gcsPaths.size() <= BQ_IMPORT_MAX_BATCH_SIZE) {
         // Directly load data into destination table when total no of input paths is loadable into BQ
         loadConfig.setSourceUris(gcsPaths);
-        loadConfig.setWriteDisposition(writeDisposition);
+        loadConfig.setWriteDisposition(writeDisposition.toUpperCase());
         loadConfig.setDestinationTable(tableRef);
 
         JobConfiguration config = new JobConfiguration();
@@ -487,7 +487,7 @@ private void loadInBatchesInTempTable(TableReference tableRef, JobConfigurationL
         .setTableId(temporaryTableName);
 
       loadConfig.setDestinationTable(temporaryTableReference);
-      loadConfig.setWriteDisposition(JobInfo.WriteDisposition.WRITE_APPEND.toString());
+      loadConfig.setWriteDisposition(JobInfo.WriteDisposition.WRITE_APPEND.name());
 
       // Split the list of files in batches 10000 (current bq load job limit) and import /append onto a temp table
       List<List<String>> gcsPathsInBatches = Lists.partition(gcsPaths, BQ_IMPORT_MAX_BATCH_SIZE);
@@ -723,7 +723,7 @@ private static TableSchema createTableSchemaFromFields(String fieldsJson) throws
     private void updateFieldDescriptions(String writeDisposition, TableReference tableRef,
                                          Map<String, String> fieldDescriptions) throws IOException {
       if (JobInfo.WriteDisposition.WRITE_TRUNCATE
-        .equals(JobInfo.WriteDisposition.valueOf(writeDisposition))) {
+        .equals(writeDisposition.toUpperCase())) {
 
         Table table = bigQueryHelper.getTable(tableRef);
         List<TableFieldSchema> tableFieldSchemas = Optional.ofNullable(table)

src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkConfig.java

@@ -198,15 +198,19 @@ public BigQuerySinkConfig(@Nullable String referenceName, String dataset, String
 
   private BigQuerySinkConfig(@Nullable String referenceName, @Nullable String project,
                              @Nullable String serviceAccountType, @Nullable String serviceFilePath,
-                             @Nullable String serviceAccountJson,
-                             @Nullable String dataset, @Nullable String table, @Nullable String location,
-                             @Nullable String cmekKey, @Nullable String bucket, @Nullable String jobLabelKeyValue) {
+                             @Nullable String serviceAccountJson, @Nullable String dataset,
+                             @Nullable String table, @Nullable String location,
+                             @Nullable String cmekKey, @Nullable String bucket,
+                             @Nullable String jobLabelKeyValue, @Nullable String writeDisposition,
+                             @Nullable Boolean truncateTable) {
     super(new BigQueryConnectorConfig(project, project, serviceAccountType,
             serviceFilePath, serviceAccountJson), dataset, cmekKey, bucket);
     this.referenceName = referenceName;
     this.table = table;
     this.location = location;
     this.jobLabelKeyValue = jobLabelKeyValue;
+    this.writeDisposition = writeDisposition;
+    this.truncateTable = truncateTable;
   }
 
   public String getTable() {
@@ -715,6 +719,8 @@ public static class Builder {
     private String location;
     private String bucket;
     private String jobLabelKeyValue;
+    private String writeDisposition;
+    private Boolean truncateTable;
 
     public BigQuerySinkConfig.Builder setReferenceName(@Nullable String referenceName) {
       this.referenceName = referenceName;
@@ -770,6 +776,16 @@ public BigQuerySinkConfig.Builder setJobLabelKeyValue(@Nullable String jobLabelK
       return this;
     }
 
+    public BigQuerySinkConfig.Builder setWriteDisposition(@Nullable String writeDisposition) {
+      this.writeDisposition = writeDisposition;
+      return this;
+    }
+
+    public BigQuerySinkConfig.Builder setTruncateTable(@Nullable Boolean truncateTable) {
+      this.truncateTable = truncateTable;
+      return this;
+    }
+
     public BigQuerySinkConfig build() {
       return new BigQuerySinkConfig(
         referenceName,
@@ -782,9 +798,10 @@ public BigQuerySinkConfig build() {
         location,
         cmekKey,
         bucket,
-        jobLabelKeyValue
+        jobLabelKeyValue,
+        writeDisposition,
+        truncateTable
       );
     }
-
   }
 }

src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkConfigTest.java

@@ -16,6 +16,7 @@
 
 package io.cdap.plugin.gcp.bigquery.sink;
 
+import com.google.cloud.bigquery.JobInfo;
 import com.google.cloud.bigquery.TimePartitioning;
 import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.cdap.etl.api.FailureCollector;
@@ -56,6 +57,27 @@ public void setup() throws NoSuchMethodException {
     arguments = new HashMap<>();
   }
 
+  @Test
+  public void testValidateWriteDisposition() {
+    BigQuerySinkConfig bigQuerySinkConfig =
+        BigQuerySinkConfig.builder()
+            .setTruncateTable(true)
+            .setWriteDisposition("WRITE_TRUNCATE_DATA")
+            .build();
+    Assert.assertEquals("WRITE_TRUNCATE_DATA", bigQuerySinkConfig.getWriteDisposition());
+
+    bigQuerySinkConfig = BigQuerySinkConfig.builder().setWriteDisposition("WRITE_APPEND").build();
+    Assert.assertEquals(bigQuerySinkConfig.getWriteDisposition(),
+        JobInfo.WriteDisposition.WRITE_APPEND.name());
+
+    bigQuerySinkConfig = BigQuerySinkConfig.builder()
+        .setTruncateTable(true)
+        .setWriteDisposition("WRITE_TRUNCATE")
+        .build();
+    Assert.assertEquals(bigQuerySinkConfig.getWriteDisposition(),
+        JobInfo.WriteDisposition.WRITE_TRUNCATE.name());
+  }
+
   @Test
   public void testValidateTimePartitioningColumnWithHourAndDate() throws
           InvocationTargetException, IllegalAccessException {

widgets/BigQueryMultiTable-batchsink.json

@@ -124,6 +124,25 @@
               "label": "False"
             }
           }
+        },
+        {
+          "name": "writeDisposition",
+          "widget-type": "radio-group",
+          "label": "Write Disposition",
+          "widget-attributes": {
+            "layout": "inline",
+            "default": "WRITE_TRUNCATE",
+            "options": [
+              {
+                "id": "WRITE_TRUNCATE",
+                "label": "WRITE_TRUNCATE"
+              },
+              {
+                "id": "WRITE_TRUNCATE_DATA",
+                "label": "WRITE_TRUNCATE_DATA"
+              }
+            ]
+          }
         }
       ]
     },
@@ -287,6 +306,18 @@
           "name": "connection"
         }
       ]
+    },
+    {
+      "name": "showWriteDisposition",
+      "condition": {
+        "expression": "truncateTable == true"
+      },
+      "show": [
+        {
+          "type": "property",
+          "name": "writeDisposition"
+        }
+      ]
     }
   ],
   "jump-config": {

widgets/BigQueryTable-batchsink.json

@@ -221,6 +221,25 @@
             }
           }
         },
+        {
+          "name": "writeDisposition",
+          "widget-type": "radio-group",
+          "label": "Write Disposition",
+          "widget-attributes": {
+            "layout": "inline",
+            "default": "WRITE_TRUNCATE",
+            "options": [
+              {
+                "id": "WRITE_TRUNCATE",
+                "label": "WRITE_TRUNCATE"
+              },
+              {
+                "id": "WRITE_TRUNCATE_DATA",
+                "label": "WRITE_TRUNCATE_DATA"
+              }
+            ]
+          }
+        },
         {
           "name": "allowSchemaRelaxation",
           "widget-type": "toggle",
@@ -542,6 +561,18 @@
           "name": "relationTableKey"
         }
       ]
+    },
+    {
+      "name": "showWriteDisposition",
+      "condition": {
+        "expression": "truncateTable == true"
+      },
+      "show": [
+        {
+          "type": "property",
+          "name": "writeDisposition"
+        }
+      ]
     }
   ],
   "jump-config": {