diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java index d12af24e407d..dd7ed0719570 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java @@ -430,7 +430,7 @@ private WriteResult expandTriggered(PCollection> inpu zeroLoadJobIdPrefixView, loadJobProjectId, WriteDisposition.WRITE_APPEND, - CreateDisposition.CREATE_NEVER, + createDisposition, maxRetryJobs, kmsKey, schemaUpdateOptions, @@ -536,7 +536,7 @@ public WriteResult expandUntriggered(PCollection> inp zeroLoadJobIdPrefixView, loadJobProjectId, WriteDisposition.WRITE_APPEND, - CreateDisposition.CREATE_NEVER, + createDisposition, maxRetryJobs, kmsKey, schemaUpdateOptions, @@ -812,6 +812,14 @@ private PCollection> writeTempTables( List> sideInputs = Lists.newArrayList(jobIdTokenView); sideInputs.addAll(dynamicDestinations.getSideInputs()); + // Temp tables need the clone source schema even when schema update options prevent matching + // the final destination's schema. + DynamicDestinations tempTableDestinations = + schemaUpdateOptions != null && !schemaUpdateOptions.isEmpty() + ? DynamicDestinationsHelpers.cloneSourceSchemaDynamicDestinations( + dynamicDestinations, bigQueryServices) + : destinationsWithMatching; + Coder, WritePartition.Result>> partitionsCoder = KvCoder.of( ShardedKeyCoder.of(NullableCoder.of(destinationCoder)), @@ -835,7 +843,7 @@ private PCollection> writeTempTables( WriteDisposition.WRITE_EMPTY, CreateDisposition.CREATE_IF_NEEDED, sideInputs, - destinationsWithMatching, + tempTableDestinations, loadJobProjectId, maxRetryJobs, ignoreUnknownValues, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index 0b8263ecb1e4..7d4ad6ffe28f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -103,6 +103,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient; import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySourceBase.ExtractResult; +import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantCloneSourceDestinations; import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantSchemaDestinations; import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantTimePartitioningClusteringDestinations; import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.SchemaFromViewDestinations; @@ -3012,6 +3013,8 @@ public enum Method { abstract @Nullable ValueProvider getJsonSchema(); + abstract @Nullable ValueProvider getJsonCloneSourceTableRef(); + abstract @Nullable ValueProvider getJsonTimePartitioning(); abstract @Nullable ValueProvider getJsonClustering(); @@ -3133,6 +3136,8 @@ abstract Builder setAvroSchemaFactory( abstract Builder setJsonSchema(ValueProvider jsonSchema); + abstract Builder setJsonCloneSourceTableRef(ValueProvider jsonCloneSourceTableRef); + abstract Builder setJsonTimePartitioning(ValueProvider jsonTimePartitioning); abstract Builder setJsonClustering(ValueProvider clustering); @@ -3473,6 +3478,39 @@ public Write withJsonSchema(ValueProvider jsonSchema) { return toBuilder().setJsonSchema(jsonSchema).build(); } + /** + * Creates missing destination tables as clones of the specified base table. + * + *

The clone source is used only if writing to a table that does not already exist and {@link + * CreateDisposition} is set to {@link CreateDisposition#CREATE_IF_NEEDED}. This option is + * mutually exclusive with {@link #withSchema(TableSchema)}, {@link #withJsonSchema(String)}, + * and {@link #withSchemaFromView(PCollectionView)}. + */ + public Write withCloneFrom(String tableSpec) { + checkArgument(tableSpec != null, "tableSpec can not be null"); + return withCloneFrom(StaticValueProvider.of(tableSpec)); + } + + /** Same as {@link #withCloneFrom(String)} but using a {@link TableReference}. */ + public Write withCloneFrom(TableReference table) { + checkArgument(table != null, "table can not be null"); + return withJsonCloneSourceTableRef( + StaticValueProvider.of(BigQueryHelpers.toJsonString(table))); + } + + /** Same as {@link #withCloneFrom(String)} but using a deferred {@link ValueProvider}. */ + public Write withCloneFrom(ValueProvider tableSpec) { + checkArgument(tableSpec != null, "tableSpec can not be null"); + return withJsonCloneSourceTableRef( + NestedValueProvider.of( + NestedValueProvider.of(tableSpec, new TableSpecToTableRef()), new TableRefToJson())); + } + + private Write withJsonCloneSourceTableRef(ValueProvider jsonCloneSourceTableRef) { + checkArgument(jsonCloneSourceTableRef != null, "jsonCloneSourceTableRef can not be null"); + return toBuilder().setJsonCloneSourceTableRef(jsonCloneSourceTableRef).build(); + } + /** * Allows the schemas for each table to be computed within the pipeline itself. * @@ -4129,8 +4167,19 @@ public WriteResult expand(PCollection input) { .collect(Collectors.toList())), "No more than one of jsonSchema, schemaFromView, or dynamicDestinations may be set"); + long activeCreationMetadataCount = + java.util.stream.Stream.of( + getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef()) + .filter(arg -> arg != null) + .count(); + checkArgument( + activeCreationMetadataCount <= 1, + "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); + // Perform some argument checks BigQueryOptions bqOptions = input.getPipeline().getOptions().as(BigQueryOptions.class); + // Resolve the token before the options are serialized to workers. + bqOptions.getCloneJobIdToken(); Write.Method method = resolveMethod(input); if (input.isBounded() == IsBounded.UNBOUNDED) { if (method == Write.Method.FILE_LOADS || method == Write.Method.STORAGE_WRITE_API) { @@ -4304,6 +4353,12 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { } dynamicDestinations = tableDestinations; } + ValueProvider jsonCloneSourceTableRef = getJsonCloneSourceTableRef(); + if (jsonCloneSourceTableRef != null) { + dynamicDestinations = + new ConstantCloneSourceDestinations<>( + (DynamicDestinations) dynamicDestinations, jsonCloneSourceTableRef); + } return expandTyped(input, dynamicDestinations); } @@ -4319,6 +4374,7 @@ private WriteResult expandTyped( getJsonSchema() != null || getDynamicDestinations() != null || getSchemaFromView() != null; + boolean hasTableCreationMetadata = hasSchema || getJsonCloneSourceTableRef() != null; Class writeProtoClass = getWriteProtosClass(); if (getUseBeamSchema()) { @@ -4347,7 +4403,7 @@ private WriteResult expandTyped( dynamicDestinations, StaticValueProvider.of(BigQueryHelpers.toJsonString(tableSchema))); } else if (writeProtoClass != null) { - if (!hasSchema) { + if (!hasTableCreationMetadata) { try { // Method.invoke takes a null receiver for a static method; that is not expressible // against the JDK's annotations. @@ -4369,8 +4425,9 @@ private WriteResult expandTyped( } else { // Require a schema if creating one or more tables. checkArgument( - getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED || hasSchema, - "CreateDisposition is CREATE_IF_NEEDED, however no schema was provided."); + getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED + || hasTableCreationMetadata, + "CreateDisposition is CREATE_IF_NEEDED, however no schema was provided and no clone source was provided."); } Coder destinationCoder; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index 9e81b675a755..0aef32c587c9 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java @@ -440,6 +440,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator transform) { if (jsonSchema != null) { fieldValues.put("json_schema", jsonSchema.get()); } + ValueProvider jsonCloneSourceTableRef = transform.getJsonCloneSourceTableRef(); + if (jsonCloneSourceTableRef != null) { + fieldValues.put("json_clone_source_table_ref", jsonCloneSourceTableRef.get()); + } ValueProvider jsonTimePartitioning = transform.getJsonTimePartitioning(); if (jsonTimePartitioning != null) { fieldValues.put("json_time_partitioning", toByteArray(jsonTimePartitioning.get())); @@ -736,6 +741,11 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { if (jsonSchema != null) { builder = builder.setJsonSchema(StaticValueProvider.of(jsonSchema)); } + String jsonCloneSourceTableRef = configRow.getString("json_clone_source_table_ref"); + if (jsonCloneSourceTableRef != null) { + builder = + builder.setJsonCloneSourceTableRef(StaticValueProvider.of(jsonCloneSourceTableRef)); + } String jsonTimePartitioning = configRow.getString("json_time_partitioning"); if (jsonTimePartitioning != null) { builder = builder.setJsonTimePartitioning(StaticValueProvider.of(jsonTimePartitioning)); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index 64382d64bb51..1fd62ae11a9a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java @@ -21,6 +21,7 @@ import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; import org.apache.beam.sdk.options.ApplicationNameOptions; import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.DefaultValueFactory; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.Hidden; import org.apache.beam.sdk.options.PipelineOptions; @@ -32,6 +33,20 @@ + "https://cloud.google.com/bigquery/what-is-bigquery for details on BigQuery.") public interface BigQueryOptions extends ApplicationNameOptions, GcpOptions, PipelineOptions, StreamingOptions { + /** Shared by workers so retries of the same clone use the same BigQuery job ID. */ + @Hidden + @Default.InstanceFactory(CloneJobIdTokenFactory.class) + String getCloneJobIdToken(); + + void setCloneJobIdToken(String token); + + class CloneJobIdTokenFactory implements DefaultValueFactory { + @Override + public String create(PipelineOptions options) { + return BigQueryHelpers.randomUUIDString(); + } + } + @Description( "Temporary dataset for BigQuery table operations. " + "Supported values are \"bigquery.googleapis.com/{dataset}\"") diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryResourceNaming.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryResourceNaming.java index 0d97b51603b5..b069f70641c6 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryResourceNaming.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryResourceNaming.java @@ -69,6 +69,7 @@ public enum JobType { LOAD, TEMP_TABLE_LOAD, COPY, + CLONE, EXPORT, QUERY, SCHEMA_UPDATE, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java index 52ac77532b1d..860a6d343820 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java @@ -27,12 +27,15 @@ import com.google.api.services.bigquery.model.BigLakeConfiguration; import com.google.api.services.bigquery.model.Clustering; import com.google.api.services.bigquery.model.EncryptionConfiguration; +import com.google.api.services.bigquery.model.JobConfigurationTableCopy; +import com.google.api.services.bigquery.model.JobReference; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableConstraints; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import com.google.api.services.bigquery.model.TimePartitioning; import io.grpc.StatusRuntimeException; +import java.io.IOException; import java.util.Collections; import java.util.Map; import java.util.Optional; @@ -42,7 +45,10 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryResourceNaming.JobType; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService; import org.apache.beam.sdk.util.FluentBackoff; import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; @@ -50,10 +56,15 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class CreateTableHelpers { + private static final Logger LOG = LoggerFactory.getLogger(CreateTableHelpers.class); + /** * The list of tables created so far, so we don't try the creation each time. * @@ -65,7 +76,8 @@ public class CreateTableHelpers { private static final FluentBackoff DEFAULT_BACKOFF_FACTORY = FluentBackoff.DEFAULT.withMaxRetries(4).withInitialBackoff(INITIAL_RPC_BACKOFF); - // When CREATE_IF_NEEDED is specified, BQ tables should be created if they do not exist. This + // When CREATE_IF_NEEDED is specified, BQ tables should be created if they do + // not exist. This // method detects // errors on table operations, and attempts to create the table if necessary. static void createTableWrapper(Callable action, Callable tryCreateTable) @@ -86,7 +98,8 @@ static void createTableWrapper(Callable action, Callable tryCreat .findAny(); lastException = (Exception) handledCause.orElseThrow(() -> e); - // TODO: Once BigQuery reliably returns a consistent error on table not found, we should + // TODO: Once BigQuery reliably returns a consistent error on table not found, + // we should // only try creating // the table on that error. boolean created = tryCreateTable.call(); @@ -103,6 +116,7 @@ static TableDestination possiblyCreateTable( TableDestination tableDestination, Supplier<@Nullable TableSchema> schemaSupplier, Supplier<@Nullable TableConstraints> tableConstraintsSupplier, + Supplier<@Nullable TableReference> cloneSourceSupplier, CreateDisposition createDisposition, @Nullable Coder tableDestinationCoder, @Nullable String kmsKey, @@ -136,7 +150,8 @@ static TableDestination possiblyCreateTable( String tableSpec = BigQueryHelpers.stripPartitionDecorator(tableDestination.getTableSpec()); if (!createdTables.contains(tableSpec)) { // Another thread may have succeeded in creating the table in the meanwhile, so - // check again. This check isn't needed for correctness, but we add it to prevent + // check again. This check isn't needed for correctness, but we add it to + // prevent // every thread from attempting a create and overwhelming our BigQuery quota. synchronized (createdTables) { if (!createdTables.contains(tableSpec)) { @@ -144,6 +159,7 @@ static TableDestination possiblyCreateTable( bigQueryOptions, schemaSupplier, tableConstraintsSupplier, + cloneSourceSupplier, tableDestination, createDisposition, tableSpec, @@ -160,6 +176,7 @@ private static void tryCreateTable( BigQueryOptions options, Supplier<@Nullable TableSchema> schemaSupplier, Supplier<@Nullable TableConstraints> tableConstraintsSupplier, + Supplier<@Nullable TableReference> cloneSourceSupplier, TableDestination tableDestination, CreateDisposition createDisposition, String tableSpec, @@ -172,60 +189,68 @@ private static void tryCreateTable( if (datasetService.getTable( tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC) == null) { - TableSchema tableSchema = schemaSupplier.get(); - @Nullable TableConstraints tableConstraints = tableConstraintsSupplier.get(); - Preconditions.checkArgumentNotNull( - tableSchema, - "Unless create disposition is %s, a schema must be specified, i.e. " - + "DynamicDestinations.getSchema() may not return null. " - + "However, create disposition is %s, and " - + " %s returned null for destination %s", - CreateDisposition.CREATE_NEVER, - createDisposition, - tableDestination); - Table table = new Table().setTableReference(tableReference).setSchema(tableSchema); - - if (tableConstraints != null) { - table = table.setTableConstraints(tableConstraints); - } + @Nullable TableReference cloneSource = cloneSourceSupplier.get(); + if (cloneSource != null) { + cloneTable( + options, datasetService, bqServices, cloneSource, tableReference, kmsKey, tableSpec); + } else { + TableSchema tableSchema = schemaSupplier.get(); + @Nullable TableConstraints tableConstraints = tableConstraintsSupplier.get(); + Preconditions.checkArgumentNotNull( + tableSchema, + "Unless create disposition is %s, a schema or clone source must be specified, i.e. " + + "DynamicDestinations.getSchema() and DynamicDestinations.getCloneSource() may " + + "not both return null. However, create disposition is %s, and " + + "no schema or clone source was returned for destination %s", + CreateDisposition.CREATE_NEVER, + createDisposition, + tableDestination); + Table table = new Table().setTableReference(tableReference).setSchema(tableSchema); - String tableDescription = tableDestination.getTableDescription(); - if (tableDescription != null) { - table = table.setDescription(tableDescription); - } + if (tableConstraints != null) { + table = table.setTableConstraints(tableConstraints); + } - TimePartitioning timePartitioning = tableDestination.getTimePartitioning(); - if (timePartitioning != null) { - table.setTimePartitioning(timePartitioning); - } + String tableDescription = tableDestination.getTableDescription(); + if (tableDescription != null) { + table = table.setDescription(tableDescription); + } - Clustering clustering = tableDestination.getClustering(); - if (clustering != null) { - table.setClustering(clustering); - } + TimePartitioning timePartitioning = tableDestination.getTimePartitioning(); + if (timePartitioning != null) { + table.setTimePartitioning(timePartitioning); + } - if (kmsKey != null) { - table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey)); - } - if (bigLakeConfiguration != null) { - TableReference ref = table.getTableReference(); - table.setBiglakeConfiguration( - new BigLakeConfiguration() - .setTableFormat( - MoreObjects.firstNonNull(bigLakeConfiguration.get("tableFormat"), "iceberg")) - .setFileFormat( - MoreObjects.firstNonNull(bigLakeConfiguration.get("fileFormat"), "parquet")) - .setConnectionId( - Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(CONNECTION_ID))) - .setStorageUri( - String.format( - "%s/%s/%s/%s", - Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(STORAGE_URI)), - ref.getProjectId(), - ref.getDatasetId(), - ref.getTableId()))); + Clustering clustering = tableDestination.getClustering(); + if (clustering != null) { + table.setClustering(clustering); + } + + if (kmsKey != null) { + table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey)); + } + if (bigLakeConfiguration != null) { + TableReference ref = table.getTableReference(); + table.setBiglakeConfiguration( + new BigLakeConfiguration() + .setTableFormat( + MoreObjects.firstNonNull( + bigLakeConfiguration.get("tableFormat"), "iceberg")) + .setFileFormat( + MoreObjects.firstNonNull(bigLakeConfiguration.get("fileFormat"), "parquet")) + .setConnectionId( + Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(CONNECTION_ID))) + .setStorageUri( + String.format( + "%s/%s/%s/%s", + Preconditions.checkArgumentNotNull( + bigLakeConfiguration.get(STORAGE_URI)), + ref.getProjectId(), + ref.getDatasetId(), + ref.getTableId()))); + } + datasetService.createTable(table); } - datasetService.createTable(table); } } catch (Exception e) { throw new RuntimeException(e); @@ -233,6 +258,131 @@ private static void tryCreateTable( createdTables.add(tableSpec); } + private static void cloneTable( + BigQueryOptions options, + DatasetService datasetService, + BigQueryServices bqServices, + TableReference cloneSource, + TableReference destination, + @Nullable String kmsKey, + String tableSpec) + throws Exception { + TableReference destinationWithDefaultProject = withDefaultProject(options, destination); + TableReference source = + withDefaultProject(options, cloneSource, destinationWithDefaultProject.getProjectId()); + JobConfigurationTableCopy copyConfig = + new JobConfigurationTableCopy() + .setSourceTables(Collections.singletonList(source)) + .setDestinationTable(destinationWithDefaultProject) + .setOperationType("CLONE") + .setWriteDisposition(WriteDisposition.WRITE_EMPTY.name()) + .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED.name()); + if (kmsKey != null) { + copyConfig.setDestinationEncryptionConfiguration( + new EncryptionConfiguration().setKmsKeyName(kmsKey)); + } + + String jobProjectId = destinationWithDefaultProject.getProjectId(); + String bqLocation = + BigQueryHelpers.getDatasetLocation( + datasetService, + destinationWithDefaultProject.getProjectId(), + destinationWithDefaultProject.getDatasetId()); + String jobIdPrefix = createCloneJobIdPrefix(options, source, destinationWithDefaultProject); + try (JobService jobService = bqServices.getJobService(options)) { + BigQueryHelpers.PendingJob cloneJob = + new BigQueryHelpers.PendingJob( + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + jobService.startCopyJob(jobRef, copyConfig); + } catch (IOException e) { + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return null; + }, + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }, + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + return jobService.getJob(jobRef); + } catch (IOException e) { + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }, + BatchLoads.DEFAULT_MAX_RETRY_JOBS, + jobIdPrefix); + cloneJob.runJob(); + new BigQueryHelpers.PendingJobManager() + .addPendingJob( + cloneJob, + job -> { + LOG.info("Created BigQuery clone table {} from {}", tableSpec, source); + return null; + }) + .waitForDone(); + } + } + + private static JobReference createJobReference(String jobId, String projectId, String location) { + return new JobReference().setProjectId(projectId).setJobId(jobId).setLocation(location); + } + + private static TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference) { + return withDefaultProject(options, tableReference, null); + } + + private static TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference, @Nullable String defaultProjectId) { + TableReference updated = tableReference.clone(); + if (Strings.isNullOrEmpty(updated.getProjectId())) { + if (defaultProjectId != null && !defaultProjectId.isEmpty()) { + updated.setProjectId(defaultProjectId); + } else { + @Nullable String projectId = options.getBigQueryProject(); + if (projectId == null || projectId.isEmpty()) { + projectId = options.getProject(); + } + if (projectId != null && !projectId.isEmpty()) { + updated.setProjectId(projectId); + } + } + } + return updated; + } + + @VisibleForTesting + static String createCloneJobIdPrefix( + BigQueryOptions options, TableReference source, TableReference destination) { + String jobName = MoreObjects.firstNonNull(options.getJobName(), "beam_bq_clone"); + String stepToken = + Hashing.murmur3_128() + .hashUnencodedChars( + BigQueryHelpers.toJsonString(source) + + BigQueryHelpers.toJsonString(destination) + + options.getCloneJobIdToken()) + .toString(); + return BigQueryResourceNaming.createJobIdPrefix(jobName, stepToken, JobType.CLONE); + } + @VisibleForTesting static void clearCreatedTables() { synchronized (createdTables) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java index 7008c049a4a5..b69bfe75d9ae 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java @@ -20,6 +20,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.api.services.bigquery.model.TableConstraints; +import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import java.util.List; import java.util.Map; @@ -117,6 +118,8 @@ public void processElement(ProcessContext context) { () -> dynamicDestinations.getSchema(dest); Supplier<@Nullable TableConstraints> tableConstraintsSupplier = () -> dynamicDestinations.getTableConstraints(dest); + Supplier<@Nullable TableReference> cloneSourceSupplier = + () -> dynamicDestinations.getCloneSource(dest); BigQueryOptions bqOptions = context.getPipelineOptions().as(BigQueryOptions.class); Lineage.getSinks() @@ -129,6 +132,7 @@ public void processElement(ProcessContext context) { tableDestination1, schemaSupplier, tableConstraintsSupplier, + cloneSourceSupplier, createDisposition, dynamicDestinations.getDestinationCoder(), kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java index 0e1ff4ea4120..93acfae19006 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java @@ -21,6 +21,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.services.bigquery.model.TableConstraints; +import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import java.io.Serializable; import java.util.List; @@ -196,6 +197,11 @@ void setSideInputAccessorFromOnWindowExpirationContext( /** Returns the table schema for the destination. */ public abstract @Nullable TableSchema getSchema(DestinationT destination); + /** Returns the base table to clone when creating the destination table. */ + public @Nullable TableReference getCloneSource(DestinationT destination) { + return null; + } + /** * Returns TableConstraints (including primary and foreign key) to be used when creating the * table. Note: this is not currently supported when using FILE_LOADS!. diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java index 4cd5d25202a3..2affd39d3a96 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java @@ -47,6 +47,7 @@ import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; @@ -183,6 +184,11 @@ public DestinationT getDestination(@Nullable ValueInSingleWindow element) { return inner.getSchema(destination); } + @Override + public @Nullable TableReference getCloneSource(DestinationT destination) { + return inner.getCloneSource(destination); + } + @Override public @Nullable TableConstraints getTableConstraints(DestinationT destination) { return inner.getTableConstraints(destination); @@ -238,6 +244,40 @@ public String toString() { } } + /** Returns the same clone source for every table. */ + static class ConstantCloneSourceDestinations + extends DelegatingDynamicDestinations { + private final ValueProvider jsonCloneSource; + private transient @Nullable TableReference cloneSource; + + ConstantCloneSourceDestinations( + DynamicDestinations inner, ValueProvider jsonCloneSource) { + super(inner); + Preconditions.checkArgumentNotNull(jsonCloneSource, "jsonCloneSource cannot be null"); + this.jsonCloneSource = jsonCloneSource; + } + + @Override + public TableReference getCloneSource(DestinationT destination) { + @Nullable TableReference cloneSource = this.cloneSource; + if (cloneSource == null) { + String jsonCloneSource = this.jsonCloneSource.get(); + checkArgument(jsonCloneSource != null, "jsonCloneSource cannot be null"); + cloneSource = BigQueryHelpers.fromJsonString(jsonCloneSource, TableReference.class); + this.cloneSource = cloneSource; + } + return cloneSource.clone(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("inner", inner) + .add("jsonCloneSource", jsonCloneSource) + .toString(); + } + } + static class ConstantTableConstraintsDestinations extends DelegatingDynamicDestinations { private final String jsonTableConstraints; @@ -433,6 +473,32 @@ static DynamicDestinations matchTableDynamicD return new MatchTableDynamicDestinations<>(inner, bqServices); } + static + DynamicDestinations cloneSourceSchemaDynamicDestinations( + DynamicDestinations inner, BigQueryServices bqServices) { + return new CloneSourceSchemaDestinations<>(inner, bqServices); + } + + /** Uses the clone source schema for temporary load tables when no schema was supplied. */ + static class CloneSourceSchemaDestinations + extends MatchTableDynamicDestinations { + private CloneSourceSchemaDestinations( + DynamicDestinations inner, BigQueryServices bqServices) { + super(inner, bqServices); + } + + @Override + public @Nullable TableSchema getSchema(DestinationT destination) { + @Nullable TableSchema schema = inner.getSchema(destination); + return schema != null ? schema : getCloneSourceSchema(destination); + } + + @Override + public TableDestination getTable(DestinationT destination) { + return inner.getTable(destination); + } + } + static class MatchTableDynamicDestinations extends DelegatingDynamicDestinations { private final BigQueryServices bqServices; @@ -444,6 +510,11 @@ private MatchTableDynamicDestinations( } private @Nullable Table getBigQueryTable(TableReference tableReference) { + return getBigQueryTable(tableReference, null); + } + + private @Nullable Table getBigQueryTable( + TableReference tableReference, @Nullable String defaultProjectId) { BackOff backoff = BackOffAdapter.toGcpBackOff( FluentBackoff.DEFAULT @@ -459,16 +530,12 @@ private MatchTableDynamicDestinations( throw new IllegalStateException("pipeline options cannot be null"); } BigQueryOptions bqOptions = options.as(BigQueryOptions.class); - if (tableReference.getProjectId() == null) { - tableReference.setProjectId( - bqOptions.getBigQueryProject() == null - ? bqOptions.getProject() - : bqOptions.getBigQueryProject()); - } + TableReference tableReferenceToFetch = + withDefaultProject(bqOptions, tableReference, defaultProjectId); try (DatasetService datasetService = bqServices.getDatasetService(bqOptions)) { - return datasetService.getTable(tableReference); + return datasetService.getTable(tableReferenceToFetch); } catch (InterruptedException | IOException e) { - LOG.info("Failed to get BigQuery table {}", tableReference); + LOG.info("Failed to get BigQuery table {}", tableReferenceToFetch); } } catch (Exception e) { throw new RuntimeException(e); @@ -480,6 +547,35 @@ private MatchTableDynamicDestinations( return null; } + private static TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference, @Nullable String defaultProjectId) { + TableReference updated = tableReference.clone(); + if (Strings.isNullOrEmpty(updated.getProjectId())) { + if (defaultProjectId != null && !defaultProjectId.isEmpty()) { + updated.setProjectId(defaultProjectId); + } else { + @Nullable String projectId = options.getBigQueryProject(); + if (projectId == null || projectId.isEmpty()) { + projectId = options.getProject(); + } + if (projectId != null && !projectId.isEmpty()) { + updated.setProjectId(projectId); + } + } + } + return updated; + } + + @Nullable TableSchema getCloneSourceSchema(DestinationT destination) { + @Nullable TableReference cloneSource = super.getCloneSource(destination); + if (cloneSource == null) { + return null; + } + TableReference destinationRef = super.getTable(destination).getTableReference(); + @Nullable Table sourceTable = getBigQueryTable(cloneSource, destinationRef.getProjectId()); + return sourceTable == null ? null : sourceTable.getSchema(); + } + /** Identical to {@link BackOffUtils#next} but without checked IOException. */ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws InterruptedException { @@ -494,16 +590,25 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) @Override public TableDestination getTable(DestinationT destination) { TableDestination wrappedDestination = super.getTable(destination); - Table existingTable = getBigQueryTable(wrappedDestination.getTableReference()); + TableReference tableReference = wrappedDestination.getTableReference(); + @Nullable String destinationProjectId = tableReference.getProjectId(); + Table existingTable = getBigQueryTable(tableReference); + Table tableToMatch = existingTable; + if (tableToMatch == null) { + @Nullable TableReference cloneSource = super.getCloneSource(destination); + if (cloneSource != null) { + tableToMatch = getBigQueryTable(cloneSource, destinationProjectId); + } + } - if (existingTable == null) { + if (tableToMatch == null) { return wrappedDestination; } else { return new TableDestination( wrappedDestination.getTableSpec(), - existingTable.getDescription(), - existingTable.getTimePartitioning(), - existingTable.getClustering()); + tableToMatch.getDescription(), + tableToMatch.getTimePartitioning(), + tableToMatch.getClustering()); } } @@ -514,10 +619,15 @@ public TableDestination getTable(DestinationT destination) { @Override public @Nullable TableSchema getSchema(DestinationT destination) { TableDestination wrappedDestination = super.getTable(destination); - @Nullable Table existingTable = getBigQueryTable(wrappedDestination.getTableReference()); + TableReference tableReference = wrappedDestination.getTableReference(); + @Nullable Table existingTable = getBigQueryTable(tableReference); if (existingTable == null || existingTable.getSchema() == null || existingTable.getSchema().isEmpty()) { + @Nullable TableSchema cloneSchema = getCloneSourceSchema(destination); + if (cloneSchema != null && !cloneSchema.isEmpty()) { + return cloneSchema; + } return super.getSchema(destination); } else { return existingTable.getSchema(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java index 8ea8d75ed17e..d7f9c8b08a5e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java @@ -1325,6 +1325,7 @@ DestinationState createDestinationState( tableDestination1, () -> dynamicDestinations.getSchema(destination), () -> dynamicDestinations.getTableConstraints(destination), + () -> dynamicDestinations.getCloneSource(destination), createDisposition, destinationCoder, kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java index 336311e3421e..5d327c31a284 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java @@ -985,6 +985,7 @@ private void processPayloads( tableDestination, () -> dynamicDestinations.getSchema(destination), () -> dynamicDestinations.getTableConstraints(destination), + () -> dynamicDestinations.getCloneSource(destination), createDisposition, destinationCoder, kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java index 536d318f0213..89da2e07b9f6 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java @@ -23,6 +23,7 @@ import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableConstraints; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import com.google.api.services.bigquery.model.TimePartitioning; @@ -31,7 +32,6 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; -import javax.annotation.Nullable; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.ValueProvider; @@ -40,9 +40,11 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -71,7 +73,7 @@ public class UpdateSchemaDestination private transient @Nullable DatasetService datasetService; private final int maxRetryJobs; private final @Nullable String kmsKey; - private @Nullable BigQueryServices.JobService jobService; + private BigQueryServices.@Nullable JobService jobService; private final Set schemaUpdateOptions; private final BigQueryIO.Write.WriteDisposition writeDisposition; private final BigQueryIO.Write.CreateDisposition createDisposition; @@ -149,6 +151,27 @@ public void processElement( DestinationT destination = entry.getKey(); TableDestination tableDestination = getTableWithDefaultProject(destination); outputs.add(KV.of(tableDestination, entry.getValue())); + @Nullable TableReference cloneSource = dynamicDestinations.getCloneSource(destination); + if (cloneSource != null + && createDisposition != BigQueryIO.Write.CreateDisposition.CREATE_NEVER) { + Supplier<@Nullable TableSchema> schemaSupplier = + () -> dynamicDestinations.getSchema(destination); + Supplier<@Nullable TableConstraints> tableConstraintsSupplier = + () -> dynamicDestinations.getTableConstraints(destination); + Supplier<@Nullable TableReference> cloneSourceSupplier = () -> cloneSource; + + CreateTableHelpers.possiblyCreateTable( + context.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, + createDisposition, + dynamicDestinations.getDestinationCoder(), + kmsKey, + bqServices, + null); + } if (pendingJobs.containsKey(destination)) { // zero load job for this destination is already set continue; @@ -230,7 +253,7 @@ public void finishBundle(FinishBundleContext context) throws Exception { jobManager.waitForDone(); } - private @Nullable BigQueryHelpers.PendingJob startZeroLoadJob( + private BigQueryHelpers.@Nullable PendingJob startZeroLoadJob( BigQueryServices.JobService jobService, DatasetService datasetService, String jobIdPrefix, @@ -272,7 +295,8 @@ public void finishBundle(FinishBundleContext context) throws Exception { LOG.warn("Failed to get table {} with {}", tableReference, e.toString()); throw new RuntimeException(e); } - // no need to update schema ahead if provided schema already matches destination schema + // no need to update schema ahead if provided schema already matches destination + // schema // or when destination schema is null (the write will set the schema) // or when provided schema is null (e.g. when using CREATE_NEVER disposition) TableSchema destinationSchema = destinationTable.getSchema(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java index a95bfc5f7c87..05c3b7046aab 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java @@ -23,6 +23,7 @@ import com.google.api.services.bigquery.model.EncryptionConfiguration; import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; +import com.google.api.services.bigquery.model.TableConstraints; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import com.google.api.services.bigquery.model.TimePartitioning; @@ -73,6 +74,7 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; @@ -198,9 +200,20 @@ public void processElement( throws Exception { dynamicDestinations.setSideInputAccessorFromProcessContext(c); DestinationT destination = c.element().getKey().getKey(); + TableDestination tableDestination = dynamicDestinations.getTable(destination); + checkArgument( + tableDestination != null, + "DynamicDestinations.getTable() may not return null, " + + "but %s returned null for destination %s", + dynamicDestinations, + destination); + @Nullable TableReference cloneSource = + tempTable ? null : dynamicDestinations.getCloneSource(destination); TableSchema tableSchema; if (firstPaneCreateDisposition == CreateDisposition.CREATE_NEVER) { tableSchema = null; + } else if (cloneSource != null) { + tableSchema = null; } else if (jsonSchemas.containsKey(destination)) { // tableSchema for the destination stored in cache (jsonSchemas) tableSchema = @@ -220,13 +233,6 @@ public void processElement( jsonSchemas.put(destination, BigQueryHelpers.toJsonString(tableSchema)); } - TableDestination tableDestination = dynamicDestinations.getTable(destination); - checkArgument( - tableDestination != null, - "DynamicDestinations.getTable() may not return null, " - + "but %s returned null for destination %s", - dynamicDestinations, - destination); boolean destinationCoderSupportsClustering = !(dynamicDestinations.getDestinationCoder() instanceof TableDestinationCoderV2); checkArgument( @@ -281,6 +287,26 @@ public void processElement( createDisposition = CreateDisposition.CREATE_IF_NEEDED; } + if (cloneSource != null && createDisposition != CreateDisposition.CREATE_NEVER) { + Supplier<@Nullable TableSchema> schemaSupplier = + () -> dynamicDestinations.getSchema(destination); + Supplier<@Nullable TableConstraints> tableConstraintsSupplier = + () -> dynamicDestinations.getTableConstraints(destination); + Supplier<@Nullable TableReference> cloneSourceSupplier = () -> cloneSource; + + CreateTableHelpers.possiblyCreateTable( + c.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, + createDisposition, + dynamicDestinations.getDestinationCoder(), + kmsKey, + bqServices, + null); + } + BigQueryHelpers.PendingJob retryJob = startLoad( getJobService(c.getPipelineOptions().as(BigQueryOptions.class)), diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java index 4391a66e27ee..0482a8752e8f 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java @@ -92,6 +92,7 @@ public class BigQueryIOTranslationTest { WRITE_TRANSFORM_SCHEMA_MAPPING.put("getUseAvroLogicalTypes", "use_avro_logical_types"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getDynamicDestinations", "dynamic_destinations"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonSchema", "json_schema"); + WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonCloneSourceTableRef", "json_clone_source_table_ref"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonTimePartitioning", "json_time_partitioning"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonClustering", "clustering"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getCreateDisposition", "create_disposition"); @@ -255,6 +256,7 @@ public void testReCreateWriteTransformFromRowTable() { .withTriggeringFrequency(org.joda.time.Duration.millis(10000)) .withWriteDisposition(WriteDisposition.WRITE_TRUNCATE) .withCreateDisposition(CreateDisposition.CREATE_NEVER) + .withCloneFrom("dummyproject:dummydataset.clone_source_table") .withClustering(testClustering) .withKmsKey("dummykmskey"); @@ -273,6 +275,9 @@ public void testReCreateWriteTransformFromRowTable() { assertEquals(WriteDisposition.WRITE_TRUNCATE, writeTransformFromRow.getWriteDisposition()); assertEquals(CreateDisposition.CREATE_NEVER, writeTransformFromRow.getCreateDisposition()); assertEquals("dummykmskey", writeTransformFromRow.getKmsKey()); + assertEquals( + writeTransform.getJsonCloneSourceTableRef().get(), + writeTransformFromRow.getJsonCloneSourceTableRef().get()); assertEquals( BigQueryHelpers.toJsonString(testClustering), writeTransformFromRow.getJsonClustering().get()); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 62e57ec848cb..a3e7e223c790 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -46,6 +46,7 @@ import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.Job; import com.google.api.services.bigquery.model.JobConfigurationLoad; +import com.google.api.services.bigquery.model.JobConfigurationTableCopy; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableCell; import com.google.api.services.bigquery.model.TableDataInsertAllResponse; @@ -180,6 +181,7 @@ import org.apache.beam.sdk.transforms.windowing.WindowFn; import org.apache.beam.sdk.transforms.windowing.WindowMappingFn; import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; @@ -346,6 +348,309 @@ void writeEmptyPCollection() { p.run(); } + @Test + public void testWriteCreatesCloneTable() throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_table"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_table"); + TableRow baseRow = new TableRow().set("name", "base"); + TableRow appendedRow = new TableRow().set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withCloneFrom("project-id:dataset-id.source_table") + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null) + .map(JobConfigurationTableCopy::getOperationType) + .collect(Collectors.toList()), + hasItem("CLONE")); + } + + @Test + public void testWriteCreatesCloneTableWithUnqualifiedCloneSourceUsesDestinationProject() + throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.source_table"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.destination_table"); + TableRow baseRow = new TableRow().set("name", "base"); + TableRow appendedRow = new TableRow().set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to("bigquery-project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withCloneFrom("dataset-id.source_table") + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null && "CLONE".equals(copy.getOperationType())) + .flatMap(copy -> copy.getSourceTables().stream()) + .collect(Collectors.toList()), + hasItem(sourceTable)); + } + + @Test + public void testWriteCreatesCloneTableWithSchemaUpdateOptionsAndMultiplePartitions() + throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_table"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_table"); + TableRow baseRow = new TableRow().set("name", "base"); + List appendedRows = + IntStream.range(0, 30) + .mapToObj(i -> new TableRow().set("name", "appended-" + i)) + .collect(Collectors.toList()); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRows)) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withSchemaUpdateOptions(EnumSet.of(SchemaUpdateOption.ALLOW_FIELD_ADDITION)) + .withMaxBytesPerPartition(1) + .withMaxFilesPerPartition(1) + .withCloneFrom("project-id:dataset-id.source_table") + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + List expectedRows = new ArrayList<>(appendedRows); + expectedRows.add(baseRow); + assertThat( + fakeDatasetService.getAllRows(destinationTable), + containsInAnyOrder(expectedRows.toArray(new TableRow[0]))); + assertTrue( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .anyMatch(copy -> copy != null && !"CLONE".equals(copy.getOperationType()))); + } + + @Test + public void testCloneJobIdStableAcrossSerializationAndUniqueAcrossRuns() throws Exception { + BigQueryOptions first = PipelineOptionsFactory.create().as(BigQueryOptions.class); + first.setJobName("same-job-name"); + BigQueryOptions second = PipelineOptionsFactory.create().as(BigQueryOptions.class); + second.setJobName("same-job-name"); + TableReference source = BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_table"); + TableReference destination = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_table"); + + String firstPrefix = CreateTableHelpers.createCloneJobIdPrefix(first, source, destination); + String restoredPrefix = + CreateTableHelpers.createCloneJobIdPrefix( + PipelineOptionsTranslation.fromJson(PipelineOptionsTranslation.toJson(first)) + .as(BigQueryOptions.class), + source, + destination); + String secondPrefix = CreateTableHelpers.createCloneJobIdPrefix(second, source, destination); + + assertEquals(firstPrefix, restoredPrefix); + assertFalse(firstPrefix.equals(secondPrefix)); + } + + @Test + public void testMatchTableDynamicDestinationsUsesDestinationProjectForCloneSourceSchema() + throws Exception { + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.source_avro_table"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + + DynamicDestinations dynamicDestinations = + new DynamicDestinations() { + @Override + public String getDestination(@Nullable ValueInSingleWindow element) { + return "destination"; + } + + @Override + public Coder getDestinationCoder() { + return StringUtf8Coder.of(); + } + + @Override + public TableDestination getTable(String destination) { + return new TableDestination( + "bigquery-project-id:dataset-id.destination_avro_table", null); + } + + @Override + public @Nullable TableSchema getSchema(String destination) { + return null; + } + + @Override + public TableReference getCloneSource(String destination) { + return BigQueryHelpers.parseTableSpec("dataset-id.source_avro_table"); + } + }; + DynamicDestinations matchingDestinations = + DynamicDestinationsHelpers.matchTableDynamicDestinations( + dynamicDestinations, fakeBqServices); + + PCollection matchedSchema = + p.apply(Create.of("destination")) + .apply( + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(ProcessContext context) { + matchingDestinations.setSideInputAccessorFromProcessContext(context); + context.output( + BigQueryHelpers.toJsonString( + checkNotNull(matchingDestinations.getSchema(context.element())))); + } + })); + + PAssert.that(matchedSchema).containsInAnyOrder(BigQueryHelpers.toJsonString(schema)); + p.run(); + } + + @Test + public void testWriteCreatesCloneTableWithDynamicDestinations() throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("bucket").setType("STRING"), + new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_bucket_a"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_bucket_a"); + TableRow baseRow = new TableRow().set("bucket", "a").set("name", "base"); + TableRow appendedRow = new TableRow().set("bucket", "a").set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to( + new DynamicDestinations() { + @Override + public String getDestination(ValueInSingleWindow element) { + return element.getValue().get("bucket").toString(); + } + + @Override + public Coder getDestinationCoder() { + return StringUtf8Coder.of(); + } + + @Override + public TableDestination getTable(String destination) { + return new TableDestination( + "project-id:dataset-id.destination_bucket_" + destination, null); + } + + @Override + public @Nullable TableSchema getSchema(String destination) { + return null; + } + + @Override + public TableReference getCloneSource(String destination) { + return BigQueryHelpers.parseTableSpec( + "project-id:dataset-id.source_bucket_" + destination); + } + }) + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null) + .map(JobConfigurationTableCopy::getOperationType) + .collect(Collectors.toList()), + hasItem("CLONE")); + } + + @Test + public void testWriteRejectsSchemaAndCloneSource() { + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage( + "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); + + p.apply(Create.empty(TableRowJsonCoder.of())) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withSchema(new TableSchema()) + .withCloneFrom("project-id:dataset-id.source_table") + .withoutValidation()); + p.run(); + } + @Test public void testWriteEmptyPCollectionGroupFilesFileLoad() throws Exception { assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); @@ -2274,6 +2579,14 @@ public void testWriteBuilderMethods() { schema, BigQueryHelpers.fromJsonString( write.withSchema(schema).getJsonSchema().get(), TableSchema.class)); + assertEquals( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_table"), + BigQueryHelpers.fromJsonString( + write + .withCloneFrom("project-id:dataset-id.source_table") + .getJsonCloneSourceTableRef() + .get(), + TableReference.class)); } @Test