From ea4e475d133b89d26982faeb97d6f922aa545c4c Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Tue, 9 Jun 2026 23:09:48 -0500 Subject: [PATCH 1/6] Add BigQueryIO clone table support --- .../beam/sdk/io/gcp/bigquery/BatchLoads.java | 4 +- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 61 ++++- .../gcp/bigquery/BigQueryResourceNaming.java | 1 + .../io/gcp/bigquery/CreateTableHelpers.java | 251 +++++++++++++----- .../sdk/io/gcp/bigquery/CreateTables.java | 1 + .../io/gcp/bigquery/DynamicDestinations.java | 6 + .../bigquery/DynamicDestinationsHelpers.java | 57 +++- .../StorageApiWriteUnshardedRecords.java | 1 + .../StorageApiWritesShardedRecords.java | 1 + .../gcp/bigquery/UpdateSchemaDestination.java | 15 ++ .../beam/sdk/io/gcp/bigquery/WriteTables.java | 33 ++- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 138 ++++++++++ 12 files changed, 488 insertions(+), 81 deletions(-) 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..bbb71deb3fc3 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, 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..3e510b86c29e 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,6 +4167,16 @@ public WriteResult expand(PCollection input) { .collect(Collectors.toList())), "No more than one of jsonSchema, schemaFromView, or dynamicDestinations may be set"); + List allTableCreationMetadataArgs = + Lists.newArrayList(getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef()); + checkArgument( + 2 + > Iterables.size( + allTableCreationMetadataArgs.stream() + .filter(Predicates.notNull()::apply) + .collect(Collectors.toList())), + "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); + // Perform some argument checks BigQueryOptions bqOptions = input.getPipeline().getOptions().as(BigQueryOptions.class); Write.Method method = resolveMethod(input); @@ -4304,6 +4352,11 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { } dynamicDestinations = tableDestinations; } + if (getJsonCloneSourceTableRef() != null) { + dynamicDestinations = + new ConstantCloneSourceDestinations<>( + (DynamicDestinations) dynamicDestinations, getJsonCloneSourceTableRef()); + } return expandTyped(input, dynamicDestinations); } @@ -4319,6 +4372,7 @@ private WriteResult expandTyped( getJsonSchema() != null || getDynamicDestinations() != null || getSchemaFromView() != null; + boolean hasTableCreationMetadata = hasSchema || getJsonCloneSourceTableRef() != null; Class writeProtoClass = getWriteProtosClass(); if (getUseBeamSchema()) { @@ -4347,7 +4401,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 +4423,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 or 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/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..448f838393a9 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,22 +56,29 @@ 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. * - *

TODO: We should put a bound on memory usage of this. Use guava cache instead. + *

+ * TODO: We should put a bound on memory usage of this. Use guava cache instead. */ private static Set createdTables = Collections.newSetFromMap(new ConcurrentHashMap<>()); private static final Duration INITIAL_RPC_BACKOFF = Duration.millis(500); - private static final FluentBackoff DEFAULT_BACKOFF_FACTORY = - FluentBackoff.DEFAULT.withMaxRetries(4).withInitialBackoff(INITIAL_RPC_BACKOFF); + 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) @@ -78,15 +91,14 @@ static void createTableWrapper(Callable action, Callable tryCreat return; } catch (Exception e) { // The Storage Write library can wrap errors in UncheckedExecutionException - Optional handledCause = - Throwables.getCausalChain(e).stream() - .filter( - cause -> - (cause instanceof ApiException || cause instanceof StatusRuntimeException)) - .findAny(); + Optional handledCause = Throwables.getCausalChain(e).stream() + .filter( + cause -> (cause instanceof ApiException || cause instanceof StatusRuntimeException)) + .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 +115,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, @@ -114,8 +127,7 @@ static TableDestination possiblyCreateTable( + "with a non-null table spec, but %s " + "has a null table spec", tableDestination); - boolean destinationCoderSupportsClustering = - !(tableDestinationCoder instanceof TableDestinationCoderV2); + boolean destinationCoderSupportsClustering = !(tableDestinationCoder instanceof TableDestinationCoderV2); checkArgument( tableDestination.getClustering() == null || destinationCoderSupportsClustering, "DynamicDestinations.getTable() may only return destinations with clustering configured" @@ -136,7 +148,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 +157,7 @@ static TableDestination possiblyCreateTable( bigQueryOptions, schemaSupplier, tableConstraintsSupplier, + cloneSourceSupplier, tableDestination, createDisposition, tableSpec, @@ -160,6 +174,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, @@ -170,62 +185,71 @@ private static void tryCreateTable( tableReference.setTableId(BigQueryHelpers.stripPartitionDecorator(tableReference.getTableId())); try (DatasetService datasetService = bqServices.getDatasetService(options)) { 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); - } + tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC) == null) { + @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 +257,103 @@ 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 source = withDefaultProject(options, cloneSource); + JobConfigurationTableCopy copyConfig = new JobConfigurationTableCopy() + .setSourceTables(Collections.singletonList(source)) + .setDestinationTable(destination) + .setOperationType("CLONE") + .setWriteDisposition(WriteDisposition.WRITE_EMPTY.name()) + .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED.name()); + if (kmsKey != null) { + copyConfig.setDestinationEncryptionConfiguration( + new EncryptionConfiguration().setKmsKeyName(kmsKey)); + } + + String jobProjectId = destination.getProjectId(); + String bqLocation = BigQueryHelpers.getDatasetLocation( + datasetService, destination.getProjectId(), destination.getDatasetId()); + String jobIdPrefix = createCloneJobIdPrefix(options, source, destination); + try (JobService jobService = bqServices.getJobService(options)) { + BigQueryHelpers.PendingJob cloneJob = new BigQueryHelpers.PendingJob( + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(jobProjectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + jobService.startCopyJob(jobRef, copyConfig); + } catch (IOException | InterruptedException e) { + throw new RuntimeException(e); + } + return null; + }, + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(jobProjectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }, + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(jobProjectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.getJob(jobRef); + } catch (IOException | InterruptedException e) { + 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 TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference) { + TableReference updated = tableReference.clone(); + if (Strings.isNullOrEmpty(updated.getProjectId())) { + updated.setProjectId( + options.getBigQueryProject() == null + ? options.getProject() + : options.getBigQueryProject()); + } + return updated; + } + + private 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)) + .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..24cdc18c4034 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 @@ -129,6 +129,7 @@ public void processElement(ProcessContext context) { tableDestination1, schemaSupplier, tableConstraintsSupplier, + () -> dynamicDestinations.getCloneSource(dest), 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..3abedc1b39a9 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 @@ -183,6 +183,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 +243,34 @@ public String toString() { } } + /** Returns the same clone source for every table. */ + static class ConstantCloneSourceDestinations + extends DelegatingDynamicDestinations { + private final ValueProvider jsonCloneSource; + + ConstantCloneSourceDestinations( + DynamicDestinations inner, ValueProvider jsonCloneSource) { + super(inner); + Preconditions.checkArgumentNotNull(jsonCloneSource, "jsonCloneSource can not be null"); + this.jsonCloneSource = jsonCloneSource; + } + + @Override + public TableReference getCloneSource(DestinationT destination) { + String jsonCloneSource = this.jsonCloneSource.get(); + checkArgument(jsonCloneSource != null, "jsonCloneSource can not be null"); + return BigQueryHelpers.fromJsonString(jsonCloneSource, TableReference.class); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("inner", inner) + .add("jsonCloneSource", jsonCloneSource) + .toString(); + } + } + static class ConstantTableConstraintsDestinations extends DelegatingDynamicDestinations { private final String jsonTableConstraints; @@ -495,15 +528,22 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) public TableDestination getTable(DestinationT destination) { TableDestination wrappedDestination = super.getTable(destination); Table existingTable = getBigQueryTable(wrappedDestination.getTableReference()); + Table tableToMatch = existingTable; + if (tableToMatch == null) { + @Nullable TableReference cloneSource = super.getCloneSource(destination); + if (cloneSource != null) { + tableToMatch = getBigQueryTable(cloneSource); + } + } - 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()); } } @@ -518,6 +558,15 @@ public TableDestination getTable(DestinationT destination) { if (existingTable == null || existingTable.getSchema() == null || existingTable.getSchema().isEmpty()) { + @Nullable TableReference cloneSource = super.getCloneSource(destination); + if (cloneSource != null) { + @Nullable Table cloneSourceTable = getBigQueryTable(cloneSource); + if (cloneSourceTable != null + && cloneSourceTable.getSchema() != null + && !cloneSourceTable.getSchema().isEmpty()) { + return cloneSourceTable.getSchema(); + } + } 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..b03425b39646 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 @@ -149,6 +149,21 @@ 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) { + CreateTableHelpers.possiblyCreateTable( + context.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + () -> dynamicDestinations.getSchema(destination), + () -> dynamicDestinations.getTableConstraints(destination), + () -> cloneSource, + createDisposition, + dynamicDestinations.getDestinationCoder(), + kmsKey, + bqServices, + null); + } if (pendingJobs.containsKey(destination)) { // zero load job for this destination is already set continue; 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..84291e984e18 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 @@ -198,9 +198,21 @@ 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 +232,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 +286,20 @@ public void processElement( createDisposition = CreateDisposition.CREATE_IF_NEEDED; } + if (cloneSource != null && createDisposition != CreateDisposition.CREATE_NEVER) { + CreateTableHelpers.possiblyCreateTable( + c.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + () -> dynamicDestinations.getSchema(destination), + () -> dynamicDestinations.getTableConstraints(destination), + () -> cloneSource, + 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/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 62e57ec848cb..c46497ad367e 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; @@ -346,6 +347,135 @@ 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 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 +2404,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 From c38a9113fd9a41b63db52f0abd428f41061a796e Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Wed, 10 Jun 2026 08:22:22 -0500 Subject: [PATCH 2/6] Fix BigQueryIO clone validation message --- .../java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 3e510b86c29e..a6c37895367a 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 @@ -4425,7 +4425,7 @@ private WriteResult expandTyped( checkArgument( getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED || hasTableCreationMetadata, - "CreateDisposition is CREATE_IF_NEEDED, however no schema or clone source was provided."); + "CreateDisposition is CREATE_IF_NEEDED, however no schema was provided and no clone source was provided."); } Coder destinationCoder; From 474c6d0379e955665256b8df662fdb4acdb33dcd Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Wed, 10 Jun 2026 12:57:56 -0500 Subject: [PATCH 3/6] Track BigQuery clone source in write transform translation --- .../sdk/io/gcp/bigquery/BigQueryIOTranslation.java | 13 +++++++++++++ .../io/gcp/bigquery/BigQueryIOTranslationTest.java | 5 +++++ 2 files changed, 18 insertions(+) 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..edbbcb1fe503 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) { ValueProvider jsonTimePartitioning = transform.getJsonTimePartitioning(); if (jsonTimePartitioning != null) { fieldValues.put("json_time_partitioning", toByteArray(jsonTimePartitioning.get())); + if (transform.getJsonCloneSourceTableRef() != null) { + fieldValues.put( + "json_clone_source_table_ref", transform.getJsonCloneSourceTableRef().get()); + } + if (transform.getJsonTimePartitioning() != null) { + fieldValues.put( + "json_time_partitioning", toByteArray(transform.getJsonTimePartitioning().get())); } ValueProvider jsonClustering = transform.getJsonClustering(); if (jsonClustering != null) { @@ -736,6 +744,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/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()); From 9974a4bbf3f524ce7d60059b9f10c08b8e0f7a12 Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Fri, 12 Jun 2026 19:22:21 -0500 Subject: [PATCH 4/6] updates based on reviews --- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 13 +- .../io/gcp/bigquery/CreateTableHelpers.java | 128 +++++---- .../sdk/io/gcp/bigquery/CreateTables.java | 5 +- .../bigquery/DynamicDestinationsHelpers.java | 63 +++-- .../gcp/bigquery/UpdateSchemaDestination.java | 249 +++++++++--------- .../beam/sdk/io/gcp/bigquery/WriteTables.java | 14 +- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 103 ++++++++ 7 files changed, 373 insertions(+), 202 deletions(-) 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 a6c37895367a..2ab069dba426 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 @@ -4167,14 +4167,13 @@ public WriteResult expand(PCollection input) { .collect(Collectors.toList())), "No more than one of jsonSchema, schemaFromView, or dynamicDestinations may be set"); - List allTableCreationMetadataArgs = - Lists.newArrayList(getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef()); + long activeCreationMetadataCount = + java.util.stream.Stream.of( + getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef()) + .filter(arg -> arg != null) + .count(); checkArgument( - 2 - > Iterables.size( - allTableCreationMetadataArgs.stream() - .filter(Predicates.notNull()::apply) - .collect(Collectors.toList())), + activeCreationMetadataCount <= 1, "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); // Perform some argument checks 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 448f838393a9..819e23397b5e 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 @@ -266,60 +266,68 @@ private static void cloneTable( @Nullable String kmsKey, String tableSpec) throws Exception { - TableReference source = withDefaultProject(options, cloneSource); - JobConfigurationTableCopy copyConfig = new JobConfigurationTableCopy() - .setSourceTables(Collections.singletonList(source)) - .setDestinationTable(destination) - .setOperationType("CLONE") - .setWriteDisposition(WriteDisposition.WRITE_EMPTY.name()) - .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED.name()); + 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 = destination.getProjectId(); - String bqLocation = BigQueryHelpers.getDatasetLocation( - datasetService, destination.getProjectId(), destination.getDatasetId()); - String jobIdPrefix = createCloneJobIdPrefix(options, source, destination); + 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 = new JobReference() - .setProjectId(jobProjectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - jobService.startCopyJob(jobRef, copyConfig); - } catch (IOException | InterruptedException e) { - throw new RuntimeException(e); - } - return null; - }, - jobId -> { - JobReference jobRef = new JobReference() - .setProjectId(jobProjectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - }, - jobId -> { - JobReference jobRef = new JobReference() - .setProjectId(jobProjectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.getJob(jobRef); - } catch (IOException | InterruptedException e) { - throw new RuntimeException(e); - } - }, - BatchLoads.DEFAULT_MAX_RETRY_JOBS, - jobIdPrefix); + 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( @@ -332,14 +340,30 @@ private static void cloneTable( } } + 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())) { - updated.setProjectId( - options.getBigQueryProject() == null - ? options.getProject() - : options.getBigQueryProject()); + 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; } 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 24cdc18c4034..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,7 +132,7 @@ public void processElement(ProcessContext context) { tableDestination1, schemaSupplier, tableConstraintsSupplier, - () -> dynamicDestinations.getCloneSource(dest), + cloneSourceSupplier, createDisposition, dynamicDestinations.getDestinationCoder(), kmsKey, 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 3abedc1b39a9..67775ce12398 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; @@ -247,19 +248,25 @@ public String toString() { 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 can not be null"); + Preconditions.checkArgumentNotNull(jsonCloneSource, "jsonCloneSource cannot be null"); this.jsonCloneSource = jsonCloneSource; } @Override public TableReference getCloneSource(DestinationT destination) { - String jsonCloneSource = this.jsonCloneSource.get(); - checkArgument(jsonCloneSource != null, "jsonCloneSource can not be null"); - return BigQueryHelpers.fromJsonString(jsonCloneSource, TableReference.class); + @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 @@ -477,6 +484,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 @@ -492,16 +504,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); @@ -513,6 +521,25 @@ 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; + } + /** Identical to {@link BackOffUtils#next} but without checked IOException. */ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws InterruptedException { @@ -527,12 +554,14 @@ 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); + tableToMatch = getBigQueryTable(cloneSource, destinationProjectId); } } @@ -554,13 +583,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 String destinationProjectId = tableReference.getProjectId(); + @Nullable Table existingTable = getBigQueryTable(tableReference); if (existingTable == null || existingTable.getSchema() == null || existingTable.getSchema().isEmpty()) { @Nullable TableReference cloneSource = super.getCloneSource(destination); if (cloneSource != null) { - @Nullable Table cloneSourceTable = getBigQueryTable(cloneSource); + @Nullable Table cloneSourceTable = getBigQueryTable(cloneSource, destinationProjectId); if (cloneSourceTable != null && cloneSourceTable.getSchema() != null && !cloneSourceTable.getSchema().isEmpty()) { 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 b03425b39646..e8ae7e621fb1 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,29 +40,37 @@ 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; /** * Update destination schema based on data that is about to be copied into it. * - *

Unlike load and query jobs, BigQuery copy jobs do not support schema field addition or - * relaxation on the destination table. This DoFn fills that gap by updating the destination table - * schemas to be compatible with the data coming from the source table so that schemaUpdateOptions - * are respected regardless of whether data is loaded directly to the destination table or loaded + *

+ * Unlike load and query jobs, BigQuery copy jobs do not support schema field + * addition or + * relaxation on the destination table. This DoFn fills that gap by updating the + * destination table + * schemas to be compatible with the data coming from the source table so that + * schemaUpdateOptions + * are respected regardless of whether data is loaded directly to the + * destination table or loaded * into temporary tables before being copied into the destination. * - *

This transform takes as input a list of KV(destination, WriteTables.Result) and emits a list - * of KV(TableDestination, WriteTables.Result) where the destination label is parsed and replaced to + *

+ * This transform takes as input a list of KV(destination, WriteTables.Result) + * and emits a list + * of KV(TableDestination, WriteTables.Result) where the destination label is + * parsed and replaced to * TableDestination objects. */ public class UpdateSchemaDestination - extends DoFn< - Iterable>, - Iterable>> { + extends DoFn>, Iterable>> { private static final Logger LOG = LoggerFactory.getLogger(UpdateSchemaDestination.class); private final BigQueryServices bqServices; @@ -71,7 +79,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,15 +157,21 @@ public void processElement( DestinationT destination = entry.getKey(); TableDestination tableDestination = getTableWithDefaultProject(destination); outputs.add(KV.of(tableDestination, entry.getValue())); - @Nullable TableReference cloneSource = dynamicDestinations.getCloneSource(destination); + @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, - () -> dynamicDestinations.getSchema(destination), - () -> dynamicDestinations.getTableConstraints(destination), - () -> cloneSource, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, createDisposition, dynamicDestinations.getDestinationCoder(), kmsKey, @@ -170,24 +184,22 @@ public void processElement( } TableSchema schema = dynamicDestinations.getSchema(destination); TableReference tableReference = tableDestination.getTableReference(); - String jobIdPrefix = - BigQueryResourceNaming.createJobIdWithDestination( - context.sideInput(zeroLoadJobIdPrefixView), - tableDestination, - 1, - context.pane().getIndex()); - BigQueryHelpers.PendingJob updateSchemaDestinationJob = - startZeroLoadJob( - getJobService(context.getPipelineOptions().as(BigQueryOptions.class)), - getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)), - jobIdPrefix, - tableReference, - tableDestination.getTimePartitioning(), - tableDestination.getClustering(), - schema, - writeDisposition, - createDisposition, - schemaUpdateOptions); + String jobIdPrefix = BigQueryResourceNaming.createJobIdWithDestination( + context.sideInput(zeroLoadJobIdPrefixView), + tableDestination, + 1, + context.pane().getIndex()); + BigQueryHelpers.PendingJob updateSchemaDestinationJob = startZeroLoadJob( + getJobService(context.getPipelineOptions().as(BigQueryOptions.class)), + getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)), + jobIdPrefix, + tableReference, + tableDestination.getTimePartitioning(), + tableDestination.getClustering(), + schema, + writeDisposition, + createDisposition, + schemaUpdateOptions); if (updateSchemaDestinationJob != null) { pendingJobs.put( destination, new PendingJobData(updateSchemaDestinationJob, tableDestination)); @@ -220,27 +232,25 @@ public void onTeardown() { @FinishBundle public void finishBundle(FinishBundleContext context) throws Exception { - DatasetService datasetService = - getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)); + DatasetService datasetService = getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)); BigQueryHelpers.PendingJobManager jobManager = new BigQueryHelpers.PendingJobManager(); for (final PendingJobData pendingJobData : pendingJobs.values()) { - jobManager = - jobManager.addPendingJob( - pendingJobData.retryJob, - j -> { - try { - if (pendingJobData.tableDestination.getTableDescription() != null) { - TableReference ref = pendingJobData.tableDestination.getTableReference(); - datasetService.patchTableDescription( - ref.clone() - .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), - pendingJobData.tableDestination.getTableDescription()); - } - } catch (IOException | InterruptedException e) { - return e; - } - return null; - }); + jobManager = jobManager.addPendingJob( + pendingJobData.retryJob, + j -> { + try { + if (pendingJobData.tableDestination.getTableDescription() != null) { + TableReference ref = pendingJobData.tableDestination.getTableReference(); + datasetService.patchTableDescription( + ref.clone() + .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), + pendingJobData.tableDestination.getTableDescription()); + } + } catch (IOException | InterruptedException e) { + return e; + } + return null; + }); } jobManager.waitForDone(); } @@ -256,22 +266,20 @@ public void finishBundle(FinishBundleContext context) throws Exception { BigQueryIO.Write.WriteDisposition writeDisposition, BigQueryIO.Write.CreateDisposition createDisposition, Set schemaUpdateOptions) { - JobConfigurationLoad loadConfig = - new JobConfigurationLoad() - .setDestinationTable(tableReference) - .setWriteDisposition(writeDisposition.name()) - .setCreateDisposition(createDisposition.name()) - .setSourceFormat("NEWLINE_DELIMITED_JSON"); + JobConfigurationLoad loadConfig = new JobConfigurationLoad() + .setDestinationTable(tableReference) + .setWriteDisposition(writeDisposition.name()) + .setCreateDisposition(createDisposition.name()) + .setSourceFormat("NEWLINE_DELIMITED_JSON"); if (schemaUpdateOptions != null) { - List options = - schemaUpdateOptions.stream() - .map(BigQueryIO.Write.SchemaUpdateOption::name) - .collect(Collectors.toList()); + List options = schemaUpdateOptions.stream() + .map(BigQueryIO.Write.SchemaUpdateOption::name) + .collect(Collectors.toList()); loadConfig.setSchemaUpdateOptions(options); } if (!loadConfig - .getWriteDisposition() - .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString()) + .getWriteDisposition() + .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString()) && !loadConfig .getWriteDisposition() .equals(BigQueryIO.Write.WriteDisposition.WRITE_APPEND.toString())) { @@ -287,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(); @@ -310,65 +319,59 @@ public void finishBundle(FinishBundleContext context) throws Exception { loadConfig.setDestinationEncryptionConfiguration( new EncryptionConfiguration().setKmsKeyName(kmsKey)); } - String projectId = - loadJobProjectId == null || loadJobProjectId.get() == null - ? tableReference.getProjectId() - : loadJobProjectId.get(); - String bqLocation = - BigQueryHelpers.getDatasetLocation( - datasetService, tableReference.getProjectId(), tableReference.getDatasetId()); + String projectId = loadJobProjectId == null || loadJobProjectId.get() == null + ? tableReference.getProjectId() + : loadJobProjectId.get(); + String bqLocation = BigQueryHelpers.getDatasetLocation( + datasetService, tableReference.getProjectId(), tableReference.getDatasetId()); - BigQueryHelpers.PendingJob retryJob = - new BigQueryHelpers.PendingJob( - // Function to load the data. - jobId -> { - JobReference jobRef = - new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - LOG.info( - "Loading zero rows using job {}, job id {} iteration {}", - tableReference, - jobRef, - jobId.getRetryIndex()); - try { - jobService.startLoadJob( - jobRef, loadConfig, new ByteArrayContent("text/plain", new byte[0])); - } catch (IOException | InterruptedException e) { - LOG.warn("Schema update load job {} failed with {}", jobRef, e.toString()); - throw new RuntimeException(e); - } - return null; - }, - // Function to poll the result of a load job. - jobId -> { - JobReference jobRef = - new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - }, - // Function to lookup a job. - jobId -> { - JobReference jobRef = - new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.getJob(jobRef); - } catch (InterruptedException | IOException e) { - throw new RuntimeException(e); - } - }, - maxRetryJobs, - jobIdPrefix); + BigQueryHelpers.PendingJob retryJob = new BigQueryHelpers.PendingJob( + // Function to load the data. + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + LOG.info( + "Loading zero rows using job {}, job id {} iteration {}", + tableReference, + jobRef, + jobId.getRetryIndex()); + try { + jobService.startLoadJob( + jobRef, loadConfig, new ByteArrayContent("text/plain", new byte[0])); + } catch (IOException | InterruptedException e) { + LOG.warn("Schema update load job {} failed with {}", jobRef, e.toString()); + throw new RuntimeException(e); + } + return null; + }, + // Function to poll the result of a load job. + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }, + // Function to lookup a job. + jobId -> { + JobReference jobRef = new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.getJob(jobRef); + } catch (InterruptedException | IOException e) { + throw new RuntimeException(e); + } + }, + maxRetryJobs, + jobIdPrefix); return retryJob; } 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 84291e984e18..970f9d5848ed 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; @@ -287,12 +289,18 @@ public void processElement( } 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, - () -> dynamicDestinations.getSchema(destination), - () -> dynamicDestinations.getTableConstraints(destination), - () -> cloneSource, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, createDisposition, dynamicDestinations.getDestinationCoder(), kmsKey, 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 c46497ad367e..f69ac688f24f 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 @@ -388,6 +388,109 @@ public void testWriteCreatesCloneTable() throws Exception { 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 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); From 3e36ed7095ad31e57c6222550cc98689b6c4aeb7 Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Tue, 27 Oct 2026 13:22:20 -0500 Subject: [PATCH 5/6] Fix BigQuery clone source translation --- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 5 +- .../gcp/bigquery/BigQueryIOTranslation.java | 11 +- .../io/gcp/bigquery/CreateTableHelpers.java | 38 +-- .../gcp/bigquery/UpdateSchemaDestination.java | 238 +++++++++--------- .../beam/sdk/io/gcp/bigquery/WriteTables.java | 3 +- 5 files changed, 150 insertions(+), 145 deletions(-) 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 2ab069dba426..4db9d237ec99 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 @@ -4351,10 +4351,11 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { } dynamicDestinations = tableDestinations; } - if (getJsonCloneSourceTableRef() != null) { + ValueProvider jsonCloneSourceTableRef = getJsonCloneSourceTableRef(); + if (jsonCloneSourceTableRef != null) { dynamicDestinations = new ConstantCloneSourceDestinations<>( - (DynamicDestinations) dynamicDestinations, getJsonCloneSourceTableRef()); + (DynamicDestinations) dynamicDestinations, jsonCloneSourceTableRef); } return expandTyped(input, dynamicDestinations); } 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 edbbcb1fe503..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 @@ -550,16 +550,13 @@ public Row toConfigRow(Write 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())); - if (transform.getJsonCloneSourceTableRef() != null) { - fieldValues.put( - "json_clone_source_table_ref", transform.getJsonCloneSourceTableRef().get()); - } - if (transform.getJsonTimePartitioning() != null) { - fieldValues.put( - "json_time_partitioning", toByteArray(transform.getJsonTimePartitioning().get())); } ValueProvider jsonClustering = transform.getJsonClustering(); if (jsonClustering != null) { 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 819e23397b5e..7bab5186b313 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 @@ -68,14 +68,13 @@ public class CreateTableHelpers { /** * The list of tables created so far, so we don't try the creation each time. * - *

- * TODO: We should put a bound on memory usage of this. Use guava cache instead. + *

TODO: We should put a bound on memory usage of this. Use guava cache instead. */ private static Set createdTables = Collections.newSetFromMap(new ConcurrentHashMap<>()); private static final Duration INITIAL_RPC_BACKOFF = Duration.millis(500); - private static final FluentBackoff DEFAULT_BACKOFF_FACTORY = FluentBackoff.DEFAULT.withMaxRetries(4) - .withInitialBackoff(INITIAL_RPC_BACKOFF); + 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 @@ -91,10 +90,12 @@ static void createTableWrapper(Callable action, Callable tryCreat return; } catch (Exception e) { // The Storage Write library can wrap errors in UncheckedExecutionException - Optional handledCause = Throwables.getCausalChain(e).stream() - .filter( - cause -> (cause instanceof ApiException || cause instanceof StatusRuntimeException)) - .findAny(); + Optional handledCause = + Throwables.getCausalChain(e).stream() + .filter( + cause -> + (cause instanceof ApiException || cause instanceof StatusRuntimeException)) + .findAny(); lastException = (Exception) handledCause.orElseThrow(() -> e); // TODO: Once BigQuery reliably returns a consistent error on table not found, @@ -127,7 +128,8 @@ static TableDestination possiblyCreateTable( + "with a non-null table spec, but %s " + "has a null table spec", tableDestination); - boolean destinationCoderSupportsClustering = !(tableDestinationCoder instanceof TableDestinationCoderV2); + boolean destinationCoderSupportsClustering = + !(tableDestinationCoder instanceof TableDestinationCoderV2); checkArgument( tableDestination.getClustering() == null || destinationCoderSupportsClustering, "DynamicDestinations.getTable() may only return destinations with clustering configured" @@ -185,16 +187,15 @@ private static void tryCreateTable( tableReference.setTableId(BigQueryHelpers.stripPartitionDecorator(tableReference.getTableId())); try (DatasetService datasetService = bqServices.getDatasetService(options)) { if (datasetService.getTable( - tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC) == null) { - @Nullable - TableReference cloneSource = cloneSourceSupplier.get(); + tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC) + == null) { + @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(); + @Nullable TableConstraints tableConstraints = tableConstraintsSupplier.get(); Preconditions.checkArgumentNotNull( tableSchema, "Unless create disposition is %s, a schema or clone source must be specified, i.e. " @@ -371,10 +372,11 @@ private static TableReference withDefaultProject( private 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)) - .toString(); + String stepToken = + Hashing.murmur3_128() + .hashUnencodedChars( + BigQueryHelpers.toJsonString(source) + BigQueryHelpers.toJsonString(destination)) + .toString(); return BigQueryResourceNaming.createJobIdPrefix(jobName, stepToken, JobType.CLONE); } 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 e8ae7e621fb1..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 @@ -51,26 +51,20 @@ /** * Update destination schema based on data that is about to be copied into it. * - *

- * Unlike load and query jobs, BigQuery copy jobs do not support schema field - * addition or - * relaxation on the destination table. This DoFn fills that gap by updating the - * destination table - * schemas to be compatible with the data coming from the source table so that - * schemaUpdateOptions - * are respected regardless of whether data is loaded directly to the - * destination table or loaded + *

Unlike load and query jobs, BigQuery copy jobs do not support schema field addition or + * relaxation on the destination table. This DoFn fills that gap by updating the destination table + * schemas to be compatible with the data coming from the source table so that schemaUpdateOptions + * are respected regardless of whether data is loaded directly to the destination table or loaded * into temporary tables before being copied into the destination. * - *

- * This transform takes as input a list of KV(destination, WriteTables.Result) - * and emits a list - * of KV(TableDestination, WriteTables.Result) where the destination label is - * parsed and replaced to + *

This transform takes as input a list of KV(destination, WriteTables.Result) and emits a list + * of KV(TableDestination, WriteTables.Result) where the destination label is parsed and replaced to * TableDestination objects. */ public class UpdateSchemaDestination - extends DoFn>, Iterable>> { + extends DoFn< + Iterable>, + Iterable>> { private static final Logger LOG = LoggerFactory.getLogger(UpdateSchemaDestination.class); private final BigQueryServices bqServices; @@ -157,13 +151,13 @@ public void processElement( DestinationT destination = entry.getKey(); TableDestination tableDestination = getTableWithDefaultProject(destination); outputs.add(KV.of(tableDestination, entry.getValue())); - @Nullable - TableReference cloneSource = dynamicDestinations.getCloneSource(destination); + @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 TableSchema> schemaSupplier = + () -> dynamicDestinations.getSchema(destination); + Supplier<@Nullable TableConstraints> tableConstraintsSupplier = + () -> dynamicDestinations.getTableConstraints(destination); Supplier<@Nullable TableReference> cloneSourceSupplier = () -> cloneSource; CreateTableHelpers.possiblyCreateTable( @@ -184,22 +178,24 @@ public void processElement( } TableSchema schema = dynamicDestinations.getSchema(destination); TableReference tableReference = tableDestination.getTableReference(); - String jobIdPrefix = BigQueryResourceNaming.createJobIdWithDestination( - context.sideInput(zeroLoadJobIdPrefixView), - tableDestination, - 1, - context.pane().getIndex()); - BigQueryHelpers.PendingJob updateSchemaDestinationJob = startZeroLoadJob( - getJobService(context.getPipelineOptions().as(BigQueryOptions.class)), - getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)), - jobIdPrefix, - tableReference, - tableDestination.getTimePartitioning(), - tableDestination.getClustering(), - schema, - writeDisposition, - createDisposition, - schemaUpdateOptions); + String jobIdPrefix = + BigQueryResourceNaming.createJobIdWithDestination( + context.sideInput(zeroLoadJobIdPrefixView), + tableDestination, + 1, + context.pane().getIndex()); + BigQueryHelpers.PendingJob updateSchemaDestinationJob = + startZeroLoadJob( + getJobService(context.getPipelineOptions().as(BigQueryOptions.class)), + getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)), + jobIdPrefix, + tableReference, + tableDestination.getTimePartitioning(), + tableDestination.getClustering(), + schema, + writeDisposition, + createDisposition, + schemaUpdateOptions); if (updateSchemaDestinationJob != null) { pendingJobs.put( destination, new PendingJobData(updateSchemaDestinationJob, tableDestination)); @@ -232,30 +228,32 @@ public void onTeardown() { @FinishBundle public void finishBundle(FinishBundleContext context) throws Exception { - DatasetService datasetService = getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)); + DatasetService datasetService = + getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class)); BigQueryHelpers.PendingJobManager jobManager = new BigQueryHelpers.PendingJobManager(); for (final PendingJobData pendingJobData : pendingJobs.values()) { - jobManager = jobManager.addPendingJob( - pendingJobData.retryJob, - j -> { - try { - if (pendingJobData.tableDestination.getTableDescription() != null) { - TableReference ref = pendingJobData.tableDestination.getTableReference(); - datasetService.patchTableDescription( - ref.clone() - .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), - pendingJobData.tableDestination.getTableDescription()); - } - } catch (IOException | InterruptedException e) { - return e; - } - return null; - }); + jobManager = + jobManager.addPendingJob( + pendingJobData.retryJob, + j -> { + try { + if (pendingJobData.tableDestination.getTableDescription() != null) { + TableReference ref = pendingJobData.tableDestination.getTableReference(); + datasetService.patchTableDescription( + ref.clone() + .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), + pendingJobData.tableDestination.getTableDescription()); + } + } catch (IOException | InterruptedException e) { + return e; + } + return null; + }); } jobManager.waitForDone(); } - private @Nullable BigQueryHelpers.PendingJob startZeroLoadJob( + private BigQueryHelpers.@Nullable PendingJob startZeroLoadJob( BigQueryServices.JobService jobService, DatasetService datasetService, String jobIdPrefix, @@ -266,20 +264,22 @@ public void finishBundle(FinishBundleContext context) throws Exception { BigQueryIO.Write.WriteDisposition writeDisposition, BigQueryIO.Write.CreateDisposition createDisposition, Set schemaUpdateOptions) { - JobConfigurationLoad loadConfig = new JobConfigurationLoad() - .setDestinationTable(tableReference) - .setWriteDisposition(writeDisposition.name()) - .setCreateDisposition(createDisposition.name()) - .setSourceFormat("NEWLINE_DELIMITED_JSON"); + JobConfigurationLoad loadConfig = + new JobConfigurationLoad() + .setDestinationTable(tableReference) + .setWriteDisposition(writeDisposition.name()) + .setCreateDisposition(createDisposition.name()) + .setSourceFormat("NEWLINE_DELIMITED_JSON"); if (schemaUpdateOptions != null) { - List options = schemaUpdateOptions.stream() - .map(BigQueryIO.Write.SchemaUpdateOption::name) - .collect(Collectors.toList()); + List options = + schemaUpdateOptions.stream() + .map(BigQueryIO.Write.SchemaUpdateOption::name) + .collect(Collectors.toList()); loadConfig.setSchemaUpdateOptions(options); } if (!loadConfig - .getWriteDisposition() - .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString()) + .getWriteDisposition() + .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString()) && !loadConfig .getWriteDisposition() .equals(BigQueryIO.Write.WriteDisposition.WRITE_APPEND.toString())) { @@ -319,59 +319,65 @@ public void finishBundle(FinishBundleContext context) throws Exception { loadConfig.setDestinationEncryptionConfiguration( new EncryptionConfiguration().setKmsKeyName(kmsKey)); } - String projectId = loadJobProjectId == null || loadJobProjectId.get() == null - ? tableReference.getProjectId() - : loadJobProjectId.get(); - String bqLocation = BigQueryHelpers.getDatasetLocation( - datasetService, tableReference.getProjectId(), tableReference.getDatasetId()); + String projectId = + loadJobProjectId == null || loadJobProjectId.get() == null + ? tableReference.getProjectId() + : loadJobProjectId.get(); + String bqLocation = + BigQueryHelpers.getDatasetLocation( + datasetService, tableReference.getProjectId(), tableReference.getDatasetId()); - BigQueryHelpers.PendingJob retryJob = new BigQueryHelpers.PendingJob( - // Function to load the data. - jobId -> { - JobReference jobRef = new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - LOG.info( - "Loading zero rows using job {}, job id {} iteration {}", - tableReference, - jobRef, - jobId.getRetryIndex()); - try { - jobService.startLoadJob( - jobRef, loadConfig, new ByteArrayContent("text/plain", new byte[0])); - } catch (IOException | InterruptedException e) { - LOG.warn("Schema update load job {} failed with {}", jobRef, e.toString()); - throw new RuntimeException(e); - } - return null; - }, - // Function to poll the result of a load job. - jobId -> { - JobReference jobRef = new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - }, - // Function to lookup a job. - jobId -> { - JobReference jobRef = new JobReference() - .setProjectId(projectId) - .setJobId(jobId.getJobId()) - .setLocation(bqLocation); - try { - return jobService.getJob(jobRef); - } catch (InterruptedException | IOException e) { - throw new RuntimeException(e); - } - }, - maxRetryJobs, - jobIdPrefix); + BigQueryHelpers.PendingJob retryJob = + new BigQueryHelpers.PendingJob( + // Function to load the data. + jobId -> { + JobReference jobRef = + new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + LOG.info( + "Loading zero rows using job {}, job id {} iteration {}", + tableReference, + jobRef, + jobId.getRetryIndex()); + try { + jobService.startLoadJob( + jobRef, loadConfig, new ByteArrayContent("text/plain", new byte[0])); + } catch (IOException | InterruptedException e) { + LOG.warn("Schema update load job {} failed with {}", jobRef, e.toString()); + throw new RuntimeException(e); + } + return null; + }, + // Function to poll the result of a load job. + jobId -> { + JobReference jobRef = + new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + }, + // Function to lookup a job. + jobId -> { + JobReference jobRef = + new JobReference() + .setProjectId(projectId) + .setJobId(jobId.getJobId()) + .setLocation(bqLocation); + try { + return jobService.getJob(jobRef); + } catch (InterruptedException | IOException e) { + throw new RuntimeException(e); + } + }, + maxRetryJobs, + jobIdPrefix); return retryJob; } 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 970f9d5848ed..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 @@ -207,8 +207,7 @@ public void processElement( + "but %s returned null for destination %s", dynamicDestinations, destination); - @Nullable - TableReference cloneSource = + @Nullable TableReference cloneSource = tempTable ? null : dynamicDestinations.getCloneSource(destination); TableSchema tableSchema; if (firstPaneCreateDisposition == CreateDisposition.CREATE_NEVER) { From 5345e74068def61f295f95b2ba2e31e3e6784744 Mon Sep 17 00:00:00 2001 From: Lalit Yadav Date: Tue, 27 Oct 2026 13:43:12 -0500 Subject: [PATCH 6/6] Fix BigQuery clone load schema and job IDs --- .../beam/sdk/io/gcp/bigquery/BatchLoads.java | 10 ++- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 2 + .../sdk/io/gcp/bigquery/BigQueryOptions.java | 15 ++++ .../io/gcp/bigquery/CreateTableHelpers.java | 7 +- .../bigquery/DynamicDestinationsHelpers.java | 48 ++++++++++--- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 72 +++++++++++++++++++ 6 files changed, 142 insertions(+), 12 deletions(-) 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 bbb71deb3fc3..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 @@ -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 4db9d237ec99..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 @@ -4178,6 +4178,8 @@ public WriteResult expand(PCollection input) { // 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) { 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/CreateTableHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java index 7bab5186b313..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 @@ -369,13 +369,16 @@ private static TableReference withDefaultProject( return updated; } - private static String createCloneJobIdPrefix( + @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)) + BigQueryHelpers.toJsonString(source) + + BigQueryHelpers.toJsonString(destination) + + options.getCloneJobIdToken()) .toString(); return BigQueryResourceNaming.createJobIdPrefix(jobName, stepToken, JobType.CLONE); } 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 67775ce12398..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 @@ -473,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; @@ -540,6 +566,16 @@ private static TableReference withDefaultProject( 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 { @@ -584,19 +620,13 @@ public TableDestination getTable(DestinationT destination) { public @Nullable TableSchema getSchema(DestinationT destination) { TableDestination wrappedDestination = super.getTable(destination); TableReference tableReference = wrappedDestination.getTableReference(); - @Nullable String destinationProjectId = tableReference.getProjectId(); @Nullable Table existingTable = getBigQueryTable(tableReference); if (existingTable == null || existingTable.getSchema() == null || existingTable.getSchema().isEmpty()) { - @Nullable TableReference cloneSource = super.getCloneSource(destination); - if (cloneSource != null) { - @Nullable Table cloneSourceTable = getBigQueryTable(cloneSource, destinationProjectId); - if (cloneSourceTable != null - && cloneSourceTable.getSchema() != null - && !cloneSourceTable.getSchema().isEmpty()) { - return cloneSourceTable.getSchema(); - } + @Nullable TableSchema cloneSchema = getCloneSourceSchema(destination); + if (cloneSchema != null && !cloneSchema.isEmpty()) { + return cloneSchema; } return super.getSchema(destination); } else { 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 f69ac688f24f..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 @@ -181,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; @@ -430,6 +431,77 @@ public void testWriteCreatesCloneTableWithUnqualifiedCloneSourceUsesDestinationP 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 {