feat: add support for table snapshot (#1320)

* feat: add table snapshot support

Allows users to generate a persistent copy of an existing table as it were at a given point in time within the same project.

* feat: add support for table snapshot

* add comment on operationType

* lint update and remove extraneous update table snapshot IT

* add cleanup for restoredTable

* update base on comments

* update based on comments

Author: stephaniewang526

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

@@ -38,6 +38,7 @@ public final class CopyJobConfiguration extends JobConfiguration {
 
   private final List<TableId> sourceTables;
   private final TableId destinationTable;
+  private final String operationType;
   private final JobInfo.CreateDisposition createDisposition;
   private final JobInfo.WriteDisposition writeDisposition;
   private final EncryptionConfiguration destinationEncryptionConfiguration;
@@ -49,6 +50,7 @@ public static final class Builder
 
     private List<TableId> sourceTables;
     private TableId destinationTable;
+    private String operationType;
     private JobInfo.CreateDisposition createDisposition;
     private JobInfo.WriteDisposition writeDisposition;
     private EncryptionConfiguration destinationEncryptionConfiguration;
@@ -63,6 +65,7 @@ private Builder(CopyJobConfiguration jobConfiguration) {
       this();
       this.sourceTables = jobConfiguration.sourceTables;
       this.destinationTable = jobConfiguration.destinationTable;
+      this.operationType = jobConfiguration.operationType;
       this.createDisposition = jobConfiguration.createDisposition;
       this.writeDisposition = jobConfiguration.writeDisposition;
       this.destinationEncryptionConfiguration = jobConfiguration.destinationEncryptionConfiguration;
@@ -74,6 +77,9 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
       this();
       JobConfigurationTableCopy copyConfigurationPb = configurationPb.getCopy();
       this.destinationTable = TableId.fromPb(copyConfigurationPb.getDestinationTable());
+      if (copyConfigurationPb.getOperationType() != null) {
+        this.operationType = copyConfigurationPb.getOperationType();
+      }
       if (copyConfigurationPb.getSourceTables() != null) {
         this.sourceTables =
             Lists.transform(copyConfigurationPb.getSourceTables(), TableId.FROM_PB_FUNCTION);
@@ -114,6 +120,15 @@ public Builder setDestinationTable(TableId destinationTable) {
       return this;
     }
 
+    /**
+     * Sets the supported operation types (COPY, SNAPSHOT or RESTORE) in table copy job. More info:
+     * https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#operationtype
+     */
+    public Builder setOperationType(String operationType) {
+      this.operationType = operationType;
+      return this;
+    }
+
     public Builder setDestinationEncryptionConfiguration(
         EncryptionConfiguration encryptionConfiguration) {
       this.destinationEncryptionConfiguration = encryptionConfiguration;
@@ -178,6 +193,7 @@ private CopyJobConfiguration(Builder builder) {
     super(builder);
     this.sourceTables = checkNotNull(builder.sourceTables);
     this.destinationTable = checkNotNull(builder.destinationTable);
+    this.operationType = builder.operationType;
     this.createDisposition = builder.createDisposition;
     this.writeDisposition = builder.writeDisposition;
     this.destinationEncryptionConfiguration = builder.destinationEncryptionConfiguration;
@@ -195,6 +211,11 @@ public TableId getDestinationTable() {
     return destinationTable;
   }
 
+  /** Returns the table copy job type */
+  public String getOperationType() {
+    return operationType;
+  }
+
   public EncryptionConfiguration getDestinationEncryptionConfiguration() {
     return destinationEncryptionConfiguration;
   }
@@ -241,6 +262,7 @@ ToStringHelper toStringHelper() {
     return super.toStringHelper()
         .add("sourceTables", sourceTables)
         .add("destinationTable", destinationTable)
+        .add("operationType", operationType)
         .add("destinationEncryptionConfiguration", destinationEncryptionConfiguration)
         .add("createDisposition", createDisposition)
         .add("writeDisposition", writeDisposition)
@@ -260,6 +282,7 @@ public int hashCode() {
         baseHashCode(),
         sourceTables,
         destinationTable,
+        operationType,
         createDisposition,
         writeDisposition,
         labels,
@@ -293,11 +316,12 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
     com.google.api.services.bigquery.model.JobConfiguration jobConfiguration =
         new com.google.api.services.bigquery.model.JobConfiguration();
     configurationPb.setDestinationTable(destinationTable.toPb());
-    if (sourceTables.size() == 1) {
-      configurationPb.setSourceTable(sourceTables.get(0).toPb());
-    } else {
+    if (sourceTables != null) {
       configurationPb.setSourceTables(Lists.transform(sourceTables, TableId.TO_PB_FUNCTION));
     }
+    if (operationType != null) {
+      configurationPb.setOperationType(operationType);
+    }
     if (createDisposition != null) {
       configurationPb.setCreateDisposition(createDisposition.toString());
     }

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

@@ -0,0 +1,112 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import com.google.api.client.util.DateTime;
+import com.google.api.core.BetaApi;
+import com.google.api.services.bigquery.model.Table;
+import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
+import javax.annotation.Nullable;
+
+@AutoValue
+@BetaApi
+public abstract class SnapshotTableDefinition extends TableDefinition {
+
+  private static final long serialVersionUID = 2113445776046717526L;
+
+  @AutoValue.Builder
+  public abstract static class Builder
+      extends TableDefinition.Builder<SnapshotTableDefinition, Builder> {
+
+    /** Reference describing the ID of the table that was snapshot. * */
+    public abstract Builder setBaseTableId(TableId baseTableId);
+
+    /**
+     * The time at which the base table was snapshot. This value is reported in the JSON response
+     * using RFC3339 format. *
+     */
+    public abstract Builder setSnapshotTime(String dateTime);
+
+    public abstract Builder setTimePartitioning(TimePartitioning timePartitioning);
+
+    public abstract Builder setRangePartitioning(RangePartitioning rangePartitioning);
+
+    public abstract Builder setClustering(Clustering clustering);
+
+    /** Creates a {@code SnapshotTableDefinition} object. */
+    public abstract SnapshotTableDefinition build();
+  }
+
+  @Nullable
+  public abstract TableId getBaseTableId();
+
+  @Nullable
+  public abstract String getSnapshotTime();
+
+  @Nullable
+  public abstract TimePartitioning getTimePartitioning();
+
+  @Nullable
+  public abstract RangePartitioning getRangePartitioning();
+
+  @Nullable
+  public abstract Clustering getClustering();
+
+  /** Returns a builder for a snapshot table definition. */
+  public static SnapshotTableDefinition.Builder newBuilder() {
+    return new AutoValue_SnapshotTableDefinition.Builder().setType(Type.SNAPSHOT);
+  }
+
+  @VisibleForTesting
+  public abstract SnapshotTableDefinition.Builder toBuilder();
+
+  @Override
+  Table toPb() {
+    Table tablePb = super.toPb();
+    com.google.api.services.bigquery.model.SnapshotDefinition snapshotDefinition =
+        new com.google.api.services.bigquery.model.SnapshotDefinition();
+    snapshotDefinition.setBaseTableReference(getBaseTableId().toPb());
+    snapshotDefinition.setSnapshotTime(DateTime.parseRfc3339(getSnapshotTime()));
+    tablePb.setSnapshotDefinition(snapshotDefinition);
+    if (getTimePartitioning() != null) {
+      tablePb.setTimePartitioning(getTimePartitioning().toPb());
+    }
+    if (getRangePartitioning() != null) {
+      tablePb.setRangePartitioning(getRangePartitioning().toPb());
+    }
+    if (getClustering() != null) {
+      tablePb.setClustering(getClustering().toPb());
+    }
+    return tablePb;
+  }
+
+  static SnapshotTableDefinition fromPb(Table tablePb) {
+    Builder builder = newBuilder().table(tablePb);
+    com.google.api.services.bigquery.model.SnapshotDefinition snapshotDefinition =
+        tablePb.getSnapshotDefinition();
+    if (snapshotDefinition != null) {
+      if (snapshotDefinition.getBaseTableReference() != null) {
+        builder.setBaseTableId(TableId.fromPb(snapshotDefinition.getBaseTableReference()));
+      }
+      if (snapshotDefinition.getSnapshotTime() != null) {
+        builder.setSnapshotTime(snapshotDefinition.getSnapshotTime().toStringRfc3339());
+      }
+    }
+    return builder.build();
+  }
+}

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

@@ -83,6 +83,8 @@ public Type apply(String constant) {
      */
     public static final Type MODEL = type.createAndRegister("MODEL");
 
+    public static final Type SNAPSHOT = type.createAndRegister("SNAPSHOT");
+
     private Type(String constant) {
       super(constant);
     }
@@ -165,6 +167,8 @@ static <T extends TableDefinition> T fromPb(Table tablePb) {
         return (T) ExternalTableDefinition.fromPb(tablePb);
       case "MODEL":
         return (T) ModelTableDefinition.fromPb(tablePb);
+      case "SNAPSHOT":
+        return (T) SnapshotTableDefinition.fromPb(tablePb);
       default:
         // never reached
         throw new IllegalArgumentException("Format " + tablePb.getType() + " is not supported");

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

@@ -113,7 +113,8 @@ public void testToPbAndFromPb() {
     assertNull(COPY_JOB_CONFIGURATION.toPb().getExtract());
     assertNull(COPY_JOB_CONFIGURATION.toPb().getLoad());
     assertNull(COPY_JOB_CONFIGURATION.toPb().getQuery());
-    assertNull(COPY_JOB_CONFIGURATION.toPb().getCopy().getSourceTables());
+    assertNull(COPY_JOB_CONFIGURATION.toPb().getCopy().getSourceTable());
+    assertNotNull(COPY_JOB_CONFIGURATION.toPb().getCopy().getSourceTables());
     assertNull(COPY_JOB_CONFIGURATION_MULTIPLE_TABLES.toPb().getCopy().getSourceTable());
     assertNotNull(COPY_JOB_CONFIGURATION.getLabels());
     assertNotNull(COPY_JOB_CONFIGURATION_MULTIPLE_TABLES.getLabels());

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+
+public class SnapshotTableDefinitionTest {
+
+  private static final TableId BASE_TABLE_ID = TableId.of("DATASET_NAME", "BASE_TABLE_NAME");
+  private static final String SNAPSHOT_TIME = "2021-05-19T11:32:26.553Z";
+  private static final SnapshotTableDefinition SNAPSHOTTABLE_DEFINITION =
+      SnapshotTableDefinition.newBuilder()
+          .setBaseTableId(BASE_TABLE_ID)
+          .setSnapshotTime(SNAPSHOT_TIME)
+          .build();
+
+  @Test
+  public void testToBuilder() {
+    compareSnapshotTableDefinition(
+        SNAPSHOTTABLE_DEFINITION, SNAPSHOTTABLE_DEFINITION.toBuilder().build());
+    SnapshotTableDefinition snapshotTableDefinition =
+        SNAPSHOTTABLE_DEFINITION.toBuilder().setSnapshotTime("2021-05-20T11:32:26.553Z").build();
+    assertEquals("2021-05-20T11:32:26.553Z", snapshotTableDefinition.getSnapshotTime());
+  }
+
+  @Test
+  public void testBuilder() {
+    assertEquals(TableDefinition.Type.SNAPSHOT, SNAPSHOTTABLE_DEFINITION.getType());
+    assertEquals(BASE_TABLE_ID, SNAPSHOTTABLE_DEFINITION.getBaseTableId());
+    assertEquals(SNAPSHOT_TIME, SNAPSHOTTABLE_DEFINITION.getSnapshotTime());
+    SnapshotTableDefinition snapshotTableDefinition =
+        SnapshotTableDefinition.newBuilder()
+            .setBaseTableId(BASE_TABLE_ID)
+            .setSnapshotTime(SNAPSHOT_TIME)
+            .build();
+    assertEquals(SNAPSHOTTABLE_DEFINITION, snapshotTableDefinition);
+  }
+
+  @Test
+  public void testToAndFromPb() {
+    SnapshotTableDefinition snapshotTableDefinition = SNAPSHOTTABLE_DEFINITION.toBuilder().build();
+    assertTrue(
+        TableDefinition.fromPb(snapshotTableDefinition.toPb()) instanceof SnapshotTableDefinition);
+    compareSnapshotTableDefinition(
+        snapshotTableDefinition,
+        TableDefinition.<SnapshotTableDefinition>fromPb(snapshotTableDefinition.toPb()));
+  }
+
+  private void compareSnapshotTableDefinition(
+      SnapshotTableDefinition expected, SnapshotTableDefinition value) {
+    assertEquals(expected, value);
+    assertEquals(expected.getBaseTableId(), value.getBaseTableId());
+    assertEquals(expected.getSnapshotTime(), value.getSnapshotTime());
+  }
+}