BigQuery: Fix useAvroLogicalTypes option in LoadJobConfiguration And WriteChannelConfiguration.

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadConfiguration.java

@@ -118,6 +118,13 @@ interface Builder {
     /** Sets the clustering specification for the destination table. */
     Builder setClustering(Clustering clustering);
 
+    /**
+     * If FormatOptions is set to AVRO, you can interpret logical types into their corresponding
+     * types (such as TIMESTAMP) instead of only using their raw types (such as INTEGER). The value
+     * may be {@code null}.
+     */
+    Builder setUseAvroLogicalTypes(Boolean useAvroLogicalTypes);
+
     LoadConfiguration build();
   }
 
@@ -204,6 +211,9 @@ interface Builder {
   /** Returns the clustering specification for the definition table. */
   Clustering getClustering();
 
+  /** Returns True/False. Indicates whether the logical type is interpreted. */
+  Boolean getUseAvroLogicalTypes();
+
   /** Returns a builder for the load configuration object. */
   Builder toBuilder();
 }

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/LoadJobConfiguration.java

@@ -49,6 +49,7 @@ public final class LoadJobConfiguration extends JobConfiguration implements Load
   private final Boolean autodetect;
   private final TimePartitioning timePartitioning;
   private final Clustering clustering;
+  private final Boolean useAvroLogicalTypes;
 
   public static final class Builder extends JobConfiguration.Builder<LoadJobConfiguration, Builder>
       implements LoadConfiguration.Builder {
@@ -68,6 +69,7 @@ public static final class Builder extends JobConfiguration.Builder<LoadJobConfig
     private Boolean autodetect;
     private TimePartitioning timePartitioning;
     private Clustering clustering;
+    private Boolean useAvroLogicalTypes;
 
     private Builder() {
       super(Type.LOAD);
@@ -90,6 +92,7 @@ private Builder(LoadJobConfiguration loadConfiguration) {
           loadConfiguration.destinationEncryptionConfiguration;
       this.timePartitioning = loadConfiguration.timePartitioning;
       this.clustering = loadConfiguration.clustering;
+      this.useAvroLogicalTypes = loadConfiguration.useAvroLogicalTypes;
     }
 
     private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -156,6 +159,7 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
         this.clustering = Clustering.fromPb(loadConfigurationPb.getClustering());
       }
       this.autodetect = loadConfigurationPb.getAutodetect();
+      this.useAvroLogicalTypes = loadConfigurationPb.getUseAvroLogicalTypes();
       if (loadConfigurationPb.getDestinationEncryptionConfiguration() != null) {
         this.destinationEncryptionConfiguration =
             new EncryptionConfiguration.Builder(
@@ -231,6 +235,12 @@ public Builder setClustering(Clustering clustering) {
       return this;
     }
 
+    @Override
+    public Builder setUseAvroLogicalTypes(Boolean useAvroLogicalTypes) {
+      this.useAvroLogicalTypes = useAvroLogicalTypes;
+      return this;
+    }
+
     /**
      * Sets the fully-qualified URIs that point to source data in Google Cloud Storage (e.g.
      * gs://bucket/path). Each URI can contain one '*' wildcard character and it must come after the
@@ -275,6 +285,7 @@ private LoadJobConfiguration(Builder builder) {
     this.destinationEncryptionConfiguration = builder.destinationEncryptionConfiguration;
     this.timePartitioning = builder.timePartitioning;
     this.clustering = builder.clustering;
+    this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
   }
 
   @Override
@@ -357,6 +368,11 @@ public Clustering getClustering() {
     return clustering;
   }
 
+  @Override
+  public Boolean getUseAvroLogicalTypes() {
+    return useAvroLogicalTypes;
+  }
+
   @Override
   public List<JobInfo.SchemaUpdateOption> getSchemaUpdateOptions() {
     return schemaUpdateOptions;
@@ -383,7 +399,8 @@ ToStringHelper toStringHelper() {
         .add("schemaUpdateOptions", schemaUpdateOptions)
         .add("autodetect", autodetect)
         .add("timePartitioning", timePartitioning)
-        .add("clustering", clustering);
+        .add("clustering", clustering)
+        .add("useAvroLogicalTypes", useAvroLogicalTypes);
   }
 
   @Override
@@ -464,6 +481,7 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
     if (clustering != null) {
       loadConfigurationPb.setClustering(clustering.toPb());
     }
+    loadConfigurationPb.setUseAvroLogicalTypes(useAvroLogicalTypes);
     return new com.google.api.services.bigquery.model.JobConfiguration()
         .setLoad(loadConfigurationPb);
   }

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java

@@ -52,6 +52,7 @@ public final class WriteChannelConfiguration implements LoadConfiguration, Seria
   private final EncryptionConfiguration destinationEncryptionConfiguration;
   private final TimePartitioning timePartitioning;
   private final Clustering clustering;
+  private final Boolean useAvroLogicalTypes;
 
   public static final class Builder implements LoadConfiguration.Builder {
 
@@ -68,6 +69,7 @@ public static final class Builder implements LoadConfiguration.Builder {
     private EncryptionConfiguration destinationEncryptionConfiguration;
     private TimePartitioning timePartitioning;
     private Clustering clustering;
+    private Boolean useAvroLogicalTypes;
 
     private Builder() {}
 
@@ -86,6 +88,7 @@ private Builder(WriteChannelConfiguration writeChannelConfiguration) {
           writeChannelConfiguration.destinationEncryptionConfiguration;
       this.timePartitioning = writeChannelConfiguration.timePartitioning;
       this.clustering = writeChannelConfiguration.clustering;
+      this.useAvroLogicalTypes = writeChannelConfiguration.useAvroLogicalTypes;
     }
 
     private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -158,6 +161,7 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
       if (loadConfigurationPb.getClustering() != null) {
         this.clustering = Clustering.fromPb(loadConfigurationPb.getClustering());
       }
+      this.useAvroLogicalTypes = loadConfigurationPb.getUseAvroLogicalTypes();
     }
 
     @Override
@@ -240,6 +244,12 @@ public Builder setClustering(Clustering clustering) {
       return this;
     }
 
+    @Override
+    public Builder setUseAvroLogicalTypes(Boolean useAvroLogicalTypes) {
+      this.useAvroLogicalTypes = useAvroLogicalTypes;
+      return this;
+    }
+
     @Override
     public WriteChannelConfiguration build() {
       return new WriteChannelConfiguration(this);
@@ -260,6 +270,7 @@ protected WriteChannelConfiguration(Builder builder) {
     this.destinationEncryptionConfiguration = builder.destinationEncryptionConfiguration;
     this.timePartitioning = builder.timePartitioning;
     this.clustering = builder.clustering;
+    this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
   }
 
   @Override
@@ -339,6 +350,11 @@ public Clustering getClustering() {
     return clustering;
   }
 
+  @Override
+  public Boolean getUseAvroLogicalTypes() {
+    return useAvroLogicalTypes;
+  }
+
   @Override
   public Builder toBuilder() {
     return new Builder(this);
@@ -358,7 +374,8 @@ MoreObjects.ToStringHelper toStringHelper() {
         .add("schemaUpdateOptions", schemaUpdateOptions)
         .add("autodetect", autodetect)
         .add("timePartitioning", timePartitioning)
-        .add("clustering", clustering);
+        .add("clustering", clustering)
+        .add("useAvroLogicalTypes", useAvroLogicalTypes);
   }
 
   @Override
@@ -387,7 +404,8 @@ public int hashCode() {
         schemaUpdateOptions,
         autodetect,
         timePartitioning,
-        clustering);
+        clustering,
+        useAvroLogicalTypes);
   }
 
   WriteChannelConfiguration setProjectId(String projectId) {
@@ -452,6 +470,7 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
     if (clustering != null) {
       loadConfigurationPb.setClustering(clustering.toPb());
     }
+    loadConfigurationPb.setUseAvroLogicalTypes(useAvroLogicalTypes);
     return new com.google.api.services.bigquery.model.JobConfiguration()
         .setLoad(loadConfigurationPb);
   }

google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java

@@ -52,6 +52,7 @@ public class LoadJobConfigurationTest {
       ImmutableList.of(SchemaUpdateOption.ALLOW_FIELD_ADDITION);
   private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA);
   private static final Boolean AUTODETECT = true;
+  private static final Boolean USERAVROLOGICALTYPES = true;
   private static final EncryptionConfiguration JOB_ENCRYPTION_CONFIGURATION =
       EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
   private static final TimePartitioning TIME_PARTITIONING = TimePartitioning.of(Type.DAY);
@@ -86,6 +87,21 @@ public class LoadJobConfigurationTest {
           .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
           .setAutodetect(AUTODETECT)
           .build();
+  private static final LoadJobConfiguration LOAD_CONFIGURATION_AVRO =
+      LoadJobConfiguration.newBuilder(TABLE_ID, SOURCE_URIS)
+          .setCreateDisposition(CREATE_DISPOSITION)
+          .setWriteDisposition(WRITE_DISPOSITION)
+          .setFormatOptions(FormatOptions.avro())
+          .setIgnoreUnknownValues(IGNORE_UNKNOWN_VALUES)
+          .setMaxBadRecords(MAX_BAD_RECORDS)
+          .setSchema(TABLE_SCHEMA)
+          .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
+          .setAutodetect(AUTODETECT)
+          .setDestinationEncryptionConfiguration(JOB_ENCRYPTION_CONFIGURATION)
+          .setTimePartitioning(TIME_PARTITIONING)
+          .setClustering(CLUSTERING)
+          .setUseAvroLogicalTypes(USERAVROLOGICALTYPES)
+          .build();
 
   @Test
   public void testToBuilder() {
@@ -109,6 +125,17 @@ public void testToBuilder() {
     assertEquals("newTable", configurationBackup.getDestinationTable().getTable());
     configurationBackup = configurationBackup.toBuilder().setDestinationTable(TABLE_ID).build();
     compareLoadJobConfiguration(LOAD_CONFIGURATION_BACKUP, configurationBackup);
+
+    compareLoadJobConfiguration(
+        LOAD_CONFIGURATION_AVRO, LOAD_CONFIGURATION_AVRO.toBuilder().build());
+    LoadJobConfiguration configurationAvro =
+        LOAD_CONFIGURATION_AVRO
+            .toBuilder()
+            .setDestinationTable(TableId.of("dataset", "newTable"))
+            .build();
+    assertEquals("newTable", configurationAvro.getDestinationTable().getTable());
+    configurationAvro = configurationAvro.toBuilder().setDestinationTable(TABLE_ID).build();
+    compareLoadJobConfiguration(LOAD_CONFIGURATION_AVRO, configurationAvro);
   }
 
   @Test
@@ -188,5 +215,6 @@ private void compareLoadJobConfiguration(
         value.getDestinationEncryptionConfiguration());
     assertEquals(expected.getTimePartitioning(), value.getTimePartitioning());
     assertEquals(expected.getClustering(), value.getClustering());
+    assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes());
   }
 }

google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java

@@ -49,6 +49,7 @@ public class WriteChannelConfigurationTest {
           .build();
   private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA);
   private static final Boolean AUTODETECT = true;
+  private static final Boolean USERAVROLOGICALTYPES = true;
   private static final List<JobInfo.SchemaUpdateOption> SCHEMA_UPDATE_OPTIONS =
       ImmutableList.of(JobInfo.SchemaUpdateOption.ALLOW_FIELD_ADDITION);
   private static final TimePartitioning TIME_PARTITIONING = TimePartitioning.of(Type.DAY);
@@ -84,6 +85,21 @@ public class WriteChannelConfigurationTest {
           .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
           .setAutodetect(AUTODETECT)
           .build();
+  private static final WriteChannelConfiguration LOAD_CONFIGURATION_AVRO =
+      WriteChannelConfiguration.newBuilder(TABLE_ID)
+          .setCreateDisposition(CREATE_DISPOSITION)
+          .setWriteDisposition(WRITE_DISPOSITION)
+          .setNullMarker(NULL_MARKER)
+          .setFormatOptions(FormatOptions.avro())
+          .setIgnoreUnknownValues(IGNORE_UNKNOWN_VALUES)
+          .setMaxBadRecords(MAX_BAD_RECORDS)
+          .setSchema(TABLE_SCHEMA)
+          .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
+          .setAutodetect(AUTODETECT)
+          .setTimePartitioning(TIME_PARTITIONING)
+          .setClustering(CLUSTERING)
+          .setUseAvroLogicalTypes(USERAVROLOGICALTYPES)
+          .build();
 
   @Test
   public void testToBuilder() {
@@ -96,6 +112,16 @@ public void testToBuilder() {
     assertEquals("newTable", configuration.getDestinationTable().getTable());
     configuration = configuration.toBuilder().setDestinationTable(TABLE_ID).build();
     compareLoadConfiguration(LOAD_CONFIGURATION_CSV, configuration);
+
+    compareLoadConfiguration(LOAD_CONFIGURATION_AVRO, LOAD_CONFIGURATION_AVRO.toBuilder().build());
+    WriteChannelConfiguration configurationAvro =
+        LOAD_CONFIGURATION_AVRO
+            .toBuilder()
+            .setDestinationTable(TableId.of("dataset", "newTable"))
+            .build();
+    assertEquals("newTable", configurationAvro.getDestinationTable().getTable());
+    configurationAvro = configurationAvro.toBuilder().setDestinationTable(TABLE_ID).build();
+    compareLoadConfiguration(LOAD_CONFIGURATION_AVRO, configurationAvro);
   }
 
   @Test
@@ -191,5 +217,6 @@ private void compareLoadConfiguration(
     assertEquals(expected.getAutodetect(), value.getAutodetect());
     assertEquals(expected.getTimePartitioning(), value.getTimePartitioning());
     assertEquals(expected.getClustering(), value.getClustering());
+    assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes());
   }
 }