diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java new file mode 100644 index 000000000000..a3bef076bd9c --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java @@ -0,0 +1,276 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import com.google.api.services.bigquery.model.TableRow; +import com.google.auto.value.AutoValue; +import com.google.cloud.bigquery.storage.v1.ProtoRows; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.util.BitSet; +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.util.ThrowingSupplier; +import org.apache.beam.sdk.values.TimestampedValue; +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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +@AutoValue +abstract class AppendRowsPacket { + abstract ProtoRows getProtoRows(); + + abstract List getTimestamps(); + + abstract List<@Nullable TableRow> getFailsafeTableRows(); + + abstract BitSet getSchemaMismatchedRows(); + + abstract Map getSchemaHashes(); + + abstract Map getUnknownFields(); + + abstract Map getOriginalPayloads(); + + // Retry deadlines for handling schema updates + abstract List getDeadlines(); + + private static final BitSet EMPTY_BIT_SET = new BitSet(0); + + static AppendRowsPacket fromStorageApiWritePayload( + PeekingIterator underlyingIterator, + long maxByteSize, + SchemaChangeDetectorHelper schemaChangeDetectorHelper, + Instant elementTimestamp, + Supplier appendClientInfoSupplier, + Function, Boolean> failedRowsHandler, + ThrowingSupplier getCurrentTableSchemaHash, + ThrowingSupplier getCurrentTableSchemaDescriptor) { + List timestamps = Lists.newArrayList(); + List<@Nullable TableRow> failsafeRows = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + long bytesSize = 0; + BitSet mismatchedRows = new BitSet(); + try { + while (underlyingIterator.hasNext()) { + // Make sure that we don't exceed the maxByteSize over multiple elements. A single + // element can exceed + // the split threshold, but in that case it should be the only element returned. + if ((bytesSize + underlyingIterator.peek().getStoragePayload().getPayload().length + > maxByteSize) + && inserts.getSerializedRowsCount() > 0) { + break; + } + StoragePayloadWithDeadline payload = underlyingIterator.next(); + StorageApiWritePayload storagePayload = payload.getStoragePayload(); + + @Nullable TableRow failsafeTableRow = null; + try { + failsafeTableRow = storagePayload.getFailsafeTableRow(); + } catch (IOException e) { + // Do nothing, table row will be generated later from row bytes + } + + // If autoUpdateSchema is set, we try to automatically merge in unknown fields. + ByteString byteString; + SchemaChangeDetectorHelper.MergePayloadResult mergeResult = + schemaChangeDetectorHelper.getMergedPayload( + storagePayload, elementTimestamp, failsafeTableRow, appendClientInfoSupplier.get()); + if (mergeResult.getKind() == SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) { + if (failedRowsHandler.apply(mergeResult.getFailed())) { + continue; + } else { + // This implies that instead of skipping failed rows, we should mark it as a mismatched + // row. + byteString = ByteString.empty(); + mismatchedRows.set(inserts.getSerializedRowsCount()); + } + } else { + byteString = mergeResult.getMerged(); + } + + if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate( + storagePayload, + byteString.toByteArray(), + getCurrentTableSchemaHash, + getCurrentTableSchemaDescriptor)) { + mismatchedRows.set(inserts.getSerializedRowsCount()); + } + + int currentIndex = inserts.getSerializedRowsCount(); + inserts.addSerializedRows(byteString); + Instant timestamp = storagePayload.getTimestamp(); + if (timestamp == null) { + timestamp = elementTimestamp; + } + timestamps.add(timestamp); + failsafeRows.add(failsafeTableRow); + if (storagePayload.getSchemaHash() != null) { + schemaHashes.put( + currentIndex, Preconditions.checkStateNotNull(storagePayload.getSchemaHash())); + } + if (storagePayload.getUnknownFields() != null) { + unknownFields.put( + currentIndex, Preconditions.checkStateNotNull(storagePayload.getUnknownFields())); + originalPayloads.put(currentIndex, storagePayload.getPayload()); + } + deadlines.add(payload.getDeadline()); + bytesSize += byteString.size(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeRows, + mismatchedRows, + schemaHashes, + unknownFields, + originalPayloads, + deadlines); + } + + AppendRowsPacket getSchemaMismatchedRowsOnly() { + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + List timestamps = Lists.newArrayList(); + List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); + if (!getSchemaMismatchedRows().isEmpty()) { + int newIndex = 0; + for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { + if (getSchemaMismatchedRows().get(i)) { + inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); + timestamps.add(getTimestamps().get(i)); + failsafeTableRows.add(getFailsafeTableRows().get(i)); + TableRow unknown = getUnknownFields().get(i); + if (unknown != null) { + unknownFields.put(newIndex, unknown); + } + byte[] originalPayload = getOriginalPayloads().get(i); + if (originalPayload != null) { + originalPayloads.put(newIndex, originalPayload); + } + deadlines.add(getDeadlines().get(i)); + byte[] schemaHash = getSchemaHashes().get(i); + if (schemaHash != null) { + schemaHashes.put(newIndex, schemaHash); + } + newIndex++; + } + } + } + BitSet allBits = new BitSet(inserts.getSerializedRowsCount()); + allBits.set(0, inserts.getSerializedRowsCount()); + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeTableRows, + allBits, + schemaHashes, + unknownFields, + originalPayloads, + deadlines); + } + + AppendRowsPacket getSchemaMatchedRowsOnly() { + if (getSchemaMismatchedRows().isEmpty()) { + return this; + } + + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + List timestamps = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); + List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); + int newIndex = 0; + for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { + if (!getSchemaMismatchedRows().get(i)) { + inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); + timestamps.add(getTimestamps().get(i)); + failsafeTableRows.add(getFailsafeTableRows().get(i)); + byte[] schemaHash = getSchemaHashes().get(i); + if (schemaHash != null) { + schemaHashes.put(newIndex, schemaHash); + } + TableRow unknown = getUnknownFields().get(i); + if (unknown != null) { + unknownFields.put(newIndex, unknown); + } + byte[] originalPayload = getOriginalPayloads().get(i); + if (originalPayload != null) { + originalPayloads.put(newIndex, originalPayload); + } + deadlines.add(getDeadlines().get(i)); + newIndex++; + } + } + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeTableRows, + EMPTY_BIT_SET, + schemaHashes, + unknownFields, + originalPayloads, + deadlines); + } + + Stream toPayloadStream() { + return IntStream.range(0, getProtoRows().getSerializedRowsCount()) + .mapToObj( + i -> { + try { + byte @Nullable [] originalBytes = getOriginalPayloads().get(i); + StorageApiWritePayload payload = + StorageApiWritePayload.of( + originalBytes != null + ? originalBytes + : getProtoRows().getSerializedRows(i).toByteArray(), + getUnknownFields().get(i), + getFailsafeTableRows().get(i)) + .withTimestamp(getTimestamps().get(i)); + byte @Nullable [] hash = getSchemaHashes().get(i); + if (hash != null) { + payload = payload.withSchemaHash(hash); + } + return StoragePayloadWithDeadline.of(payload, getDeadlines().get(i)); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } +} 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..1360f58c6df5 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 @@ -108,6 +108,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.PassThroughThenCleanup.ContextContainer; import org.apache.beam.sdk.io.gcp.bigquery.RowWriterFactory.OutputType; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -2496,6 +2497,7 @@ public static Write write() { .setAutoSharding(false) .setPropagateSuccessful(true) .setAutoSchemaUpdate(false) + .setAutoSchemaUpdateStrictTimeout(null) .setDeterministicRecordIdFn(null) .setMaxRetryJobs(1000) .setPropagateSuccessfulStorageApiWrites(false) @@ -2777,6 +2779,8 @@ public enum Method { abstract boolean getAutoSchemaUpdate(); + abstract @Nullable Duration getAutoSchemaUpdateStrictTimeout(); + abstract @Nullable Class getWriteProtosClass(); abstract boolean getDirectWriteProtos(); @@ -2897,6 +2901,8 @@ abstract Builder setDefaultMissingValueInterpretation( abstract Builder setAutoSchemaUpdate(boolean autoSchemaUpdate); + abstract Builder setAutoSchemaUpdateStrictTimeout(@Nullable Duration timeout); + abstract Builder setWriteProtosClass(@Nullable Class clazz); abstract Builder setDirectWriteProtos(boolean direct); @@ -3562,11 +3568,46 @@ public Write withSuccessfulInsertsPropagation(boolean propagateSuccessful) { * If true, enables automatically detecting BigQuery table schema updates. Table schema updates * are usually noticed within several minutes. Only supported when using one of the STORAGE_API * insert methods. + * + *

Rows that contain new columns will only have the new columns sent to BigQuery after the + * new table schema has been observed. Until then, only the known columns will be sent to + * BigQuery. + * + *

Note that this option detects table schema updates performed on the table, usually by an + * external process. If you want Beam to update the table schema for you, please see {@link + * #withSchemaUpdateOptions}. */ public Write withAutoSchemaUpdate(boolean autoSchemaUpdate) { return toBuilder().setAutoSchemaUpdate(autoSchemaUpdate).build(); } + /** + * If true, enables automatically detecting BigQuery table schema updates. Table schema updates + * are usually noticed within several minutes. Only supported when using one of the STORAGE_API + * insert methods. + * + *

This mode ensures consistent writes - rows with new schemas will not be sent to BigQuery + * until we have observed the new table schema. This comes with a few caveats: - Detecting + * unknown columns requires extra parsing. Some increase in CPU usage may be noticed. - Rows + * with unknown columns will be retried until a new schema is observed. This may temporarily + * block inserts of rows into BigQuery whenever a schema update happens. + * + *

The timeout parameter specifies how long to wait until we see the new table schema. If + * more than the specified time goes by before a matching table schema is seen, the row will be + * sent to the failedRows output collection. + * + *

Note that this option detects table schema updates performed on the table, usually by an + * external process. If you want Beam to update the table schema for you, please see {@link + * #withSchemaUpdateOptions}. + */ + public Write withAutoSchemaUpdateConsistent( + boolean autoSchemaUpdate, Duration waitForSchemaTimeout) { + return toBuilder() + .setAutoSchemaUpdate(autoSchemaUpdate) + .setAutoSchemaUpdateStrictTimeout(waitForSchemaTimeout) + .build(); + } + /* * Provides a function which can serve as a source of deterministic unique ids for each record * to be written, replacing the unique ids generated with the default scheme. When used with @@ -4090,6 +4131,14 @@ private WriteResult continueExpandTyped( DynamicDestinations dynamicDestinations, RowWriterFactory rowWriterFactory, Write.Method method) { + if (getAutoSchemaUpdateStrictTimeout() != null) { + checkArgument( + method == Method.STORAGE_API_AT_LEAST_ONCE || method == Method.STORAGE_WRITE_API, + "Auto update schema only supported when using storage write API"); + checkArgument( + input.getPipeline().getOptions().as(StreamingOptions.class).isStreaming(), + "auto update schema only supported on streaming pipelines"); + } if (method == Write.Method.STREAMING_INSERTS) { checkArgument( getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE, @@ -4294,6 +4343,10 @@ private WriteResult continueExpandTyped( RowWriterFactory.TableRowWriterFactory tableRowWriterFactory = (RowWriterFactory.TableRowWriterFactory) rowWriterFactory; // Fallback behavior: convert to JSON TableRows and convert those into Beam TableRows. + @Nullable Set schemaUpdateOptions = getSchemaUpdateOptions(); + boolean useSchemaUpdatingTableRow = + (schemaUpdateOptions != null && !schemaUpdateOptions.isEmpty()) + || (getAutoSchemaUpdate() && getAutoSchemaUpdateStrictTimeout() != null); storageApiDynamicDestinations = new StorageApiDynamicDestinationsTableRow<>( dynamicDestinations, @@ -4303,9 +4356,7 @@ private WriteResult continueExpandTyped( getCreateDisposition(), getIgnoreUnknownValues(), getAutoSchemaUpdate(), - getSchemaUpdateOptions() == null - ? Collections.emptySet() - : getSchemaUpdateOptions()); + useSchemaUpdatingTableRow); } int numShards = getStorageApiNumStreams(bqOptions); @@ -4328,6 +4379,7 @@ private WriteResult continueExpandTyped( method == Method.STORAGE_API_AT_LEAST_ONCE, enableAutoSharding, getAutoSchemaUpdate(), + getAutoSchemaUpdateStrictTimeout(), getIgnoreUnknownValues(), getPropagateSuccessfulStorageApiWrites(), getPropagateSuccessfulStorageApiWritesPredicate(), 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..ee490b227242 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 @@ -438,6 +438,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator transform) { fieldValues.put("use_beam_schema", transform.getUseBeamSchema()); fieldValues.put("auto_sharding", transform.getAutoSharding()); fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate()); + if (transform.getAutoSchemaUpdateStrictTimeout() != null) { + fieldValues.put( + "auto_schema_update_strict_timeout_ms", + transform.getAutoSchemaUpdateStrictTimeout().getMillis()); + } if (transform.getWriteProtosClass() != null) { fieldValues.put("write_protos_class", toByteArray(transform.getWriteProtosClass())); } @@ -870,6 +876,15 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { if (autoSchemaUpdate != null) { builder = builder.setAutoSchemaUpdate(autoSchemaUpdate); } + Long autoSchemaUpdateStrictTimeoutMillis = + configRow.getInt64("auto_schema_update_strict_timeout_ms"); + if (autoSchemaUpdateStrictTimeoutMillis != null) { + builder = + builder + .setAutoSchemaUpdate(true) + .setAutoSchemaUpdateStrictTimeout( + org.joda.time.Duration.millis(autoSchemaUpdateStrictTimeoutMillis)); + } byte[] writeProtosClasses = configRow.getBytes("write_protos_class"); if (writeProtosClasses != null) { builder = 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 face2ef5841a..da8526bd9948 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 @@ -260,4 +260,24 @@ public interface BigQueryOptions Integer getSchemaUpgradeBufferingShards(); void setSchemaUpgradeBufferingShards(Integer value); + + @Description("How long to retry locally before buffering when a schema mismatch is detected.") + @Default.Integer(5000) + Integer getStorageApiMismatchLocalRetryTimeMilliSec(); + + void setStorageApiMismatchLocalRetryTimeMilliSec(Integer value); + + @Description("The retry time in milliseconds when a schema mismatch is detected.") + @Default.Integer(60000) + Integer getStorageApiMismatchRetryTimeMilliSec(); + + void setStorageApiMismatchRetryTimeMilliSec(Integer value); + + @Description( + "If a pipeline is drained while waiting for a BigQuery schema update, we will wait this long for the " + + "schema update before sending the rows to the failed-rows collection.") + @Default.Integer(300000) + Integer getStorageApiMismatchDrainRetryTimeMilliSec(); + + void setStorageApiMismatchDrainRetryTimeMilliSec(Integer value); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java index aec011cebb54..acb0bd5f6946 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java @@ -51,6 +51,7 @@ public class BigQuerySinkMetrics { public static final String OK = Status.Code.OK.toString(); static final String INTERNAL = "INTERNAL"; public static final String PAYLOAD_TOO_LARGE = "PayloadTooLarge"; + public static final String SCHEMA_MISMATCHED = "SchemaMismatched"; // Base Metric names private static final String RPC_REQUESTS = "RpcRequestsCount"; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java new file mode 100644 index 000000000000..a8202744642b --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -0,0 +1,355 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import com.google.api.services.bigquery.model.TableRow; +import java.nio.ByteBuffer; +import java.time.Instant; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; +import java.util.stream.StreamSupport; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.state.TimerSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.ShardedKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +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.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; + +class BufferMismatchedRows + extends PTransform< + PCollection>, PCollectionTuple> { + private final Coder failedRowsCoder; + private final Coder successfulRowsCoder; + private final Coder destinationCoder; + private final StorageApiDynamicDestinations dynamicDestinations; + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl + writeDoFn; + private final TupleTag failedRowsTag; + private final @Nullable TupleTag successfulRowsTag; + // This output is effectively ignored, since we only support this code path for + // StorageApiWriteRecordsInconsistent. + private final TupleTag> finalizeTag = new TupleTag<>("finalizeTag"); + private static final int NUM_DEFAULT_SHARDS = 20; + + public BufferMismatchedRows( + Coder failedRowsCoder, + Coder successfulRowsCoder, + Coder destinationCoder, + StorageApiDynamicDestinations dynamicDestinations, + StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn, + TupleTag failedRowsTag, + @Nullable TupleTag successfulRowsTag) { + this.failedRowsCoder = failedRowsCoder; + this.successfulRowsCoder = successfulRowsCoder; + this.destinationCoder = destinationCoder; + this.dynamicDestinations = dynamicDestinations; + this.writeDoFn = writeDoFn; + this.failedRowsTag = failedRowsTag; + this.successfulRowsTag = successfulRowsTag; + } + + @Override + public PCollectionTuple expand(PCollection> input) { + // Append records to the Storage API streams. + TupleTagList tupleTagList = TupleTagList.of(failedRowsTag); + if (successfulRowsTag != null) { + tupleTagList = tupleTagList.and(successfulRowsTag); + } + + PCollectionTuple result = + input + .apply( + "addShard", + ParDo.of( + new DoFn< + KV, + KV, StoragePayloadWithDeadline>>() { + int shardNumber; + + @Setup + public void setup() { + shardNumber = ThreadLocalRandom.current().nextInt(NUM_DEFAULT_SHARDS); + } + + @ProcessElement + public void process( + @Element KV element, + OutputReceiver, StoragePayloadWithDeadline>> + o) { + ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); + buffer.putInt(++shardNumber % NUM_DEFAULT_SHARDS); + o.output( + KV.of( + ShardedKey.of(element.getKey(), buffer.array()), + element.getValue())); + } + })) + .setCoder( + KvCoder.of( + ShardedKey.Coder.of(destinationCoder), StoragePayloadWithDeadline.Coder.of())) + .apply( + "bufferMismatchedRows", + ParDo.of(new BufferingDoFn(writeDoFn)) + .withOutputTags(finalizeTag, tupleTagList) + .withSideInputs(dynamicDestinations.getSideInputs())); + + result.get(failedRowsTag).setCoder(failedRowsCoder); + if (successfulRowsTag != null) { + result.get(successfulRowsTag).setCoder(successfulRowsCoder); + } + return result; + } + + class BufferingDoFn + extends DoFn, StoragePayloadWithDeadline>, KV> { + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl + writeDoFn; + + @StateId("mismatchedRows") + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(StoragePayloadWithDeadline.Coder.of()); + + @TimerId("retryMismatchedRowsTimer") + private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + + @StateId("currentMismatchedRowTimerValue") + private final StateSpec> currentMismatchedRowTimerValueSpec = + StateSpecs.value(); + + @StateId("minPendingTimestamp") + private final StateSpec> minPendingTimestampSpec = StateSpecs.value(); + + private final Counter rowsSentToFailedRowsCollection = + Metrics.counter(BufferMismatchedRows.BufferingDoFn.class, "rowsSentToFailedRowsCollection"); + + public BufferingDoFn( + StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn) { + this.writeDoFn = writeDoFn; + } + + @ProcessElement + public void process( + PipelineOptions pipelineOptions, + ProcessContext processContext, + @Element KV, StoragePayloadWithDeadline> element, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @TimerId("retryMismatchedRowsTimer") Timer retryTimer, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromProcessContext(processContext); + TableDestination tableDestination = dynamicDestinations.getTable(element.getKey().getKey()); + + Duration timerRetryDuration = + Duration.millis( + pipelineOptions.as(BigQueryOptions.class).getStorageApiMismatchRetryTimeMilliSec()); + SchemaChangeDetectorHelper.bufferMismatchedRows( + Collections.singleton(element.getValue()), + mismatchedRowsBag, + retryTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + null, + rowsSentToFailedRowsCollection, + timerRetryDuration); + } + + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(Long.MAX_VALUE); + } + + @OnTimer("retryMismatchedRowsTimer") + public void onTimer( + OnTimerContext context, + @Key ShardedKey shardedDestination, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, + @TimerId("retryMismatchedRowsTimer") Timer retryTimer, + PipelineOptions pipelineOptions, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnTimerContext(context); + writeDoFn.startBundle(); + + mismatchedRowsBag.readLater(); + currentTimerValue.readLater(); + minPendingTimestamp.readLater(); + + TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey()); + StorageApiDynamicDestinations.MessageConverter messageConverter = + writeDoFn.messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + writeDoFn.getDatasetService(pipelineOptions), + writeDoFn.getWriteStreamService(pipelineOptions)); + messageConverter.updateSchemaFromTable(); + + List>> mismatchedRowsList = + Lists.newArrayList(); + for (StoragePayloadWithDeadline row : mismatchedRowsBag.read()) { + Iterable> mismatchedRows = + writeDoFn.processElement( + pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRowsList.add(mismatchedRows); + } + } + // Once we're done, delegate to finishBundle to finish things. + Iterable> mismatchedDestRows = + writeDoFn.finishBundle( + o.get(failedRowsTag), + successfulRowsTag != null ? o.get(successfulRowsTag) : null, + o.get(finalizeTag), + null); + if (!Iterables.isEmpty(mismatchedDestRows)) { + mismatchedRowsList.add(mismatchedDestRows); + } + + mismatchedRowsBag.clear(); + currentTimerValue.clear(); + minPendingTimestamp.clear(); + if (!mismatchedRowsList.isEmpty()) { + AppendClientInfo appendClientInfo = + AppendClientInfo.of( + messageConverter.getTableSchema(), + messageConverter.getDescriptor(writeDoFn.usesCdc), + AutoCloseable::close); + + Iterable mismatchedRows = + () -> + StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false) + .map(KV::getValue) + .iterator(); + + Duration timerRetryDuration = + Duration.millis( + pipelineOptions.as(BigQueryOptions.class).getStorageApiMismatchRetryTimeMilliSec()); + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRows, + mismatchedRowsBag, + retryTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + appendClientInfo, + rowsSentToFailedRowsCollection, + timerRetryDuration); + } + } + + @OnWindowExpiration + public void onWindowExpiration( + OnWindowExpirationContext context, + @Key ShardedKey shardedDestination, + @Timestamp org.joda.time.Instant elementTs, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + PipelineOptions pipelineOptions, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + writeDoFn.messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + writeDoFn.getDatasetService(pipelineOptions), + writeDoFn.getWriteStreamService(pipelineOptions)); + messageConverter.updateSchemaFromTable(); + + java.time.Duration waitTime = + java.time.Duration.ofMillis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchDrainRetryTimeMilliSec()); + + Iterable bufferedRows = mismatchedRowsBag.read(); + Instant start = Instant.now(); + while (!Iterables.isEmpty(bufferedRows) && start.plus(waitTime).isAfter(Instant.now())) { + writeDoFn.startBundle(); + List>> mismatchedRowsList = + Lists.newArrayList(); + for (StoragePayloadWithDeadline row : bufferedRows) { + Iterable> mismatchedRows = + writeDoFn.processElement( + pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRowsList.add(mismatchedRows); + } + } + // Once we're done, delegate to finishBundle to finish things. + Iterable> mismatchedDestRows = + writeDoFn.finishBundle( + o.get(failedRowsTag), + successfulRowsTag != null ? o.get(successfulRowsTag) : null, + o.get(finalizeTag), + null); + if (!Iterables.isEmpty(mismatchedDestRows)) { + mismatchedRowsList.add(mismatchedDestRows); + } + + bufferedRows = + () -> + StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false) + .map(KV::getValue) + .iterator(); + } + writeDoFn.writeFailedRows( + shardedDestination.getKey(), + bufferedRows, + "Timed out waiting for table schema update in OnWindowExpiration", + pipelineOptions, + elementTs, + o.get(failedRowsTag)); + } + + @Teardown + public void onTeardown() { + writeDoFn.teardown(); + } + } +} 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..f40918391200 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 @@ -82,6 +82,33 @@ interface SideInputAccessor { private transient @Nullable SideInputAccessor sideInputAccessor; private transient @Nullable PipelineOptions options; + static class SideInputAccessorViaOnTimerContext implements SideInputAccessor { + private DoFn.OnTimerContext onTimerContext; + + public SideInputAccessorViaOnTimerContext(DoFn.OnTimerContext onTimerContext) { + this.onTimerContext = onTimerContext; + } + + @Override + public SideInputT sideInput(PCollectionView view) { + return onTimerContext.sideInput(view); + } + } + + static class SideInputAccessorViaOnWindowExpirationContext implements SideInputAccessor { + private DoFn.OnWindowExpirationContext onWindowExpirationContext; + + public SideInputAccessorViaOnWindowExpirationContext( + DoFn.OnWindowExpirationContext onWindowExpirationContext) { + this.onWindowExpirationContext = onWindowExpirationContext; + } + + @Override + public SideInputT sideInput(PCollectionView view) { + return onWindowExpirationContext.sideInput(view); + } + } + static class SideInputAccessorViaProcessContext implements SideInputAccessor { private DoFn.ProcessContext processContext; @@ -129,6 +156,17 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { this.options = context.getPipelineOptions(); } + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + this.sideInputAccessor = new SideInputAccessorViaOnTimerContext(context); + this.options = context.getPipelineOptions(); + } + + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + this.sideInputAccessor = new SideInputAccessorViaOnWindowExpirationContext(context); + this.options = context.getPipelineOptions(); + } + /** * Returns an object that represents at a high level which table is being written to. May not * return null. 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..14a7ef3de660 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 @@ -215,6 +215,19 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { inner.setSideInputAccessorFromProcessContext(context); } + @Override + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + super.setSideInputAccessorFromOnTimerContext(context); + inner.setSideInputAccessorFromOnTimerContext(context); + } + + @Override + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + super.setSideInputAccessorFromOnWindowExpirationContext(context); + inner.setSideInputAccessorFromOnWindowExpirationContext(context); + } + @Override public String toString() { return MoreObjects.toStringHelper(this).add("inner", inner).toString(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java new file mode 100644 index 000000000000..517b319a8c14 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java @@ -0,0 +1,220 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.auto.value.AutoOneOf; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.util.Optional; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.util.ThrowingSupplier; +import org.apache.beam.sdk.values.TimestampedValue; +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.Predicates; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +class SchemaChangeDetectorHelper { + @AutoOneOf(MergePayloadResult.Kind.class) + abstract static class MergePayloadResult { + enum Kind { + MERGED, + FAILED + }; + + abstract Kind getKind(); + + abstract ByteString getMerged(); + + abstract TimestampedValue getFailed(); + + static MergePayloadResult ofMerged(ByteString merged) { + return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.merged(merged); + } + + static MergePayloadResult ofFailed(TimestampedValue failed) { + return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.failed(failed); + } + } + + private final boolean autoUpdateSchema; + private final boolean ignoreUnknownValues; + private final TableReference tableReference; + private final boolean ignoreSchemaHashes; + + public SchemaChangeDetectorHelper( + boolean autoUpdateSchema, + boolean ignoreUnknownValues, + TableReference tableReference, + boolean ignoreSchemaHashes) { + this.autoUpdateSchema = autoUpdateSchema; + this.ignoreUnknownValues = ignoreUnknownValues; + this.tableReference = tableReference; + this.ignoreSchemaHashes = ignoreSchemaHashes; + } + + Optional checkUpdatedSchema( + TableSchema currentSchema, + String streamName, + BigQueryServices.WriteStreamService writeStreamService) { + if (autoUpdateSchema) { + @Nullable TableSchema streamSchema = writeStreamService.getWriteStreamSchema(streamName); + if (streamSchema != null) { + return TableSchemaUpdateUtils.getUpdatedSchema(currentSchema, streamSchema); + } + } + return Optional.empty(); + } + + Optional checkResponseForUpdatedSchema( + TableSchema currentSchema, BigQueryServices.StreamAppendClient streamAppendClient) { + @Nullable TableSchema updatedSchemaReturned = streamAppendClient.getUpdatedSchema(); + + // Update the table schema and clear the append client. + if (updatedSchemaReturned != null) { + return TableSchemaUpdateUtils.getUpdatedSchema(currentSchema, updatedSchemaReturned); + } + return Optional.empty(); + } + + MergePayloadResult getMergedPayload( + StorageApiWritePayload payload, + Instant elementTimestamp, + @Nullable TableRow failsafeTableRow, + AppendClientInfo appendClientInfo) + throws IOException { + ByteString byteString = ByteString.copyFrom(payload.getPayload()); + + if (autoUpdateSchema) { + @Nullable TableRow unknownFields = payload.getUnknownFields(); + + if (unknownFields != null && !unknownFields.isEmpty()) { + try { + // Protocol buffer serialization format supports concatenation. We serialize any new + // "known" fields into a proto and concatenate to the existing proto. + byteString = + Preconditions.checkStateNotNull(appendClientInfo) + .mergeNewFields(byteString, unknownFields, ignoreUnknownValues); + } catch (TableRowToStorageApiProto.SchemaConversionException e) { + // This generally implies that ignoreUnknownValues=false and there were still + // unknown values here. + // Reconstitute the TableRow and send it to the failed-rows consumer. + TableRow tableRow = + failsafeTableRow != null + ? failsafeTableRow + : appendClientInfo.toTableRow(byteString, Predicates.alwaysTrue()); + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError(tableRow, e.toString(), tableReference); + Instant timestamp = MoreObjects.firstNonNull(payload.getTimestamp(), elementTimestamp); + return MergePayloadResult.ofFailed(TimestampedValue.of(error, timestamp)); + } + } + } + return MergePayloadResult.ofMerged(byteString); + } + + boolean isPayloadSchemaOutOfDate( + StorageApiWritePayload payload, + byte[] mergedPayload, + ThrowingSupplier schemaHash, + ThrowingSupplier schemaDescriptor) + throws Exception { + if (ignoreSchemaHashes) { + if (payload.getUnknownFields() != null + && UpgradeTableSchema.missingUnknownField( + Preconditions.checkStateNotNull(payload.getUnknownFields()), schemaDescriptor)) { + return true; + } + // We currently rely on getting an append failure from Vortex if there are + // missing required + // fields. However + // we should consider explicitly checking here in the future. + } else { + return UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), () -> mergedPayload, schemaHash, schemaDescriptor); + } + return false; + } + + static void bufferMismatchedRows( + Iterable rows, + BagState bufferedBag, + Timer retryRowsTimers, + ValueState currentTimerValue, + ValueState minPendingTimestamp, + TableDestination tableDestination, + DoFn.OutputReceiver failedRowsReceiver, + @Nullable AppendClientInfo appendClientInfo, + Counter rowsSentToFailedRowsCollection, + Duration retryPeriod) + throws IOException { + org.joda.time.Instant minTimestamp = + org.joda.time.Instant.ofEpochMilli( + MoreObjects.firstNonNull( + minPendingTimestamp.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + + for (StoragePayloadWithDeadline row : rows) { + if (appendClientInfo == null + || row.getDeadline().isAfter(retryRowsTimers.getCurrentRelativeTime())) { + bufferedBag.add(row); + org.joda.time.@Nullable Instant timestamp = row.getStoragePayload().getTimestamp(); + if (timestamp != null && timestamp.isBefore(minTimestamp)) { + minTimestamp = timestamp; + } + } else { + @Nullable TableRow failedRow = row.getStoragePayload().getFailsafeTableRow(); + if (failedRow == null) { + ByteString rowBytes = ByteString.copyFrom(row.getStoragePayload().getPayload()); + failedRow = appendClientInfo.toTableRow(rowBytes, Predicates.alwaysTrue()); + } + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, "Mismatched schema", tableDestination.getTableReference()); + failedRowsReceiver.outputWithTimestamp( + error, Preconditions.checkStateNotNull(row.getStoragePayload().getTimestamp())); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.SCHEMA_MISMATCHED, + tableDestination.getShortTableUrn()) + .inc(1); + } + } + minPendingTimestamp.write(minTimestamp.getMillis()); + + long targetTs = + MoreObjects.firstNonNull( + currentTimerValue.read(), + retryRowsTimers.getCurrentRelativeTime().getMillis() + retryPeriod.getMillis()); + retryRowsTimers + .withOutputTimestamp(minTimestamp) + .set(org.joda.time.Instant.ofEpochMilli(targetTs)); + currentTimerValue.write(targetTs); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java index b65b11ef3194..456b557d61c4 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java @@ -165,6 +165,7 @@ public Duration getAllowedTimestampSkew() { @OnTimer("pollTimer") public void onPollTimer( + OnTimerContext context, @Key ShardedKey key, PipelineOptions pipelineOptions, @StateId("bufferedElements") BagState> bag, @@ -174,6 +175,8 @@ public void onPollTimer( BoundedWindow window, MultiOutputReceiver o) throws Exception { + convertMessagesDoFn.getDynamicDestinations().setSideInputAccessorFromOnTimerContext(context); + if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o)) { timerTs.clear(); } else { @@ -188,12 +191,17 @@ public void onPollTimer( @OnWindowExpiration public void onWindowExpiration( + OnWindowExpirationContext context, @Key ShardedKey key, PipelineOptions pipelineOptions, @StateId("bufferedElements") BagState> bag, @StateId("minBufferedTimestamp") CombiningState minBufferedTimestamp, MultiOutputReceiver o) throws Exception { + convertMessagesDoFn + .getDynamicDestinations() + .setSideInputAccessorFromOnWindowExpirationContext(context); + // This can happen on test completion or drain. We can't set any more timers in window // expiration, so we just have to loop until the schema is updated. BackOff backoff = diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java index 26eb6d55d208..6dc321be2e0a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java @@ -17,23 +17,15 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.sdk.io.gcp.bigquery.UpgradeTableSchema.isPayloadSchemaOutOfDate; - -import com.google.api.services.bigquery.model.TableRow; -import com.google.auto.value.AutoValue; -import com.google.cloud.bigquery.storage.v1.ProtoRows; -import com.google.protobuf.ByteString; import com.google.protobuf.Descriptors; -import java.io.IOException; import java.util.Iterator; -import java.util.List; import java.util.NoSuchElementException; -import java.util.function.BiConsumer; import java.util.function.Function; +import java.util.function.Supplier; +import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.util.ThrowingSupplier; import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; -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.PeekingIterator; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -43,147 +35,74 @@ * parameter controls how many rows are batched into a single ProtoRows object before we move on to * the next one. */ -class SplittingIterable implements Iterable { - @AutoValue - abstract static class Value { - abstract ProtoRows getProtoRows(); - - abstract List getTimestamps(); - - abstract List<@Nullable TableRow> getFailsafeTableRows(); - - abstract boolean getSchemaMismatchSeen(); - } - - interface ConcatFields { - ByteString concat(ByteString bytes, TableRow tableRows) - throws TableRowToStorageApiProto.SchemaConversionException; - } - - private final Iterable underlying; +class SplittingIterable implements Iterable { + private final Iterable underlying; private final long splitSize; - private final ConcatFields concatProtoAndTableRow; - private final Function protoToTableRow; - private final BiConsumer, String> failedRowsConsumer; + private final Function, Boolean> + failedRowsHandler; private final ThrowingSupplier getCurrentTableSchemaHash; private final ThrowingSupplier getCurrentTableSchemaDescriptor; - private final boolean autoUpdateSchema; - private final Instant elementsTimestamp; + private final Instant elementTimestamp; + private final Supplier appendClientSupplier; + private final SchemaChangeDetectorHelper schemaChangeDetectorHelper; public SplittingIterable( - Iterable underlying, + Iterable underlying, long splitSize, - ConcatFields concatProtoAndTableRow, - Function protoToTableRow, - BiConsumer, String> failedRowsConsumer, + Function, Boolean> failedRowsHandler, ThrowingSupplier getCurrentTableSchemaHash, ThrowingSupplier getCurrentTableSchemaDescriptor, - boolean autoUpdateSchema, - Instant elementsTimestamp) { + Instant elementTimestamp, + Supplier appendClientSupplier, + SchemaChangeDetectorHelper schemaChangeDetectorHelper) { this.underlying = underlying; this.splitSize = splitSize; - this.concatProtoAndTableRow = concatProtoAndTableRow; - this.protoToTableRow = protoToTableRow; - this.failedRowsConsumer = failedRowsConsumer; + this.failedRowsHandler = failedRowsHandler; this.getCurrentTableSchemaHash = getCurrentTableSchemaHash; this.getCurrentTableSchemaDescriptor = getCurrentTableSchemaDescriptor; - this.autoUpdateSchema = autoUpdateSchema; - this.elementsTimestamp = elementsTimestamp; + this.elementTimestamp = elementTimestamp; + this.appendClientSupplier = appendClientSupplier; + this.schemaChangeDetectorHelper = schemaChangeDetectorHelper; } @Override - public Iterator iterator() { - return new Iterator() { - final PeekingIterator underlyingIterator = + public Iterator iterator() { + return new Iterator() { + final PeekingIterator underlyingIterator = Iterators.peekingIterator(underlying.iterator()); + @Nullable AppendRowsPacket cachedNext = null; @Override public boolean hasNext() { - return underlyingIterator.hasNext(); + if (cachedNext == null && underlyingIterator.hasNext()) { + cachedNext = calculateNext(); + } + return cachedNext != null; } @Override - public Value next() { + public AppendRowsPacket next() { if (!hasNext()) { throw new NoSuchElementException(); } + AppendRowsPacket result = Preconditions.checkStateNotNull(cachedNext); + cachedNext = null; + return result; + } - List timestamps = Lists.newArrayList(); - List<@Nullable TableRow> failsafeRows = Lists.newArrayList(); - ProtoRows.Builder inserts = ProtoRows.newBuilder(); - long bytesSize = 0; - - boolean schemaMismatchSeen = false; - try { - while (underlyingIterator.hasNext()) { - // Make sure that we don't exceed the split-size length over multiple elements. A single - // element can exceed - // the split threshold, but in that case it should be the only element returned. - if ((bytesSize + underlyingIterator.peek().getPayload().length > splitSize) - && inserts.getSerializedRowsCount() > 0) { - break; - } - StorageApiWritePayload payload = underlyingIterator.next(); - schemaMismatchSeen = - schemaMismatchSeen - || isPayloadSchemaOutOfDate( - payload, getCurrentTableSchemaHash, getCurrentTableSchemaDescriptor); - - ByteString byteString = ByteString.copyFrom(payload.getPayload()); - @Nullable TableRow failsafeTableRow = null; - try { - failsafeTableRow = payload.getFailsafeTableRow(); - } catch (IOException e) { - // Do nothing, table row will be generated later from row bytes - } - if (autoUpdateSchema) { - @Nullable TableRow unknownFields = payload.getUnknownFields(); - if (unknownFields != null && !unknownFields.isEmpty()) { - // Protocol buffer serialization format supports concatenation. We serialize any new - // "known" fields - // into a proto and concatenate to the existing proto. - - try { - byteString = concatProtoAndTableRow.concat(byteString, unknownFields); - } catch (TableRowToStorageApiProto.SchemaConversionException e) { - // This generally implies that ignoreUnknownValues=false and there were still - // unknown values here. - // Reconstitute the TableRow and send it to the failed-rows consumer. - TableRow tableRow = - failsafeTableRow != null - ? failsafeTableRow - : protoToTableRow.apply(byteString); - // TODO(24926, reuvenlax): We need to merge the unknown fields in! Currently we - // only execute this - // codepath when ignoreUnknownFields==true, so we should never hit this codepath. - // However once - // 24926 is fixed, we need to merge the unknownFields back into the main row - // before outputting to the - // failed-rows consumer. - Instant timestamp = payload.getTimestamp(); - if (timestamp == null) { - timestamp = elementsTimestamp; - } - failedRowsConsumer.accept(TimestampedValue.of(tableRow, timestamp), e.toString()); - continue; - } - } - } - inserts.addSerializedRows(byteString); - Instant timestamp = payload.getTimestamp(); - if (timestamp == null) { - timestamp = elementsTimestamp; - } - timestamps.add(timestamp); - failsafeRows.add(failsafeTableRow); - bytesSize += byteString.size(); - } - } catch (Exception e) { - throw new RuntimeException(e); - } - return new AutoValue_SplittingIterable_Value( - inserts.build(), timestamps, failsafeRows, schemaMismatchSeen); + public @Nullable AppendRowsPacket calculateNext() { + AppendRowsPacket value = + AppendRowsPacket.fromStorageApiWritePayload( + underlyingIterator, + splitSize, + schemaChangeDetectorHelper, + elementTimestamp, + appendClientSupplier, + failedRowsHandler, + getCurrentTableSchemaHash, + getCurrentTableSchemaDescriptor); + return value.getProtoRows().getSerializedRowsCount() == 0 ? null : value; } }; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java index d6981858a4b6..1c21260b864c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.DescriptorProtos; import java.io.IOException; import javax.annotation.Nullable; @@ -29,7 +30,7 @@ abstract class StorageApiDynamicDestinations extends DynamicDestinationsHelpers.DelegatingDynamicDestinations { public interface MessageConverter { - com.google.cloud.bigquery.storage.v1.TableSchema getTableSchema(); + TableSchema getTableSchema(); DescriptorProtos.DescriptorProto getDescriptor(boolean includeCdcColumns) throws Exception; @@ -42,6 +43,8 @@ StorageApiWritePayload toMessage( TableRow toFailsafeTableRow(T element); void updateSchemaFromTable() throws IOException, InterruptedException; + + default void updateSchema(TableSchema schema) {} } StorageApiDynamicDestinations(DynamicDestinations inner) { @@ -60,4 +63,17 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { super.setSideInputAccessorFromProcessContext(context); inner.setSideInputAccessorFromProcessContext(context); } + + @Override + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + super.setSideInputAccessorFromOnTimerContext(context); + inner.setSideInputAccessorFromOnTimerContext(context); + } + + @Override + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + super.setSideInputAccessorFromOnWindowExpirationContext(context); + inner.setSideInputAccessorFromOnWindowExpirationContext(context); + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java index 1710d32689c9..ac6aaa2846fb 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java @@ -24,7 +24,6 @@ import com.google.protobuf.Descriptors.Descriptor; import com.google.protobuf.Message; import java.io.IOException; -import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; @@ -48,7 +47,7 @@ public class StorageApiDynamicDestinationsTableRow schemaUpdateOptions; + private final boolean useSchemaUpdatingTableRow; private static final TableSchemaCache SCHEMA_CACHE = new TableSchemaCache(Duration.standardSeconds(1)); @@ -64,7 +63,7 @@ public class StorageApiDynamicDestinationsTableRow schemaUpdateOptions) { + boolean useSchemaUpdatingTableRow) { super(inner); this.formatFunction = formatFunction; this.formatRecordOnFailureFunction = formatRecordOnFailureFunction; @@ -72,7 +71,7 @@ public class StorageApiDynamicDestinationsTableRow getMessageConverter( } }; - return schemaUpdateOptions.isEmpty() + return !useSchemaUpdatingTableRow ? getConverter.apply(getSchema(destination)) : new SchemaUpgradingTableRowConverter( getConverter, options, datasetService, writeStreamService); @@ -143,6 +142,7 @@ public StorageApiWritePayload toMessage( converter.toMessage(element, rowMutationInformation, collectedExceptions); // Set the schema hash on the payload so the next transform knows whether it has an // out-of-date schema. + // TODO: Don't set this for autoUpdateSchema, as it's ignored. payload = payload.toBuilder().setSchemaHash(converter.getSchemaHash()).build(); return payload; @@ -153,9 +153,16 @@ public void updateSchemaFromTable() throws IOException, InterruptedException { SCHEMA_CACHE.refreshSchema( delegate.get().tableReference, datasetService, writeStreamService, bigQueryOptions); // Recycle the internal MessageConverter so that we pick up the new schema from the cache. + // TODO: only do this if the schema has changed. this.delegate.set(getConverter.apply(null)); } + @Override + public void updateSchema(com.google.cloud.bigquery.storage.v1.TableSchema schema) { + TableSchema modelTableSchema = TableRowToStorageApiProto.protoSchemaToTableSchema(schema); + this.delegate.set(getConverter.apply(modelTableSchema)); + } + @Override public TableRow toFailsafeTableRow(T element) { return delegate.get().toFailsafeTableRow(element); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java index 007bba5c6cdf..5bd5a6d0dd38 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java @@ -75,6 +75,7 @@ public class StorageApiLoads private final boolean allowInconsistentWrites; private final boolean allowAutosharding; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; private final boolean ignoreUnknownValues; private final boolean usesCdc; @@ -99,6 +100,7 @@ public StorageApiLoads( boolean allowInconsistentWrites, boolean allowAutosharding, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, boolean propagateSuccessfulStorageApiWrites, Predicate propagateSuccessfulStorageApiWritesPredicate, @@ -120,6 +122,7 @@ public StorageApiLoads( this.allowInconsistentWrites = allowInconsistentWrites; this.allowAutosharding = allowAutosharding; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; this.ignoreUnknownValues = ignoreUnknownValues; if (propagateSuccessfulStorageApiWrites) { this.successfulWrittenRowsTag = new TupleTag<>("successfulPublishedRowsTag"); @@ -195,7 +198,9 @@ public WriteResult expandInconsistent( successfulRowsPredicate, BigQueryStorageApiInsertErrorCoder.of(), TableRowJsonCoder.of(), + destinationCoder, autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, createDisposition, kmsKey, @@ -296,9 +301,11 @@ public WriteResult expandTriggered( successfulWrittenRowsTag, successfulRowsPredicate, autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, defaultMissingValueInterpretation, - bigLakeConfiguration)); + bigLakeConfiguration, + hasSchemaUpdateOptions)); PCollection insertErrors = PCollectionList.of(convertMessagesResult.get(failedRowsTag)) @@ -399,6 +406,7 @@ public WriteResult expandUntriggered( BigQueryStorageApiInsertErrorCoder.of(), TableRowJsonCoder.of(), autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, createDisposition, kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java index 3f00ed67a8a5..280906d08975 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java @@ -83,6 +83,22 @@ static StorageApiWritePayload of( .build(); } + static StorageApiWritePayload of( + byte[] payload, + @Nullable Instant timestamp, + @Nullable byte[] unknownFieldsPayload, + @Nullable byte[] failsafeTableRowPayload, + @Nullable byte[] schemaHash) + throws IOException { + return new AutoValue_StorageApiWritePayload.Builder() + .setPayload(payload) + .setTimestamp(timestamp) + .setUnknownFieldsPayload(unknownFieldsPayload) + .setFailsafeTableRowPayload(failsafeTableRowPayload) + .setSchemaHash(schemaHash) + .build(); + } + public StorageApiWritePayload withTimestamp(Instant instant) { return toBuilder().setTimestamp(instant).build(); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java index 58bbed8ba5a9..3e32249b0e4b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java @@ -23,13 +23,18 @@ import java.util.function.Predicate; import javax.annotation.Nullable; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.joda.time.Duration; /** * A transform to write sharded records to BigQuery using the Storage API. This transform uses the @@ -38,7 +43,7 @@ * writes, use {@link StorageApiWritesShardedRecords} or {@link StorageApiWriteUnshardedRecords}. */ @SuppressWarnings("FutureReturnValueIgnored") -public class StorageApiWriteRecordsInconsistent +public class StorageApiWriteRecordsInconsistent extends PTransform>, PCollectionTuple> { private final StorageApiDynamicDestinations dynamicDestinations; private final BigQueryServices bqServices; @@ -47,10 +52,15 @@ public class StorageApiWriteRecordsInconsistent private final Predicate successfulRowsPredicate; + // This output is effectively ignored, since this code path uses the default stream which is never + // finalized. private final TupleTag> finalizeTag = new TupleTag<>("finalizeTag"); private final Coder failedRowsCoder; private final Coder successfulRowsCoder; + private final Coder destinationCoder; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; + private final @Nullable TupleTag> mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; @@ -66,7 +76,9 @@ public StorageApiWriteRecordsInconsistent( Predicate successfulRowsPredicate, Coder failedRowsCoder, Coder successfulRowsCoder, + Coder destinationCoder, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, @@ -78,9 +90,16 @@ public StorageApiWriteRecordsInconsistent( this.failedRowsTag = failedRowsTag; this.failedRowsCoder = failedRowsCoder; this.successfulRowsCoder = successfulRowsCoder; + this.destinationCoder = destinationCoder; this.successfulRowsTag = successfulRowsTag; this.successfulRowsPredicate = successfulRowsPredicate; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; + if (autoUpdateSchemaStrictTimeout != null) { + this.mismatchedRowsTag = new TupleTag<>("mismatchedRowsTag"); + } else { + this.mismatchedRowsTag = null; + } this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -98,36 +117,88 @@ public PCollectionTuple expand(PCollection fnImpl = + new StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<>( + operationName, + dynamicDestinations, + bqServices, + true, + bigQueryOptions.getStorageApiAppendThresholdBytes(), + bigQueryOptions.getStorageApiAppendThresholdRecordCount(), + bigQueryOptions.getNumStorageWriteApiStreamAppendClients(), + finalizeTag, + failedRowsTag, + successfulRowsTag, + successfulRowsPredicate, + autoUpdateSchema, + autoUpdateSchemaStrictTimeout, + mismatchedRowsTag, + ignoreUnknownValues, + createDisposition, + kmsKey, + usesCdc, + defaultMissingValueInterpretation, + bigQueryOptions.getStorageWriteApiMaxRetries(), + bigLakeConfiguration, + bigQueryOptions.getStorageApiMismatchLocalRetryTimeMilliSec()); + StorageApiWriteUnshardedRecords.WriteRecordsDoFn fn = + new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>(fnImpl); + ; PCollectionTuple result = input.apply( "Write Records", - ParDo.of( - new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>( - operationName, - dynamicDestinations, - bqServices, - true, - bigQueryOptions.getStorageApiAppendThresholdBytes(), - bigQueryOptions.getStorageApiAppendThresholdRecordCount(), - bigQueryOptions.getNumStorageWriteApiStreamAppendClients(), - finalizeTag, - failedRowsTag, - successfulRowsTag, - successfulRowsPredicate, - autoUpdateSchema, - ignoreUnknownValues, - createDisposition, - kmsKey, - usesCdc, - defaultMissingValueInterpretation, - bigQueryOptions.getStorageWriteApiMaxRetries(), - bigLakeConfiguration)) + ParDo.of(fn) .withOutputTags(finalizeTag, tupleTagList) .withSideInputs(dynamicDestinations.getSideInputs())); result.get(failedRowsTag).setCoder(failedRowsCoder); if (successfulRowsTag != null) { result.get(successfulRowsTag).setCoder(successfulRowsCoder); } - return result; + + @Nullable PCollectionTuple mismatchedResult = null; + if (mismatchedRowsTag != null) { + PCollection> mismatchedRows = + result.get(mismatchedRowsTag); + mismatchedRows.setCoder(KvCoder.of(destinationCoder, StoragePayloadWithDeadline.Coder.of())); + mismatchedResult = + mismatchedRows.apply( + "bufferMismatched", + new BufferMismatchedRows<>( + failedRowsCoder, + successfulRowsCoder, + destinationCoder, + dynamicDestinations, + fnImpl, + failedRowsTag, + successfulRowsTag)); + mismatchedResult.get(failedRowsTag).setCoder(failedRowsCoder); + if (successfulRowsTag != null) { + mismatchedResult.get(successfulRowsTag).setCoder(successfulRowsCoder); + } + } + + if (mismatchedResult != null) { + // We need to merge in any results from BufferMismatchRows. + PCollection flattenedErrors = + PCollectionList.of(result.get(failedRowsTag)) + .and(mismatchedResult.get(failedRowsTag)) + .apply("flattenErrors", Flatten.pCollections()); + + PCollectionTuple flattenedResult = PCollectionTuple.of(failedRowsTag, flattenedErrors); + if (successfulRowsTag != null) { + PCollection flattenedSuccesses = + PCollectionList.of(result.get(successfulRowsTag)) + .and(mismatchedResult.get(successfulRowsTag)) + .apply("flattenSucesses", Flatten.pCollections()); + flattenedResult = flattenedResult.and(successfulRowsTag, flattenedSuccesses); + } + return flattenedResult; + } else { + return result; + } } } 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..86573e34f265 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 @@ -17,9 +17,7 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.sdk.io.gcp.bigquery.UpgradeTableSchema.isPayloadSchemaOutOfDate; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.core.ApiFuture; @@ -35,11 +33,12 @@ import com.google.protobuf.DescriptorProtos; import com.google.protobuf.Descriptors.Descriptor; import com.google.protobuf.Descriptors.DescriptorValidationException; -import com.google.protobuf.DynamicMessage; import io.grpc.Status; import java.io.IOException; +import java.io.Serializable; import java.time.Instant; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -49,7 +48,10 @@ import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiConsumer; +import java.util.function.Function; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; @@ -81,14 +83,18 @@ import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.WindowedValues; +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.Predicates; 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.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; @@ -138,6 +144,7 @@ public StorageApiWriteUnshardedRecords( Coder failedRowsCoder, Coder successfulRowsCoder, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, @@ -152,6 +159,7 @@ public StorageApiWriteUnshardedRecords( this.failedRowsCoder = failedRowsCoder; this.successfulRowsCoder = successfulRowsCoder; this.autoUpdateSchema = autoUpdateSchema; + checkState(autoUpdateSchemaStrictTimeout == null, "Not supported in this configuration"); this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -171,30 +179,35 @@ public PCollectionTuple expand(PCollection( - operationName, - dynamicDestinations, - bqServices, - false, - options.getStorageApiAppendThresholdBytes(), - options.getStorageApiAppendThresholdRecordCount(), - options.getNumStorageWriteApiStreamAppendClients(), - finalizeTag, - failedRowsTag, - successfulRowsTag, - successfulRowsPredicate, - autoUpdateSchema, - ignoreUnknownValues, - createDisposition, - kmsKey, - usesCdc, - defaultMissingValueInterpretation, - options.getStorageWriteApiMaxRetries(), - bigLakeConfiguration)) + new WriteRecordsDoFnImpl<>( + operationName, + dynamicDestinations, + bqServices, + false, + options.getStorageApiAppendThresholdBytes(), + options.getStorageApiAppendThresholdRecordCount(), + options.getNumStorageWriteApiStreamAppendClients(), + finalizeTag, + failedRowsTag, + successfulRowsTag, + successfulRowsPredicate, + autoUpdateSchema, + null, + null, + ignoreUnknownValues, + createDisposition, + kmsKey, + usesCdc, + defaultMissingValueInterpretation, + options.getStorageWriteApiMaxRetries(), + bigLakeConfiguration, + options.getStorageApiMismatchLocalRetryTimeMilliSec()))) .withOutputTags(finalizeTag, tupleTagList) .withSideInputs(dynamicDestinations.getSideInputs())); @@ -215,6 +228,155 @@ public PCollectionTuple expand(PCollection extends DoFn, KV> { + WriteRecordsDoFnImpl writeRecordsDoFnImpl; + + WriteRecordsDoFn(WriteRecordsDoFnImpl writeRecordsDoFnImpl) { + this.writeRecordsDoFnImpl = writeRecordsDoFnImpl; + } + + @StartBundle + public void startBundle() throws IOException { + writeRecordsDoFnImpl.startBundle(); + } + + @FinishBundle + public void finishBundle(final FinishBundleContext context) throws Exception { + OutputReceiver failedRowsReceiver = + value -> + WindowedValues.builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + writeRecordsDoFnImpl.failedRowsTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + + OutputReceiver successfulRowsReceiver = null; + if (writeRecordsDoFnImpl.successfulRowsTag != null) { + successfulRowsReceiver = + makeSuccessfulRowsReceiver(context, writeRecordsDoFnImpl.successfulRowsTag); + } + + OutputReceiver> finalizeOutputRecever = + value -> + WindowedValues.>builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + writeRecordsDoFnImpl.finalizeTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + Iterable> mismatchedRows = + writeRecordsDoFnImpl.finishBundle( + failedRowsReceiver, + successfulRowsReceiver, + finalizeOutputRecever, + org.joda.time.Instant.now()); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull( + m.getValue().getStoragePayload().getTimestamp(), + GlobalWindow.INSTANCE.maxTimestamp()); + context.output( + Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag), + m, + ts, + GlobalWindow.INSTANCE); + }); + } + } + + private OutputReceiver makeSuccessfulRowsReceiver( + FinishBundleContext context, TupleTag successfulRowsTag) { + return new OutputReceiver() { + @Override + public OutputBuilder builder(TableRow value) { + return WindowedValues.builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + successfulRowsTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + } + }; + } + + @ProcessElement + public void process( + ProcessContext c, + PipelineOptions pipelineOptions, + @Element KV element, + @Timestamp org.joda.time.Instant elementTs, + MultiOutputReceiver o) + throws Exception { + writeRecordsDoFnImpl.dynamicDestinations.setSideInputAccessorFromProcessContext(c); + org.joda.time.Instant startTimeForLocalRetry = org.joda.time.Instant.now(); + org.joda.time.Instant targetDeadline = + (writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout != null) + ? startTimeForLocalRetry.plus(writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout) + : BoundedWindow.TIMESTAMP_MAX_VALUE; + + Iterable> mismatchedRows = + writeRecordsDoFnImpl.processElement( + pipelineOptions, + KV.of( + element.getKey(), + StoragePayloadWithDeadline.of(element.getValue(), targetDeadline)), + startTimeForLocalRetry, + o); + if (!Iterables.isEmpty(mismatchedRows)) { + final OutputReceiver> mismatchedReceiver = + o.get(Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag)); + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull( + m.getValue().getStoragePayload().getTimestamp(), elementTs); + mismatchedReceiver.outputWithTimestamp(m, ts); + }); + } + } + + @Teardown + public void teardown() { + writeRecordsDoFnImpl.teardown(); + } + + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()); + } + } + + static class WriteRecordsDoFnImpl + implements Serializable { private final Counter forcedFlushes = Metrics.counter(WriteRecordsDoFn.class, "forcedFlushes"); private final TupleTag> finalizeTag; private final TupleTag failedRowsTag; @@ -222,10 +384,13 @@ static class WriteRecordsDoFn private final Predicate successfulRowsPredicate; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; + private final @Nullable TupleTag> + mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; - private final boolean usesCdc; + final boolean usesCdc; private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation; static class AppendRowsContext extends RetryManager.Operation.Context { @@ -250,14 +415,14 @@ public AppendRowsContext( } class DestinationState { + private final DestinationT destination; private final TableDestination tableDestination; private final String tableUrn; private final String shortTableUrn; private String streamName = ""; private @Nullable AppendClientInfo appendClientInfo = null; private long currentOffset = 0; - private List pendingMessages; - private List pendingTimestamps; + private List pendingMessages; private transient @Nullable WriteStreamService maybeWriteStreamService; private final Counter recordsAppended = Metrics.counter(WriteRecordsDoFn.class, "recordsAppended"); @@ -279,8 +444,10 @@ class DestinationState { private final long maxRequestSize; private final boolean includeCdcColumns; + private final SchemaChangeDetectorHelper schemaChangeDetectorHelper; public DestinationState( + DestinationT destination, TableDestination tableDestination, String tableUrn, String shortTableUrn, @@ -293,11 +460,11 @@ public DestinationState( Callable tryCreateTable, boolean includeCdcColumns) throws Exception { + this.destination = destination; this.tableDestination = tableDestination; this.tableUrn = tableUrn; this.shortTableUrn = shortTableUrn; this.pendingMessages = Lists.newArrayList(); - this.pendingTimestamps = Lists.newArrayList(); this.maybeWriteStreamService = writeStreamService; this.useDefaultStream = useDefaultStream; this.messageConverter = messageConverter; @@ -309,6 +476,12 @@ public DestinationState( if (includeCdcColumns) { checkState(useDefaultStream); } + this.schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + autoUpdateSchema, + ignoreUnknownValues, + tableDestination.getTableReference(), + autoUpdateSchemaStrictTimeout != null); } public TableDestination getTableDestination() { @@ -409,22 +582,15 @@ SchemaAndDescriptor getCurrentTableSchema(String stream, @Nullable TableSchema u AtomicBoolean updated = new AtomicBoolean(); CreateTableHelpers.createTableWrapper( () -> { - if (autoUpdateSchema) { - @Nullable - TableSchema streamSchema = - Preconditions.checkStateNotNull(maybeWriteStreamService) - .getWriteStreamSchema(streamName); - if (streamSchema != null) { - Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema( - messageConverter.getTableSchema(), streamSchema); - if (newSchema.isPresent()) { - currentSchema.set(newSchema.get()); - updated.set(true); - LOG.debug( - "Fetched updated schema for table {}:\n\t{}", tableUrn, newSchema.get()); - } - } + Optional newSchema = + schemaChangeDetectorHelper.checkUpdatedSchema( + messageConverter.getTableSchema(), + streamName, + Preconditions.checkStateNotNull(maybeWriteStreamService)); + if (newSchema.isPresent()) { + currentSchema.set(newSchema.get()); + updated.set(true); + LOG.debug("Fetched updated schema for table {}:\n\t{}", tableUrn, newSchema.get()); } return null; }, @@ -484,76 +650,53 @@ void invalidateAppendClient(boolean invalidateCache) { } } - void addMessage( - StorageApiWritePayload payload, - org.joda.time.Instant elementTs, - OutputReceiver failedRowsReceiver) - throws Exception { - maybeTickleCache(); - - @Nullable ByteString mergedPayloadBytes = null; - if (autoUpdateSchema) { - if (appendClientInfo == null) { - appendClientInfo = getAppendClientInfo(true, null); + void writeFailedRows( + Iterable failedRows, + String msg, + org.joda.time.Instant defaultTs, + DoFn.OutputReceiver failedRowsReceiver) + throws IOException { + + for (StoragePayloadWithDeadline mismatchedRow : failedRows) { + TableRow failedRow = mismatchedRow.getStoragePayload().getFailsafeTableRow(); + if (failedRow == null) { + AppendClientInfo aci = getAppendClientInfo(true, null); + failedRow = + aci.toTableRow( + ByteString.copyFrom(mismatchedRow.getStoragePayload().getPayload()), + Predicates.alwaysTrue()); } - @Nullable TableRow unknownFields = payload.getUnknownFields(); - if (unknownFields != null && !unknownFields.isEmpty()) { - // check if unknownFields contains repeated struct, merge - // otherwise use concat - mergedPayloadBytes = ByteString.copyFrom(payload.getPayload()); - try { - mergedPayloadBytes = - Preconditions.checkStateNotNull(appendClientInfo) - .mergeNewFields(mergedPayloadBytes, unknownFields, ignoreUnknownValues); - } catch (TableRowToStorageApiProto.SchemaConversionException e) { - @Nullable TableRow tableRow = payload.getFailsafeTableRow(); - if (tableRow == null) { - tableRow = - checkNotNull(appendClientInfo) - .toTableRow(mergedPayloadBytes, Predicates.alwaysTrue()); - } - // TODO(24926, reuvenlax): We need to merge the unknown fields in! Currently we only - // execute this - // codepath when ignoreUnknownFields==true, so we should never hit this codepath. - // However once - // 24926 is fixed, we need to merge the unknownFields back into the main row before - // outputting to the - // failed-rows consumer. - org.joda.time.Instant timestamp = payload.getTimestamp(); - rowsSentToFailedRowsCollection.inc(); - failedRowsReceiver.outputWithTimestamp( - new BigQueryStorageApiInsertError( - tableRow, e.toString(), tableDestination.getTableReference()), - timestamp != null ? timestamp : elementTs); - return; - } - } - } - byte[] byteArray = - mergedPayloadBytes != null ? mergedPayloadBytes.toByteArray() : payload.getPayload(); - StorageApiWritePayload pending = - StorageApiWritePayload.of(byteArray, null, payload.getFailsafeTableRow()); - byte[] schemaHash = payload.getSchemaHash(); - if (schemaHash != null) { - pending = pending.withSchemaHash(schemaHash); + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, msg, tableDestination.getTableReference()); + org.joda.time.Instant ts = + MoreObjects.firstNonNull(mismatchedRow.getStoragePayload().getTimestamp(), defaultTs); + failedRowsReceiver.outputWithTimestamp(error, ts); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, + tableDestination.getShortTableUrn()) + .inc(1); } - pendingMessages.add(pending); + } - org.joda.time.Instant timestamp = payload.getTimestamp(); - pendingTimestamps.add(timestamp != null ? timestamp : elementTs); + void addMessage(StoragePayloadWithDeadline payload) throws Exception { + maybeTickleCache(); + pendingMessages.add(payload); } long flush( RetryManager retryManager, - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver) + org.joda.time.@Nullable Instant startTimeForLocalRetry, + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + BiConsumer bufferMismatchedRow) throws Exception { if (pendingMessages.isEmpty()) { return 0; } - final ProtoRows.Builder insertsBuilder = ProtoRows.newBuilder(); - List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); boolean schemaMismatchSeen; BackOff backoff = FluentBackoff.DEFAULT @@ -564,26 +707,42 @@ long flush( BigQuerySinkMetrics.throttledTimeCounter( BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM)) .backoff(); + + final Duration initialMismatchRetryTime = + Duration.millis(WriteRecordsDoFnImpl.this.initialMismatchRetryTimeMilliSec); + + AppendRowsPacket rowsToProcess; + Iterable payloadsToIterate = pendingMessages; do { - insertsBuilder.clear(); - failsafeTableRows.clear(); - schemaMismatchSeen = false; - for (StorageApiWritePayload payload : pendingMessages) { - schemaMismatchSeen = - isPayloadSchemaOutOfDate( - payload, - () -> getAppendClientInfo(true, null).getTableSchemaHash(), - () -> - TableRowToStorageApiProto.wrapDescriptorProto( - messageConverter.getDescriptor(includeCdcColumns))); - if (schemaMismatchSeen) { - break; - } + final AppendRowsPacket processingRows = + buildInserts(payloadsToIterate, failedRowsReceiver); + rowsToProcess = processingRows; - insertsBuilder.addSerializedRows(ByteString.copyFrom(payload.getPayload())); - failsafeTableRows.add(payload.getFailsafeTableRow()); - } + schemaMismatchSeen = !processingRows.getSchemaMismatchedRows().isEmpty(); if (schemaMismatchSeen) { + if (autoUpdateSchemaStrictTimeout != null) { + if (startTimeForLocalRetry == null + || startTimeForLocalRetry + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { + // Local retry time has expired! Pull out the mismatched rows so that we can buffer + // them for later retrying. + // Buffer those rows. + processingRows + .getSchemaMismatchedRowsOnly() + .toPayloadStream() + .forEach(row -> bufferMismatchedRow.accept(destination, row)); + + // Process the good rows! + Iterable goodRows = + () -> processingRows.getSchemaMatchedRowsOnly().toPayloadStream().iterator(); + + rowsToProcess = buildInserts(goodRows, failedRowsReceiver); + checkState(rowsToProcess.getSchemaMismatchedRows().isEmpty()); + break; + } + } + LOG.info("Schema out of date: refreshing table schema for {}.", tableUrn); // Refresh our view of the schema and try again.. this.messageConverter.updateSchemaFromTable(); @@ -596,12 +755,15 @@ long flush( } while (schemaMismatchSeen && BackOffUtils.next(Sleeper.DEFAULT, backoff)); pendingMessages.clear(); - final ProtoRows inserts = insertsBuilder.build(); - List insertTimestamps = pendingTimestamps; + + if (rowsToProcess.getProtoRows().getSerializedRowsCount() == 0) { + // All rows either failed or were buffered with schema mismatches. + return 0; + } // Handle the case where the request is too large. - if (inserts.getSerializedSize() >= maxRequestSize) { - if (inserts.getSerializedRowsCount() > 1) { + if (rowsToProcess.getProtoRows().getSerializedSize() >= maxRequestSize) { + if (rowsToProcess.getProtoRows().getSerializedRowsCount() > 1) { // TODO(reuvenlax): Is it worth trying to handle this case by splitting the protoRows? // Given that we split // the ProtoRows iterable at 2MB and the max request size is 10MB, this scenario seems @@ -611,21 +773,14 @@ long flush( + "This is unexpected. All rows in the request will be sent to the failed-rows PCollection.", maxRequestSize); } - for (int i = 0; i < inserts.getSerializedRowsCount(); ++i) { - @Nullable TableRow failedRow = failsafeTableRows.get(i); + for (int i = 0; i < rowsToProcess.getProtoRows().getSerializedRowsCount(); ++i) { + @Nullable TableRow failedRow = rowsToProcess.getFailsafeTableRows().get(i); if (failedRow == null) { - ByteString rowBytes = inserts.getSerializedRows(i); + ByteString rowBytes = rowsToProcess.getProtoRows().getSerializedRows(i); AppendClientInfo aci = getAppendClientInfo(true, null); - failedRow = - TableRowToStorageApiProto.tableRowFromMessage( - aci.getSchemaInformation(), - DynamicMessage.parseFrom( - TableRowToStorageApiProto.wrapDescriptorProto(aci.getDescriptor()), - rowBytes), - true, - successfulRowsPredicate); + failedRow = aci.toTableRow(rowBytes, successfulRowsPredicate); } - org.joda.time.Instant timestamp = insertTimestamps.get(i); + org.joda.time.Instant timestamp = rowsToProcess.getTimestamps().get(i); failedRowsReceiver.outputWithTimestamp( new BigQueryStorageApiInsertError( failedRow, @@ -633,7 +788,7 @@ long flush( tableDestination.getTableReference()), timestamp); } - int numRowsFailed = inserts.getSerializedRowsCount(); + int numRowsFailed = rowsToProcess.getProtoRows().getSerializedRowsCount(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.FAILED, BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, @@ -647,10 +802,14 @@ long flush( if (!this.useDefaultStream) { getOrCreateStreamName(); // Force creation of the stream before we get offsets. offset = this.currentOffset; - this.currentOffset += inserts.getSerializedRowsCount(); + this.currentOffset += rowsToProcess.getProtoRows().getSerializedRowsCount(); } AppendRowsContext appendRowsContext = - new AppendRowsContext(offset, inserts, insertTimestamps, failsafeTableRows); + new AppendRowsContext( + offset, + rowsToProcess.getProtoRows(), + rowsToProcess.getTimestamps(), + rowsToProcess.getFailsafeTableRows()); retryManager.addOperation( c -> { @@ -703,15 +862,7 @@ long flush( ByteString protoBytes = failedContext.protoRows.getSerializedRows(failedIndex); AppendClientInfo aci = Preconditions.checkStateNotNull(this.appendClientInfo); - failedRow = - TableRowToStorageApiProto.tableRowFromMessage( - aci.getSchemaInformation(), - DynamicMessage.parseFrom( - TableRowToStorageApiProto.wrapDescriptorProto( - aci.getDescriptor()), - protoBytes), - true, - Predicates.alwaysTrue()); + failedRow = aci.toTableRow(protoBytes, Predicates.alwaysTrue()); } element = new BigQueryStorageApiInsertError( @@ -894,12 +1045,8 @@ long flush( ByteString rowBytes = c.protoRows.getSerializedRowsList().get(i); try { TableRow row = - TableRowToStorageApiProto.tableRowFromMessage( - Preconditions.checkStateNotNull(appendClientInfo) - .getSchemaInformation(), - DynamicMessage.parseFrom(descriptor, rowBytes), - true, - successfulRowsPredicate); + Preconditions.checkStateNotNull(appendClientInfo) + .toTableRow(rowBytes, successfulRowsPredicate); org.joda.time.Instant timestamp = c.timestamps.get(i); successfulRowsReceiver.outputWithTimestamp(row, timestamp); } catch (Exception e) { @@ -911,7 +1058,42 @@ long flush( }, appendRowsContext); maybeTickleCache(); - return inserts.getSerializedRowsCount(); + return rowsToProcess.getProtoRows().getSerializedRowsCount(); + } + + private AppendRowsPacket buildInserts( + Iterable payloads, + DoFn.OutputReceiver failedRowsReceiver) { + Supplier appendClientInfoSupplier = + () -> { + if (appendClientInfo == null) { + appendClientInfo = getAppendClientInfo(true, null); + } + return appendClientInfo; + }; + + Function, Boolean> failedRowsHandler = + (autoUpdateSchemaStrictTimeout == null) + ? tv -> { + rowsSentToFailedRowsCollection.inc(); + failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp()); + return true; + } + : tv -> false; + + PeekingIterator peekingIterator = + Iterators.peekingIterator(payloads.iterator()); + return AppendRowsPacket.fromStorageApiWritePayload( + peekingIterator, + Long.MAX_VALUE, + schemaChangeDetectorHelper, + BoundedWindow.TIMESTAMP_MAX_VALUE, + appendClientInfoSupplier, + failedRowsHandler, + () -> getAppendClientInfo(true, null).getTableSchemaHash(), + () -> + TableRowToStorageApiProto.wrapDescriptorProto( + messageConverter.getDescriptor(includeCdcColumns))); } String retrieveErrorDetails(Iterable failedContext) { @@ -930,16 +1112,15 @@ String retrieveErrorDetails(Iterable failedContext) { void postFlush() { // If we got a response indicating an updated schema, recreate the client. - if (this.appendClientInfo != null && autoUpdateSchema) { + if (this.appendClientInfo != null + && autoUpdateSchema + && autoUpdateSchemaStrictTimeout == null) { @Nullable StreamAppendClient streamAppendClient = appendClientInfo.getStreamAppendClient(); - @Nullable - TableSchema updatedTableSchemaReturned = - (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null; - if (updatedTableSchemaReturned != null) { + if (streamAppendClient != null) { Optional updatedTableSchema = - TableSchemaUpdateUtils.getUpdatedSchema( - this.messageConverter.getTableSchema(), updatedTableSchemaReturned); + schemaChangeDetectorHelper.checkResponseForUpdatedSchema( + this.messageConverter.getTableSchema(), streamAppendClient); if (updatedTableSchema.isPresent()) { invalidateAppendClient(false); // TODO: This overwrites whatever is in the cache which can cause races between @@ -956,7 +1137,7 @@ void postFlush() { } private @Nullable Map destinations = Maps.newHashMap(); - private final TwoLevelMessageConverterCache messageConverters; + final TwoLevelMessageConverterCache messageConverters; private transient @Nullable DatasetService maybeDatasetService; private transient @Nullable WriteStreamService maybeWriteStreamService; private int numPendingRecords = 0; @@ -964,13 +1145,14 @@ void postFlush() { private final int flushThresholdBytes; private final int flushThresholdCount; private final int maxRetries; - private final StorageApiDynamicDestinations dynamicDestinations; + final StorageApiDynamicDestinations dynamicDestinations; private final BigQueryServices bqServices; private final boolean useDefaultStream; private int streamAppendClientCount; private final @Nullable Map bigLakeConfiguration; + private final int initialMismatchRetryTimeMilliSec; - WriteRecordsDoFn( + WriteRecordsDoFnImpl( String operationName, StorageApiDynamicDestinations dynamicDestinations, BigQueryServices bqServices, @@ -983,13 +1165,16 @@ void postFlush() { @Nullable TupleTag successfulRowsTag, Predicate successfulRowsPredicate, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, + @Nullable TupleTag> mismatchedRowsTag, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, boolean usesCdc, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, int maxRetries, - @Nullable Map bigLakeConfiguration) { + @Nullable Map bigLakeConfiguration, + int initialMismatchRetryTimeMilliSec) { this.messageConverters = new TwoLevelMessageConverterCache<>(operationName); this.dynamicDestinations = dynamicDestinations; this.bqServices = bqServices; @@ -1002,6 +1187,8 @@ void postFlush() { this.successfulRowsTag = successfulRowsTag; this.successfulRowsPredicate = successfulRowsPredicate; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; + this.mismatchedRowsTag = mismatchedRowsTag; this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -1009,6 +1196,7 @@ void postFlush() { this.defaultMissingValueInterpretation = defaultMissingValueInterpretation; this.maxRetries = maxRetries; this.bigLakeConfiguration = bigLakeConfiguration; + this.initialMismatchRetryTimeMilliSec = initialMismatchRetryTimeMilliSec; } boolean shouldFlush(int recordBytes) { @@ -1017,9 +1205,10 @@ boolean shouldFlush(int recordBytes) { && numPendingRecords > 0); } - void flushIfNecessary( - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver, + Iterable> flushIfNecessary( + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry, int recordBytes) throws Exception { if (shouldFlush(recordBytes)) { @@ -1027,14 +1216,17 @@ void flushIfNecessary( // Too much memory being used. Flush the state and wait for it to drain out. // TODO(reuvenlax): Consider waiting for memory usage to drop instead of waiting for all the // appends to finish. - flushAll(failedRowsReceiver, successfulRowsReceiver); + return flushAll(failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry); } + return Collections.emptyList(); } - void flushAll( - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver) + Iterable> flushAll( + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry) throws Exception { + List> mismatchedRows = Lists.newArrayList(); List> retryManagers = Lists.newArrayListWithCapacity(Preconditions.checkStateNotNull(destinations).size()); long numRowsWritten = 0; @@ -1049,7 +1241,12 @@ void flushAll( BigQuerySinkMetrics.RpcMethod.APPEND_ROWS)); retryManagers.add(retryManager); numRowsWritten += - destinationState.flush(retryManager, failedRowsReceiver, successfulRowsReceiver); + destinationState.flush( + retryManager, + startTimeForLocalRetry, + failedRowsReceiver, + successfulRowsReceiver, + (d, m) -> mismatchedRows.add(KV.of(d, m))); retryManager.run(false); } if (numRowsWritten > 0) { @@ -1067,9 +1264,10 @@ void flushAll( } numPendingRecords = 0; numPendingRecordBytes = 0; + return mismatchedRows; } - private DatasetService initializeDatasetService(PipelineOptions pipelineOptions) { + DatasetService getDatasetService(PipelineOptions pipelineOptions) { if (maybeDatasetService == null) { maybeDatasetService = bqServices.getDatasetService(pipelineOptions.as(BigQueryOptions.class)); @@ -1077,7 +1275,7 @@ private DatasetService initializeDatasetService(PipelineOptions pipelineOptions) return maybeDatasetService; } - private WriteStreamService initializeWriteStreamService(PipelineOptions pipelineOptions) { + WriteStreamService getWriteStreamService(PipelineOptions pipelineOptions) { if (maybeWriteStreamService == null) { maybeWriteStreamService = bqServices.getWriteStreamService(pipelineOptions.as(BigQueryOptions.class)); @@ -1085,7 +1283,6 @@ private WriteStreamService initializeWriteStreamService(PipelineOptions pipeline return maybeWriteStreamService; } - @StartBundle public void startBundle() throws IOException { destinations = Maps.newHashMap(); numPendingRecords = 0; @@ -1093,7 +1290,7 @@ public void startBundle() throws IOException { } DestinationState createDestinationState( - ProcessContext c, + PipelineOptions options, DestinationT destination, boolean useCdc, DatasetService datasetService, @@ -1110,7 +1307,7 @@ DestinationState createDestinationState( Callable tryCreateTable = () -> { CreateTableHelpers.possiblyCreateTable( - c.getPipelineOptions().as(BigQueryOptions.class), + options.as(BigQueryOptions.class), tableDestination1, () -> dynamicDestinations.getSchema(destination), () -> dynamicDestinations.getTableConstraints(destination), @@ -1126,12 +1323,9 @@ DestinationState createDestinationState( try { messageConverter = messageConverters.get( - destination, - dynamicDestinations, - c.getPipelineOptions(), - datasetService, - writeStreamService); + destination, dynamicDestinations, options, datasetService, writeStreamService); return new DestinationState( + destination, tableDestination1, tableDestination1.getTableUrn(bigQueryOptions), tableDestination1.getShortTableUrn(), @@ -1148,25 +1342,46 @@ DestinationState createDestinationState( } } - @ProcessElement - public void process( - ProcessContext c, + void writeFailedRows( + DestinationT destination, + Iterable failedRows, + String msg, PipelineOptions pipelineOptions, - @Element KV element, - @Timestamp org.joda.time.Instant elementTs, - MultiOutputReceiver o) + org.joda.time.Instant defaultTs, + DoFn.OutputReceiver failedRowsReceiver) + throws IOException { + DatasetService initializedDatasetService = getDatasetService(pipelineOptions); + WriteStreamService initializedWriteStreamService = getWriteStreamService(pipelineOptions); + DestinationState state = + Preconditions.checkStateNotNull(destinations) + .computeIfAbsent( + destination, + d -> + createDestinationState( + pipelineOptions, + destination, + usesCdc, + initializedDatasetService, + initializedWriteStreamService, + pipelineOptions.as(BigQueryOptions.class))); + state.writeFailedRows(failedRows, msg, defaultTs, failedRowsReceiver); + } + + Iterable> processElement( + PipelineOptions pipelineOptions, + KV element, + org.joda.time.@Nullable Instant startTimeForLocalRetry, + DoFn.MultiOutputReceiver o) throws Exception { - DatasetService initializedDatasetService = initializeDatasetService(pipelineOptions); - WriteStreamService initializedWriteStreamService = - initializeWriteStreamService(pipelineOptions); - dynamicDestinations.setSideInputAccessorFromProcessContext(c); + DatasetService initializedDatasetService = getDatasetService(pipelineOptions); + WriteStreamService initializedWriteStreamService = getWriteStreamService(pipelineOptions); DestinationState state = Preconditions.checkStateNotNull(destinations) .computeIfAbsent( element.getKey(), destination -> createDestinationState( - c, + pipelineOptions, destination, usesCdc, initializedDatasetService, @@ -1179,92 +1394,44 @@ public void process( state.getTableDestination().getTableReference(), pipelineOptions.as(BigQueryOptions.class))); - OutputReceiver failedRowsReceiver = o.get(failedRowsTag); - @Nullable - OutputReceiver successfulRowsReceiver = + DoFn.OutputReceiver failedRowsReceiver = o.get(failedRowsTag); + DoFn.@Nullable OutputReceiver successfulRowsReceiver = (successfulRowsTag != null) ? o.get(successfulRowsTag) : null; - int recordBytes = element.getValue().getPayload().length; - flushIfNecessary(failedRowsReceiver, successfulRowsReceiver, recordBytes); - state.addMessage(element.getValue(), elementTs, failedRowsReceiver); + int recordBytes = element.getValue().getStoragePayload().getPayload().length; + Iterable> mismatchedRows = + flushIfNecessary( + failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry, recordBytes); + state.addMessage(element.getValue()); ++numPendingRecords; numPendingRecordBytes += recordBytes; - } - private OutputReceiver makeSuccessfulRowsreceiver( - FinishBundleContext context, TupleTag successfulRowsTag) { - return new OutputReceiver() { - @Override - public OutputBuilder builder(TableRow value) { - return WindowedValues.builder() - .setValue(value) - .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) - .setWindow(GlobalWindow.INSTANCE) - .setPaneInfo(PaneInfo.NO_FIRING) - .setReceiver( - windowedValue -> { - for (BoundedWindow window : windowedValue.getWindows()) { - context.output( - successfulRowsTag, - windowedValue.getValue(), - windowedValue.getTimestamp(), - window); - } - }); - } - }; + return mismatchedRows; } - @FinishBundle - public void finishBundle(FinishBundleContext context) throws Exception { - - OutputReceiver failedRowsReceiver = - new OutputReceiver() { - @Override - public OutputBuilder builder( - BigQueryStorageApiInsertError value) { - return WindowedValues.builder() - .setValue(value) - .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) - .setWindow(GlobalWindow.INSTANCE) - .setPaneInfo(PaneInfo.NO_FIRING) - .setReceiver( - windowedValue -> { - for (BoundedWindow window : windowedValue.getWindows()) { - context.output( - failedRowsTag, - windowedValue.getValue(), - windowedValue.getTimestamp(), - window); - } - }); - } - }; - - @Nullable OutputReceiver successfulRowsReceiver = null; - if (successfulRowsTag != null) { - successfulRowsReceiver = makeSuccessfulRowsreceiver(context, successfulRowsTag); - } + public Iterable> finishBundle( + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + DoFn.OutputReceiver> finalizeOutputReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry) + throws Exception { - flushAll(failedRowsReceiver, successfulRowsReceiver); + Iterable> mismatchedRows = + flushAll(failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry); final Map destinations = Preconditions.checkStateNotNull(this.destinations); for (DestinationState state : destinations.values()) { if (!useDefaultStream && !Strings.isNullOrEmpty(state.streamName)) { - context.output( - finalizeTag, - KV.of(state.tableUrn, state.streamName), - GlobalWindow.INSTANCE.maxTimestamp(), - GlobalWindow.INSTANCE); + finalizeOutputReceiver.output(KV.of(state.tableUrn, state.streamName)); } state.teardown(); } destinations.clear(); this.destinations = null; + return mismatchedRows; } - @Teardown public void teardown() { destinations = null; try { @@ -1280,10 +1447,5 @@ public void teardown() { throw new RuntimeException(e); } } - - @Override - public Duration getAllowedTimestampSkew() { - return Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()); - } } } 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..82c396d76c26 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 @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -33,6 +34,7 @@ import com.google.cloud.bigquery.storage.v1.WriteStream; import com.google.protobuf.ByteString; import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; import io.grpc.Status; import io.grpc.Status.Code; import java.io.IOException; @@ -48,10 +50,13 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.extensions.protobuf.ProtoCoder; +import org.apache.beam.sdk.function.ThrowingConsumer; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StreamAppendClient; @@ -66,6 +71,7 @@ import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.schemas.SchemaRegistry; +import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.TimeDomain; @@ -128,6 +134,7 @@ public class StorageApiWritesShardedRecords destinationCoder; private final Coder failedRowsCoder; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; private final boolean ignoreUnknownValues; private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation; private final @Nullable Map bigLakeConfiguration; @@ -140,6 +147,7 @@ public class StorageApiWritesShardedRecords succussfulRowsCoder; private final TupleTag> flushTag = new TupleTag<>("flushTag"); + private final boolean managedSchemaUpdate; private static final AppendClientCache>> APPEND_CLIENTS = new AppendClientCache<>(Duration.standardMinutes(5)); @@ -160,9 +168,11 @@ public StorageApiWritesShardedRecords( @Nullable TupleTag successfulRowsTag, Predicate successfulRowsPredicate, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, - @Nullable Map bigLakeConfiguration) { + @Nullable Map bigLakeConfiguration, + boolean managedSchemaUpdate) { this.dynamicDestinations = dynamicDestinations; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -174,9 +184,11 @@ public StorageApiWritesShardedRecords( this.successfulRowsPredicate = successfulRowsPredicate; this.succussfulRowsCoder = successfulRowsCoder; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; this.ignoreUnknownValues = ignoreUnknownValues; this.defaultMissingValueInterpretation = defaultMissingValueInterpretation; this.bigLakeConfiguration = bigLakeConfiguration; + this.managedSchemaUpdate = managedSchemaUpdate; } @Override @@ -344,11 +356,25 @@ class WriteRecordsDoFn private final StateSpec> streamOffsetSpec = StateSpecs.value(); @StateId("updatedSchema") - private final StateSpec> updatedSchema = + private final StateSpec> updatedSchemaSpe = StateSpecs.value(ProtoCoder.of(TableSchema.class)); + @StateId("mismatchedRows") + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(StoragePayloadWithDeadline.Coder.of()); + + @TimerId("retryMismatchedRowsTimer") + private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + + @StateId("currentMismatchedRowTimerValue") + private final StateSpec> currentMismatchedRowTimerValueSpec = + StateSpecs.value(); + + @StateId("minPendingTimestamp") + private final StateSpec> minPendingTimestampSpec = StateSpecs.value(); + @TimerId("idleTimer") - private final TimerSpec idleTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + private final TimerSpec idleTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); private final Duration streamIdleTime; private final long splitSize; @@ -379,7 +405,7 @@ String getOrCreateStream( String tableId, ValueState streamName, ValueState streamOffset, - Timer streamIdleTimer, + @Nullable Timer streamIdleTimer, WriteStreamService writeStreamService, Callable tryCreateTable) { try { @@ -404,7 +430,9 @@ String getOrCreateStream( stream.set(streamValue); } // Reset the idle timer. - streamIdleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + if (streamIdleTimer != null) { + streamIdleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + } return stream.get(); } catch (Exception e) { @@ -475,8 +503,7 @@ public void close() { } AppendClientInfo get() { - org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState( - valid); + checkState(valid); return appendClientInfo; } @@ -487,7 +514,7 @@ StreamAppendClient getStreamAppendClient() { private CreateRetryManagerResult createRetryManager( ShardedKey key, - Iterable messages, + Iterable messages, Function, ApiFuture> runOperation, Function>, RetryType> onError, Consumer> onSuccess, @@ -504,8 +531,8 @@ private CreateRetryManagerResult createRetryManager( List> failedRows = Lists.newArrayList(); int recordsAppended = 0; List histogramUpdates = Lists.newArrayList(); - for (SplittingIterable.Value splitValue : messages) { - if (splitValue.getSchemaMismatchSeen()) { + for (AppendRowsPacket splitValue : messages) { + if (!splitValue.getSchemaMismatchedRows().isEmpty()) { return CreateRetryManagerResult.schemaMismatch(); } // Handle the case of a row that is too large. @@ -767,18 +794,16 @@ public void process( final PipelineOptions pipelineOptions, @Element KV, Iterable> element, @Timestamp org.joda.time.Instant elementTs, - final @AlwaysFetched @StateId("streamName") ValueState streamName, - final @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, - final @StateId("updatedSchema") ValueState updatedSchema, + @AlwaysFetched @StateId("streamName") ValueState streamName, + @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @StateId("updatedSchema") ValueState updatedSchema, @TimerId("idleTimer") Timer idleTimer, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @TimerId("retryMismatchedRowsTimer") Timer mismatchedRowsRetryTimer, + @StateId("currentMismatchedRowTimerValue") ValueState mismatchedRowsRetryTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, final MultiOutputReceiver o) throws Exception { - BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class); - - if (autoUpdateSchema) { - updatedSchema.readLater(); - } - dynamicDestinations.setSideInputAccessorFromProcessContext(c); TableDestination tableDestination = destinations.computeIfAbsent( @@ -793,11 +818,104 @@ public void process( dest); return tableDestination1; }); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + element.getKey().getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + + ThrowingConsumer> processMismatchedRows = + mismatchedRows -> { + if (!Iterables.isEmpty(mismatchedRows)) { + AppendClientInfo info = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + Duration timerRetryDuration = + Duration.millis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchRetryTimeMilliSec()); + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRows, + mismatchedRowsBag, + mismatchedRowsRetryTimer, + mismatchedRowsRetryTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + info, + rowsSentToFailedRowsCollection, + timerRetryDuration); + } + }; + + org.joda.time.Instant now = org.joda.time.Instant.now(); + org.joda.time.Instant targetDeadline = + (autoUpdateSchemaStrictTimeout != null) + ? now.plus(autoUpdateSchemaStrictTimeout) + : BoundedWindow.TIMESTAMP_MAX_VALUE; + + Iterable withDeadlines = + () -> + StreamSupport.stream(element.getValue().spliterator(), false) + .map(p -> StoragePayloadWithDeadline.of(p, targetDeadline)) + .iterator(); + processPayloads( + pipelineOptions, + element.getKey(), + tableDestination, + messageConverter, + withDeadlines, + now, + elementTs, + streamName, + streamOffset, + updatedSchema, + idleTimer, + o, + processMismatchedRows); + } + + private void processPayloads( + final PipelineOptions pipelineOptions, + ShardedKey shardedDestination, + TableDestination tableDestination, + StorageApiDynamicDestinations.MessageConverter messageConverter, + Iterable element, + org.joda.time.@Nullable Instant startTimeForLocalRetry, + org.joda.time.Instant elementTs, + final ValueState streamName, + final ValueState streamOffset, + final ValueState updatedSchema, + @Nullable Timer idleTimer, + final MultiOutputReceiver o, + ThrowingConsumer> processMismatchedRows) + throws Exception { + BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class); + + if (autoUpdateSchema) { + updatedSchema.readLater(); + } + + final DestinationT destination = shardedDestination.getKey(); + final String tableId = tableDestination.getTableUrn(bigQueryOptions); final String shortTableId = tableDestination.getShortTableUrn(); final TableReference tableReference = tableDestination.getTableReference(); final DatasetService datasetService = getDatasetService(pipelineOptions); final WriteStreamService writeStreamService = getWriteStreamService(pipelineOptions); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + autoUpdateSchema, + ignoreUnknownValues, + tableReference, + autoUpdateSchemaStrictTimeout != null); Lineage.getSinks() .add( @@ -808,12 +926,11 @@ public void process( Coder destinationCoder = dynamicDestinations.getDestinationCoder(); Callable tryCreateTable = () -> { - DestinationT dest = element.getKey().getKey(); CreateTableHelpers.possiblyCreateTable( - c.getPipelineOptions().as(BigQueryOptions.class), + bigQueryOptions, tableDestination, - () -> dynamicDestinations.getSchema(dest), - () -> dynamicDestinations.getTableConstraints(dest), + () -> dynamicDestinations.getSchema(destination), + () -> dynamicDestinations.getTableConstraints(destination), createDisposition, destinationCoder, kmsKey, @@ -827,20 +944,16 @@ public void process( getOrCreateStream( tableId, streamName, streamOffset, idleTimer, writeStreamService, tryCreateTable); - StorageApiDynamicDestinations.MessageConverter messageConverter = - messageConverters.get( - element.getKey().getKey(), - dynamicDestinations, - pipelineOptions, - datasetService, - writeStreamService); Callable getAppendClientInfo = () -> { @Nullable TableSchema tableSchema; DescriptorProtos.DescriptorProto descriptor; + TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null; - if (autoUpdateSchema && updatedSchemaValue != null) { + if (autoUpdateSchema + && updatedSchemaValue != null + && autoUpdateSchemaStrictTimeout == null) { // This means that Vortex has told us in the past that the table schema has been // updated. We should use // this updated schema instead of the initial schema from the messageConverter. @@ -849,23 +962,23 @@ public void process( TableRowToStorageApiProto.descriptorSchemaFromTableSchema( tableSchema, true, false); } else { - // Start off with the base schema. As we get notified of schema updates, we - // will update the descriptor. + if (autoUpdateSchemaStrictTimeout != null && updatedSchemaValue != null) { + Optional updated = + TableSchemaUpdateUtils.getUpdatedSchema( + messageConverter.getTableSchema(), updatedSchemaValue); + updated.ifPresent(messageConverter::updateSchema); + } + tableSchema = messageConverter.getTableSchema(); descriptor = messageConverter.getDescriptor(false); - if (autoUpdateSchema) { + if (autoUpdateSchema && autoUpdateSchemaStrictTimeout == null) { // A StreamWriter ignores table schema updates that happen prior to its creation. // So before creating a StreamWriter below, we fetch the table schema to check if we // missed an update. If so, use the new schema instead of the base schema. - // TODO: There's still a race here! - @Nullable - TableSchema streamSchema = - MoreObjects.firstNonNull( - writeStreamService.getWriteStreamSchema(getOrCreateStream.get()), - TableSchema.getDefaultInstance()); Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema(tableSchema, streamSchema); + schemaChangeDetectorHelper.checkUpdatedSchema( + tableSchema, getOrCreateStream.get(), writeStreamService); if (newSchema.isPresent()) { tableSchema = newSchema.get(); @@ -895,8 +1008,9 @@ public void process( // cache could in // theory evict the object during execution and we want a pin held throughout the execution of // this function. + Iterable mismatchedRows = Collections.emptyList(); try (AppendClientHolder appendClientHolder = - new AppendClientHolder(element.getKey(), getAppendClientInfo)) { + new AppendClientHolder(shardedDestination, getAppendClientInfo)) { String currentStream = getOrCreateStream.get(); if (!currentStream.equals(appendClientHolder.get().getStreamName())) { // Cached append client is inconsistent with persisted state. Throw away cached item and @@ -981,66 +1095,127 @@ public void process( BigQuerySinkMetrics.throttledTimeCounter( BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM)) .backoff(); + CreateRetryManagerResult createRetryManagerResult; + final Duration initialMismatchRetryTime = + Duration.millis( + bigQueryOptions.getStorageApiMismatchLocalRetryTimeMilliSec()); // Retry locally + Iterable payloadsToIterate = element; do { // Each ProtoRows object contains at most 1MB of rows. // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely // if // already proto or already schema. - Iterable messages = - new SplittingIterable( - element.getValue(), - splitSize, - // Unknown field merger - (bytes, tableRow) -> - appendClientHolder.get().mergeNewFields(bytes, tableRow, ignoreUnknownValues), - // Convert back to TableRow - bytes -> appendClientHolder.get().toTableRow(bytes, Predicates.alwaysTrue()), - // Failed rows consumer - (failedRow, errorMessage) -> { + + // If either managedSchemaUpdate==true (implying that the schema will be updated by + // another branch of this + // same pipeline) or autoUpdateSchemaStrictTimeout != null (implying that we want to + // strictly buffer records + // until the schema is updated or the timeout expires), we want the failedRowsHandler to + // be a noop. + Function, Boolean> failedRowsHandler = + (autoUpdateSchemaStrictTimeout == null) && !managedSchemaUpdate + ? (TimestampedValue error) -> { o.get(failedRowsTag) - .outputWithTimestamp( - new BigQueryStorageApiInsertError( - failedRow.getValue(), errorMessage, tableReference), - failedRow.getTimestamp()); + .outputWithTimestamp(error.getValue(), error.getTimestamp()); rowsSentToFailedRowsCollection.inc(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.FAILED, BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, shortTableId) .inc(1); - }, + return true; + } + : e -> false; + final byte[] currentTableSchemaHash = appendClientHolder.get().getTableSchemaHash(); + final Descriptors.Descriptor currentDescriptorProto = + TableRowToStorageApiProto.wrapDescriptorProto(messageConverter.getDescriptor(false)); + final Iterable messages = + new SplittingIterable( + payloadsToIterate, + splitSize, + failedRowsHandler, // Get the currently-known TableSchema hash - () -> appendClientHolder.get().getTableSchemaHash(), - () -> - TableRowToStorageApiProto.wrapDescriptorProto( - messageConverter.getDescriptor(false)), - autoUpdateSchema, - elementTs); + () -> currentTableSchemaHash, + () -> currentDescriptorProto, + elementTs, + appendClientHolder::get, + schemaChangeDetectorHelper); + Iterable messagesToProcess = messages; createRetryManagerResult = createRetryManager( - element.getKey(), - messages, + shardedDestination, + messagesToProcess, runOperation, onError, onSuccess, appendClientHolder.get(), tableReference); + if (createRetryManagerResult.getSchemaMismatchSeen()) { + if (autoUpdateSchemaStrictTimeout != null) { + if (startTimeForLocalRetry == null + || startTimeForLocalRetry + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { + // Local retry time has expired! Pull out the mismatched rows so that we can buffer + // them for later + // retrying. + + mismatchedRows = + () -> + StreamSupport.stream(messages.spliterator(), false) + .map(AppendRowsPacket::getSchemaMismatchedRowsOnly) + .flatMap(AppendRowsPacket::toPayloadStream) + .iterator(); + + // Continue processing only the messages that matched the schema. + messagesToProcess = + () -> + StreamSupport.stream(messages.spliterator(), false) + .map(AppendRowsPacket::getSchemaMatchedRowsOnly) + .iterator(); + + createRetryManagerResult = + createRetryManager( + shardedDestination, + messagesToProcess, + runOperation, + onError, + onSuccess, + appendClientHolder.get(), + tableReference); + checkState(!createRetryManagerResult.getSchemaMismatchSeen()); + // Break out of the loop and start processing. + break; + } + } + // TODO: The call to updateSchemaFromTable will throttle the DoFn (both because of the // RPC // call and because // the cache has a delay on refresh). We should update throttling counters here as well. LOG.info("Schema out of date: refreshing table schema for {}", tableId); + LOG.info( + "DEBUG: updatedSchemaValue={}, autoUpdateSchemaStrictTimeout={}, autoUpdateSchema={}", + updatedSchemaValue, + autoUpdateSchemaStrictTimeout, + autoUpdateSchema); // Force the message converter to get the schema again from the table. messageConverter.updateSchemaFromTable(); + if (autoUpdateSchemaStrictTimeout != null) { + updatedSchema.write(messageConverter.getTableSchema()); + } // Close all RPC clients that were opened with the old descriptor. Clear the cache, // forcing us to create a new append client with the updated descriptor. appendClientHolder.invalidateAndReset(); } } while (createRetryManagerResult.getSchemaMismatchSeen() && BackOffUtils.next(Sleeper.DEFAULT, backoff)); + if (createRetryManagerResult.getSchemaMismatchSeen()) { + throw new RuntimeException("Timed out waiting for schema update"); + } // Output any rows that failed along they way. createRetryManagerResult @@ -1070,21 +1245,19 @@ public void process( appendSplitDistribution.update(numAppends); if (autoUpdateSchema) { - @Nullable StreamAppendClient streamAppendClient = appendClientHolder.getStreamAppendClient(); TableSchema originalSchema = appendClientHolder.get().getTableSchema(); - @Nullable - TableSchema updatedSchemaReturned = - (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null; - // Update the table schema and clear the append client. - if (updatedSchemaReturned != null) { - Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema(originalSchema, updatedSchemaReturned); - if (newSchema.isPresent()) { - APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey())); - LOG.debug( - "Fetched updated schema for table {}:\n\t{}", tableId, updatedSchemaReturned); + Optional newSchema = + schemaChangeDetectorHelper.checkResponseForUpdatedSchema( + originalSchema, streamAppendClient); + if (newSchema.isPresent()) { + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(shardedDestination)); + LOG.debug("Fetched updated schema for table {}:\n\t{}", tableId, newSchema.get()); + if (autoUpdateSchemaStrictTimeout != null) { + messageConverter.updateSchemaFromTable(); + updatedSchema.write(messageConverter.getTableSchema()); + } else { updatedSchema.write(newSchema.get()); } } @@ -1093,8 +1266,120 @@ public void process( java.time.Duration timeElapsed = java.time.Duration.between(now, Instant.now()); appendLatencyDistribution.update(timeElapsed.toMillis()); } + processMismatchedRows.accept(mismatchedRows); } - idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + if (idleTimer != null) { + idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + } + } + + @OnTimer("retryMismatchedRowsTimer") + public void onMismatchedRowsTimer( + OnTimerContext context, + PipelineOptions pipelineOptions, + @Key ShardedKey shardedDestination, + @Timestamp org.joda.time.Instant elementTs, + @AlwaysFetched @StateId("streamName") ValueState streamName, + @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @StateId("updatedSchema") ValueState updatedSchema, + @TimerId("idleTimer") Timer idleTimer, + MultiOutputReceiver o, + @TimerId("retryMismatchedRowsTimer") Timer retryRowsTimer, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnTimerContext(context); + + mismatchedRowsBag.readLater(); + currentTimerValue.readLater(); + minPendingTimestamp.readLater(); + + TableDestination tableDestination = + destinations.computeIfAbsent( + shardedDestination.getKey(), + dest -> { + TableDestination tableDestination1 = dynamicDestinations.getTable(dest); + checkArgument( + tableDestination1 != null, + "DynamicDestinations.getTable() may not return null, " + + "but %s returned null for destination %s", + dynamicDestinations, + dest); + return tableDestination1; + }); + + Iterable payloads = mismatchedRowsBag.read(); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + LOG.info( + "Schema out of date: refreshing table schema for {}", + tableDestination.getShortTableUrn()); + // Force the message converter to get the schema again from the table. + messageConverter.updateSchemaFromTable(); + updatedSchema.write(messageConverter.getTableSchema()); + + // Close all RPC clients that were opened with the old descriptor. Clear the cache, + // forcing us to create a new append client with the updated descriptor. + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(shardedDestination)); + + // Try to reprocess all of these rows. + ThrowingConsumer> processMismatchedRows = + mismatchedRows -> { + // TODO: Try to refactor so that we don't have to materialize this list. + List mismatchedRowsList = + StreamSupport.stream(mismatchedRows.spliterator(), false) + .collect(Collectors.toList()); + // Rebuffer the ones that are still not succeeding. + mismatchedRowsBag.clear(); + currentTimerValue.clear(); + minPendingTimestamp.clear(); + if (!Iterables.isEmpty(mismatchedRowsList)) { + AppendClientInfo info = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + Duration timerRetryDuration = + Duration.millis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchRetryTimeMilliSec()); + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRowsList, + mismatchedRowsBag, + retryRowsTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + info, + rowsSentToFailedRowsCollection, + timerRetryDuration); + } + }; + + processPayloads( + pipelineOptions, + shardedDestination, + tableDestination, + messageConverter, + payloads, + null, + elementTs, + streamName, + streamOffset, + updatedSchema, + idleTimer, + o, + processMismatchedRows); } // called by the idleTimer and window-expiration handlers. @@ -1138,11 +1423,103 @@ public void onTimer( @OnWindowExpiration public void onWindowExpiration( + PipelineOptions pipelineOptions, @Key ShardedKey key, + @Timestamp org.joda.time.Instant elementTs, + @StateId("updatedSchema") ValueState updatedSchema, @AlwaysFetched @StateId("streamName") ValueState streamName, @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @AlwaysFetched @StateId("mismatchedRows") + BagState mismatchedRowsBag, MultiOutputReceiver o, - BoundedWindow window) { + BoundedWindow window) + throws Exception { + TableDestination tableDestination = + destinations.computeIfAbsent( + key.getKey(), + dest -> { + TableDestination tableDestination1 = dynamicDestinations.getTable(dest); + checkArgument( + tableDestination1 != null, + "DynamicDestinations.getTable() may not return null, " + + "but %s returned null for destination %s", + dynamicDestinations, + dest); + return tableDestination1; + }); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + key.getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + + java.time.Duration waitTime = + java.time.Duration.ofMillis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchDrainRetryTimeMilliSec()); + AtomicReference> mismatchedRows = + new AtomicReference<>(mismatchedRowsBag.read()); + Instant start = Instant.now(); + while (!Iterables.isEmpty(mismatchedRows.get()) + && start.plus(waitTime).isAfter(Instant.now())) { + messageConverter.updateSchemaFromTable(); + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(key)); + processPayloads( + pipelineOptions, + key, + tableDestination, + messageConverter, + mismatchedRows.get(), + null, + elementTs, + streamName, + streamOffset, + updatedSchema, + null, + o, + mismatchedRows::set); + } + + if (mismatchedRows.get() != null) { + // At this point, there's no more waiting. Output the remaining elements to the failed-rows + // collection. + AppendClientInfo appendClientInfo = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + for (StoragePayloadWithDeadline mismatchedRow : mismatchedRows.get()) { + TableRow failedRow = mismatchedRow.getStoragePayload().getFailsafeTableRow(); + if (failedRow == null) { + failedRow = + appendClientInfo.toTableRow( + ByteString.copyFrom(mismatchedRow.getStoragePayload().getPayload()), + Predicates.alwaysTrue()); + } + + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, + "Timed out waiting for table schema update in OnWindowExpiration", + tableDestination.getTableReference()); + org.joda.time.Instant ts = + MoreObjects.firstNonNull(mismatchedRow.getStoragePayload().getTimestamp(), elementTs); + o.get(failedRowsTag).outputWithTimestamp(error, ts); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, + tableDestination.getShortTableUrn()) + .inc(1); + } + ; + } + // Window is done - usually because the pipeline has been drained. Make sure to clean up // streams so that they are not leaked. finalizeStream(streamName, streamOffset, key, o, window.maxTimestamp()); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java new file mode 100644 index 000000000000..8000986405fc --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import com.google.auto.value.AutoValue; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +@DefaultSchema(AutoValueSchema.class) +@AutoValue +abstract class StoragePayloadWithDeadline { + abstract StorageApiWritePayload getStoragePayload(); + + abstract org.joda.time.Instant getDeadline(); + + static StoragePayloadWithDeadline of( + StorageApiWritePayload payload, org.joda.time.Instant deadline) { + return new AutoValue_StoragePayloadWithDeadline(payload, deadline); + } + + // Schemas give us a coder, however there are still some limitations to storing schema objects + // inside of state + // variables (mostly involving the Dataflow runner and update). Therefore we use a custom coder. + static class Coder extends CustomCoder { + static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of(); + static final InstantCoder INSTANT_CODER = InstantCoder.of(); + static final NullableCoder NULLABLE_BYTE_ARRAY_CODER = + NullableCoder.of(BYTE_ARRAY_CODER); + static final NullableCoder NULLABLE_INSTANT_CODER = + NullableCoder.of(InstantCoder.of()); + + static StoragePayloadWithDeadline.Coder of() { + return new Coder(); + } + + @Override + public void encode(StoragePayloadWithDeadline value, OutputStream outStream) + throws CoderException, IOException { + BYTE_ARRAY_CODER.encode(value.getStoragePayload().getPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode( + value.getStoragePayload().getUnknownFieldsPayload(), outStream); + NULLABLE_INSTANT_CODER.encode(value.getStoragePayload().getTimestamp(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode( + value.getStoragePayload().getFailsafeTableRowPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode(value.getStoragePayload().getSchemaHash(), outStream); + INSTANT_CODER.encode(value.getDeadline(), outStream); + } + + @Override + public StoragePayloadWithDeadline decode(InputStream inStream) + throws CoderException, IOException { + byte[] innerPayload = BYTE_ARRAY_CODER.decode(inStream); + byte @Nullable [] unknownFieldsPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + @Nullable Instant timestamp = NULLABLE_INSTANT_CODER.decode(inStream); + byte @Nullable [] failsafeTableRowPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + byte @Nullable [] schemaHash = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + Instant deadline = INSTANT_CODER.decode(inStream); + + StorageApiWritePayload payload = + StorageApiWritePayload.of( + innerPayload, timestamp, unknownFieldsPayload, failsafeTableRowPayload, schemaHash); + return new AutoValue_StoragePayloadWithDeadline(payload, deadline); + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java index ba72bb8682fd..898c486baa7f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java @@ -1193,6 +1193,7 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) if (!collectedExceptions.isEmpty()) { return null; } + try { return builder.build(); } catch (Exception e) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java index 6c3c028b6d0d..a6cc916b7802 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java @@ -17,18 +17,22 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; +import com.google.api.services.bigquery.model.TableCell; +import com.google.cloud.bigquery.storage.v1.BigQuerySchemaUtil; import com.google.cloud.bigquery.storage.v1.TableFieldSchema; import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.Descriptors; import com.google.protobuf.DynamicMessage; import com.google.protobuf.Message; import com.google.protobuf.UnknownFieldSet; +import java.util.AbstractMap; import java.util.Arrays; import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Supplier; import org.apache.beam.sdk.util.ThrowingSupplier; 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; @@ -247,26 +251,31 @@ private static TableFieldSchema mergeField(TableFieldSchema f1, TableFieldSchema } public static boolean isPayloadSchemaOutOfDate( - StorageApiWritePayload payload, + byte @Nullable [] payloadSchemaHash, + Supplier payloadSupplier, ThrowingSupplier schemaHash, ThrowingSupplier schemaDescriptor) throws Exception { - byte @Nullable [] payloadSchemaHash = payload.getSchemaHash(); if (payloadSchemaHash != null) { // Schema hash is only included in the payload if schema update options are set. HashCode lhs = HashCode.fromBytes(payloadSchemaHash); HashCode rhs = HashCode.fromBytes(schemaHash.get()); if (!lhs.equals(rhs)) { - DynamicMessage msg = - DynamicMessage.newBuilder(schemaDescriptor.get()) - .mergeFrom(payload.getPayload()) - .buildPartial(); - return !msg.isInitialized() || hasUnknownFields(msg); + return isPayloadSchemaOutOfDate(payloadSupplier, schemaDescriptor); } } return false; } + public static boolean isPayloadSchemaOutOfDate( + Supplier payloadSupplier, ThrowingSupplier schemaDescriptor) + throws Exception { + Descriptors.Descriptor descriptor = schemaDescriptor.get(); + byte[] payload = payloadSupplier.get(); + DynamicMessage msg = DynamicMessage.newBuilder(descriptor).mergeFrom(payload).buildPartial(); + return !msg.isInitialized() || hasUnknownFields(msg); + } + private static boolean hasUnknownFields(Message message) { if (message == null) { return false; @@ -305,4 +314,88 @@ private static boolean hasUnknownFields(Message message) { // If we reach here, neither this message nor its descendants have unknown fields return false; } + + public static boolean missingUnknownField( + AbstractMap unknownFields, + ThrowingSupplier schemaDescriptor) + throws Exception { + @Nullable Object fValue = unknownFields.get("f"); + if (fValue instanceof List) { + List cells = (List) fValue; + return missingUnknownField(cells, schemaDescriptor.get()); + } else { + return missingUnknownField(unknownFields, schemaDescriptor.get()); + } + } + + public static boolean missingUnknownField( + List unknownFields, Descriptors.Descriptor schemaDescriptor) { + for (int i = 0; i < unknownFields.size(); i++) { + Object cell = unknownFields.get(i); + Object value; + if (cell instanceof TableCell) { + value = ((TableCell) cell).getV(); + } else if (cell instanceof Map) { + value = ((Map) cell).get("v"); + } else { + value = cell; + } + if (i >= schemaDescriptor.getFields().size()) { + return true; + } + Descriptors.FieldDescriptor fieldDescriptor = schemaDescriptor.getFields().get(i); + if (missingUnknownFieldObject(value, fieldDescriptor)) { + return true; + } + } + return false; + } + + public static boolean missingUnknownField( + AbstractMap unknownFields, Descriptors.Descriptor schemaDescriptor) { + for (Map.Entry entry : unknownFields.entrySet()) { + String key = entry.getKey().toLowerCase(); + String protoFieldName = + BigQuerySchemaUtil.isProtoCompatible(key) + ? key + : BigQuerySchemaUtil.generatePlaceholderFieldName(key); + + Descriptors.FieldDescriptor fieldDescriptor = + schemaDescriptor.findFieldByName(protoFieldName); + if (fieldDescriptor == null) { + return true; + } + Object value = entry.getValue(); + if (value == null) { + continue; + } + if (missingUnknownFieldObject(value, fieldDescriptor)) { + return true; + } + } + return false; + } + + private static boolean missingUnknownFieldObject( + @Nullable Object value, Descriptors.FieldDescriptor fieldDescriptor) { + if (value == null) { + return false; + } + + if (fieldDescriptor.getType() != Descriptors.FieldDescriptor.Type.MESSAGE) { + return false; + } + if (value instanceof List) { + for (Object element : (List) value) { + if (missingUnknownFieldObject(element, fieldDescriptor)) { + return true; + } + } + return false; + } else if (value instanceof AbstractMap) { + return missingUnknownField( + (AbstractMap) value, fieldDescriptor.getMessageType()); + } + return false; + } } 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..4391a66e27ee 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 @@ -132,6 +132,8 @@ public class BigQueryIOTranslationTest { WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSharding", "auto_sharding"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getPropagateSuccessful", "propagate_successful"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSchemaUpdate", "auto_schema_update"); + WRITE_TRANSFORM_SCHEMA_MAPPING.put( + "getAutoSchemaUpdateStrictTimeout", "auto_schema_update_strict_timeout_ms"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getWriteProtosClass", "write_protos_class"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getDirectWriteProtos", "direct_write_protos"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getDeterministicRecordIdFn", "deterministic_record_id_fn"); 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..c1cd2a010c52 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 @@ -136,6 +136,7 @@ import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.schemas.JavaFieldSchema; import org.apache.beam.sdk.schemas.Schema; @@ -2301,12 +2302,22 @@ public TableRow apply(Long input) { @Test public void testUpdateTableSchemaUseSet() throws Exception { - updateTableSchemaTest(true); + updateTableSchemaTest(true, false); } @Test public void testUpdateTableSchemaUseSetF() throws Exception { - updateTableSchemaTest(false); + updateTableSchemaTest(false, false); + } + + @Test + public void testUpdateTableSchemaConsistentUseSet() throws Exception { + updateTableSchemaTest(true, true); + } + + @Test + public void testUpdateTableSchemaConsistentUseSetF() throws Exception { + updateTableSchemaTest(false, true); } @Test @@ -2396,13 +2407,18 @@ public void onTimer(@Key String tableSpec) throws IOException { } } - public void updateTableSchemaTest(boolean useSet) throws Exception { + public void updateTableSchemaTest(boolean useSet, boolean withConsistentSchemaUpdate) + throws Exception { assumeTrue(useStreaming); assumeTrue(useStorageApi); + if (withConsistentSchemaUpdate) { + p.getOptions().as(StreamingOptions.class).setStreaming(true); + } // Make sure that GroupIntoBatches does not buffer data. p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchLocalRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; @@ -2417,15 +2433,12 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("req").setType("STRING").setMode("REQUIRED"))); - // Add new fields to the update schema. Also reorder some existing fields to validate that we - // handle update - // field reordering correctly. TableSchema tableSchemaUpdated = new TableSchema() .setFields( ImmutableList.of( - new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("number").setType("INTEGER"), + new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("req").setType("STRING"), new TableFieldSchema().setName("double_number").setType("INTEGER"), new TableFieldSchema().setName("12_special_name").setType("STRING"))); @@ -2476,6 +2489,8 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { for (long i = 6; i < 10; i++) { testStream = testStream.addElements(i); } + // Expire the buffering timer so the elements get processed + testStream = testStream.advanceProcessingTime(Duration.standardMinutes(2)); PCollection tableRows = p.apply(testStream.advanceWatermarkToInfinity()) @@ -2488,25 +2503,41 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { Duration.standardSeconds(5), tableSchemaUpdated, fakeDatasetService))) .setCoder(TableRowJsonCoder.of()); - tableRows.apply( + BigQueryIO.Write write = BigQueryIO.writeTableRows() .to(tableRef) .withMethod(method) .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER) .ignoreUnknownValues() - .withAutoSchemaUpdate(true) .withTestServices(fakeBqServices) - .withoutValidation()); + .withoutValidation(); + if (withConsistentSchemaUpdate) { + write = write.withAutoSchemaUpdateConsistent(true, Duration.standardHours(1)); + } else { + write = write.withAutoSchemaUpdate(true); + } + tableRows.apply(write); p.run(); Iterable expectedDroppedValues = - LongStream.range(0, 6) + (withConsistentSchemaUpdate ? LongStream.empty() : LongStream.range(0, 6)) .mapToObj(getRowSet) .map(tr -> filterUnknownValues(tr, tableSchema.getFields())) .collect(Collectors.toList()); Iterable expectedFullValues = - LongStream.range(6, 10).mapToObj(getRowSet).collect(Collectors.toList()); + (withConsistentSchemaUpdate ? LongStream.range(0, 10) : LongStream.range(6, 10)) + .mapToObj(getRowSet) + .collect(Collectors.toList()); + System.err.println( + "GOT " + + fakeDatasetService.getAllRows( + tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId())); + System.err.println( + "WANT " + + Arrays.toString( + Iterables.toArray( + Iterables.concat(expectedDroppedValues, expectedFullValues), TableRow.class))); assertThat( fakeDatasetService.getAllRows( tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()), @@ -2524,6 +2555,7 @@ public void testAutoPatchTableSchemaTest() throws Exception { p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); p.getOptions().as(BigQueryOptions.class).setSchemaUpgradeBufferingShards(2); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchLocalRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java new file mode 100644 index 000000000000..8ba90ad01d6d --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java @@ -0,0 +1,305 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableFieldSchema; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.ByteString; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import java.nio.charset.StandardCharsets; +import java.util.Optional; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SchemaChangeDetectorHelperTest { + + private TableReference tableReference; + private BigQueryServices.WriteStreamService mockWriteStreamService; + private BigQueryServices.StreamAppendClient mockStreamAppendClient; + private AppendClientInfo mockAppendClientInfo; + + @Before + public void setUp() { + tableReference = new TableReference().setProjectId("p").setDatasetId("d").setTableId("t"); + mockWriteStreamService = mock(BigQueryServices.WriteStreamService.class); + mockStreamAppendClient = mock(BigQueryServices.StreamAppendClient.class); + mockAppendClientInfo = mock(AppendClientInfo.class); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateTrue_hasUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + TableSchema newSchema = + TableSchema.newBuilder() + .addFields(TableFieldSchema.newBuilder().setName("new_field").build()) + .build(); + + when(mockWriteStreamService.getWriteStreamSchema("test_stream")).thenReturn(newSchema); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertTrue(updated.isPresent()); + assertEquals(1, updated.get().getFieldsCount()); + assertEquals("new_field", updated.get().getFields(0).getName()); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateFalse() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertFalse(updated.isPresent()); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateTrue_noUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + when(mockWriteStreamService.getWriteStreamSchema("test_stream")).thenReturn(null); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertFalse(updated.isPresent()); + } + + @Test + public void testCheckResponseForUpdatedSchema_hasUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + TableSchema newSchema = + TableSchema.newBuilder() + .addFields(TableFieldSchema.newBuilder().setName("new_field").build()) + .build(); + + when(mockStreamAppendClient.getUpdatedSchema()).thenReturn(newSchema); + + Optional updated = + helper.checkResponseForUpdatedSchema(currentSchema, mockStreamAppendClient); + + assertTrue(updated.isPresent()); + assertEquals(1, updated.get().getFieldsCount()); + assertEquals("new_field", updated.get().getFields(0).getName()); + } + + @Test + public void testCheckResponseForUpdatedSchema_noUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + when(mockStreamAppendClient.getUpdatedSchema()).thenReturn(null); + + Optional updated = + helper.checkResponseForUpdatedSchema(currentSchema, mockStreamAppendClient); + + assertFalse(updated.isPresent()); + } + + @Test + public void testGetMergedPayload_autoUpdateFalse() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, new TableRow().set("foo", "bar"), null); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {1, 2, 3}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_emptyUnknownFields() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + StorageApiWritePayload payload = StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {1, 2, 3}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_mergeSuccess() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + ByteString mergedBytes = ByteString.copyFrom(new byte[] {4, 5, 6}); + when(mockAppendClientInfo.mergeNewFields(any(ByteString.class), eq(unknownFields), eq(false))) + .thenReturn(mergedBytes); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {4, 5, 6}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_mergeFailure() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + when(mockAppendClientInfo.mergeNewFields(any(ByteString.class), eq(unknownFields), eq(false))) + .thenThrow(new TableRowToStorageApiProto.SchemaDoesntMatchException("conversion error")); + TableRow expectedFailsafe = new TableRow().set("failsafe", "true"); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), expectedFailsafe, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED, result.getKind()); + TimestampedValue failed = result.getFailed(); + assertEquals( + "org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto$SchemaDoesntMatchException: conversion error", + failed.getValue().getErrorMessage()); + assertEquals(expectedFailsafe, failed.getValue().getRow()); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesTrue_hasMissingUnknownField() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, true); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> new byte[0], () -> descriptor); + + assertTrue(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesTrue_noUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, true); + StorageApiWritePayload payload = StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> new byte[0], () -> null); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMatch() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] hash = "hash".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null).withSchemaHash(hash); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> hash, () -> null); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMismatch_noUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] payloadHash = "hash1".getBytes(StandardCharsets.UTF_8); + byte[] currentHash = "hash2".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(payloadHash); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + // DynamicMessage will parse new byte[0] cleanly, and we have no unknown fields. + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> currentHash, () -> descriptor); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMismatch_hasUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] payloadHash = "hash1".getBytes(StandardCharsets.UTF_8); + byte[] currentHash = "hash2".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(payloadHash); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + // To have unknown fields, we provide some valid protobuf bytes that contain a tag not in + // descriptor. + // Tag 1 (varint) = 1 << 3 | 0 = 8. Value 1. + byte[] payloadWithUnknown = new byte[] {8, 1}; + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate( + payload, payloadWithUnknown, () -> currentHash, () -> descriptor); + + assertTrue(outOfDate); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java new file mode 100644 index 000000000000..5bc7aea1b3c1 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java @@ -0,0 +1,260 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.DescriptorProtos; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SplittingIterableTest { + + private AppendClientInfo createAppendClientInfo() throws Exception { + TableSchema tableSchema = TableSchema.newBuilder().build(); + DescriptorProtos.DescriptorProto descriptor = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + return AppendClientInfo.of(tableSchema, descriptor, client -> {}); + } + + @Test + public void testBatchingBySplitSize() throws Exception { + List payloads = new ArrayList<>(); + // Payload of 10 bytes each + for (int i = 0; i < 5; i++) { + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[10], null, null) + .withTimestamp(Instant.ofEpochMilli(i)); + payloads.add(StoragePayloadWithDeadline.of(payload, null)); + } + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, false, new TableReference(), false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + // Split size 25 means 2 payloads (20 bytes) per batch, then 2, then 1. + SplittingIterable iterable = + new SplittingIterable( + payloads, + 25, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + + assertTrue(it.hasNext()); + AppendRowsPacket batch1 = it.next(); + assertEquals(2, batch1.getProtoRows().getSerializedRowsCount()); + assertEquals(Instant.ofEpochMilli(0), batch1.getTimestamps().get(0)); + assertEquals(Instant.ofEpochMilli(1), batch1.getTimestamps().get(1)); + + assertTrue(it.hasNext()); + AppendRowsPacket batch2 = it.next(); + assertEquals(2, batch2.getProtoRows().getSerializedRowsCount()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch3 = it.next(); + assertEquals(1, batch3.getProtoRows().getSerializedRowsCount()); + + assertFalse(it.hasNext()); + assertTrue(failedRows.isEmpty()); + } + + @Test + public void testLargeElementExceedingSplitSize() throws Exception { + List payloads = new ArrayList<>(); + // Payload of 10 bytes, 100 bytes, 10 bytes + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[10], null, null), null)); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[100], null, null), null)); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[10], null, null), null)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, false, new TableReference(), false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + // Split size 25 + SplittingIterable iterable = + new SplittingIterable( + payloads, + 25, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + + assertTrue(it.hasNext()); + AppendRowsPacket batch1 = it.next(); + assertEquals(1, batch1.getProtoRows().getSerializedRowsCount()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch2 = it.next(); + assertEquals(1, batch2.getProtoRows().getSerializedRowsCount()); + assertEquals(100, batch2.getProtoRows().getSerializedRows(0).size()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch3 = it.next(); + assertEquals(1, batch3.getProtoRows().getSerializedRowsCount()); + + assertFalse(it.hasNext()); + } + + @Test + public void testSchemaMismatchedAndMatchedRows() throws Exception { + List payloads = new ArrayList<>(); + byte[] hash1 = "currentHash".getBytes(StandardCharsets.UTF_8); + byte[] hash2 = "oldHash".getBytes(StandardCharsets.UTF_8); + + TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); + + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1), null)); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null).withSchemaHash(hash2), + null)); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1), null)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, true, new TableReference(), true); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + SplittingIterable iterable = + new SplittingIterable( + payloads, + 100, + failedRows::add, + () -> "currentHash".getBytes(StandardCharsets.UTF_8), // Matches hash1 + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + assertTrue(it.hasNext()); + AppendRowsPacket batch = it.next(); + + assertEquals(3, batch.getProtoRows().getSerializedRowsCount()); + + // Check getting only mismatched rows + AppendRowsPacket mismatched = batch.getSchemaMismatchedRowsOnly(); + assertEquals(1, mismatched.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], mismatched.getProtoRows().getSerializedRows(0).toByteArray()); + assertEquals(unknownFieldsRow, mismatched.getUnknownFields().get(0)); + assertArrayEquals(new byte[0], mismatched.getOriginalPayloads().get(0)); + assertArrayEquals(hash2, mismatched.getSchemaHashes().get(0)); + + // Check getting only matched rows + AppendRowsPacket matched = batch.getSchemaMatchedRowsOnly(); + assertEquals(2, matched.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], matched.getProtoRows().getSerializedRows(0).toByteArray()); + assertArrayEquals(new byte[0], matched.getProtoRows().getSerializedRows(1).toByteArray()); + assertArrayEquals(hash1, matched.getSchemaHashes().get(0)); + assertArrayEquals(hash1, matched.getSchemaHashes().get(1)); + assertNull(matched.getUnknownFields().get(0)); + assertNull(matched.getOriginalPayloads().get(0)); + + // Reconstruct stream + List reconstructed = + batch + .toPayloadStream() + .map(StoragePayloadWithDeadline::getStoragePayload) + .collect(Collectors.toList()); + assertEquals(3, reconstructed.size()); + assertArrayEquals(new byte[0], reconstructed.get(0).getPayload()); + assertArrayEquals(new byte[0], reconstructed.get(1).getPayload()); + assertArrayEquals(new byte[0], reconstructed.get(2).getPayload()); + } + + @Test + public void testFailedRows() throws Exception { + List payloads = new ArrayList<>(); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[0], null, null), null)); + // Provide unknown fields, so auto update schema tries to merge and fails + TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null), null)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + true, // autoUpdateSchema + false, // ignoreUnknownValues (this will trigger SchemaConversionException) + new TableReference().setTableId("test_table"), + false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + SplittingIterable iterable = + new SplittingIterable( + payloads, + 100, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + assertTrue(it.hasNext()); + AppendRowsPacket batch = it.next(); + + // Only 1 row successfully added to batch + assertEquals(1, batch.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], batch.getProtoRows().getSerializedRows(0).toByteArray()); + + // 1 row failed + assertEquals(1, failedRows.size()); + BigQueryStorageApiInsertError error = failedRows.get(0).getValue(); + assertEquals("test_table", error.getTable().getTableId()); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java index 3cb9897ada52..21e694fa2591 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java @@ -49,6 +49,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition; import org.apache.beam.sdk.io.gcp.testing.BigqueryClient; import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.ValueState; @@ -358,7 +359,10 @@ private static TableSchema makeTableSchemaFromTypes( } private void runStreamingPipelineWithSchemaChange( - Write.Method method, boolean useAutoSchemaUpdate, boolean useIgnoreUnknownValues) + Write.Method method, + boolean useAutoSchemaUpdate, + boolean consistentAutoUpdate, + boolean useIgnoreUnknownValues) throws Exception { Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); // Set threshold bytes to 0 so that the stream attempts to fetch an updated schema after each @@ -367,6 +371,8 @@ private void runStreamingPipelineWithSchemaChange( // Limit parallelism so that all streams recognize the new schema in an expected short amount // of time (before we start writing rows with updated schema) p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(TOTAL_NUM_STREAMS); + p.getOptions().as(StreamingOptions.class).setStreaming(true); + // Need to manually enable streaming engine for legacy dataflow runner ExperimentalOptions.addExperiment( p.getOptions().as(ExperimentalOptions.class), GcpOptions.STREAMING_ENGINE_EXPERIMENT); @@ -374,7 +380,11 @@ private void runStreamingPipelineWithSchemaChange( if (p.getOptions().getRunner().getName().contains("DataflowRunner")) { assumeTrue( "Skipping in favor of more relevant test case and to avoid timing issues", - !changeTableSchema && useInputSchema && useAutoSchemaUpdate); + consistentAutoUpdate || (!changeTableSchema && useInputSchema && useAutoSchemaUpdate)); + } + if (consistentAutoUpdate) { + assumeTrue(changeTableSchema); + assumeFalse(useAutoSchemaUpdate); } List fieldNamesOrigin = new ArrayList(Arrays.asList(FIELDS)); @@ -402,6 +412,7 @@ private void runStreamingPipelineWithSchemaChange( BigQueryIO.writeTableRows() .to(tableSpec) .withAutoSchemaUpdate(useAutoSchemaUpdate) + .withAutoSchemaUpdateConsistent(consistentAutoUpdate, Duration.standardMinutes(5)) .withMethod(method) .withCreateDisposition(CreateDisposition.CREATE_NEVER) .withWriteDisposition(WriteDisposition.WRITE_APPEND); @@ -413,7 +424,8 @@ private void runStreamingPipelineWithSchemaChange( } // We give a healthy waiting period between each element to give Storage API streams a chance to // recognize the new schema. Apply on relevant tests. - boolean waitLonger = changeTableSchema && (useAutoSchemaUpdate || !useInputSchema); + boolean waitLonger = + changeTableSchema && (useAutoSchemaUpdate || !useInputSchema) && !consistentAutoUpdate; if (method == Write.Method.STORAGE_WRITE_API) { write = write.withTriggeringFrequency( @@ -465,7 +477,7 @@ private void runStreamingPipelineWithSchemaChange( PROJECT, BIG_QUERY_DATASET_ID, ImmutableMap.of(tableId, updatedSchema)))); } WriteResult result = rows.apply("Stream to BigQuery", write); - if (useIgnoreUnknownValues) { + if (useIgnoreUnknownValues || consistentAutoUpdate) { // We ignore the extra fields, so no rows should have been sent to DLQ PAssert.that("Check DLQ is empty", result.getFailedStorageApiInserts()).empty(); } else { @@ -478,11 +490,12 @@ private void runStreamingPipelineWithSchemaChange( p.run().waitUntilFinish(); // Check row completeness, non-duplication, and that schema update works as intended. - int expectedCount = useIgnoreUnknownValues ? TOTAL_N : ORIGINAL_N; - boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API) ? true : false; + int expectedCount = (useIgnoreUnknownValues || consistentAutoUpdate) ? TOTAL_N : ORIGINAL_N; + boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API); checkRowCompleteness(tableSpec, expectedCount, checkNoDuplication); - if (useIgnoreUnknownValues) { - checkRowsWithUpdatedSchema(tableSpec, extraField, useAutoSchemaUpdate); + if (useIgnoreUnknownValues || consistentAutoUpdate) { + checkRowsWithUpdatedSchema( + tableSpec, extraField, useAutoSchemaUpdate || consistentAutoUpdate); } } @@ -553,7 +566,6 @@ public void checkRowsWithUpdatedSchema( List actualRows = BQ_CLIENT.queryUnflattened( String.format("SELECT * FROM [%s]", tableSpec), PROJECT, true, false, bigQueryLocation); - for (TableRow row : actualRows) { // Rows written to the table should not have the extra field if // 1. The row has original schema @@ -579,39 +591,53 @@ public void testExactlyOnce() throws Exception { Write.Method.STORAGE_WRITE_API, /** autoSchemaUpdate */ false, + false, /** ignoreUnknownvalues */ false); } @Test public void testExactlyOnceWithIgnoreUnknownValues() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, false, true); } @Test public void testExactlyOnceWithAutoSchemaUpdate() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, false, true); + } + + @Test + public void testExactlyOnceWithAutoSchemaUpdateConsistent() throws Exception { + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, true, true); } @Test public void testAtLeastOnce() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false); + runStreamingPipelineWithSchemaChange( + Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false, false); } @Test public void testAtLeastOnceWithIgnoreUnknownValues() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, true); + runStreamingPipelineWithSchemaChange( + Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false, true); } @Test public void testAtLeastOnceWithAutoSchemaUpdate() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, false, true); + } + + @Test + public void testAtLeastOnceWithAutoSchemaUpdateConsistent() throws Exception { + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, true, true); } public void runDynamicDestinationsWithAutoSchemaUpdate(boolean useAtLeastOnce) throws Exception { Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); // 0 threshold so that the stream tries fetching an updated schema after each append p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(0); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchRetryTimeMilliSec(20); // Total streams per destination p.getOptions() .as(BigQueryOptions.class) diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java index 13724ab83e61..2e1da866b260 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java @@ -22,6 +22,8 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import com.google.api.services.bigquery.model.TableCell; +import com.google.api.services.bigquery.model.TableRow; import com.google.cloud.bigquery.storage.v1.TableFieldSchema; import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.DescriptorProtos; @@ -29,6 +31,7 @@ import com.google.protobuf.DynamicMessage; import com.google.protobuf.UnknownFieldSet; import java.nio.charset.StandardCharsets; +import java.util.Arrays; import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; @@ -325,14 +328,19 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { // 1. Missing hash in payload StorageApiWritePayload payloadNoHash = StorageApiWritePayload.of(new byte[0], null, null); - assertFalse(UpgradeTableSchema.isPayloadSchemaOutOfDate(payloadNoHash, () -> hash1, null)); + assertFalse( + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payloadNoHash.getSchemaHash(), payloadNoHash::getPayload, () -> hash1, null)); // 2. Equal hash - assertFalse(UpgradeTableSchema.isPayloadSchemaOutOfDate(payload, () -> hash1, null)); + assertFalse( + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), payload::getPayload, () -> hash1, null)); // 3. Different hash, but message doesn't have unknown fields (initialized and empty) assertFalse( - UpgradeTableSchema.isPayloadSchemaOutOfDate(payload, () -> hash2, () -> descriptor)); + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), payload::getPayload, () -> hash2, () -> descriptor)); // 4. Different hash with unknown fields. DynamicMessage unknownFieldSet = @@ -348,7 +356,10 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { StorageApiWritePayload.of(unknownFieldSet.toByteArray(), null, null).withSchemaHash(hash1); assertTrue( UpgradeTableSchema.isPayloadSchemaOutOfDate( - payloadWithUnknown, () -> hash2, () -> descriptor)); + payloadWithUnknown.getSchemaHash(), + payloadWithUnknown::getPayload, + () -> hash2, + () -> descriptor)); // 5. Different hash with missing required fields. DynamicMessage missingField = DynamicMessage.newBuilder(descriptor).buildPartial(); @@ -356,6 +367,110 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { StorageApiWritePayload.of(missingField.toByteArray(), null, null).withSchemaHash(hash1); assertTrue( UpgradeTableSchema.isPayloadSchemaOutOfDate( - payloadMissingField, () -> hash2, () -> descriptor)); + payloadMissingField.getSchemaHash(), + payloadMissingField::getPayload, + () -> hash2, + () -> descriptor)); + } + + @Test + public void testMissingUnknownField_TableRowWithF() throws Exception { + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + TableRow row = new TableRow(); + TableCell cell1 = new TableCell(); + cell1.setV("value1"); + row.set("f", Arrays.asList(cell1)); + + assertFalse(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + + TableCell cell2 = new TableCell(); + cell2.setV("value2"); + row.set("f", Arrays.asList(cell1, cell2)); + + assertTrue(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + } + + @Test + public void testMissingUnknownField_TableRow() throws Exception { + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + TableRow row = new TableRow(); + row.set("field1", "value1"); + + assertFalse(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + + row.set("field2", "value2"); + + assertTrue(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + } + + @Test + public void testMissingUnknownField_NestedObject() throws Exception { + DescriptorProtos.DescriptorProto nestedProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("NestedMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("nested1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + + DescriptorProtos.DescriptorProto parentProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_MESSAGE) + .setTypeName("NestedMessage") + .build()) + .build(); + + DescriptorProtos.FileDescriptorProto fileProto = + DescriptorProtos.FileDescriptorProto.newBuilder() + .addMessageType(nestedProto) + .addMessageType(parentProto) + .build(); + + Descriptors.FileDescriptor fileDescriptor = + Descriptors.FileDescriptor.buildFrom(fileProto, new Descriptors.FileDescriptor[0]); + Descriptors.Descriptor descriptor = fileDescriptor.findMessageTypeByName("TestMessage"); + + TableRow nestedRow = new TableRow(); + nestedRow.set("nested1", "value1"); + + TableRow parentRow = new TableRow(); + parentRow.set("field1", nestedRow); + + assertFalse(UpgradeTableSchema.missingUnknownField(parentRow, () -> descriptor)); + + nestedRow.set("nested2", "value2"); + assertTrue(UpgradeTableSchema.missingUnknownField(parentRow, () -> descriptor)); } }