feat: add support for Hive partitioning options when creating external tables (#235) (#236)

ekanth · web-flow · commit ccc2bb3de28a · 2020-03-20T16:51:15.000-04:00
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalTableDefinition.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalTableDefinition.java
@@ -135,6 +135,14 @@ public Builder setFormatOptions(FormatOptions formatOptions) {
     /** Sets the table schema. */
     public abstract Builder setSchema(Schema schema);
 
+    /** Sets the table Hive partitioning options. */
+    public Builder setHivePartitioningOptions(HivePartitioningOptions hivePartitioningOptions) {
+      return setHivePartitioningOptionsInner(hivePartitioningOptions);
+    };
+
+    abstract Builder setHivePartitioningOptionsInner(
+        HivePartitioningOptions hivePartitioningOptions);
+
     /** Creates an {@code ExternalTableDefinition} object. */
     @Override
     public abstract ExternalTableDefinition build();
@@ -212,6 +220,19 @@ public <F extends FormatOptions> F getFormatOptions() {
   @Nullable
   public abstract Boolean getAutodetect();
 
+  /**
+   * [Experimental] Returns the HivePartitioningOptions when the data layout follows Hive
+   * partitioning convention
+   */
+  @SuppressWarnings("unchecked")
+  @Nullable
+  public HivePartitioningOptions getHivePartitioningOptions() {
+    return getHivePartitioningOptionsInner();
+  }
+
+  @Nullable
+  abstract HivePartitioningOptions getHivePartitioningOptionsInner();
+
   /** Returns a builder for the {@code ExternalTableDefinition} object. */
   public abstract Builder toBuilder();
 
@@ -257,6 +278,9 @@ com.google.api.services.bigquery.model.ExternalDataConfiguration toExternalDataC
     if (getAutodetect() != null) {
       externalConfigurationPb.setAutodetect(getAutodetect());
     }
+    if (getHivePartitioningOptions() != null) {
+      externalConfigurationPb.setHivePartitioningOptions(getHivePartitioningOptions().toPb());
+    }
     return externalConfigurationPb;
   }
 
@@ -405,6 +429,10 @@ static ExternalTableDefinition fromPb(Table tablePb) {
       }
       builder.setMaxBadRecords(externalDataConfiguration.getMaxBadRecords());
       builder.setAutodetect(externalDataConfiguration.getAutodetect());
+      if (externalDataConfiguration.getHivePartitioningOptions() != null) {
+        builder.setHivePartitioningOptions(
+            HivePartitioningOptions.fromPb(externalDataConfiguration.getHivePartitioningOptions()));
+      }
     }
     return builder.build();
   }
@@ -444,6 +472,10 @@ static ExternalTableDefinition fromExternalDataConfiguration(
     if (externalDataConfiguration.getAutodetect() != null) {
       builder.setAutodetect(externalDataConfiguration.getAutodetect());
     }
+    if (externalDataConfiguration.getHivePartitioningOptions() != null) {
+      builder.setHivePartitioningOptions(
+          HivePartitioningOptions.fromPb(externalDataConfiguration.getHivePartitioningOptions()));
+    }
     return builder.build();
   }
 }
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/HivePartitioningOptions.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/HivePartitioningOptions.java
@@ -103,6 +103,16 @@ public String toString() {
         .toString();
   }
 
+  @Override
+  public boolean equals(Object obj) {
+    return obj == this
+        || obj != null
+            && obj.getClass().equals(HivePartitioningOptions.class)
+            && Objects.equals(mode, ((HivePartitioningOptions) obj).getMode())
+            && Objects.equals(
+                sourceUriPrefix, ((HivePartitioningOptions) obj).getSourceUriPrefix());
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(mode, sourceUriPrefix);
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExternalTableDefinitionTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExternalTableDefinitionTest.java
@@ -46,12 +46,18 @@ public class ExternalTableDefinitionTest {
   private static final String COMPRESSION = "GZIP";
   private static final Boolean AUTODETECT = true;
   private static final CsvOptions CSV_OPTIONS = CsvOptions.newBuilder().build();
+  private static final HivePartitioningOptions HIVE_PARTITIONING_OPTIONS =
+      HivePartitioningOptions.newBuilder()
+          .setMode("AUTO")
+          .setSourceUriPrefix(SOURCE_URIS.get(0))
+          .build();
   private static final ExternalTableDefinition EXTERNAL_TABLE_DEFINITION =
       ExternalTableDefinition.newBuilder(SOURCE_URIS, TABLE_SCHEMA, CSV_OPTIONS)
           .setCompression(COMPRESSION)
           .setIgnoreUnknownValues(IGNORE_UNKNOWN_VALUES)
           .setMaxBadRecords(MAX_BAD_RECORDS)
           .setAutodetect(AUTODETECT)
+          .setHivePartitioningOptions(HIVE_PARTITIONING_OPTIONS)
           .build();
 
   @Test
@@ -83,6 +89,7 @@ public void testBuilder() {
     assertEquals(TABLE_SCHEMA, EXTERNAL_TABLE_DEFINITION.getSchema());
     assertEquals(SOURCE_URIS, EXTERNAL_TABLE_DEFINITION.getSourceUris());
     assertEquals(AUTODETECT, EXTERNAL_TABLE_DEFINITION.getAutodetect());
+    assertEquals(HIVE_PARTITIONING_OPTIONS, EXTERNAL_TABLE_DEFINITION.getHivePartitioningOptions());
   }
 
   @Test
@@ -107,5 +114,6 @@ private void compareExternalTableDefinition(
     assertEquals(expected.getSourceUris(), value.getSourceUris());
     assertEquals(expected.hashCode(), value.hashCode());
     assertEquals(expected.getAutodetect(), value.getAutodetect());
+    assertEquals(expected.getHivePartitioningOptions(), value.getHivePartitioningOptions());
   }
 }
