diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TimePartitioning.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TimePartitioning.java index 3359736ca1a0..659c902b56b8 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TimePartitioning.java +++ b/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 Partitioned Tables */ -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(); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TimePartitioningTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TimePartitioningTest.java index 037ca4625419..75488afceb27 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TimePartitioningTest.java +++ b/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()); }