feat: add support for table snapshot

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,12 @@ public Builder setDestinationTable(TableId destinationTable) {
       return this;
     }
 
+    /** Sets the supported operation types in table copy job. */
+    public Builder setOperationType(String operationType) {
+      this.operationType = operationType;
+      return this;
+    }
+
     public Builder setDestinationEncryptionConfiguration(
         EncryptionConfiguration encryptionConfiguration) {
       this.destinationEncryptionConfiguration = encryptionConfiguration;
@@ -178,6 +190,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 +208,11 @@ public TableId getDestinationTable() {
     return destinationTable;
   }
 
+  /** Returns the table copy job type */
+  public String getOperationType() {
+    return operationType;
+  }
+
   public EncryptionConfiguration getDestinationEncryptionConfiguration() {
     return destinationEncryptionConfiguration;
   }
@@ -241,6 +259,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 +279,7 @@ public int hashCode() {
         baseHashCode(),
         sourceTables,
         destinationTable,
+        operationType,
         createDisposition,
         writeDisposition,
         labels,
@@ -293,11 +313,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/SnapshotDefinition.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 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 SnapshotDefinition extends TableDefinition {
+
+  private static final long serialVersionUID = 2113445776046717526L;
+
+  @AutoValue.Builder
+  public abstract static class Builder
+      extends TableDefinition.Builder<SnapshotDefinition, Builder> {
+
+    public abstract Builder setBaseTableId(TableId baseTableId);
+
+    public abstract Builder setSnapshotTime(DateTime dateTime);
+
+    /** Creates a {@code SnapshotDefinition} object. */
+    public abstract SnapshotDefinition build();
+  }
+
+  @Nullable
+  public abstract TableId getBaseTableId();
+
+  @Nullable
+  public abstract DateTime getSnapshotTime();
+
+  /** Returns a builder for a snapshot table definition. */
+  public static SnapshotDefinition.Builder newBuilder() {
+    return new AutoValue_SnapshotDefinition.Builder().setType(Type.SNAPSHOT);
+  }
+
+  @VisibleForTesting
+  public abstract SnapshotDefinition.Builder toBuilder();
+
+  static SnapshotDefinition 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());
+      }
+    }
+    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) SnapshotDefinition.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/SnapshotDefinitionTest.java

@@ -0,0 +1,57 @@
+/*
+ * 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 com.google.api.client.util.DateTime;
+import org.junit.Test;
+
+public class SnapshotDefinitionTest {
+
+  private static final TableId BASE_TABLE_ID = TableId.of("DATASET_NAME", "BASE_TABLE_NAME");
+  private static final DateTime SNAPSHOT_TIME = new DateTime("2021-05-18");
+  private static final SnapshotDefinition SNAPSHOT_DEFINITION =
+      SnapshotDefinition.newBuilder()
+          .setBaseTableId(BASE_TABLE_ID)
+          .setSnapshotTime(SNAPSHOT_TIME)
+          .build();
+
+  @Test
+  public void testToBuilder() {
+    compareSnapshotDefinition(SNAPSHOT_DEFINITION, SNAPSHOT_DEFINITION.toBuilder().build());
+  }
+
+  @Test
+  public void testBuilder() {
+    assertEquals(TableDefinition.Type.SNAPSHOT, SNAPSHOT_DEFINITION.getType());
+    assertEquals(BASE_TABLE_ID, SNAPSHOT_DEFINITION.getBaseTableId());
+    assertEquals(SNAPSHOT_TIME, SNAPSHOT_DEFINITION.getSnapshotTime());
+    SnapshotDefinition snapshotDefinition =
+        SnapshotDefinition.newBuilder()
+            .setBaseTableId(BASE_TABLE_ID)
+            .setSnapshotTime(SNAPSHOT_TIME)
+            .build();
+    assertEquals(SNAPSHOT_DEFINITION, snapshotDefinition);
+  }
+
+  private void compareSnapshotDefinition(SnapshotDefinition expected, SnapshotDefinition value) {
+    assertEquals(expected, value);
+    assertEquals(expected.getBaseTableId(), value.getBaseTableId());
+    assertEquals(expected.getSnapshotTime(), value.getSnapshotTime());
+  }
+}

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -2589,6 +2589,68 @@ public void testCopyJob() throws InterruptedException, TimeoutException {
     assertTrue(remoteTable.delete());
   }
 
+  @Test
+  public void testSnapshotCopyJob() throws InterruptedException {
+    String sourceTableName = "test_copy_job_base_table";
+    String snapshotTableName = "test_snapshot_table";
+    // Create source table
+    TableId sourceTableId = TableId.of(DATASET, sourceTableName);
+    StandardTableDefinition tableDefinition = StandardTableDefinition.of(TABLE_SCHEMA);
+    TableInfo tableInfo = TableInfo.of(sourceTableId, tableDefinition);
+    Table createdTable = bigquery.create(tableInfo);
+    assertNotNull(createdTable);
+
+    // Create snapshot table using source table as the base table
+    TableId snapshotTableId = TableId.of(DATASET, snapshotTableName);
+    CopyJobConfiguration snapshotConfiguration =
+        CopyJobConfiguration.newBuilder(snapshotTableId, sourceTableId)
+            .setOperationType("SNAPSHOT")
+            .build();
+    Job createdJob = bigquery.create(JobInfo.of(snapshotConfiguration));
+    CopyJobConfiguration createdConfiguration = createdJob.getConfiguration();
+    assertNotNull(createdConfiguration.getSourceTables());
+    assertNotNull(createdConfiguration.getOperationType());
+    assertNotNull(createdConfiguration.getDestinationTable());
+    Job completedJob = createdJob.waitFor();
+    assertNull(completedJob.getStatus().getError());
+    Table snapshotTable = bigquery.getTable(DATASET, snapshotTableName);
+    assertNotNull(snapshotTable);
+    assertEquals(snapshotTableId.getDataset(), snapshotTable.getTableId().getDataset());
+    assertEquals(snapshotTableName, snapshotTable.getTableId().getTable());
+    assertEquals(TABLE_SCHEMA, snapshotTable.getDefinition().getSchema());
+
+    // Update snapshot table with some description
+    Table updatedSnapshotTable =
+        bigquery.update(
+            snapshotTable.toBuilder().setDescription("This is a snapshot table").build());
+    assertEquals(updatedSnapshotTable.getDescription(), "This is a snapshot table");
+    assertEquals(snapshotTableName, snapshotTable.getTableId().getTable());
+
+    // Restore base table to a new table
+    String restoredTableName = "test_restore_table";
+    TableId restoredTableId = TableId.of(DATASET, restoredTableName);
+    CopyJobConfiguration restoreConfiguration =
+        CopyJobConfiguration.newBuilder(restoredTableId, snapshotTableId)
+            .setOperationType("RESTORE")
+            .build();
+    Job createdRestoreJob = bigquery.create(JobInfo.of(restoreConfiguration));
+    CopyJobConfiguration createdRestoreConfiguration = createdRestoreJob.getConfiguration();
+    assertNotNull(createdRestoreConfiguration.getSourceTables());
+    assertNotNull(createdRestoreConfiguration.getOperationType());
+    assertNotNull(createdRestoreConfiguration.getDestinationTable());
+    Job completedRestoreJob = createdRestoreJob.waitFor();
+    assertNull(completedRestoreJob.getStatus().getError());
+    Table restoredTable = bigquery.getTable(DATASET, restoredTableName);
+    assertNotNull(restoredTable);
+    assertEquals(restoredTableId.getDataset(), restoredTable.getTableId().getDataset());
+    assertEquals(restoredTableName, restoredTable.getTableId().getTable());
+    assertEquals(TABLE_SCHEMA, restoredTable.getDefinition().getSchema());
+
+    // Clean up
+    assertTrue(createdTable.delete());
+    assertTrue(snapshotTable.delete());
+  }
+
   @Test
   public void testCopyJobWithLabels() throws InterruptedException {
     String sourceTableName = "test_copy_job_source_table_label";