Add support for new api features in BigQuery TimePartitioning.

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TimePartitioning.java

@@ -18,8 +18,11 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.api.core.BetaApi;
+import com.google.auto.value.AutoValue;
 import com.google.common.base.MoreObjects;
 
+import javax.annotation.Nullable;
 import java.io.Serializable;
 import java.util.Objects;
 
@@ -30,13 +33,11 @@
  *
  * @see <a href="https://cloud.google.com/bigquery/docs/partitioned-tables">Partitioned Tables</a>
  */
-public final class TimePartitioning implements Serializable {
+@AutoValue
+public abstract class TimePartitioning implements Serializable {
 
   private static final long serialVersionUID = -8565064035346940951L;
 
-  private final Type type;
-  private final Long expirationMs;
-
   /**
    * The type of time partitioning. Currently, the only type supported is {@code DAY}, which will
    * generate one partition per day based on data loading time.
@@ -49,47 +50,66 @@ public enum Type {
     DAY
   }
 
-  private TimePartitioning(Type type, Long expirationMs) {
-    this.type = checkNotNull(type);
-    this.expirationMs = expirationMs;
+  TimePartitioning() {
+    // Users cannot extend this, but AutoValue can.
   }
 
-
   /**
    * Returns the time partitioning type. Currently, the only type supported is {@link Type#DAY},
    * which will generate one partition per day based on data loading time.
    */
-  public Type getType() {
-    return type;
-  }
+  public abstract Type getType();
 
 
   /**
    * Returns the number of milliseconds for which to keep the storage for a partition. When expired,
    * the storage for the partition is reclaimed.
    */
-  public Long getExpirationMs() {
-    return expirationMs;
-  }
+  @Nullable
+  public abstract Long getExpirationMs();
 
-  @Override
-  public String toString() {
-    return MoreObjects.toStringHelper(this)
-        .add("type", type)
-        .add("expirationMs", expirationMs)
-        .toString();
-  }
 
-  @Override
-  public int hashCode() {
-    return Objects.hash(type, expirationMs);
+  /**
+   * If not set, the table is partitioned by pseudo column '_PARTITIONTIME'; if set, the table is
+   * partitioned by this field.
+   */
+  @BetaApi
+  @Nullable
+  public abstract String getField();
+
+
+  /**
+   * If set to true, queries over this table require a partition filter (that can be used for
+   * partition elimination) to be specified.
+   */
+  @BetaApi
+  @Nullable
+  public abstract Boolean getRequirePartitionFilter();
+
+  public abstract Builder toBuilder();
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+    abstract Builder setType(Type type);
+
+    public abstract Builder setExpirationMs(Long expirationMs);
+
+    @BetaApi
+    public abstract Builder setRequirePartitionFilter(Boolean requirePartitionFilter);
+
+    @BetaApi
+    public abstract Builder setField(String field);
+
+    public abstract TimePartitioning build();
   }
 
-  @Override
-  public boolean equals(Object obj) {
-    return obj == this
-        || obj instanceof TimePartitioning
-        && Objects.equals(toPb(), ((TimePartitioning) obj).toPb());
+  /**
+   * Returns a {@code TimePartitioning} object given the time partitioning type. Currently, the only
+   * type supported is {@link Type#DAY}, which will generate one partition per day based on data
+   * loading time.
+   */
+  public static Builder newBuilder(Type type) {
+    return new AutoValue_TimePartitioning.Builder().setType(type);
   }
 
   /**
@@ -98,7 +118,7 @@ public boolean equals(Object obj) {
    * loading time.
    */
   public static TimePartitioning of(Type type) {
-    return new TimePartitioning(type, null);
+    return newBuilder(type).build();
   }
 
   /**
@@ -110,20 +130,25 @@ public static TimePartitioning of(Type type) {
    * @param expirationMs the number of milliseconds for which to keep the storage for a partition
    */
   public static TimePartitioning of(Type type, long expirationMs) {
-    return new TimePartitioning(type, expirationMs);
+    return newBuilder(type).setExpirationMs(expirationMs).build();
   }
 
   com.google.api.services.bigquery.model.TimePartitioning toPb() {
     com.google.api.services.bigquery.model.TimePartitioning partitioningPb =
         new com.google.api.services.bigquery.model.TimePartitioning();
-    partitioningPb.setType(type.name());
-    partitioningPb.setExpirationMs(expirationMs);
+    partitioningPb.setType(getType().name());
+    partitioningPb.setExpirationMs(getExpirationMs());
+    partitioningPb.setRequirePartitionFilter(getRequirePartitionFilter());
+    partitioningPb.setField(getField());
     return partitioningPb;
   }
 
   static TimePartitioning fromPb(
       com.google.api.services.bigquery.model.TimePartitioning partitioningPb) {
-    return new TimePartitioning(
-        Type.valueOf(partitioningPb.getType()), partitioningPb.getExpirationMs());
+    return newBuilder(Type.valueOf(partitioningPb.getType()))
+            .setExpirationMs(partitioningPb.getExpirationMs())
+            .setField(partitioningPb.getField())
+            .setRequirePartitionFilter(partitioningPb.getRequirePartitionFilter())
+            .build();
   }
 }

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TimePartitioningTest.java

@@ -29,8 +29,14 @@ public class TimePartitioningTest {
 
   private static final Type TYPE = Type.DAY;
   private static final long EXPIRATION_MS = 42;
+  private static final boolean REQUIRE_PARTITION_FILTER = false;
+  private static final String FIELD = "field";
   private static final TimePartitioning TIME_PARTITIONING =
-      TimePartitioning.of(TYPE, EXPIRATION_MS);
+      TimePartitioning.newBuilder(TYPE)
+              .setExpirationMs(EXPIRATION_MS)
+              .setRequirePartitionFilter(REQUIRE_PARTITION_FILTER)
+              .setField(FIELD)
+              .build();
 
   @Rule
   public ExpectedException thrown = ExpectedException.none();
@@ -39,11 +45,26 @@ public class TimePartitioningTest {
   public void testOf() {
     assertEquals(TYPE, TIME_PARTITIONING.getType());
     assertEquals(EXPIRATION_MS, TIME_PARTITIONING.getExpirationMs().longValue());
+    assertEquals(REQUIRE_PARTITION_FILTER, TIME_PARTITIONING.getRequirePartitionFilter());
+    assertEquals(FIELD, TIME_PARTITIONING.getField());
     TimePartitioning partitioning = TimePartitioning.of(TYPE);
     assertEquals(TYPE, partitioning.getType());
     assertNull(partitioning.getExpirationMs());
   }
 
+  @Test
+  public void testBuilder() {
+    TimePartitioning partitioning = TimePartitioning.newBuilder(TYPE).build();
+    assertEquals(TYPE, partitioning.getType());
+    assertNull(partitioning.getExpirationMs());
+    assertNull(partitioning.getRequirePartitionFilter());
+    assertNull(partitioning.getField());
+    partitioning = TimePartitioning.newBuilder(TYPE).setExpirationMs(100L).build();
+    assertEquals(TYPE, partitioning.getType());
+    assertEquals(100, (long) partitioning.getExpirationMs());
+    assertNull(partitioning.getRequirePartitionFilter());
+    assertNull(partitioning.getField());
+  }
 
   @Test
   public void testTypeOf_Npe() {
@@ -68,6 +89,8 @@ private void compareTimePartitioning(TimePartitioning expected, TimePartitioning
     assertEquals(expected, value);
     assertEquals(expected.getType(), value.getType());
     assertEquals(expected.getExpirationMs(), value.getExpirationMs());
+    assertEquals(expected.getRequirePartitionFilter(), value.getRequirePartitionFilter());
+    assertEquals(expected.getField(), value.getField());
     assertEquals(expected.hashCode(), value.hashCode());
     assertEquals(expected.toString(), value.toString());
   }