From d783292f4bd2d971235d8afa6318976856529cb5 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Mon, 28 Jun 2021 14:53:45 -0400 Subject: [PATCH] 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 --- .../cloud/bigquery/CopyJobConfiguration.java | 30 ++++- .../bigquery/SnapshotTableDefinition.java | 112 ++++++++++++++++++ .../cloud/bigquery/TableDefinition.java | 4 + .../bigquery/CopyJobConfigurationTest.java | 3 +- .../bigquery/SnapshotTableDefinitionTest.java | 72 +++++++++++ .../cloud/bigquery/it/ITBigQueryTest.java | 84 +++++++++++++ 6 files changed, 301 insertions(+), 4 deletions(-) create mode 100644 google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/SnapshotTableDefinition.java create mode 100644 google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SnapshotTableDefinitionTest.java diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java index d0e15e49e..0b11bf7be 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/CopyJobConfiguration.java +++ b/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 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 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()); } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/SnapshotTableDefinition.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/SnapshotTableDefinition.java new file mode 100644 index 000000000..ad0aeb0ce --- /dev/null +++ b/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 { + + /** 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(); + } +} diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableDefinition.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableDefinition.java index 6babd4e6f..1fa902498 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableDefinition.java +++ b/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 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"); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java index be62bc1d6..dfe76adaf 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/CopyJobConfigurationTest.java +++ b/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()); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SnapshotTableDefinitionTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SnapshotTableDefinitionTest.java new file mode 100644 index 000000000..c739bcf5a --- /dev/null +++ b/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.fromPb(snapshotTableDefinition.toPb())); + } + + private void compareSnapshotTableDefinition( + SnapshotTableDefinition expected, SnapshotTableDefinition value) { + assertEquals(expected, value); + assertEquals(expected.getBaseTableId(), value.getBaseTableId()); + assertEquals(expected.getSnapshotTime(), value.getSnapshotTime()); + } +} diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 4d7b30bef..a0135fb9e 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -87,6 +87,7 @@ import com.google.cloud.bigquery.RoutineId; import com.google.cloud.bigquery.RoutineInfo; import com.google.cloud.bigquery.Schema; +import com.google.cloud.bigquery.SnapshotTableDefinition; import com.google.cloud.bigquery.StandardSQLDataType; import com.google.cloud.bigquery.StandardSQLField; import com.google.cloud.bigquery.StandardSQLTableType; @@ -2655,6 +2656,89 @@ public void testCopyJob() throws InterruptedException, TimeoutException { assertTrue(remoteTable.delete()); } + @Test + public void testSnapshotTableCopyJob() throws InterruptedException { + String sourceTableName = "test_copy_job_base_table"; + String ddlTableName = TABLE_ID_DDL.getTable(); + // this creates a snapshot table at specified snapshotTime + String snapshotTableName = String.format("test_snapshot_table"); + // Create source table with some data in it + String ddlQuery = + String.format( + "CREATE OR REPLACE TABLE %s (" + + "TimestampField TIMESTAMP OPTIONS(description='TimestampDescription'), " + + "StringField STRING OPTIONS(description='StringDescription'), " + + "BooleanField BOOLEAN OPTIONS(description='BooleanDescription') " + + ") AS SELECT * FROM %s", + sourceTableName, ddlTableName); + QueryJobConfiguration ddlConfig = + QueryJobConfiguration.newBuilder(ddlQuery).setDefaultDataset(DatasetId.of(DATASET)).build(); + TableId sourceTableId = TableId.of(DATASET, sourceTableName); + TableResult result = bigquery.query(ddlConfig); + assertEquals(DDL_TABLE_SCHEMA, result.getSchema()); + Table remoteTable = bigquery.getTable(DATASET, sourceTableName); + assertNotNull(remoteTable); + // StandardTableDefinition tableDefinition = StandardTableDefinition.of(TABLE_SCHEMA); + // TableInfo tableInfo = TableInfo.of(sourceTableId, tableDefinition); + // Table createdTable = bigquery.create(tableInfo); + + // 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()); + assertTrue(snapshotTable.getDefinition() instanceof SnapshotTableDefinition); + assertEquals(DDL_TABLE_SCHEMA, snapshotTable.getDefinition().getSchema()); + assertNotNull(((SnapshotTableDefinition) snapshotTable.getDefinition()).getSnapshotTime()); + assertEquals( + sourceTableName, + ((SnapshotTableDefinition) snapshotTable.getDefinition()).getBaseTableId().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(); + assertEquals( + restoreConfiguration.getSourceTables().get(0).getTable(), + createdRestoreConfiguration.getSourceTables().get(0).getTable()); + assertEquals( + restoreConfiguration.getOperationType(), createdRestoreConfiguration.getOperationType()); + assertEquals( + restoreConfiguration.getDestinationTable().getTable(), + createdRestoreConfiguration.getDestinationTable().getTable()); + 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(DDL_TABLE_SCHEMA, restoredTable.getDefinition().getSchema()); + assertEquals(snapshotTable.getNumBytes(), restoredTable.getNumBytes()); + assertEquals(snapshotTable.getNumRows(), restoredTable.getNumRows()); + + // Clean up + assertTrue(remoteTable.delete()); + assertTrue(restoredTable.delete()); + assertTrue(snapshotTable.delete()); + } + @Test public void testCopyJobWithLabels() throws InterruptedException { String sourceTableName = "test_copy_job_source_table_label";