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 252e55d34c07..68a0dff22081 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 @@ -429,7 +429,7 @@ private WriteResult expandTriggered(PCollection> inpu zeroLoadJobIdPrefixView, loadJobProjectId, WriteDisposition.WRITE_APPEND, - CreateDisposition.CREATE_NEVER, + createDisposition, maxRetryJobs, kmsKey, schemaUpdateOptions, @@ -535,7 +535,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 b222b358f547..b1b49f99cdc7 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 @@ -100,6 +100,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; @@ -2705,6 +2706,8 @@ public enum Method { abstract @Nullable ValueProvider getJsonSchema(); + abstract @Nullable ValueProvider getJsonCloneSourceTableRef(); + abstract @Nullable ValueProvider getJsonTimePartitioning(); abstract @Nullable ValueProvider getJsonClustering(); @@ -2824,6 +2827,8 @@ abstract Builder setAvroSchemaFactory( abstract Builder setJsonSchema(ValueProvider jsonSchema); + abstract Builder setJsonCloneSourceTableRef(ValueProvider jsonCloneSourceTableRef); + abstract Builder setJsonTimePartitioning(ValueProvider jsonTimePartitioning); abstract Builder setJsonClustering(ValueProvider clustering); @@ -3162,6 +3167,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. * @@ -3771,6 +3809,15 @@ public WriteResult expand(PCollection input) { .collect(Collectors.toList())), "No more than one of jsonSchema, schemaFromView, or dynamicDestinations may be set"); + long activeCreationMetadataCount = + java.util.stream.Stream.of( + getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef()) + .filter(arg -> arg != null) + .count(); + checkArgument( + activeCreationMetadataCount <= 1, + "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); + // Perform some argument checks BigQueryOptions bqOptions = input.getPipeline().getOptions().as(BigQueryOptions.class); Write.Method method = resolveMethod(input); @@ -3940,6 +3987,11 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { new TableConstraints.PrimaryKey().setColumns(getPrimaryKey()))); } } + if (getJsonCloneSourceTableRef() != null) { + dynamicDestinations = + new ConstantCloneSourceDestinations<>( + (DynamicDestinations) dynamicDestinations, getJsonCloneSourceTableRef()); + } return expandTyped(input, dynamicDestinations); } @@ -3955,6 +4007,7 @@ private WriteResult expandTyped( getJsonSchema() != null || getDynamicDestinations() != null || getSchemaFromView() != null; + boolean hasTableCreationMetadata = hasSchema || getJsonCloneSourceTableRef() != null; Class writeProtoClass = getWriteProtosClass(); if (getUseBeamSchema()) { @@ -3983,7 +4036,7 @@ private WriteResult expandTyped( dynamicDestinations, StaticValueProvider.of(BigQueryHelpers.toJsonString(tableSchema))); } else if (writeProtoClass != null) { - if (!hasSchema) { + if (!hasTableCreationMetadata) { try { @SuppressWarnings({"unchecked", "nullness"}) Descriptors.Descriptor descriptor = @@ -4005,8 +4058,9 @@ private WriteResult expandTyped( } else { // Require a schema if creating one or more tables. checkArgument( - getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED || hasSchema, - "CreateDisposition is CREATE_IF_NEEDED, however no schema was provided."); + getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED + || hasTableCreationMetadata, + "CreateDisposition is CREATE_IF_NEEDED, however no schema was provided and no clone source was provided."); } Coder destinationCoder; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index dd59939726bf..d8dcfd78888b 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 @@ -396,6 +396,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator transform) { if (transform.getJsonSchema() != null) { fieldValues.put("json_schema", transform.getJsonSchema().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())); @@ -668,6 +673,11 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { if (jsonSchema != null) { builder = builder.setJsonSchema(StaticValueProvider.of(jsonSchema)); } + String jsonCloneSourceTableRef = configRow.getString("json_clone_source_table_ref"); + if (jsonCloneSourceTableRef != null) { + builder = + builder.setJsonCloneSourceTableRef(StaticValueProvider.of(jsonCloneSourceTableRef)); + } String jsonTimePartitioning = configRow.getString("json_time_partitioning"); if (jsonTimePartitioning != null) { builder = builder.setJsonTimePartitioning(StaticValueProvider.of(jsonTimePartitioning)); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/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 7c428917503f..d6d5d67b245e 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.Set; @@ -41,17 +44,25 @@ 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; 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.base.Supplier; +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. * @@ -93,6 +104,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, @@ -134,6 +146,7 @@ static TableDestination possiblyCreateTable( bigQueryOptions, schemaSupplier, tableConstraintsSupplier, + cloneSourceSupplier, tableDestination, createDisposition, tableSpec, @@ -150,6 +163,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, @@ -162,60 +176,68 @@ private static void tryCreateTable( if (datasetService.getTable( tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC) == null) { - TableSchema tableSchema = schemaSupplier.get(); - @Nullable TableConstraints tableConstraints = tableConstraintsSupplier.get(); - Preconditions.checkArgumentNotNull( - tableSchema, - "Unless create disposition is %s, a schema must be specified, i.e. " - + "DynamicDestinations.getSchema() may not return null. " - + "However, create disposition is %s, and " - + " %s returned null for destination %s", - CreateDisposition.CREATE_NEVER, - createDisposition, - tableDestination); - Table table = new Table().setTableReference(tableReference).setSchema(tableSchema); - - if (tableConstraints != null) { - table = table.setTableConstraints(tableConstraints); - } + @Nullable TableReference cloneSource = cloneSourceSupplier.get(); + if (cloneSource != null) { + cloneTable( + options, datasetService, bqServices, cloneSource, tableReference, kmsKey, tableSpec); + } else { + TableSchema tableSchema = schemaSupplier.get(); + @Nullable TableConstraints tableConstraints = tableConstraintsSupplier.get(); + Preconditions.checkArgumentNotNull( + tableSchema, + "Unless create disposition is %s, a schema or clone source must be specified, i.e. " + + "DynamicDestinations.getSchema() and DynamicDestinations.getCloneSource() may " + + "not both return null. However, create disposition is %s, and " + + "no schema or clone source was returned for destination %s", + CreateDisposition.CREATE_NEVER, + createDisposition, + tableDestination); + Table table = new Table().setTableReference(tableReference).setSchema(tableSchema); - String tableDescription = tableDestination.getTableDescription(); - if (tableDescription != null) { - table = table.setDescription(tableDescription); - } + if (tableConstraints != null) { + table = table.setTableConstraints(tableConstraints); + } - TimePartitioning timePartitioning = tableDestination.getTimePartitioning(); - if (timePartitioning != null) { - table.setTimePartitioning(timePartitioning); - } + String tableDescription = tableDestination.getTableDescription(); + if (tableDescription != null) { + table = table.setDescription(tableDescription); + } - Clustering clustering = tableDestination.getClustering(); - if (clustering != null) { - table.setClustering(clustering); - } + TimePartitioning timePartitioning = tableDestination.getTimePartitioning(); + if (timePartitioning != null) { + table.setTimePartitioning(timePartitioning); + } - if (kmsKey != null) { - table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey)); - } - if (bigLakeConfiguration != null) { - TableReference ref = table.getTableReference(); - table.setBiglakeConfiguration( - new BigLakeConfiguration() - .setTableFormat( - MoreObjects.firstNonNull(bigLakeConfiguration.get("tableFormat"), "iceberg")) - .setFileFormat( - MoreObjects.firstNonNull(bigLakeConfiguration.get("fileFormat"), "parquet")) - .setConnectionId( - Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(CONNECTION_ID))) - .setStorageUri( - String.format( - "%s/%s/%s/%s", - Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(STORAGE_URI)), - ref.getProjectId(), - ref.getDatasetId(), - ref.getTableId()))); + Clustering clustering = tableDestination.getClustering(); + if (clustering != null) { + table.setClustering(clustering); + } + + if (kmsKey != null) { + table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey)); + } + if (bigLakeConfiguration != null) { + TableReference ref = table.getTableReference(); + table.setBiglakeConfiguration( + new BigLakeConfiguration() + .setTableFormat( + MoreObjects.firstNonNull( + bigLakeConfiguration.get("tableFormat"), "iceberg")) + .setFileFormat( + MoreObjects.firstNonNull(bigLakeConfiguration.get("fileFormat"), "parquet")) + .setConnectionId( + Preconditions.checkArgumentNotNull(bigLakeConfiguration.get(CONNECTION_ID))) + .setStorageUri( + String.format( + "%s/%s/%s/%s", + Preconditions.checkArgumentNotNull( + bigLakeConfiguration.get(STORAGE_URI)), + ref.getProjectId(), + ref.getDatasetId(), + ref.getTableId()))); + } + datasetService.createTable(table); } - datasetService.createTable(table); } } catch (Exception e) { throw new RuntimeException(e); @@ -223,6 +245,128 @@ private static void tryCreateTable( createdTables.add(tableSpec); } + private static void cloneTable( + BigQueryOptions options, + DatasetService datasetService, + BigQueryServices bqServices, + TableReference cloneSource, + TableReference destination, + @Nullable String kmsKey, + String tableSpec) + throws Exception { + TableReference destinationWithDefaultProject = withDefaultProject(options, destination); + TableReference source = + withDefaultProject(options, cloneSource, destinationWithDefaultProject.getProjectId()); + JobConfigurationTableCopy copyConfig = + new JobConfigurationTableCopy() + .setSourceTables(Collections.singletonList(source)) + .setDestinationTable(destinationWithDefaultProject) + .setOperationType("CLONE") + .setWriteDisposition(WriteDisposition.WRITE_EMPTY.name()) + .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED.name()); + if (kmsKey != null) { + copyConfig.setDestinationEncryptionConfiguration( + new EncryptionConfiguration().setKmsKeyName(kmsKey)); + } + + String jobProjectId = destinationWithDefaultProject.getProjectId(); + String bqLocation = + BigQueryHelpers.getDatasetLocation( + datasetService, + destinationWithDefaultProject.getProjectId(), + destinationWithDefaultProject.getDatasetId()); + String jobIdPrefix = createCloneJobIdPrefix(options, source, destinationWithDefaultProject); + try (JobService jobService = bqServices.getJobService(options)) { + BigQueryHelpers.PendingJob cloneJob = + new BigQueryHelpers.PendingJob( + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + jobService.startCopyJob(jobRef, copyConfig); + } catch (IOException e) { + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return null; + }, + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + return jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }, + jobId -> { + JobReference jobRef = + createJobReference(jobId.getJobId(), jobProjectId, bqLocation); + try { + return jobService.getJob(jobRef); + } catch (IOException e) { + throw new RuntimeException(e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }, + BatchLoads.DEFAULT_MAX_RETRY_JOBS, + jobIdPrefix); + cloneJob.runJob(); + new BigQueryHelpers.PendingJobManager() + .addPendingJob( + cloneJob, + job -> { + LOG.info("Created BigQuery clone table {} from {}", tableSpec, source); + return null; + }) + .waitForDone(); + } + } + + private static JobReference createJobReference(String jobId, String projectId, String location) { + return new JobReference().setProjectId(projectId).setJobId(jobId).setLocation(location); + } + + private static TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference) { + return withDefaultProject(options, tableReference, null); + } + + private static TableReference withDefaultProject( + BigQueryOptions options, TableReference tableReference, @Nullable String defaultProjectId) { + TableReference updated = tableReference.clone(); + if (Strings.isNullOrEmpty(updated.getProjectId())) { + if (defaultProjectId != null && !defaultProjectId.isEmpty()) { + updated.setProjectId(defaultProjectId); + } else { + @Nullable String projectId = options.getBigQueryProject(); + if (projectId == null || projectId.isEmpty()) { + projectId = options.getProject(); + } + if (projectId != null && !projectId.isEmpty()) { + updated.setProjectId(projectId); + } + } + } + return updated; + } + + 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..b69bfe75d9ae 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTables.java @@ -20,6 +20,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.api.services.bigquery.model.TableConstraints; +import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import java.util.List; import java.util.Map; @@ -117,6 +118,8 @@ public void processElement(ProcessContext context) { () -> dynamicDestinations.getSchema(dest); Supplier<@Nullable TableConstraints> tableConstraintsSupplier = () -> dynamicDestinations.getTableConstraints(dest); + Supplier<@Nullable TableReference> cloneSourceSupplier = + () -> dynamicDestinations.getCloneSource(dest); BigQueryOptions bqOptions = context.getPipelineOptions().as(BigQueryOptions.class); Lineage.getSinks() @@ -129,6 +132,7 @@ public void processElement(ProcessContext context) { tableDestination1, schemaSupplier, tableConstraintsSupplier, + cloneSourceSupplier, createDisposition, dynamicDestinations.getDestinationCoder(), kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java index 105da60c75b1..a04a77d92e86 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; @@ -159,6 +160,11 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { /** 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 6370244c268c..a66ab1770717 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; @@ -179,6 +180,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); @@ -221,6 +227,40 @@ public String toString() { } } + /** Returns the same clone source for every table. */ + static class ConstantCloneSourceDestinations + extends DelegatingDynamicDestinations { + private final ValueProvider jsonCloneSource; + private transient @Nullable TableReference cloneSource; + + ConstantCloneSourceDestinations( + DynamicDestinations inner, ValueProvider jsonCloneSource) { + super(inner); + Preconditions.checkArgumentNotNull(jsonCloneSource, "jsonCloneSource cannot be null"); + this.jsonCloneSource = jsonCloneSource; + } + + @Override + public TableReference getCloneSource(DestinationT destination) { + @Nullable TableReference cloneSource = this.cloneSource; + if (cloneSource == null) { + String jsonCloneSource = this.jsonCloneSource.get(); + checkArgument(jsonCloneSource != null, "jsonCloneSource cannot be null"); + cloneSource = BigQueryHelpers.fromJsonString(jsonCloneSource, TableReference.class); + this.cloneSource = cloneSource; + } + return cloneSource.clone(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("inner", inner) + .add("jsonCloneSource", jsonCloneSource) + .toString(); + } + } + static class ConstantTableConstraintsDestinations extends DelegatingDynamicDestinations { private final String jsonTableConstraints; @@ -427,6 +467,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 @@ -442,16 +487,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); @@ -463,6 +504,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 { @@ -477,16 +537,25 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) @Override public TableDestination getTable(DestinationT destination) { TableDestination wrappedDestination = super.getTable(destination); - Table existingTable = getBigQueryTable(wrappedDestination.getTableReference()); + TableReference tableReference = wrappedDestination.getTableReference(); + @Nullable String destinationProjectId = tableReference.getProjectId(); + Table existingTable = getBigQueryTable(tableReference); + Table tableToMatch = existingTable; + if (tableToMatch == null) { + @Nullable TableReference cloneSource = super.getCloneSource(destination); + if (cloneSource != null) { + tableToMatch = getBigQueryTable(cloneSource, destinationProjectId); + } + } - if (existingTable == null) { + if (tableToMatch == null) { return wrappedDestination; } else { return new TableDestination( wrappedDestination.getTableSpec(), - existingTable.getDescription(), - existingTable.getTimePartitioning(), - existingTable.getClustering()); + tableToMatch.getDescription(), + tableToMatch.getTimePartitioning(), + tableToMatch.getClustering()); } } @@ -497,10 +566,21 @@ 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, destinationProjectId); + 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 2dfc8b2f1c00..40f245e75f94 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 @@ -1114,6 +1114,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 b644d7aa752c..fda0c52d36bc 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 @@ -814,6 +814,7 @@ public void process( tableDestination, () -> dynamicDestinations.getSchema(dest), () -> dynamicDestinations.getTableConstraints(dest), + () -> dynamicDestinations.getCloneSource(dest), 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 b9cd554b578b..f485b88030d7 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,8 +40,10 @@ 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.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -71,7 +73,7 @@ public class UpdateSchemaDestination private transient @Nullable DatasetService datasetService; private final int maxRetryJobs; private final @Nullable String kmsKey; - private @Nullable BigQueryServices.JobService jobService; + private BigQueryServices.@Nullable JobService jobService; private final Set schemaUpdateOptions; private final BigQueryIO.Write.WriteDisposition writeDisposition; private final BigQueryIO.Write.CreateDisposition createDisposition; @@ -148,6 +150,27 @@ public void processElement( DestinationT destination = entry.getKey(); TableDestination tableDestination = getTableWithDefaultProject(destination); outputs.add(KV.of(tableDestination, entry.getValue())); + @Nullable TableReference cloneSource = dynamicDestinations.getCloneSource(destination); + if (cloneSource != null + && createDisposition != BigQueryIO.Write.CreateDisposition.CREATE_NEVER) { + Supplier<@Nullable TableSchema> schemaSupplier = + () -> dynamicDestinations.getSchema(destination); + Supplier<@Nullable TableConstraints> tableConstraintsSupplier = + () -> dynamicDestinations.getTableConstraints(destination); + Supplier<@Nullable TableReference> cloneSourceSupplier = () -> cloneSource; + + CreateTableHelpers.possiblyCreateTable( + context.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, + createDisposition, + dynamicDestinations.getDestinationCoder(), + kmsKey, + bqServices, + null); + } if (pendingJobs.containsKey(destination)) { // zero load job for this destination is already set continue; 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 fb9ad2e8d0f5..3329731603fa 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; @@ -197,9 +199,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 = @@ -219,13 +233,6 @@ public void processElement( jsonSchemas.put(destination, BigQueryHelpers.toJsonString(tableSchema)); } - TableDestination tableDestination = dynamicDestinations.getTable(destination); - checkArgument( - tableDestination != null, - "DynamicDestinations.getTable() may not return null, " - + "but %s returned null for destination %s", - dynamicDestinations, - destination); boolean destinationCoderSupportsClustering = !(dynamicDestinations.getDestinationCoder() instanceof TableDestinationCoderV2); checkArgument( @@ -280,6 +287,26 @@ public void processElement( createDisposition = CreateDisposition.CREATE_IF_NEEDED; } + if (cloneSource != null && createDisposition != CreateDisposition.CREATE_NEVER) { + Supplier<@Nullable TableSchema> schemaSupplier = + () -> dynamicDestinations.getSchema(destination); + Supplier<@Nullable TableConstraints> tableConstraintsSupplier = + () -> dynamicDestinations.getTableConstraints(destination); + Supplier<@Nullable TableReference> cloneSourceSupplier = () -> cloneSource; + + CreateTableHelpers.possiblyCreateTable( + c.getPipelineOptions().as(BigQueryOptions.class), + tableDestination, + schemaSupplier, + tableConstraintsSupplier, + cloneSourceSupplier, + createDisposition, + dynamicDestinations.getDestinationCoder(), + kmsKey, + bqServices, + null); + } + BigQueryHelpers.PendingJob retryJob = startLoad( getJobService(c.getPipelineOptions().as(BigQueryOptions.class)), diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java index de63120c93cc..1ae3a395e964 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"); @@ -253,6 +254,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"); @@ -271,6 +273,9 @@ public void testReCreateWriteTransformFromRowTable() { assertEquals(WriteDisposition.WRITE_TRUNCATE, writeTransformFromRow.getWriteDisposition()); assertEquals(CreateDisposition.CREATE_NEVER, writeTransformFromRow.getCreateDisposition()); assertEquals("dummykmskey", writeTransformFromRow.getKmsKey()); + assertEquals( + writeTransform.getJsonCloneSourceTableRef().get(), + writeTransformFromRow.getJsonCloneSourceTableRef().get()); assertEquals( BigQueryHelpers.toJsonString(testClustering), writeTransformFromRow.getJsonClustering().get()); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 601ed71473ed..7ef56303dc02 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; @@ -350,6 +351,238 @@ void writeEmptyPCollection() { p.run(); } + @Test + public void testWriteCreatesCloneTable() throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_table"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_table"); + TableRow baseRow = new TableRow().set("name", "base"); + TableRow appendedRow = new TableRow().set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withCloneFrom("project-id:dataset-id.source_table") + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null) + .map(JobConfigurationTableCopy::getOperationType) + .collect(Collectors.toList()), + hasItem("CLONE")); + } + + @Test + public void testWriteCreatesCloneTableWithUnqualifiedCloneSourceUsesDestinationProject() + throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.source_table"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.destination_table"); + TableRow baseRow = new TableRow().set("name", "base"); + TableRow appendedRow = new TableRow().set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to("bigquery-project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withCloneFrom("dataset-id.source_table") + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null && "CLONE".equals(copy.getOperationType())) + .flatMap(copy -> copy.getSourceTables().stream()) + .collect(Collectors.toList()), + hasItem(sourceTable)); + } + + @Test + public void testMatchTableDynamicDestinationsUsesDestinationProjectForCloneSourceSchema() + throws Exception { + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset-id.source_avro_table"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + + DynamicDestinations dynamicDestinations = + new DynamicDestinations() { + @Override + public String getDestination(@Nullable ValueInSingleWindow element) { + return "destination"; + } + + @Override + public Coder getDestinationCoder() { + return StringUtf8Coder.of(); + } + + @Override + public TableDestination getTable(String destination) { + return new TableDestination( + "bigquery-project-id:dataset-id.destination_avro_table", null); + } + + @Override + public @Nullable TableSchema getSchema(String destination) { + return null; + } + + @Override + public TableReference getCloneSource(String destination) { + return BigQueryHelpers.parseTableSpec("dataset-id.source_avro_table"); + } + }; + DynamicDestinations matchingDestinations = + DynamicDestinationsHelpers.matchTableDynamicDestinations( + dynamicDestinations, fakeBqServices); + + PCollection matchedSchema = + p.apply(Create.of("destination")) + .apply( + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(ProcessContext context) { + matchingDestinations.setSideInputAccessorFromProcessContext(context); + context.output( + BigQueryHelpers.toJsonString( + checkNotNull(matchingDestinations.getSchema(context.element())))); + } + })); + + PAssert.that(matchedSchema).containsInAnyOrder(BigQueryHelpers.toJsonString(schema)); + p.run(); + } + + @Test + public void testWriteCreatesCloneTableWithDynamicDestinations() throws Exception { + assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); + + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("bucket").setType("STRING"), + new TableFieldSchema().setName("name").setType("STRING"))); + TableReference sourceTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.source_bucket_a"); + TableReference destinationTable = + BigQueryHelpers.parseTableSpec("project-id:dataset-id.destination_bucket_a"); + TableRow baseRow = new TableRow().set("bucket", "a").set("name", "base"); + TableRow appendedRow = new TableRow().set("bucket", "a").set("name", "appended"); + fakeDatasetService.createTable(new Table().setTableReference(sourceTable).setSchema(schema)); + fakeDatasetService.insertAll(sourceTable, ImmutableList.of(baseRow), null); + + p.apply(Create.of(appendedRow)) + .apply( + BigQueryIO.writeTableRows() + .to( + new DynamicDestinations() { + @Override + public String getDestination(ValueInSingleWindow element) { + return element.getValue().get("bucket").toString(); + } + + @Override + public Coder getDestinationCoder() { + return StringUtf8Coder.of(); + } + + @Override + public TableDestination getTable(String destination) { + return new TableDestination( + "project-id:dataset-id.destination_bucket_" + destination, null); + } + + @Override + public @Nullable TableSchema getSchema(String destination) { + return null; + } + + @Override + public TableReference getCloneSource(String destination) { + return BigQueryHelpers.parseTableSpec( + "project-id:dataset-id.source_bucket_" + destination); + } + }) + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withoutValidation()); + p.run(); + + assertEquals(schema, fakeDatasetService.getTable(destinationTable).getSchema()); + assertThat( + fakeDatasetService.getAllRows(destinationTable), containsInAnyOrder(baseRow, appendedRow)); + assertThat( + fakeJobService.getAllJobs().stream() + .map(Job::getConfiguration) + .map(configuration -> configuration.getCopy()) + .filter(copy -> copy != null) + .map(JobConfigurationTableCopy::getOperationType) + .collect(Collectors.toList()), + hasItem("CLONE")); + } + + @Test + public void testWriteRejectsSchemaAndCloneSource() { + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage( + "No more than one of jsonSchema, schemaFromView, or cloneSource may be set"); + + p.apply(Create.empty(TableRowJsonCoder.of())) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.destination_table") + .withTestServices(fakeBqServices) + .withSchema(new TableSchema()) + .withCloneFrom("project-id:dataset-id.source_table") + .withoutValidation()); + p.run(); + } + @Test public void testWriteEmptyPCollectionGroupFilesFileLoad() throws Exception { assumeFalse(useStorageApi || useStorageApiApproximate || useStreaming); @@ -2157,6 +2390,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