You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by re...@apache.org on 2026/09/24 18:57:55 UTC
(beam) branch master updated: Merge pull request #38854 from reuvenlax/schema_update_handle_unknown_fields
This is an automated email from the ASF dual-hosted git repository.
reuvenlax pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0793f1dd10f Merge pull request #38854 from reuvenlax/schema_update_handle_unknown_fields
0793f1dd10f is described below
commit 0793f1dd10fe6fec58b3bd2d32ddbc3a4f6983bb
Author: Reuven Lax <re...@google.com>
AuthorDate: Thu Sep 24 11:57:47 2026 -0700
Merge pull request #38854 from reuvenlax/schema_update_handle_unknown_fields
Allow strongly-consistent schema updates in BigQueryIO
---
.../beam/sdk/io/gcp/bigquery/AppendClientInfo.java | 27 +-
.../beam/sdk/io/gcp/bigquery/AppendRowsPacket.java | 155 +++++
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 63 +-
.../sdk/io/gcp/bigquery/BigQueryIOTranslation.java | 15 +
.../beam/sdk/io/gcp/bigquery/BigQueryOptions.java | 20 +
.../sdk/io/gcp/bigquery/BigQuerySinkMetrics.java | 1 +
.../sdk/io/gcp/bigquery/BufferMismatchedRows.java | 501 ++++++++++++++
.../sdk/io/gcp/bigquery/DynamicDestinations.java | 38 ++
.../gcp/bigquery/DynamicDestinationsHelpers.java | 13 +
.../gcp/bigquery/SchemaChangeDetectorHelper.java | 236 +++++++
.../sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java | 85 ++-
.../sdk/io/gcp/bigquery/SplittingIterable.java | 175 ++---
.../bigquery/StorageApiDynamicDestinations.java | 18 +-
.../StorageApiDynamicDestinationsTableRow.java | 30 +-
.../beam/sdk/io/gcp/bigquery/StorageApiLoads.java | 10 +-
.../io/gcp/bigquery/StorageApiWritePayload.java | 16 +
.../StorageApiWriteRecordsInconsistent.java | 117 +++-
.../bigquery/StorageApiWriteUnshardedRecords.java | 740 +++++++++++++--------
.../bigquery/StorageApiWritesShardedRecords.java | 662 +++++++++++++++---
.../gcp/bigquery/StoragePayloadWithDeadline.java | 90 +++
.../io/gcp/bigquery/TableRowToStorageApiProto.java | 48 +-
.../sdk/io/gcp/bigquery/UpgradeTableSchema.java | 120 +++-
.../sdk/io/gcp/bigquery/AppendRowsPacketTest.java | 268 ++++++++
.../io/gcp/bigquery/BigQueryIOTranslationTest.java | 2 +
.../sdk/io/gcp/bigquery/BigQueryIOWriteTest.java | 70 +-
.../SchemaChangeDetectorHelperBufferingTest.java | 431 ++++++++++++
.../bigquery/SchemaChangeDetectorHelperTest.java | 307 +++++++++
.../sdk/io/gcp/bigquery/SplittingIterableTest.java | 242 +++++++
.../StorageApiSchemaMismatchDrainTest.java | 641 ++++++++++++++++++
.../bigquery/StorageApiSinkSchemaUpdateITBase.java | 345 +++++++---
...torageApiSinkSchemaUpdateWithInputSchemaIT.java | 5 +-
...ageApiSinkSchemaUpdateWithoutInputSchemaIT.java | 5 +-
.../io/gcp/bigquery/UpgradeTableSchemaTest.java | 212 +++++-
33 files changed, 5020 insertions(+), 688 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java
index 4adae9ccea0..fa6e9794a0f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java
@@ -191,21 +191,30 @@ abstract class AppendClientInfo {
public ByteString mergeNewFields(
ByteString payloadBytes, TableRow unknownFields, boolean ignoreUnknownValues)
throws TableRowToStorageApiProto.SchemaConversionException {
- return TableRowToStorageApiProto.mergeNewFields(
- payloadBytes,
- getDescriptor(),
- getTableSchema(),
- getSchemaInformation(),
- unknownFields,
- ignoreUnknownValues);
+ try {
+ return TableRowToStorageApiProto.mergeNewFields(
+ payloadBytes,
+ getWrappedDescriptor(),
+ getDescriptorIgnoreRequired(),
+ getSchemaInformation(),
+ unknownFields,
+ ignoreUnknownValues);
+ } catch (Descriptors.DescriptorValidationException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Memoized
+ public Descriptors.Descriptor getWrappedDescriptor()
+ throws Descriptors.DescriptorValidationException {
+ return TableRowToStorageApiProto.wrapDescriptorProto(getDescriptor());
}
public TableRow toTableRow(ByteString protoBytes, Predicate<String> includeField) {
try {
return TableRowToStorageApiProto.tableRowFromMessage(
getSchemaInformation(),
- DynamicMessage.parseFrom(
- TableRowToStorageApiProto.wrapDescriptorProto(getDescriptor()), protoBytes),
+ DynamicMessage.parseFrom(getWrappedDescriptor(), protoBytes),
true,
includeField);
} catch (Exception e) {
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 00000000000..341ce5e5b46
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java
@@ -0,0 +1,155 @@
+/*
+ * 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 java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Function;
+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.PeekingIterator;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+
+@AutoValue
+abstract class AppendRowsPacket {
+ abstract ProtoRows getProtoRows();
+
+ abstract List<Instant> getTimestamps();
+
+ abstract List<@Nullable TableRow> getFailsafeTableRows();
+
+ abstract List<StoragePayloadWithDeadline> getSchemaMismatchedRows();
+
+ /**
+ * Creates a packet.
+ *
+ * <p>Exists so that the generated AutoValue subclass is never named outside this file;
+ * ErrorProne's AutoValueSubclassLeaked check forbids that, which would otherwise keep tests from
+ * constructing a packet directly.
+ */
+ static AppendRowsPacket create(
+ ProtoRows protoRows,
+ List<Instant> timestamps,
+ List<@Nullable TableRow> failsafeTableRows,
+ List<StoragePayloadWithDeadline> schemaMismatchedRows) {
+ return new AutoValue_AppendRowsPacket(
+ protoRows, timestamps, failsafeTableRows, schemaMismatchedRows);
+ }
+
+ static AppendRowsPacket fromStorageApiWritePayload(
+ PeekingIterator<StoragePayloadWithDeadline> underlyingIterator,
+ long maxByteSize,
+ SchemaChangeDetectorHelper schemaChangeDetectorHelper,
+ Instant elementTimestamp,
+ AppendClientInfo appendClientInfo,
+ Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean> failedRowsHandler) {
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeRows = Lists.newArrayList();
+ List<StoragePayloadWithDeadline> mismatchedRows = Lists.newArrayList();
+ ProtoRows.Builder inserts = ProtoRows.newBuilder();
+ long bytesSize = 0;
+ 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 || !mismatchedRows.isEmpty())) {
+ 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 = ByteString.empty();
+ SchemaChangeDetectorHelper.MergePayloadResult mergeResult =
+ schemaChangeDetectorHelper.getMergedPayload(
+ storagePayload, elementTimestamp, failsafeTableRow, appendClientInfo);
+ boolean mismatched = false;
+ 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. Note that this path can only happen if there are unknown fields, as otherwise
+ // the call to
+ // getMergedPayload is a noop.
+ mismatched = true;
+ }
+ } else {
+ byteString = mergeResult.getMerged();
+ }
+
+ if (!mismatched
+ && schemaChangeDetectorHelper.isPayloadSchemaOutOfDate(
+ storagePayload,
+ byteString,
+ appendClientInfo::getTableSchemaHash,
+ appendClientInfo::getWrappedDescriptor)) {
+ mismatched = true;
+ }
+
+ Instant timestamp = storagePayload.getTimestamp();
+ if (timestamp == null) {
+ timestamp = elementTimestamp;
+ }
+
+ if (mismatched) {
+ if (storagePayload.getTimestamp() == null) {
+ payload =
+ StoragePayloadWithDeadline.of(
+ storagePayload.withTimestamp(timestamp), payload.getDeadline());
+ }
+ mismatchedRows.add(payload);
+ bytesSize += storagePayload.getPayload().length;
+ continue;
+ }
+
+ inserts.addSerializedRows(byteString);
+ timestamps.add(timestamp);
+ failsafeRows.add(failsafeTableRow);
+ bytesSize += byteString.size();
+ }
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+
+ return create(inserts.build(), timestamps, failsafeRows, mismatchedRows);
+ }
+
+ AppendRowsPacket getSchemaMatchedRowsOnly() {
+ if (getSchemaMismatchedRows().isEmpty()) {
+ return this;
+ }
+ return create(getProtoRows(), getTimestamps(), getFailsafeTableRows(), Collections.emptyList());
+ }
+}
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 b8f5b0a6682..0cefea4639f 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
@@ -110,6 +110,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.PassThroughThenCleanup.CleanupOperati
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;
@@ -2562,6 +2563,7 @@ public class BigQueryIO {
.setAutoSharding(false)
.setPropagateSuccessful(true)
.setAutoSchemaUpdate(false)
+ .setAutoSchemaUpdateStrictTimeout(null)
.setDeterministicRecordIdFn(null)
.setMaxRetryJobs(1000)
.setPropagateSuccessfulStorageApiWrites(false)
@@ -2844,6 +2846,8 @@ public class BigQueryIO {
abstract boolean getAutoSchemaUpdate();
+ abstract @Nullable Duration getAutoSchemaUpdateStrictTimeout();
+
abstract @Nullable Class<T> getWriteProtosClass();
abstract boolean getDirectWriteProtos();
@@ -2964,6 +2968,8 @@ public class BigQueryIO {
abstract Builder<T> setAutoSchemaUpdate(boolean autoSchemaUpdate);
+ abstract Builder<T> setAutoSchemaUpdateStrictTimeout(@Nullable Duration timeout);
+
abstract Builder<T> setWriteProtosClass(@Nullable Class<T> clazz);
abstract Builder<T> setDirectWriteProtos(boolean direct);
@@ -3629,9 +3635,50 @@ public class BigQueryIO {
* 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.
+ *
+ * <p>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.
+ *
+ * <p>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<T> withAutoSchemaUpdate(boolean autoSchemaUpdate) {
- return toBuilder().setAutoSchemaUpdate(autoSchemaUpdate).build();
+ return toBuilder()
+ .setAutoSchemaUpdate(autoSchemaUpdate)
+ .setAutoSchemaUpdateStrictTimeout(null)
+ .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.
+ *
+ * <p>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.
+ *
+ * <p>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.
+ *
+ * <p>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<T> withAutoSchemaUpdateConsistent(
+ boolean autoSchemaUpdate, Duration waitForSchemaTimeout) {
+ return toBuilder()
+ .setAutoSchemaUpdate(autoSchemaUpdate)
+ // Never leave a strict timeout configured when auto schema update is disabled -
+ // that would half-enable the consistent code path (buffering enabled, unknown-field
+ // capture and stream-schema refresh disabled).
+ .setAutoSchemaUpdateStrictTimeout(autoSchemaUpdate ? waitForSchemaTimeout : null)
+ .build();
}
/*
@@ -4170,6 +4217,14 @@ public class BigQueryIO {
DynamicDestinations<T, DestinationT> dynamicDestinations,
RowWriterFactory<T, DestinationT> 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,
@@ -4376,6 +4431,9 @@ public class BigQueryIO {
RowWriterFactory.TableRowWriterFactory<T, DestinationT> tableRowWriterFactory =
(RowWriterFactory.TableRowWriterFactory<T, DestinationT>) rowWriterFactory;
// Fallback behavior: convert to JSON TableRows and convert those into Beam TableRows.
+ boolean useSchemaUpdatingTableRow =
+ (schemaUpdateOptions != null && !schemaUpdateOptions.isEmpty())
+ || (getAutoSchemaUpdate() && getAutoSchemaUpdateStrictTimeout() != null);
storageApiDynamicDestinations =
new StorageApiDynamicDestinationsTableRow<>(
dynamicDestinations,
@@ -4385,7 +4443,7 @@ public class BigQueryIO {
getCreateDisposition(),
getIgnoreUnknownValues(),
getAutoSchemaUpdate(),
- schemaUpdateOptions == null ? Collections.emptySet() : schemaUpdateOptions);
+ useSchemaUpdatingTableRow);
}
int numShards = getStorageApiNumStreams(bqOptions);
@@ -4408,6 +4466,7 @@ public class BigQueryIO {
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 ee29189a065..9e81b675a75 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
@@ -482,6 +482,7 @@ public class BigQueryIOTranslation {
.addNullableByteArrayField("row_mutation_information_fn")
.addNullableByteArrayField("bad_record_error_handler")
.addNullableByteArrayField("bad_record_router")
+ .addNullableInt64Field("auto_schema_update_strict_timeout_ms")
.build();
public static final String BIGQUERY_WRITE_TRANSFORM_URN =
@@ -642,6 +643,11 @@ public class BigQueryIOTranslation {
fieldValues.put("use_beam_schema", transform.getUseBeamSchema());
fieldValues.put("auto_sharding", transform.getAutoSharding());
fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate());
+ org.joda.time.Duration autoSchemaUpdateTimeout = transform.getAutoSchemaUpdateStrictTimeout();
+ if (autoSchemaUpdateTimeout != null) {
+ fieldValues.put(
+ "auto_schema_update_strict_timeout_ms", autoSchemaUpdateTimeout.getMillis());
+ }
Class<?> writeProtosClass = transform.getWriteProtosClass();
if (writeProtosClass != null) {
fieldValues.put("write_protos_class", toByteArray(writeProtosClass));
@@ -932,6 +938,15 @@ public class BigQueryIOTranslation {
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) {
// NOTE: this deserializes 'defaultMissingValueInterpretationsBytes' rather than
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 face2ef5841..da8526bd994 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 aec011cebb5..acb0bd5f694 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 00000000000..139c73608f5
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java
@@ -0,0 +1,501 @@
+/*
+ * 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.client.util.BackOff;
+import com.google.api.client.util.BackOffUtils;
+import com.google.api.client.util.ExponentialBackOff;
+import com.google.api.services.bigquery.model.TableRow;
+import java.nio.ByteBuffer;
+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.transforms.windowing.BoundedWindow;
+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.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
+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;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class BufferMismatchedRows<DestinationT extends @NonNull Object, ElementT>
+ extends PTransform<
+ PCollection<KV<DestinationT, StoragePayloadWithDeadline>>, PCollectionTuple> {
+ private static final Logger LOG = LoggerFactory.getLogger(BufferMismatchedRows.class);
+ private final Coder<BigQueryStorageApiInsertError> failedRowsCoder;
+ private final Coder<TableRow> successfulRowsCoder;
+ private final Coder<DestinationT> destinationCoder;
+ private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
+ private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT>
+ writeDoFn;
+ private final TupleTag<BigQueryStorageApiInsertError> failedRowsTag;
+ private final @Nullable TupleTag<TableRow> successfulRowsTag;
+ // This output is effectively ignored, since we only support this code path for
+ // StorageApiWriteRecordsInconsistent.
+ private final TupleTag<KV<String, String>> finalizeTag = new TupleTag<>("finalizeTag");
+ private static final int NUM_DEFAULT_SHARDS = 20;
+ private static final byte[][] SHARD_KEYS = new byte[NUM_DEFAULT_SHARDS][];
+
+ // Precalculate the shard keys.
+ static {
+ for (int i = 0; i < NUM_DEFAULT_SHARDS; i++) {
+ SHARD_KEYS[i] = ByteBuffer.allocate(Integer.BYTES).putInt(i).array();
+ }
+ }
+
+ private static final String DRAIN_TIMEOUT_MESSAGE =
+ "Timed out waiting for table schema update in OnWindowExpiration";
+
+ public BufferMismatchedRows(
+ Coder<BigQueryStorageApiInsertError> failedRowsCoder,
+ Coder<TableRow> successfulRowsCoder,
+ Coder<DestinationT> destinationCoder,
+ StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations,
+ StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT> writeDoFn,
+ TupleTag<BigQueryStorageApiInsertError> failedRowsTag,
+ @Nullable TupleTag<TableRow> 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<KV<DestinationT, StoragePayloadWithDeadline>> 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<DestinationT, StoragePayloadWithDeadline>,
+ KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline>>() {
+ long shardNumber;
+
+ @Setup
+ public void setup() {
+ shardNumber = ThreadLocalRandom.current().nextLong(NUM_DEFAULT_SHARDS);
+ }
+
+ @ProcessElement
+ public void process(
+ @Element KV<DestinationT, StoragePayloadWithDeadline> element,
+ OutputReceiver<KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline>>
+ o) {
+ byte[] shardKey = SHARD_KEYS[(int) (++shardNumber % NUM_DEFAULT_SHARDS)];
+ o.output(
+ KV.of(ShardedKey.of(element.getKey(), shardKey), 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<KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline>, KV<String, String>> {
+ private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT>
+ writeDoFn;
+
+ @StateId("mismatchedRows")
+ private final StateSpec<BagState<StoragePayloadWithDeadline>> mismatchedRowsSpec =
+ StateSpecs.bag(StoragePayloadWithDeadline.Coder.of());
+
+ @TimerId("retryMismatchedRowsTimer")
+ private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ @StateId("currentMismatchedRowTimerValue")
+ private final StateSpec<ValueState<Long>> currentMismatchedRowTimerValueSpec =
+ StateSpecs.value();
+
+ @StateId("minPendingTimestamp")
+ private final StateSpec<ValueState<Long>> minPendingTimestampSpec = StateSpecs.value();
+
+ private final Counter rowsSentToFailedRowsCollection =
+ Metrics.counter(BufferMismatchedRows.BufferingDoFn.class, "rowsSentToFailedRowsCollection");
+
+ // Bounds the total time this DoFn instance spends blocking in onWindowExpiration,
+ // per window, across every key it handles. Window expiration fires once per key and
+ // runs synchronously on the worker thread, so a per-key budget multiplies by the
+ // number of keys and can exceed the runner's bundle execution limit during a drain.
+ //
+ // Deliberately an instance field rather than static: a static cache would share one
+ // budget across every BigQuery sink in the worker JVM.
+ //
+ // Keyed by window so a budget cannot leak from one window into the next. Today the
+ // sink is always rewindowed into the global window by StorageApiLoads, so in practice
+ // there is a single entry and expiration fires once, but nothing here enforces that
+ // and a stale deadline would silently skip the retry and dead-letter the rows.
+ //
+ // transient, and initialized lazily rather than inline: Java deserialization does not
+ // run field initializers, so an inline initializer would leave this null on workers.
+ private transient @Nullable Cache<BoundedWindow, Instant> drainDeadlines = null;
+
+ private Cache<BoundedWindow, Instant> getDrainDeadlines() {
+ Cache<BoundedWindow, Instant> cache = drainDeadlines;
+ if (cache == null) {
+ cache =
+ CacheBuilder.newBuilder().expireAfterAccess(java.time.Duration.ofMinutes(30)).build();
+ drainDeadlines = cache;
+ }
+ return cache;
+ }
+
+ public BufferingDoFn(
+ StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT> writeDoFn) {
+ this.writeDoFn = writeDoFn;
+ }
+
+ @ProcessElement
+ public void process(
+ PipelineOptions pipelineOptions,
+ ProcessContext processContext,
+ @Element KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline> element,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
+ @TimerId("retryMismatchedRowsTimer") Timer retryTimer,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long> currentTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
+ @Timestamp Instant originalElementTimestamp,
+ 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,
+ originalElementTimestamp);
+ }
+
+ @Override
+ public Duration getAllowedTimestampSkew() {
+ return Duration.millis(Long.MAX_VALUE);
+ }
+
+ @OnTimer("retryMismatchedRowsTimer")
+ public void onTimer(
+ OnTimerContext context,
+ @Key ShardedKey<DestinationT> shardedDestination,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long> currentTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
+ @TimerId("retryMismatchedRowsTimer") Timer retryTimer,
+ @Timestamp org.joda.time.Instant outputTimestamp,
+ 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();
+ writeDoFn.invalidateAllAppendClientsForTable(
+ tableDestination.getTableUrn(pipelineOptions.as(BigQueryOptions.class)));
+
+ // TODO: An optimization would be to detect whether a schema change has been observed (e.g. by
+ // storing a schema
+ // hash in state) and to skip reading the bag in this case.
+ List<Iterable<KV<DestinationT, StoragePayloadWithDeadline>>> mismatchedRowsList =
+ Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : mismatchedRowsBag.read()) {
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> 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<KV<DestinationT, StoragePayloadWithDeadline>> 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<StoragePayloadWithDeadline> 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,
+ outputTimestamp);
+ }
+ }
+
+ @OnWindowExpiration
+ public void onWindowExpiration(
+ OnWindowExpirationContext context,
+ @Key ShardedKey<DestinationT> shardedDestination,
+ @Timestamp org.joda.time.Instant elementTs,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
+ PipelineOptions pipelineOptions,
+ BoundedWindow window,
+ MultiOutputReceiver o)
+ throws Exception {
+ Iterable<StoragePayloadWithDeadline> bufferedRows = mismatchedRowsBag.read();
+ if (Iterables.isEmpty(bufferedRows)) {
+ return;
+ }
+
+ dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context);
+
+ StorageApiDynamicDestinations.MessageConverter<?> messageConverter =
+ writeDoFn.messageConverters.get(
+ shardedDestination.getKey(),
+ dynamicDestinations,
+ pipelineOptions,
+ writeDoFn.getDatasetService(pipelineOptions),
+ writeDoFn.getWriteStreamService(pipelineOptions));
+ messageConverter.updateSchemaFromTable();
+ TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey());
+ writeDoFn.invalidateAllAppendClientsForTable(
+ tableDestination.getTableUrn(pipelineOptions.as(BigQueryOptions.class)));
+
+ java.time.Duration waitTime =
+ java.time.Duration.ofMillis(
+ pipelineOptions
+ .as(BigQueryOptions.class)
+ .getStorageApiMismatchDrainRetryTimeMilliSec());
+
+ // Shared across every key in this window; see drainDeadlines.
+ // TODO: This needs to use the internal clock instead Instant.now(), as the row deadlines are
+ // based on the
+ // internal clock. Add support for accessing the internal clock in OnWindowExpiration.
+ Instant drainDeadline =
+ getDrainDeadlines()
+ .get(window, () -> Instant.now().plus(Duration.millis(waitTime.toMillis())));
+ LOG.info(
+ "Draining buffered schema-mismatched rows for table {}, waiting until {} for schema update.",
+ tableDestination.getShortTableUrn(),
+ drainDeadline);
+ BackOff backoff =
+ new ExponentialBackOff.Builder()
+ .setMaxElapsedTimeMillis((int) waitTime.toMillis())
+ .build();
+ // Rows whose deadline has passed are written to the failed-rows collection as soon as they
+ // are discovered, rather than accumulated into a list that lives until the drain finishes.
+ // The drain can run for many rounds, and a row that has already been dead-lettered has no
+ // further use here, so holding it only raises the peak heap of a handler whose failure
+ // prevents the drain from ever completing.
+ //
+ // One exception: writeFailedRows needs an open bundle, because finishBundle nulls out the
+ // destination map it looks the row's table up in. Rows that finishBundle itself reports as
+ // mismatched therefore cannot be written where they are discovered, and are carried to the
+ // next round's bundle (or to the final flush below). That carry is bounded by a single
+ // append batch -- finishBundle only flushes whatever processElement left pending -- and not
+ // by the size of the buffer, and it is emptied on every round.
+ List<StoragePayloadWithDeadline> expiredDuringFlush = Lists.newArrayList();
+ do {
+ writeDoFn.startBundle();
+ if (!expiredDuringFlush.isEmpty()) {
+ writeDoFn.writeFailedRows(
+ shardedDestination.getKey(),
+ expiredDuringFlush,
+ DRAIN_TIMEOUT_MESSAGE,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ pipelineOptions,
+ elementTs,
+ o.get(failedRowsTag));
+ // Safe to clear: writeFailedRows consumes the iterable before returning.
+ expiredDuringFlush.clear();
+ }
+
+ // Read the clock once per round so that every row in the round is judged against the same
+ // instant, and so that the reading is taken before the round's work rather than after it.
+ // This check is hard to test with DirectRunner. Ideally we should use the injected clock,
+ // but I'm not sure how to access it from OnWindowExpiration (it's usually accessed via
+ // the Timer object).
+ Instant now = Instant.now();
+ List<StoragePayloadWithDeadline> stillPending = Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : bufferedRows) {
+ // Partitioned inline, while the bundle that produced these rows is still open. Both
+ // halves must be derived from the same source: previously the pending rows were taken
+ // from the whole mismatchedRowsList while the expired rows were taken only from
+ // mismatchedDestRows, so a row returned by processElement whose deadline had passed was
+ // filtered out of the pending set without ever being recorded as expired, and was
+ // therefore never written to the failed-rows collection.
+ for (KV<DestinationT, StoragePayloadWithDeadline> kv :
+ writeDoFn.processElement(
+ pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o)) {
+ StoragePayloadWithDeadline mismatchedRow = kv.getValue();
+ if (mismatchedRow.getDeadline().isAfter(now)) {
+ stillPending.add(mismatchedRow);
+ } else {
+ // Written and dropped. An expired row never enters stillPending, so no later round
+ // and no final flush can emit it a second time.
+ writeDoFn.writeFailedRows(
+ shardedDestination.getKey(),
+ Collections.singletonList(mismatchedRow),
+ DRAIN_TIMEOUT_MESSAGE,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ pipelineOptions,
+ elementTs,
+ o.get(failedRowsTag));
+ }
+ }
+ }
+
+ // Once we're done, delegate to finishBundle to finish things. Its mismatched rows are
+ // partitioned the same way, except that the bundle is closed by the time they are
+ // returned, so the expired half has to wait for the next open bundle.
+ for (KV<DestinationT, StoragePayloadWithDeadline> kv :
+ writeDoFn.finishBundle(
+ o.get(failedRowsTag),
+ successfulRowsTag != null ? o.get(successfulRowsTag) : null,
+ o.get(finalizeTag),
+ null)) {
+ StoragePayloadWithDeadline mismatchedRow = kv.getValue();
+ if (mismatchedRow.getDeadline().isAfter(now)) {
+ stillPending.add(mismatchedRow);
+ } else {
+ expiredDuringFlush.add(mismatchedRow);
+ }
+ }
+ bufferedRows = stillPending;
+
+ // Checked before sleeping so we never sleep past the shared deadline.
+ } while (!Iterables.isEmpty(bufferedRows)
+ && Instant.now().isBefore(drainDeadline)
+ && BackOffUtils.next(com.google.api.client.util.Sleeper.DEFAULT, backoff));
+
+ // Whichever way the loop exited, expiredDuringFlush holds at most the last round's tail, and
+ // is written exactly once: either here, or by the next round that never came.
+ if (!Iterables.isEmpty(bufferedRows) || !expiredDuringFlush.isEmpty()) {
+ LOG.warn(
+ "Timed out waiting for table schema update during drain for table {}; routing remaining buffered rows to failed-rows collection.",
+ tableDestination.getShortTableUrn());
+ writeDoFn.startBundle();
+ writeDoFn.writeFailedRows(
+ shardedDestination.getKey(),
+ Iterables.concat(bufferedRows, expiredDuringFlush),
+ DRAIN_TIMEOUT_MESSAGE,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ pipelineOptions,
+ elementTs,
+ o.get(failedRowsTag));
+ writeDoFn.finishBundle(
+ o.get(failedRowsTag),
+ successfulRowsTag != null ? o.get(successfulRowsTag) : null,
+ o.get(finalizeTag),
+ null);
+ }
+ }
+
+ @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 74f998da81d..0e1ff4ea412 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 @@ public abstract class DynamicDestinations<T, DestinationT> implements Serializab
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> SideInputT sideInput(PCollectionView<SideInputT> 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> SideInputT sideInput(PCollectionView<SideInputT> view) {
+ return onWindowExpirationContext.sideInput(view);
+ }
+ }
+
static class SideInputAccessorViaProcessContext implements SideInputAccessor {
private DoFn<?, ?>.ProcessContext processContext;
@@ -128,6 +155,17 @@ public abstract class DynamicDestinations<T, DestinationT> implements Serializab
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 31898696049..4cd5d25202a 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
@@ -219,6 +219,19 @@ class DynamicDestinationsHelpers {
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 00000000000..301a0e09d34
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java
@@ -0,0 +1,236 @@
+/*
+ * 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.GlobalWindow;
+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<BigQueryStorageApiInsertError> getFailed();
+
+ static MergePayloadResult ofMerged(ByteString merged) {
+ return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.merged(merged);
+ }
+
+ static MergePayloadResult ofFailed(TimestampedValue<BigQueryStorageApiInsertError> failed) {
+ return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.failed(failed);
+ }
+ }
+
+ private final boolean autoUpdateSchema;
+ private final boolean ignoreUnknownValues;
+ private final TableReference tableReference;
+ private final boolean checkPayloadUnknownFields;
+
+ public SchemaChangeDetectorHelper(
+ boolean autoUpdateSchema,
+ boolean ignoreUnknownValues,
+ TableReference tableReference,
+ boolean checkPayloadUnknownFields) {
+ this.autoUpdateSchema = autoUpdateSchema;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.tableReference = tableReference;
+ this.checkPayloadUnknownFields = checkPayloadUnknownFields;
+ }
+
+ Optional<TableSchema> 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<TableSchema> 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,
+ ByteString mergedPayload,
+ ThrowingSupplier<byte[]> schemaHash,
+ ThrowingSupplier<Descriptors.Descriptor> schemaDescriptor)
+ throws Exception {
+ if (checkPayloadUnknownFields) {
+ 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.
+ }
+ return UpgradeTableSchema.isPayloadSchemaOutOfDate(
+ payload.getSchemaHash(), mergedPayload, schemaHash, schemaDescriptor);
+ }
+
+ static void bufferMismatchedRows(
+ Iterable<StoragePayloadWithDeadline> rows,
+ BagState<StoragePayloadWithDeadline> bufferedBag,
+ Timer retryRowsTimers,
+ ValueState<Long> currentTimerValue,
+ ValueState<Long> minPendingTimestamp,
+ TableDestination tableDestination,
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
+ @Nullable AppendClientInfo appendClientInfo,
+ Counter rowsSentToFailedRowsCollection,
+ Duration retryPeriod,
+ Instant originaElementTimestamp)
+ throws IOException {
+ minPendingTimestamp.readLater();
+ currentTimerValue.readLater();
+
+ org.joda.time.Instant minTimestamp =
+ org.joda.time.Instant.ofEpochMilli(
+ MoreObjects.firstNonNull(
+ minPendingTimestamp.read(), GlobalWindow.INSTANCE.maxTimestamp().getMillis()));
+
+ boolean modified = false;
+ for (StoragePayloadWithDeadline row : rows) {
+ if (appendClientInfo == null
+ || row.getDeadline().isAfter(retryRowsTimers.getCurrentRelativeTime())) {
+ bufferedBag.add(row);
+ org.joda.time.Instant timestamp =
+ MoreObjects.firstNonNull(
+ row.getStoragePayload().getTimestamp(), originaElementTimestamp);
+ if (timestamp.isBefore(minTimestamp)) {
+ minTimestamp = timestamp;
+ }
+ modified = true;
+ } else {
+ @Nullable TableRow failedRow = row.getStoragePayload().getFailsafeTableRow();
+ if (failedRow == null) {
+ ByteString rowBytes = ByteString.copyFrom(row.getStoragePayload().getPayload());
+ failedRow = appendClientInfo.toTableRow(rowBytes, Predicates.alwaysTrue());
+ }
+ @Nullable TableRow unknownFields = row.getStoragePayload().getUnknownFields();
+ String errorMessage =
+ (unknownFields != null && !unknownFields.isEmpty())
+ ? "Timed out waiting for table schema update. Unknown fields: "
+ + unknownFields.keySet()
+ : "Timed out waiting for table schema update (Mismatched schema)";
+ BigQueryStorageApiInsertError error =
+ new BigQueryStorageApiInsertError(
+ failedRow, errorMessage, tableDestination.getTableReference());
+ org.joda.time.Instant timestamp =
+ MoreObjects.firstNonNull(
+ row.getStoragePayload().getTimestamp(), originaElementTimestamp);
+ failedRowsReceiver.outputWithTimestamp(error, timestamp);
+ rowsSentToFailedRowsCollection.inc();
+ BigQuerySinkMetrics.appendRowsRowStatusCounter(
+ BigQuerySinkMetrics.RowStatus.FAILED,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ tableDestination.getShortTableUrn())
+ .inc(1);
+ }
+ }
+
+ if (modified) {
+ 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 b65b11ef319..87a7c3842f7 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
@@ -21,7 +21,6 @@ import com.google.api.client.util.BackOff;
import com.google.api.client.util.BackOffUtils;
import com.google.api.client.util.ExponentialBackOff;
import java.util.List;
-import java.util.concurrent.TimeUnit;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.BagState;
@@ -39,12 +38,16 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.ShardedKey;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
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;
import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* This is a stateful DoFn that buffers elements that triggered table schema update. Once the table
@@ -58,6 +61,7 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
extends DoFn<
KV<ShardedKey<DestinationT>, @Nullable ElementT>,
KV<DestinationT, StorageApiWritePayload>> {
+ private static final Logger LOG = LoggerFactory.getLogger(SchemaUpdateHoldingFn.class);
private static final Duration POLL_DURATION = Duration.standardSeconds(1);
@StateId("bufferedElements")
@@ -74,6 +78,17 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
private final ConvertMessagesDoFn<DestinationT, ElementT> convertMessagesDoFn;
+ private transient @Nullable Cache<BoundedWindow, Instant> drainDeadlines = null;
+
+ private Cache<BoundedWindow, Instant> getDrainDeadlines() {
+ Cache<BoundedWindow, Instant> cache = drainDeadlines;
+ if (cache == null) {
+ cache = CacheBuilder.newBuilder().expireAfterAccess(java.time.Duration.ofMinutes(30)).build();
+ drainDeadlines = cache;
+ }
+ return cache;
+ }
+
public SchemaUpdateHoldingFn(
Coder<ElementT> elementCoder,
ConvertMessagesDoFn<DestinationT, ElementT> convertMessagesDoFn) {
@@ -140,7 +155,12 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
} else {
// This means that the table schema was recently updated. Try to flush the pending elements.
if (tryFlushBuffer(
- element.getKey().getKey(), context.getPipelineOptions(), bag, minBufferedTimestamp, o)) {
+ element.getKey().getKey(),
+ context.getPipelineOptions(),
+ bag,
+ minBufferedTimestamp,
+ o,
+ true)) {
// Nothing left in buffer. clear timer.
pollTimer.clear();
timerTs.clear();
@@ -165,6 +185,7 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
@OnTimer("pollTimer")
public void onPollTimer(
+ OnTimerContext context,
@Key ShardedKey<DestinationT> key,
PipelineOptions pipelineOptions,
@StateId("bufferedElements") BagState<TimestampedValue<ElementT>> bag,
@@ -174,7 +195,9 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
BoundedWindow window,
MultiOutputReceiver o)
throws Exception {
- if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o)) {
+ convertMessagesDoFn.getDynamicDestinations().setSideInputAccessorFromOnTimerContext(context);
+
+ if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o, true)) {
timerTs.clear();
} else {
// We still have buffered elements. Make sure that the polling timer keeps looping.
@@ -188,24 +211,58 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
@OnWindowExpiration
public void onWindowExpiration(
+ OnWindowExpirationContext context,
@Key ShardedKey<DestinationT> key,
PipelineOptions pipelineOptions,
@StateId("bufferedElements") BagState<TimestampedValue<ElementT>> bag,
@StateId("minBufferedTimestamp") CombiningState<Long, long[], Long> minBufferedTimestamp,
+ BoundedWindow window,
MultiOutputReceiver o)
throws Exception {
+ minBufferedTimestamp.readLater();
+ bag.readLater();
+ if (Iterables.isEmpty(bag.read())) {
+ return;
+ }
+
+ convertMessagesDoFn
+ .getDynamicDestinations()
+ .setSideInputAccessorFromOnWindowExpirationContext(context);
+
+ java.time.Duration waitTime =
+ java.time.Duration.ofMillis(
+ pipelineOptions
+ .as(BigQueryOptions.class)
+ .getStorageApiMismatchDrainRetryTimeMilliSec());
+ // Shared across every key in this window; see drainDeadlines.
+ // TODO: This needs to use the internal clock instead Instant.now(), as the row deadlines are
+ // based on the
+ // internal clock. Add support for accessing the internal clock in OnWindowExpiration.
+ Instant drainDeadline =
+ getDrainDeadlines()
+ .get(window, () -> Instant.now().plus(Duration.millis(waitTime.toMillis())));
+ LOG.info(
+ "Draining buffered schema-mismatched rows for destination {}, waiting until {} for schema update.",
+ key.getKey(),
+ drainDeadline);
+
// 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 =
- new ExponentialBackOff.Builder()
- .setMaxElapsedTimeMillis((int) TimeUnit.SECONDS.toMillis(10))
- .build();
+ new ExponentialBackOff.Builder().setMaxElapsedTimeMillis((int) waitTime.toMillis()).build();
do {
- if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o)) {
+ if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o, true)) {
return;
}
- } while (BackOffUtils.next(com.google.api.client.util.Sleeper.DEFAULT, backoff));
- throw new RuntimeException("Failed to flush elements on window expiration!");
+ } while (Instant.now().isBefore(drainDeadline)
+ && BackOffUtils.next(com.google.api.client.util.Sleeper.DEFAULT, backoff));
+
+ // Drain deadline expired: route any remaining buffered elements that still cannot be converted
+ // to the dead-letter collection (failedWritesTag).
+ LOG.warn(
+ "Timed out waiting for table schema update during drain for destination {}; routing remaining buffered rows to failed-rows collection.",
+ key.getKey());
+ tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o, false);
}
// Returns true if the buffer is completely flushed.
@@ -214,7 +271,8 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
PipelineOptions pipelineOptions,
@StateId("bufferedElements") BagState<TimestampedValue<ElementT>> bag,
@StateId("minBufferedTimestamp") CombiningState<Long, long[], Long> minBufferedTimestamp,
- MultiOutputReceiver o)
+ MultiOutputReceiver o,
+ boolean collectSchemaErrors)
throws Exception {
// Force an update of the MessageConverter schema.
StorageApiDynamicDestinations.MessageConverter<ElementT> messageConverter =
@@ -236,8 +294,13 @@ public class SchemaUpdateHoldingFn<DestinationT extends @NonNull Object, Element
bag.read(),
e -> TimestampedValue.of(KV.of(destination, e.getValue()), e.getTimestamp()));
+ // if collectSchemaErrors==false, convertMessagesDoFn should route failed row conversion to the
+ // failed-row
+ // PCollection.
TableRowToStorageApiProto.ErrorCollector errorCollector =
- UpgradeTableSchema.newErrorCollector();
+ collectSchemaErrors
+ ? UpgradeTableSchema.newErrorCollector()
+ : TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT;
Iterable<TimestampedValue<KV<DestinationT, ElementT>>> unProcessed =
convertMessagesDoFn.handleProcessElements(
messageConverter, kvBagElements, o, errorCollector);
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 26eb6d55d20..ebcfc2f143c 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,12 @@
*/
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 org.apache.beam.sdk.util.ThrowingSupplier;
+import org.apache.beam.sdk.util.Preconditions;
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 +32,69 @@ import org.joda.time.Instant;
* parameter controls how many rows are batched into a single ProtoRows object before we move on to
* the next one.
*/
-class SplittingIterable implements Iterable<SplittingIterable.Value> {
- @AutoValue
- abstract static class Value {
- abstract ProtoRows getProtoRows();
-
- abstract List<Instant> getTimestamps();
-
- abstract List<@Nullable TableRow> getFailsafeTableRows();
-
- abstract boolean getSchemaMismatchSeen();
- }
-
- interface ConcatFields {
- ByteString concat(ByteString bytes, TableRow tableRows)
- throws TableRowToStorageApiProto.SchemaConversionException;
- }
-
- private final Iterable<StorageApiWritePayload> underlying;
+class SplittingIterable implements Iterable<AppendRowsPacket> {
+ private final Iterable<StoragePayloadWithDeadline> underlying;
private final long splitSize;
- private final ConcatFields concatProtoAndTableRow;
- private final Function<ByteString, TableRow> protoToTableRow;
- private final BiConsumer<TimestampedValue<TableRow>, String> failedRowsConsumer;
- private final ThrowingSupplier<byte[]> getCurrentTableSchemaHash;
- private final ThrowingSupplier<Descriptors.Descriptor> getCurrentTableSchemaDescriptor;
- private final boolean autoUpdateSchema;
- private final Instant elementsTimestamp;
+ private final Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean>
+ failedRowsHandler;
+ private final Instant elementTimestamp;
+ private final AppendClientInfo appendClientInfo;
+ private final SchemaChangeDetectorHelper schemaChangeDetectorHelper;
public SplittingIterable(
- Iterable<StorageApiWritePayload> underlying,
+ Iterable<StoragePayloadWithDeadline> underlying,
long splitSize,
- ConcatFields concatProtoAndTableRow,
- Function<ByteString, TableRow> protoToTableRow,
- BiConsumer<TimestampedValue<TableRow>, String> failedRowsConsumer,
- ThrowingSupplier<byte[]> getCurrentTableSchemaHash,
- ThrowingSupplier<Descriptors.Descriptor> getCurrentTableSchemaDescriptor,
- boolean autoUpdateSchema,
- Instant elementsTimestamp) {
+ Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean> failedRowsHandler,
+ Instant elementTimestamp,
+ AppendClientInfo appendClientInfo,
+ SchemaChangeDetectorHelper schemaChangeDetectorHelper) {
this.underlying = underlying;
this.splitSize = splitSize;
- this.concatProtoAndTableRow = concatProtoAndTableRow;
- this.protoToTableRow = protoToTableRow;
- this.failedRowsConsumer = failedRowsConsumer;
- this.getCurrentTableSchemaHash = getCurrentTableSchemaHash;
- this.getCurrentTableSchemaDescriptor = getCurrentTableSchemaDescriptor;
- this.autoUpdateSchema = autoUpdateSchema;
- this.elementsTimestamp = elementsTimestamp;
+ this.failedRowsHandler = failedRowsHandler;
+ this.elementTimestamp = elementTimestamp;
+ this.appendClientInfo = appendClientInfo;
+ this.schemaChangeDetectorHelper = schemaChangeDetectorHelper;
}
@Override
- public Iterator<Value> iterator() {
- return new Iterator<Value>() {
- final PeekingIterator<StorageApiWritePayload> underlyingIterator =
+ public Iterator<AppendRowsPacket> iterator() {
+ return new Iterator<AppendRowsPacket>() {
+ final PeekingIterator<StoragePayloadWithDeadline> 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<Instant> 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,
+ appendClientInfo,
+ failedRowsHandler);
+ return value.getProtoRows().getSerializedRowsCount() == 0
+ && value.getSchemaMismatchedRows().isEmpty()
+ ? 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 d6981858a4b..1c21260b864 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 @@ import org.apache.beam.sdk.transforms.DoFn;
abstract class StorageApiDynamicDestinations<T, DestinationT>
extends DynamicDestinationsHelpers.DelegatingDynamicDestinations<T, DestinationT> {
public interface MessageConverter<T> {
- com.google.cloud.bigquery.storage.v1.TableSchema getTableSchema();
+ TableSchema getTableSchema();
DescriptorProtos.DescriptorProto getDescriptor(boolean includeCdcColumns) throws Exception;
@@ -42,6 +43,8 @@ abstract class StorageApiDynamicDestinations<T, DestinationT>
TableRow toFailsafeTableRow(T element);
void updateSchemaFromTable() throws IOException, InterruptedException;
+
+ default void updateSchema(TableSchema schema) {}
}
StorageApiDynamicDestinations(DynamicDestinations<T, DestinationT> inner) {
@@ -60,4 +63,17 @@ abstract class StorageApiDynamicDestinations<T, DestinationT>
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 86307eb7017..4dc5987c779 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.DescriptorProtos;
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<T, DestinationT extends @NonN
private final CreateDisposition createDisposition;
private final boolean ignoreUnknownValues;
private final boolean autoSchemaUpdates;
- private final Set<BigQueryIO.Write.SchemaUpdateOption> schemaUpdateOptions;
+ private final boolean useSchemaUpdatingTableRow;
private static final TableSchemaCache SCHEMA_CACHE =
new TableSchemaCache(Duration.standardSeconds(1));
@@ -64,7 +63,7 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
CreateDisposition createDisposition,
boolean ignoreUnknownValues,
boolean autoSchemaUpdates,
- Set<BigQueryIO.Write.SchemaUpdateOption> schemaUpdateOptions) {
+ boolean useSchemaUpdatingTableRow) {
super(inner);
this.formatFunction = formatFunction;
this.formatRecordOnFailureFunction = formatRecordOnFailureFunction;
@@ -72,7 +71,7 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
this.createDisposition = createDisposition;
this.ignoreUnknownValues = ignoreUnknownValues;
this.autoSchemaUpdates = autoSchemaUpdates;
- this.schemaUpdateOptions = schemaUpdateOptions;
+ this.useSchemaUpdatingTableRow = useSchemaUpdatingTableRow;
}
static void clearSchemaCache() throws ExecutionException, InterruptedException {
@@ -94,7 +93,7 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
}
};
- return schemaUpdateOptions.isEmpty()
+ return !useSchemaUpdatingTableRow
? getConverter.apply(getSchema(destination))
: new SchemaUpgradingTableRowConverter(
getConverter, options, datasetService, writeStreamService);
@@ -153,9 +152,16 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
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);
@@ -181,6 +187,15 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
// If the table already exists, then try and fetch the schema from the existing
// table.
localTableSchema = SCHEMA_CACHE.getSchema(tableReference, datasetService);
+ if (localTableSchema == null && createDisposition != CreateDisposition.CREATE_NEVER) {
+ localTableSchema = getSchema(destination);
+ if (localTableSchema != null) {
+ localTableSchema =
+ MoreObjects.firstNonNull(
+ SCHEMA_CACHE.putSchemaIfAbsent(tableReference, localTableSchema),
+ localTableSchema);
+ }
+ }
if (localTableSchema == null) {
if (createDisposition == CreateDisposition.CREATE_NEVER) {
throw new RuntimeException(
@@ -283,6 +298,11 @@ public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonN
changeType,
changeSequenceNum,
collectedExceptions);
+
+ if (unknownFields != null && unknownFields.isEmpty()) {
+ unknownFields = null;
+ }
+
return StorageApiWritePayload.of(
msg == null ? new byte[0] : msg.toByteArray(),
unknownFields,
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 38c104ab969..5adfbc60324 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
@@ -76,6 +76,7 @@ public class StorageApiLoads<DestinationT, ElementT>
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;
@@ -100,6 +101,7 @@ public class StorageApiLoads<DestinationT, ElementT>
boolean allowInconsistentWrites,
boolean allowAutosharding,
boolean autoUpdateSchema,
+ @Nullable Duration autoUpdateSchemaStrictTimeout,
boolean ignoreUnknownValues,
boolean propagateSuccessfulStorageApiWrites,
Predicate<String> propagateSuccessfulStorageApiWritesPredicate,
@@ -121,6 +123,7 @@ public class StorageApiLoads<DestinationT, ElementT>
this.allowInconsistentWrites = allowInconsistentWrites;
this.allowAutosharding = allowAutosharding;
this.autoUpdateSchema = autoUpdateSchema;
+ this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout;
this.ignoreUnknownValues = ignoreUnknownValues;
if (propagateSuccessfulStorageApiWrites) {
this.successfulWrittenRowsTag = new TupleTag<>("successfulPublishedRowsTag");
@@ -196,7 +199,9 @@ public class StorageApiLoads<DestinationT, ElementT>
successfulRowsPredicate,
BigQueryStorageApiInsertErrorCoder.of(),
TableRowJsonCoder.of(),
+ destinationCoder,
autoUpdateSchema,
+ autoUpdateSchemaStrictTimeout,
ignoreUnknownValues,
createDisposition,
kmsKey,
@@ -301,9 +306,11 @@ public class StorageApiLoads<DestinationT, ElementT>
successfulWrittenRowsTag,
successfulRowsPredicate,
autoUpdateSchema,
+ autoUpdateSchemaStrictTimeout,
ignoreUnknownValues,
defaultMissingValueInterpretation,
- bigLakeConfiguration));
+ bigLakeConfiguration,
+ hasSchemaUpdateOptions));
PCollection<BigQueryStorageApiInsertError> insertErrors =
PCollectionList.of(convertMessagesResult.get(failedRowsTag))
@@ -404,6 +411,7 @@ public class StorageApiLoads<DestinationT, ElementT>
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 87ed4518bb0..6e54c3da404 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
@@ -82,6 +82,22 @@ public abstract class StorageApiWritePayload {
.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 58bbed8ba5a..3e32249b0e4 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.Map;
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 @@ import org.apache.beam.sdk.values.TupleTagList;
* writes, use {@link StorageApiWritesShardedRecords} or {@link StorageApiWriteUnshardedRecords}.
*/
@SuppressWarnings("FutureReturnValueIgnored")
-public class StorageApiWriteRecordsInconsistent<DestinationT, ElementT>
+public class StorageApiWriteRecordsInconsistent<DestinationT extends @NonNull Object, ElementT>
extends PTransform<PCollection<KV<DestinationT, StorageApiWritePayload>>, PCollectionTuple> {
private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
private final BigQueryServices bqServices;
@@ -47,10 +52,15 @@ public class StorageApiWriteRecordsInconsistent<DestinationT, ElementT>
private final Predicate<String> successfulRowsPredicate;
+ // This output is effectively ignored, since this code path uses the default stream which is never
+ // finalized.
private final TupleTag<KV<String, String>> finalizeTag = new TupleTag<>("finalizeTag");
private final Coder<BigQueryStorageApiInsertError> failedRowsCoder;
private final Coder<TableRow> successfulRowsCoder;
+ private final Coder<DestinationT> destinationCoder;
private final boolean autoUpdateSchema;
+ private final @Nullable Duration autoUpdateSchemaStrictTimeout;
+ private final @Nullable TupleTag<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRowsTag;
private final boolean ignoreUnknownValues;
private final BigQueryIO.Write.CreateDisposition createDisposition;
private final @Nullable String kmsKey;
@@ -66,7 +76,9 @@ public class StorageApiWriteRecordsInconsistent<DestinationT, ElementT>
Predicate<String> successfulRowsPredicate,
Coder<BigQueryStorageApiInsertError> failedRowsCoder,
Coder<TableRow> successfulRowsCoder,
+ Coder<DestinationT> destinationCoder,
boolean autoUpdateSchema,
+ @Nullable Duration autoUpdateSchemaStrictTimeout,
boolean ignoreUnknownValues,
BigQueryIO.Write.CreateDisposition createDisposition,
@Nullable String kmsKey,
@@ -78,9 +90,16 @@ public class StorageApiWriteRecordsInconsistent<DestinationT, ElementT>
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 class StorageApiWriteRecordsInconsistent<DestinationT, ElementT>
if (successfulRowsTag != null) {
tupleTagList = tupleTagList.and(successfulRowsTag);
}
+ if (mismatchedRowsTag != null) {
+ tupleTagList = tupleTagList.and(mismatchedRowsTag);
+ }
+
+ StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT> 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<DestinationT, ElementT> 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<KV<DestinationT, StoragePayloadWithDeadline>> 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<BigQueryStorageApiInsertError> flattenedErrors =
+ PCollectionList.of(result.get(failedRowsTag))
+ .and(mismatchedResult.get(failedRowsTag))
+ .apply("flattenErrors", Flatten.pCollections());
+
+ PCollectionTuple flattenedResult = PCollectionTuple.of(failedRowsTag, flattenedErrors);
+ if (successfulRowsTag != null) {
+ PCollection<TableRow> 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 170de2ed2c5..33ed612d29c 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.ByteString;
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,6 +48,8 @@ import java.util.Set;
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.stream.Collectors;
import java.util.stream.StreamSupport;
@@ -81,14 +82,18 @@ import org.apache.beam.sdk.values.KV;
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;
@@ -122,7 +127,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation;
private final @Nullable Map<String, String> bigLakeConfiguration;
- private static final AppendClientCache<String> APPEND_CLIENTS =
+ static final AppendClientCache<String> APPEND_CLIENTS =
new AppendClientCache<>(Duration.standardMinutes(15));
static void clearCache() {
@@ -138,6 +143,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
Coder<BigQueryStorageApiInsertError> failedRowsCoder,
Coder<TableRow> successfulRowsCoder,
boolean autoUpdateSchema,
+ @Nullable Duration autoUpdateSchemaStrictTimeout,
boolean ignoreUnknownValues,
BigQueryIO.Write.CreateDisposition createDisposition,
@Nullable String kmsKey,
@@ -152,6 +158,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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 +178,35 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
if (successfulRowsTag != null) {
tupleTagList = tupleTagList.and(successfulRowsTag);
}
+
PCollectionTuple writeResults =
input.apply(
"Write Records",
ParDo.of(
new WriteRecordsDoFn<>(
- 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 +227,158 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
static class WriteRecordsDoFn<DestinationT extends @NonNull Object, ElementT>
extends DoFn<KV<DestinationT, StorageApiWritePayload>, KV<String, String>> {
+ WriteRecordsDoFnImpl<DestinationT, ElementT> writeRecordsDoFnImpl;
+
+ WriteRecordsDoFn(WriteRecordsDoFnImpl<DestinationT, ElementT> writeRecordsDoFnImpl) {
+ this.writeRecordsDoFnImpl = writeRecordsDoFnImpl;
+ }
+
+ @StartBundle
+ public void startBundle() throws IOException {
+ writeRecordsDoFnImpl.startBundle();
+ }
+
+ @FinishBundle
+ public void finishBundle(final FinishBundleContext context) throws Exception {
+ OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver =
+ value ->
+ WindowedValues.<BigQueryStorageApiInsertError>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<TableRow> successfulRowsReceiver = null;
+ if (writeRecordsDoFnImpl.successfulRowsTag != null) {
+ successfulRowsReceiver =
+ makeSuccessfulRowsReceiver(context, writeRecordsDoFnImpl.successfulRowsTag);
+ }
+
+ OutputReceiver<KV<String, String>> finalizeOutputRecever =
+ value ->
+ WindowedValues.<KV<String, String>>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<KV<DestinationT, StoragePayloadWithDeadline>> 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<TableRow> makeSuccessfulRowsReceiver(
+ FinishBundleContext context, TupleTag<TableRow> successfulRowsTag) {
+ return new OutputReceiver<TableRow>() {
+ @Override
+ public OutputBuilder<TableRow> builder(TableRow value) {
+ return WindowedValues.<TableRow>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<DestinationT, StorageApiWritePayload> 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();
+ // TODO: Once the processing-time clock is accessible in @OnWindowExpiration, compute
+ // targetDeadline using the processing-time clock instead of Instant.now() so it matches
+ // Timer.getCurrentRelativeTime() in TestStream.
+ org.joda.time.Instant targetDeadline =
+ (writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout != null)
+ ? startTimeForLocalRetry.plus(writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout)
+ : BoundedWindow.TIMESTAMP_MAX_VALUE;
+
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRows =
+ writeRecordsDoFnImpl.processElement(
+ pipelineOptions,
+ KV.of(
+ element.getKey(),
+ StoragePayloadWithDeadline.of(element.getValue(), targetDeadline)),
+ startTimeForLocalRetry,
+ o);
+ if (!Iterables.isEmpty(mismatchedRows)) {
+ final OutputReceiver<KV<DestinationT, StoragePayloadWithDeadline>> 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<DestinationT extends @NonNull Object, ElementT>
+ implements Serializable {
private final Counter forcedFlushes = Metrics.counter(WriteRecordsDoFn.class, "forcedFlushes");
private final TupleTag<KV<String, String>> finalizeTag;
private final TupleTag<BigQueryStorageApiInsertError> failedRowsTag;
@@ -222,10 +386,13 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
private final Predicate<String> successfulRowsPredicate;
private final boolean autoUpdateSchema;
+ private final @Nullable Duration autoUpdateSchemaStrictTimeout;
+ private final @Nullable TupleTag<KV<DestinationT, StoragePayloadWithDeadline>>
+ 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<AppendRowsResponse> {
@@ -250,14 +417,14 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
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<StorageApiWritePayload> pendingMessages;
- private List<org.joda.time.Instant> pendingTimestamps;
+ private List<StoragePayloadWithDeadline> pendingMessages;
private transient @Nullable WriteStreamService maybeWriteStreamService;
private final Counter recordsAppended =
Metrics.counter(WriteRecordsDoFn.class, "recordsAppended");
@@ -279,8 +446,10 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
private final long maxRequestSize;
private final boolean includeCdcColumns;
+ private final SchemaChangeDetectorHelper schemaChangeDetectorHelper;
public DestinationState(
+ DestinationT destination,
TableDestination tableDestination,
String tableUrn,
String shortTableUrn,
@@ -293,11 +462,11 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
Callable<Boolean> 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 +478,12 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
if (includeCdcColumns) {
checkState(useDefaultStream);
}
+ this.schemaChangeDetectorHelper =
+ new SchemaChangeDetectorHelper(
+ autoUpdateSchema,
+ ignoreUnknownValues,
+ tableDestination.getTableReference(),
+ autoUpdateSchemaStrictTimeout != null);
}
public TableDestination getTableDestination() {
@@ -323,18 +498,13 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
String getDefaultStreamName() {
- return BigQueryHelpers.stripPartitionDecorator(tableUrn) + "/streams/_default";
+ return WriteRecordsDoFnImpl.this.getDefaultStreamName(tableUrn);
}
String getStreamAppendClientCacheEntryKey() {
- if (useDefaultStream) {
- String defaultStreamKey = getDefaultStreamName() + "-client" + clientNumber;
- // The storage write API doesn't currently allow both inserts and updates/deletes on the
- // same connection.
- // Once this limitation is removed, we can remove this.
- return includeCdcColumns ? defaultStreamKey + "-cdc" : defaultStreamKey;
- }
- return this.streamName;
+ return useDefaultStream
+ ? getDefaultStreamAppendClientCacheEntryKey(tableUrn, clientNumber, includeCdcColumns)
+ : this.streamName;
}
String getOrCreateStreamName() throws Exception {
@@ -409,21 +579,15 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
AtomicBoolean updated = new AtomicBoolean();
CreateTableHelpers.createTableWrapper(
() -> {
- if (autoUpdateSchema) {
- @Nullable TableSchema streamSchema =
- Preconditions.checkStateNotNull(maybeWriteStreamService)
- .getWriteStreamSchema(streamName);
- if (streamSchema != null) {
- Optional<TableSchema> 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<TableSchema> 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;
},
@@ -483,76 +647,54 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
}
- void addMessage(
- StorageApiWritePayload payload,
- org.joda.time.Instant elementTs,
- OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver)
- throws Exception {
- maybeTickleCache();
-
- @Nullable ByteString mergedPayloadBytes = null;
- if (autoUpdateSchema) {
- if (appendClientInfo == null) {
- appendClientInfo = getAppendClientInfo(true, null);
- }
- @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;
- }
+ void writeFailedRows(
+ Iterable<StoragePayloadWithDeadline> failedRows,
+ String msg,
+ String statusMetric,
+ org.joda.time.Instant defaultTs,
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> 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());
}
+ 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,
+ statusMetric,
+ tableDestination.getShortTableUrn())
+ .inc(1);
}
- 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);
- }
- 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<AppendRowsResponse, AppendRowsContext> retryManager,
- OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
- @Nullable OutputReceiver<TableRow> successfulRowsReceiver)
+ org.joda.time.@Nullable Instant startTimeForLocalRetry,
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
+ DoFn.@Nullable OutputReceiver<TableRow> successfulRowsReceiver,
+ BiConsumer<DestinationT, StoragePayloadWithDeadline> 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
@@ -563,26 +705,38 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
BigQuerySinkMetrics.throttledTimeCounter(
BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM))
.backoff();
+
+ final Duration initialMismatchRetryTime =
+ Duration.millis(WriteRecordsDoFnImpl.this.initialMismatchRetryTimeMilliSec);
+
+ AppendRowsPacket rowsToProcess;
+ Iterable<StoragePayloadWithDeadline> 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)
+ .isAfter(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
+ .getSchemaMismatchedRows()
+ .forEach(row -> bufferMismatchedRow.accept(destination, row));
+
+ // Process the good rows!
+ rowsToProcess = processingRows.getSchemaMatchedRowsOnly();
+ 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();
@@ -595,12 +749,15 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
} while (schemaMismatchSeen && BackOffUtils.next(Sleeper.DEFAULT, backoff));
pendingMessages.clear();
- final ProtoRows inserts = insertsBuilder.build();
- List<org.joda.time.Instant> 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
@@ -610,21 +767,14 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
+ "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,
@@ -632,7 +782,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
tableDestination.getTableReference()),
timestamp);
}
- int numRowsFailed = inserts.getSerializedRowsCount();
+ int numRowsFailed = rowsToProcess.getProtoRows().getSerializedRowsCount();
BigQuerySinkMetrics.appendRowsRowStatusCounter(
BigQuerySinkMetrics.RowStatus.FAILED,
BigQuerySinkMetrics.PAYLOAD_TOO_LARGE,
@@ -646,10 +796,14 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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 -> {
@@ -702,15 +856,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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(
@@ -889,8 +1035,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
Descriptor descriptor = null;
try {
descriptor =
- TableRowToStorageApiProto.wrapDescriptorProto(
- Preconditions.checkStateNotNull(appendClientInfo).getDescriptor());
+ Preconditions.checkStateNotNull(appendClientInfo).getWrappedDescriptor();
} catch (DescriptorValidationException e) {
LOG.warn(
"Failure getting proto descriptor. Successful output will not be produced.",
@@ -901,12 +1046,8 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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) {
@@ -918,7 +1059,35 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
},
appendRowsContext);
maybeTickleCache();
- return inserts.getSerializedRowsCount();
+ return rowsToProcess.getProtoRows().getSerializedRowsCount();
+ }
+
+ private AppendRowsPacket buildInserts(
+ Iterable<StoragePayloadWithDeadline> payloads,
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver) {
+ Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean> failedRowsHandler =
+ (autoUpdateSchemaStrictTimeout == null)
+ ? tv -> {
+ rowsSentToFailedRowsCollection.inc();
+ BigQuerySinkMetrics.appendRowsRowStatusCounter(
+ BigQuerySinkMetrics.RowStatus.FAILED,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ shortTableUrn)
+ .inc(1);
+ failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp());
+ return true;
+ }
+ : tv -> false;
+
+ PeekingIterator<StoragePayloadWithDeadline> peekingIterator =
+ Iterators.peekingIterator(payloads.iterator());
+ return AppendRowsPacket.fromStorageApiWritePayload(
+ peekingIterator,
+ Long.MAX_VALUE,
+ schemaChangeDetectorHelper,
+ GlobalWindow.INSTANCE.maxTimestamp(),
+ getAppendClientInfo(true, null),
+ failedRowsHandler);
}
String retrieveErrorDetails(Iterable<AppendRowsContext> failedContext) {
@@ -937,15 +1106,15 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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<TableSchema> 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
@@ -961,8 +1130,8 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
}
- private @Nullable Map<DestinationT, DestinationState> destinations = Maps.newHashMap();
- private final TwoLevelMessageConverterCache<DestinationT, ElementT> messageConverters;
+ @Nullable Map<DestinationT, DestinationState> destinations = Maps.newHashMap();
+ final TwoLevelMessageConverterCache<DestinationT, ElementT> messageConverters;
private transient @Nullable DatasetService maybeDatasetService;
private transient @Nullable WriteStreamService maybeWriteStreamService;
private int numPendingRecords = 0;
@@ -970,13 +1139,14 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
private final int flushThresholdBytes;
private final int flushThresholdCount;
private final int maxRetries;
- private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
+ final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
private final BigQueryServices bqServices;
private final boolean useDefaultStream;
private int streamAppendClientCount;
private final @Nullable Map<String, String> bigLakeConfiguration;
+ private final int initialMismatchRetryTimeMilliSec;
- WriteRecordsDoFn(
+ WriteRecordsDoFnImpl(
String operationName,
StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations,
BigQueryServices bqServices,
@@ -989,13 +1159,16 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
@Nullable TupleTag<TableRow> successfulRowsTag,
Predicate<String> successfulRowsPredicate,
boolean autoUpdateSchema,
+ @Nullable Duration autoUpdateSchemaStrictTimeout,
+ @Nullable TupleTag<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRowsTag,
boolean ignoreUnknownValues,
BigQueryIO.Write.CreateDisposition createDisposition,
@Nullable String kmsKey,
boolean usesCdc,
AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation,
int maxRetries,
- @Nullable Map<String, String> bigLakeConfiguration) {
+ @Nullable Map<String, String> bigLakeConfiguration,
+ int initialMismatchRetryTimeMilliSec) {
this.messageConverters = new TwoLevelMessageConverterCache<>(operationName);
this.dynamicDestinations = dynamicDestinations;
this.bqServices = bqServices;
@@ -1008,6 +1181,8 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
this.successfulRowsTag = successfulRowsTag;
this.successfulRowsPredicate = successfulRowsPredicate;
this.autoUpdateSchema = autoUpdateSchema;
+ this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout;
+ this.mismatchedRowsTag = mismatchedRowsTag;
this.ignoreUnknownValues = ignoreUnknownValues;
this.createDisposition = createDisposition;
this.kmsKey = kmsKey;
@@ -1015,6 +1190,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
this.defaultMissingValueInterpretation = defaultMissingValueInterpretation;
this.maxRetries = maxRetries;
this.bigLakeConfiguration = bigLakeConfiguration;
+ this.initialMismatchRetryTimeMilliSec = initialMismatchRetryTimeMilliSec;
}
boolean shouldFlush(int recordBytes) {
@@ -1023,9 +1199,10 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
&& numPendingRecords > 0);
}
- void flushIfNecessary(
- OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
- @Nullable OutputReceiver<TableRow> successfulRowsReceiver,
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> flushIfNecessary(
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
+ DoFn.@Nullable OutputReceiver<TableRow> successfulRowsReceiver,
+ org.joda.time.@Nullable Instant startTimeForLocalRetry,
int recordBytes)
throws Exception {
if (shouldFlush(recordBytes)) {
@@ -1033,14 +1210,17 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
// 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<BigQueryStorageApiInsertError> failedRowsReceiver,
- @Nullable OutputReceiver<TableRow> successfulRowsReceiver)
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> flushAll(
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
+ DoFn.@Nullable OutputReceiver<TableRow> successfulRowsReceiver,
+ org.joda.time.@Nullable Instant startTimeForLocalRetry)
throws Exception {
+ List<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRows = Lists.newArrayList();
List<RetryManager<AppendRowsResponse, AppendRowsContext>> retryManagers =
Lists.newArrayListWithCapacity(Preconditions.checkStateNotNull(destinations).size());
long numRowsWritten = 0;
@@ -1055,7 +1235,12 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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) {
@@ -1073,9 +1258,10 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
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));
@@ -1083,7 +1269,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
return maybeDatasetService;
}
- private WriteStreamService initializeWriteStreamService(PipelineOptions pipelineOptions) {
+ WriteStreamService getWriteStreamService(PipelineOptions pipelineOptions) {
if (maybeWriteStreamService == null) {
maybeWriteStreamService =
bqServices.getWriteStreamService(pipelineOptions.as(BigQueryOptions.class));
@@ -1091,7 +1277,6 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
return maybeWriteStreamService;
}
- @StartBundle
public void startBundle() throws IOException {
destinations = Maps.newHashMap();
numPendingRecords = 0;
@@ -1099,7 +1284,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
DestinationState createDestinationState(
- ProcessContext c,
+ PipelineOptions options,
DestinationT destination,
boolean useCdc,
DatasetService datasetService,
@@ -1116,7 +1301,7 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
Callable<Boolean> tryCreateTable =
() -> {
CreateTableHelpers.possiblyCreateTable(
- c.getPipelineOptions().as(BigQueryOptions.class),
+ options.as(BigQueryOptions.class),
tableDestination1,
() -> dynamicDestinations.getSchema(destination),
() -> dynamicDestinations.getTableConstraints(destination),
@@ -1132,12 +1317,9 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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(),
@@ -1154,25 +1336,67 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
}
}
- @ProcessElement
- public void process(
- ProcessContext c,
+ void writeFailedRows(
+ DestinationT destination,
+ Iterable<StoragePayloadWithDeadline> failedRows,
+ String msg,
+ String statusMetric,
PipelineOptions pipelineOptions,
- @Element KV<DestinationT, StorageApiWritePayload> element,
- @Timestamp org.joda.time.Instant elementTs,
- MultiOutputReceiver o)
+ org.joda.time.Instant defaultTs,
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> 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, statusMetric, defaultTs, failedRowsReceiver);
+ }
+
+ private String getDefaultStreamName(String tableUrn) {
+ return BigQueryHelpers.stripPartitionDecorator(tableUrn) + "/streams/_default";
+ }
+
+ private String getDefaultStreamAppendClientCacheEntryKey(
+ String tableUrn, int clientNumber, boolean includeCdcColumns) {
+ String defaultStreamKey = getDefaultStreamName(tableUrn) + "-client" + clientNumber;
+ // The storage write API doesn't currently allow both inserts and updates/deletes on the
+ // same connection.
+ // Once this limitation is removed, we can remove this.
+ return includeCdcColumns ? defaultStreamKey + "-cdc" : defaultStreamKey;
+ }
+
+ void invalidateAllAppendClientsForTable(String tableUrn) {
+ for (int i = 0; i < streamAppendClientCount; ++i) {
+ String key = getDefaultStreamAppendClientCacheEntryKey(tableUrn, i, usesCdc);
+ APPEND_CLIENTS.invalidate(key);
+ }
+ }
+
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> processElement(
+ PipelineOptions pipelineOptions,
+ KV<DestinationT, StoragePayloadWithDeadline> 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,
@@ -1185,91 +1409,50 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
state.getTableDestination().getTableReference(),
pipelineOptions.as(BigQueryOptions.class)));
- OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver = o.get(failedRowsTag);
- @Nullable OutputReceiver<TableRow> successfulRowsReceiver =
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver = o.get(failedRowsTag);
+ DoFn.@Nullable OutputReceiver<TableRow> 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<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRows =
+ flushIfNecessary(
+ failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry, recordBytes);
+ state.addMessage(element.getValue());
++numPendingRecords;
numPendingRecordBytes += recordBytes;
- }
- private OutputReceiver<TableRow> makeSuccessfulRowsreceiver(
- FinishBundleContext context, TupleTag<TableRow> successfulRowsTag) {
- return new OutputReceiver<TableRow>() {
- @Override
- public OutputBuilder<TableRow> builder(TableRow value) {
- return WindowedValues.<TableRow>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<BigQueryStorageApiInsertError> failedRowsReceiver =
- new OutputReceiver<BigQueryStorageApiInsertError>() {
- @Override
- public OutputBuilder<BigQueryStorageApiInsertError> builder(
- BigQueryStorageApiInsertError value) {
- return WindowedValues.<BigQueryStorageApiInsertError>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<TableRow> successfulRowsReceiver = null;
- if (successfulRowsTag != null) {
- successfulRowsReceiver = makeSuccessfulRowsreceiver(context, successfulRowsTag);
- }
+ public Iterable<KV<DestinationT, StoragePayloadWithDeadline>> finishBundle(
+ DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver,
+ DoFn.@Nullable OutputReceiver<TableRow> successfulRowsReceiver,
+ DoFn.OutputReceiver<KV<String, String>> finalizeOutputReceiver,
+ org.joda.time.@Nullable Instant startTimeForLocalRetry)
+ throws Exception {
- flushAll(failedRowsReceiver, successfulRowsReceiver);
-
- final Map<DestinationT, DestinationState> 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);
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> mismatchedRows;
+ try {
+ mismatchedRows =
+ flushAll(failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry);
+
+ final Map<DestinationT, DestinationState> destinations =
+ Preconditions.checkStateNotNull(this.destinations);
+ for (DestinationState state : destinations.values()) {
+ if (!useDefaultStream && !Strings.isNullOrEmpty(state.streamName)) {
+ finalizeOutputReceiver.output(KV.of(state.tableUrn, state.streamName));
+ }
}
- state.teardown();
+ } finally {
+ final Map<DestinationT, DestinationState> destinations =
+ Preconditions.checkStateNotNull(this.destinations);
+ destinations.values().forEach(DestinationState::teardown);
+ destinations.clear();
+ this.destinations = null;
}
- destinations.clear();
- this.destinations = null;
+ return mismatchedRows;
}
- @Teardown
public void teardown() {
destinations = null;
try {
@@ -1285,10 +1468,5 @@ public class StorageApiWriteUnshardedRecords<DestinationT, ElementT>
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 51be28fedf8..e615182154c 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,7 +18,9 @@
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.client.util.ExponentialBackOff;
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.api.services.bigquery.model.TableReference;
@@ -48,10 +50,13 @@ import java.util.function.Consumer;
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.options.PipelineOptions;
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;
@@ -98,6 +104,8 @@ import org.apache.beam.sdk.values.TypeDescriptor;
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.cache.Cache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
@@ -128,6 +136,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
private final Coder<DestinationT> destinationCoder;
private final Coder<BigQueryStorageApiInsertError> failedRowsCoder;
private final boolean autoUpdateSchema;
+ private final @Nullable Duration autoUpdateSchemaStrictTimeout;
private final boolean ignoreUnknownValues;
private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation;
private final @Nullable Map<String, String> bigLakeConfiguration;
@@ -140,6 +149,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
private final Coder<TableRow> succussfulRowsCoder;
private final TupleTag<KV<String, Operation>> flushTag = new TupleTag<>("flushTag");
+ private final boolean managedSchemaUpdate;
private static final AppendClientCache<KV<String, ShardedKey<?>>> APPEND_CLIENTS =
new AppendClientCache<>(Duration.standardMinutes(5));
@@ -160,9 +170,11 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
@Nullable TupleTag<TableRow> successfulRowsTag,
Predicate<String> successfulRowsPredicate,
boolean autoUpdateSchema,
+ @Nullable Duration autoUpdateSchemaStrictTimeout,
boolean ignoreUnknownValues,
AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation,
- @Nullable Map<String, String> bigLakeConfiguration) {
+ @Nullable Map<String, String> bigLakeConfiguration,
+ boolean managedSchemaUpdate) {
this.dynamicDestinations = dynamicDestinations;
this.createDisposition = createDisposition;
this.kmsKey = kmsKey;
@@ -174,9 +186,11 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
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
@@ -336,6 +350,33 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
private transient @Nullable WriteStreamService writeStreamServiceInternal = null;
+ // Bounds the total time this DoFn instance spends blocking in onWindowExpiration,
+ // per window, across every key it handles. Window expiration fires once per key and
+ // runs synchronously on the worker thread, so a per-key budget multiplies by the
+ // number of keys and can exceed the runner's bundle execution limit during a drain.
+ //
+ // Deliberately an instance field rather than static: a static cache would share one
+ // budget across every BigQuery sink in the worker JVM.
+ //
+ // Keyed by window so a budget cannot leak from one window into the next. Today the
+ // sink is always rewindowed into the global window by StorageApiLoads, so in practice
+ // there is a single entry and expiration fires once, but nothing here enforces that
+ // and a stale deadline would silently skip the retry and dead-letter the rows.
+ //
+ // transient, and initialized lazily rather than inline: Java deserialization does not
+ // run field initializers, so an inline initializer would leave this null on workers.
+ private transient @Nullable Cache<BoundedWindow, org.joda.time.Instant> drainDeadlines = null;
+
+ private Cache<BoundedWindow, org.joda.time.Instant> getDrainDeadlines() {
+ Cache<BoundedWindow, org.joda.time.Instant> cache = drainDeadlines;
+ if (cache == null) {
+ cache =
+ CacheBuilder.newBuilder().expireAfterAccess(java.time.Duration.ofMinutes(30)).build();
+ drainDeadlines = cache;
+ }
+ return cache;
+ }
+
// Stores the current stream for this key.
@StateId("streamName")
private final StateSpec<ValueState<String>> streamNameSpec = StateSpecs.value();
@@ -345,11 +386,25 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
private final StateSpec<ValueState<Long>> streamOffsetSpec = StateSpecs.value();
@StateId("updatedSchema")
- private final StateSpec<ValueState<TableSchema>> updatedSchema =
+ private final StateSpec<ValueState<TableSchema>> updatedSchemaSpec =
StateSpecs.value(ProtoCoder.of(TableSchema.class));
+ @StateId("mismatchedRows")
+ private final StateSpec<BagState<StoragePayloadWithDeadline>> mismatchedRowsSpec =
+ StateSpecs.bag(StoragePayloadWithDeadline.Coder.of());
+
+ @TimerId("retryMismatchedRowsTimer")
+ private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ @StateId("currentMismatchedRowTimerValue")
+ private final StateSpec<ValueState<Long>> currentMismatchedRowTimerValueSpec =
+ StateSpecs.value();
+
+ @StateId("minPendingTimestamp")
+ private final StateSpec<ValueState<Long>> 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;
@@ -380,7 +435,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
String tableId,
ValueState<String> streamName,
ValueState<Long> streamOffset,
- Timer streamIdleTimer,
+ @Nullable Timer streamIdleTimer,
WriteStreamService writeStreamService,
Callable<Boolean> tryCreateTable) {
try {
@@ -405,7 +460,9 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
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) {
@@ -476,8 +533,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
}
AppendClientInfo get() {
- org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState(
- valid);
+ checkState(valid);
return appendClientInfo;
}
@@ -488,7 +544,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
private CreateRetryManagerResult<DestinationT> createRetryManager(
ShardedKey<DestinationT> key,
- Iterable<SplittingIterable.Value> messages,
+ Iterable<AppendRowsPacket> messages,
Function<AppendRowsContext<DestinationT>, ApiFuture<AppendRowsResponse>> runOperation,
Function<Iterable<AppendRowsContext<DestinationT>>, RetryType> onError,
Consumer<AppendRowsContext<DestinationT>> onSuccess,
@@ -505,10 +561,13 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
List<TimestampedValue<BigQueryStorageApiInsertError>> failedRows = Lists.newArrayList();
int recordsAppended = 0;
List<Integer> histogramUpdates = Lists.newArrayList();
- for (SplittingIterable.Value splitValue : messages) {
- if (splitValue.getSchemaMismatchSeen()) {
+ for (AppendRowsPacket splitValue : messages) {
+ if (!splitValue.getSchemaMismatchedRows().isEmpty()) {
return CreateRetryManagerResult.schemaMismatch();
}
+ if (splitValue.getProtoRows().getSerializedRowsCount() == 0) {
+ continue;
+ }
// Handle the case of a row that is too large.
if (splitValue.getProtoRows().getSerializedSize() >= maxRequestSize) {
if (splitValue.getProtoRows().getSerializedRowsCount() > 1) {
@@ -768,18 +827,16 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
final PipelineOptions pipelineOptions,
@Element KV<ShardedKey<DestinationT>, Iterable<StorageApiWritePayload>> element,
@Timestamp org.joda.time.Instant elementTs,
- final @AlwaysFetched @StateId("streamName") ValueState<String> streamName,
- final @AlwaysFetched @StateId("streamOffset") ValueState<Long> streamOffset,
- final @StateId("updatedSchema") ValueState<TableSchema> updatedSchema,
+ @AlwaysFetched @StateId("streamName") ValueState<String> streamName,
+ @AlwaysFetched @StateId("streamOffset") ValueState<Long> streamOffset,
+ @StateId("updatedSchema") ValueState<TableSchema> updatedSchema,
@TimerId("idleTimer") Timer idleTimer,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
+ @TimerId("retryMismatchedRowsTimer") Timer mismatchedRowsRetryTimer,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long> mismatchedRowsRetryTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
final MultiOutputReceiver o)
throws Exception {
- BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class);
-
- if (autoUpdateSchema) {
- updatedSchema.readLater();
- }
-
dynamicDestinations.setSideInputAccessorFromProcessContext(c);
TableDestination tableDestination =
destinations.computeIfAbsent(
@@ -794,11 +851,108 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
dest);
return tableDestination1;
});
+
+ StorageApiDynamicDestinations.MessageConverter<?> messageConverter =
+ messageConverters.get(
+ element.getKey().getKey(),
+ dynamicDestinations,
+ pipelineOptions,
+ getDatasetService(pipelineOptions),
+ getWriteStreamService(pipelineOptions));
+
+ ThrowingConsumer<Exception, Iterable<StoragePayloadWithDeadline>> 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,
+ elementTs);
+ }
+ };
+
+ org.joda.time.Instant now = org.joda.time.Instant.now();
+ // TODO: Once the processing-time clock is accessible in @OnWindowExpiration, compute
+ // targetDeadline from retryRowsTimers.getCurrentRelativeTime() instead of Instant.now() so
+ // it matches Timer.getCurrentRelativeTime() in TestStream.
+ org.joda.time.Instant targetDeadline =
+ (autoUpdateSchemaStrictTimeout != null)
+ ? now.plus(autoUpdateSchemaStrictTimeout)
+ : BoundedWindow.TIMESTAMP_MAX_VALUE;
+
+ Iterable<StoragePayloadWithDeadline> 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<DestinationT> shardedDestination,
+ TableDestination tableDestination,
+ StorageApiDynamicDestinations.MessageConverter<?> messageConverter,
+ Iterable<StoragePayloadWithDeadline> element,
+ org.joda.time.@Nullable Instant startTimeForLocalRetry,
+ org.joda.time.Instant elementTs,
+ final ValueState<String> streamName,
+ final ValueState<Long> streamOffset,
+ final ValueState<TableSchema> updatedSchema,
+ @Nullable Timer idleTimer,
+ final MultiOutputReceiver o,
+ ThrowingConsumer<Exception, Iterable<StoragePayloadWithDeadline>> 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(
@@ -809,12 +963,11 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
Coder<DestinationT> destinationCoder = dynamicDestinations.getDestinationCoder();
Callable<Boolean> 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,
@@ -828,20 +981,16 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
getOrCreateStream(
tableId, streamName, streamOffset, idleTimer, writeStreamService, tryCreateTable);
- StorageApiDynamicDestinations.MessageConverter<?> messageConverter =
- messageConverters.get(
- element.getKey().getKey(),
- dynamicDestinations,
- pipelineOptions,
- datasetService,
- writeStreamService);
Callable<AppendClientInfo> 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.
@@ -850,22 +999,23 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
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<TableSchema> 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<TableSchema> newSchema =
- TableSchemaUpdateUtils.getUpdatedSchema(tableSchema, streamSchema);
+ schemaChangeDetectorHelper.checkUpdatedSchema(
+ tableSchema, getOrCreateStream.get(), writeStreamService);
if (newSchema.isPresent()) {
tableSchema = newSchema.get();
@@ -895,8 +1045,9 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
// cache could in
// theory evict the object during execution and we want a pin held throughout the execution of
// this function.
+ Iterable<StoragePayloadWithDeadline> 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,53 +1132,110 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
BigQuerySinkMetrics.throttledTimeCounter(
BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM))
.backoff();
+
CreateRetryManagerResult<DestinationT> createRetryManagerResult;
+ final Duration initialMismatchRetryTime =
+ Duration.millis(
+ bigQueryOptions.getStorageApiMismatchLocalRetryTimeMilliSec()); // Retry locally
+ Iterable<StoragePayloadWithDeadline> 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<SplittingIterable.Value> 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<TimestampedValue<BigQueryStorageApiInsertError>, Boolean> failedRowsHandler =
+ (autoUpdateSchemaStrictTimeout == null) && !managedSchemaUpdate
+ ? (TimestampedValue<BigQueryStorageApiInsertError> 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,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
shortTableId)
.inc(1);
- },
- // Get the currently-known TableSchema hash
- () -> appendClientHolder.get().getTableSchemaHash(),
- () ->
- TableRowToStorageApiProto.wrapDescriptorProto(
- messageConverter.getDescriptor(false)),
- autoUpdateSchema,
- elementTs);
+ return true;
+ }
+ : e -> false;
+ // Check payload staleness against the descriptor that will actually be used to
+ // serialize the row, i.e. the one belonging to this shard's append client.
+ //
+ // The messageConverter is shared by every shard writing to this destination, but each
+ // shard owns its own append client and its own "updatedSchema" state. When one shard
+ // detects a schema change it upgrades the shared converter, but the other shards only
+ // rebuild their append clients once their own state says the schema changed. Comparing
+ // against the converter therefore reports "no missing unknown field" for a shard that
+ // is still holding a pre-upgrade append client, so the row is treated as matching and
+ // the new column is silently dropped during serialization.
+ //
+ // Captured per iteration of the do-while loop so that messagesToProcess and
+ // mismatchedRows always evaluate against the exact same AppendClientInfo snapshot even
+ // if retryManager.run(true) later resets appendClientHolder.
+ final Iterable<AppendRowsPacket> messages =
+ new SplittingIterable(
+ payloadsToIterate,
+ splitSize,
+ failedRowsHandler,
+ elementTs,
+ appendClientHolder.get(),
+ schemaChangeDetectorHelper);
+ Iterable<AppendRowsPacket> 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)
+ .isAfter(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)
+ .flatMap(packet -> packet.getSchemaMismatchedRows().stream())
+ .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
@@ -1035,12 +1243,18 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
LOG.info("Schema out of date: refreshing table schema for {}", tableId);
// 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
@@ -1088,20 +1302,19 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
appendSplitDistribution.update(numAppends);
if (autoUpdateSchema) {
- @Nullable StreamAppendClient streamAppendClient =
- appendClientHolder.getStreamAppendClient();
+ 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<TableSchema> 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<TableSchema> 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());
}
}
@@ -1110,8 +1323,121 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
java.time.Duration timeElapsed = java.time.Duration.between(now, Instant.now());
appendLatencyDistribution.update(timeElapsed.toMillis());
}
+ processMismatchedRows.accept(mismatchedRows);
+ }
+ if (idleTimer != null) {
+ idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative();
}
- idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative();
+ }
+
+ @OnTimer("retryMismatchedRowsTimer")
+ public void onMismatchedRowsTimer(
+ OnTimerContext context,
+ PipelineOptions pipelineOptions,
+ @Key ShardedKey<DestinationT> shardedDestination,
+ @Timestamp org.joda.time.Instant elementTs,
+ @AlwaysFetched @StateId("streamName") ValueState<String> streamName,
+ @AlwaysFetched @StateId("streamOffset") ValueState<Long> streamOffset,
+ @StateId("updatedSchema") ValueState<TableSchema> updatedSchema,
+ @TimerId("idleTimer") Timer idleTimer,
+ MultiOutputReceiver o,
+ @TimerId("retryMismatchedRowsTimer") Timer retryRowsTimer,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long> currentTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> 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<StoragePayloadWithDeadline> 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<Exception, Iterable<StoragePayloadWithDeadline>> processMismatchedRows =
+ mismatchedRows -> {
+ // TODO: Try to refactor so that we don't have to materialize this list.
+ List<StoragePayloadWithDeadline> 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,
+ elementTs);
+ }
+ };
+
+ processPayloads(
+ pipelineOptions,
+ shardedDestination,
+ tableDestination,
+ messageConverter,
+ payloads,
+ null,
+ elementTs,
+ streamName,
+ streamOffset,
+ updatedSchema,
+ idleTimer,
+ o,
+ processMismatchedRows);
}
// called by the idleTimer and window-expiration handlers.
@@ -1153,13 +1479,187 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
streamsIdle.inc();
}
+ /**
+ * Builds the converter used to render a drained row as a {@link TableRow} when the row carries
+ * no failsafe representation of its own. Rebuilt after each {@code updateSchemaFromTable} call,
+ * since that may change the descriptor.
+ */
+ private AppendClientInfo drainFailureConverter(
+ StorageApiDynamicDestinations.MessageConverter<?> messageConverter) throws Exception {
+ return AppendClientInfo.of(
+ Preconditions.checkStateNotNull(messageConverter.getTableSchema()),
+ messageConverter.getDescriptor(false),
+ AutoCloseable::close);
+ }
+
+ /**
+ * Routes one row that could not be written before the drain finished to the failed-rows
+ * collection.
+ *
+ * <p>Callers own the exactly-once property: a row passed here must not also be carried into a
+ * later retry round or into the final flush.
+ */
+ private void writeDrainFailure(
+ StoragePayloadWithDeadline row,
+ AppendClientInfo appendClientInfo,
+ TableDestination tableDestination,
+ org.joda.time.Instant elementTs,
+ MultiOutputReceiver o)
+ throws Exception {
+ TableRow failedRow = row.getStoragePayload().getFailsafeTableRow();
+ if (failedRow == null) {
+ failedRow =
+ appendClientInfo.toTableRow(
+ ByteString.copyFrom(row.getStoragePayload().getPayload()), Predicates.alwaysTrue());
+ }
+ BigQueryStorageApiInsertError error =
+ new BigQueryStorageApiInsertError(
+ failedRow,
+ "Timed out waiting for table schema update in OnWindowExpiration",
+ tableDestination.getTableReference());
+ o.get(failedRowsTag)
+ .outputWithTimestamp(
+ error, MoreObjects.firstNonNull(row.getStoragePayload().getTimestamp(), elementTs));
+ rowsSentToFailedRowsCollection.inc();
+ BigQuerySinkMetrics.appendRowsRowStatusCounter(
+ BigQuerySinkMetrics.RowStatus.FAILED,
+ BigQuerySinkMetrics.SCHEMA_MISMATCHED,
+ tableDestination.getShortTableUrn())
+ .inc(1);
+ }
+
@OnWindowExpiration
public void onWindowExpiration(
+ OnWindowExpirationContext context,
+ PipelineOptions pipelineOptions,
@Key ShardedKey<DestinationT> key,
+ @Timestamp org.joda.time.Instant elementTs,
+ @StateId("updatedSchema") ValueState<TableSchema> updatedSchema,
@AlwaysFetched @StateId("streamName") ValueState<String> streamName,
@AlwaysFetched @StateId("streamOffset") ValueState<Long> streamOffset,
+ @AlwaysFetched @StateId("mismatchedRows")
+ BagState<StoragePayloadWithDeadline> mismatchedRowsBag,
MultiOutputReceiver o,
- BoundedWindow window) {
+ BoundedWindow window)
+ throws Exception {
+ dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context);
+
+ 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<Iterable<StoragePayloadWithDeadline>> mismatchedRows =
+ new AtomicReference<>(mismatchedRowsBag.read());
+ if (mismatchedRows.get() != null && !Iterables.isEmpty(mismatchedRows.get())) {
+ // Shared across every key in this window; see drainDeadlines.
+ // TODO: This needs to use the internal clock instead Instant.now(), as the row deadlines
+ // are based on the
+ // internal clock. Add support for accessing the internal clock in OnWindowExpiration.
+ org.joda.time.Instant drainDeadline =
+ getDrainDeadlines()
+ .get(
+ window,
+ () -> org.joda.time.Instant.now().plus(Duration.millis(waitTime.toMillis())));
+ LOG.info(
+ "Draining buffered schema-mismatched rows for table {}, waiting until {} for schema update.",
+ tableDestination.getShortTableUrn(),
+ drainDeadline);
+ com.google.api.client.util.BackOff backoff =
+ new ExponentialBackOff.Builder()
+ .setMaxElapsedTimeMillis((int) waitTime.toMillis())
+ .build();
+ do {
+ messageConverter.updateSchemaFromTable();
+ APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(key));
+ final AppendClientInfo roundConverter = drainFailureConverter(messageConverter);
+ processPayloads(
+ pipelineOptions,
+ key,
+ tableDestination,
+ messageConverter,
+ mismatchedRows.get(),
+ null,
+ elementTs,
+ streamName,
+ streamOffset,
+ updatedSchema,
+ null,
+ o,
+ // Partition in a single streaming pass, inside the consumer rather than after
+ // processPayloads returns. What processPayloads hands back is a lazy view over its
+ // SplittingIterable, which resolves rows through the AppendClientHolder that
+ // processPayloads closes on its way out, so this is the only point at which it can
+ // be read at all.
+ //
+ // Splitting here also avoids snapshotting the whole round into a list purely to walk
+ // it a second time. Rows that can still be retried are the only thing retained; rows
+ // whose deadline has passed go straight to the failed-rows collection instead of
+ // accumulating in a list that lives for the rest of the drain.
+ //
+ // Every row is emitted at most once. processPayloads invokes this consumer exactly
+ // once per call, after its internal retry loop and before closing the holder; this
+ // makes a single pass over "rows" with no isEmpty() pre-check that would re-iterate;
+ // each row takes exactly one branch; and a row handed to writeDrainFailure is dropped
+ // rather than carried, so it can reappear neither in a later round nor in the final
+ // flush below.
+ rows -> {
+ org.joda.time.Instant now = org.joda.time.Instant.now();
+ List<StoragePayloadWithDeadline> stillPending = Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : rows) {
+ // This check is hard to test with DirectRunner. Ideally we should use the
+ // injected clock, but I'm not sure how to access it from OnWindowExpiration
+ // (it's usually accessed via the Timer object).
+ if (row.getDeadline().isAfter(now)) {
+ stillPending.add(row);
+ } else {
+ writeDrainFailure(row, roundConverter, tableDestination, elementTs, o);
+ }
+ }
+ mismatchedRows.set(stillPending);
+ });
+ // Checked before sleeping so we never sleep past the shared deadline.
+ } while (!Iterables.isEmpty(mismatchedRows.get())
+ && org.joda.time.Instant.now().isBefore(drainDeadline)
+ && com.google.api.client.util.BackOffUtils.next(
+ com.google.api.client.util.Sleeper.DEFAULT, backoff));
+
+ // Whatever is still pending once the drain budget is gone. Rows whose deadline lapsed
+ // during a round were already written to the failed-rows collection at that point and were
+ // not carried forward, so this set cannot overlap with them.
+ if (!Iterables.isEmpty(mismatchedRows.get())) {
+ LOG.warn(
+ "Timed out waiting for table schema update during drain for table {}; routing remaining buffered rows to failed-rows collection.",
+ tableDestination.getShortTableUrn());
+ AppendClientInfo finalConverter = drainFailureConverter(messageConverter);
+ for (StoragePayloadWithDeadline row : mismatchedRows.get()) {
+ writeDrainFailure(row, finalConverter, tableDestination, elementTs, o);
+ }
+ }
+ }
+
// 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 00000000000..aefc7fd8690
--- /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.AtomicCoder;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.CoderException;
+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 AtomicCoder<StoragePayloadWithDeadline> {
+ static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of();
+ static final InstantCoder INSTANT_CODER = InstantCoder.of();
+ static final NullableCoder<byte[]> NULLABLE_BYTE_ARRAY_CODER =
+ NullableCoder.of(BYTE_ARRAY_CODER);
+ static final NullableCoder<Instant> 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 ec1daed48b9..687796ed5b4 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
@@ -552,6 +552,9 @@ public class TableRowToStorageApiProto {
hashCodes.add(tableSchemaHash(name, tableFieldSchema.getFieldsList()));
}
}
+ if (hashCodes.isEmpty()) {
+ return SCHEMA_HASH_FUNCTION.hashInt(0);
+ }
return Hashing.combineOrdered(hashCodes);
}
@@ -1092,7 +1095,7 @@ public class TableRowToStorageApiProto {
@Nullable Object fValue = tableRow.get("f");
if (fValue instanceof List) {
if (descriptor == null) {
- // This only happens if we are recursively finding unknwon field names. We don't support
+ // This only happens if we are recursively finding unknown field names. We don't support
// this for list cells
// (all we have is field position in that case) so just bail out.
return null;
@@ -1141,6 +1144,8 @@ public class TableRowToStorageApiProto {
return null;
}
TableRow localUnknownFields = Preconditions.checkStateNotNull(unknownFields);
+ // TODO(reuvenlax): We should probably assume a List instead of a TableRow if the
+ // field is repeated.
@Nullable TableRow nested =
(TableRow) localUnknownFields.getF().get(finalIndex).getV();
if (nested == null) {
@@ -1195,6 +1200,7 @@ public class TableRowToStorageApiProto {
if (!collectedExceptions.isEmpty()) {
return null;
}
+
try {
return builder.build();
} catch (Exception e) {
@@ -1465,6 +1471,31 @@ public class TableRowToStorageApiProto {
// nothing to do here
return tableRowProto;
}
+ try {
+ return mergeNewFields(
+ tableRowProto,
+ wrapDescriptorProto(descriptorProto),
+ TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, false, false),
+ schemaInformation,
+ unknownFields,
+ ignoreUnknownValues);
+ } catch (DescriptorValidationException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ public static ByteString mergeNewFields(
+ ByteString tableRowProto,
+ Descriptor descriptor,
+ Descriptor descriptorIgnoreRequired,
+ SchemaInformation schemaInformation,
+ TableRow unknownFields,
+ boolean ignoreUnknownValues)
+ throws TableRowToStorageApiProto.SchemaConversionException {
+ if (unknownFields == null || unknownFields.isEmpty()) {
+ // nothing to do here
+ return tableRowProto;
+ }
// check if unknownFields contains repeated struct, merge
boolean hasRepeatedStruct =
unknownFields.entrySet().stream()
@@ -1474,13 +1505,6 @@ public class TableRowToStorageApiProto {
&& !((List<?>) entry.getValue()).isEmpty()
&& ((List<?>) entry.getValue()).get(0) instanceof TableRow);
if (!hasRepeatedStruct) {
- Descriptor descriptorIgnoreRequired = null;
- try {
- descriptorIgnoreRequired =
- TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, false, false);
- } catch (DescriptorValidationException e) {
- throw new RuntimeException(e);
- }
ByteString unknownFieldsProto =
Preconditions.checkArgumentNotNull(
messageFromTableRow(
@@ -1497,13 +1521,7 @@ public class TableRowToStorageApiProto {
return tableRowProto.concat(unknownFieldsProto);
}
- DynamicMessage message = null;
- Descriptor descriptor = null;
- try {
- descriptor = wrapDescriptorProto(descriptorProto);
- } catch (DescriptorValidationException e) {
- throw new RuntimeException(e);
- }
+ DynamicMessage message;
try {
message = DynamicMessage.parseFrom(descriptor, tableRowProto);
} catch (InvalidProtocolBufferException 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 899c3169bba..2cb5fb7a6cb 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,12 +17,16 @@
*/
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.ByteString;
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;
@@ -33,7 +37,6 @@ 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;
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.HashCode;
import org.checkerframework.checker.nullness.qual.Nullable;
/** Helper functions for SchemaUpdateOptions. */
@@ -246,26 +249,28 @@ public class UpgradeTableSchema {
}
public static boolean isPayloadSchemaOutOfDate(
- StorageApiWritePayload payload,
+ byte @Nullable [] payloadSchemaHash,
+ ByteString payload,
ThrowingSupplier<byte[]> schemaHash,
ThrowingSupplier<Descriptors.Descriptor> 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);
+ if (!Arrays.equals(payloadSchemaHash, schemaHash.get())) {
+ return isPayloadSchemaOutOfDate(payload, schemaDescriptor);
}
}
return false;
}
+ public static boolean isPayloadSchemaOutOfDate(
+ ByteString payload, ThrowingSupplier<Descriptors.Descriptor> schemaDescriptor)
+ throws Exception {
+ Descriptors.Descriptor descriptor = schemaDescriptor.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;
@@ -304,4 +309,97 @@ public class UpgradeTableSchema {
// If we reach here, neither this message nor its descendants have unknown fields
return false;
}
+
+ public static boolean missingUnknownField(
+ AbstractMap<String, Object> unknownFields,
+ ThrowingSupplier<Descriptors.Descriptor> schemaDescriptor)
+ throws Exception {
+ 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<String, Object> unknownFields, Descriptors.Descriptor schemaDescriptor) {
+ @Nullable Object fValue = unknownFields.get("f");
+ if (fValue instanceof List) {
+ return missingUnknownField((List<?>) fValue, schemaDescriptor);
+ }
+ for (Map.Entry<String, Object> 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 TableCell) {
+ value = ((TableCell) value).getV();
+ } else if (value instanceof Map
+ && ((Map<?, ?>) value).size() == 1
+ && ((Map<?, ?>) value).containsKey("v")
+ && fieldDescriptor.getMessageType().findFieldByName("v") == null) {
+ value = ((Map<?, ?>) value).get("v");
+ }
+ if (value == null) {
+ 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<String, Object>) value, fieldDescriptor.getMessageType());
+ }
+ return false;
+ }
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacketTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacketTest.java
new file mode 100644
index 00000000000..1b1e64dd029
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacketTest.java
@@ -0,0 +1,268 @@
+/*
+ * 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.assertNull;
+import static org.junit.Assert.assertSame;
+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.ProtoRows;
+import com.google.cloud.bigquery.storage.v1.TableSchema;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.DescriptorProtos;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.BitSet;
+import java.util.List;
+import java.util.stream.Collectors;
+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.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link AppendRowsPacket}, focused on partitioning matched vs. mismatched rows. */
+@RunWith(JUnit4.class)
+public class AppendRowsPacketTest {
+
+ private static boolean hasHash(int i) {
+ return i % 2 == 0;
+ }
+
+ private static boolean hasUnknown(int i) {
+ return i % 3 == 0;
+ }
+
+ // Some rows have no failsafe row, matching the @Nullable element type in production.
+ private static boolean hasFailsafe(int i) {
+ return i % 4 != 3;
+ }
+
+ private static ByteString serializedFor(int i) {
+ return ByteString.copyFromUtf8("row" + i);
+ }
+
+ private static byte[] hashFor(int i) {
+ return ("hash" + i).getBytes(StandardCharsets.UTF_8);
+ }
+
+ private static byte[] originalFor(int i) {
+ return ("original" + i).getBytes(StandardCharsets.UTF_8);
+ }
+
+ private static TableRow unknownFor(int i) {
+ return new TableRow().set("unknown", "u" + i);
+ }
+
+ private static @Nullable TableRow failsafeFor(int i) {
+ return hasFailsafe(i) ? new TableRow().set("failsafe", "f" + i) : null;
+ }
+
+ private static Instant timestampFor(int i) {
+ return new Instant(1_000 + i);
+ }
+
+ private static Instant deadlineFor(int i) {
+ return new Instant(2_000 + i);
+ }
+
+ private static StoragePayloadWithDeadline payloadFor(int i) {
+ try {
+ StorageApiWritePayload payload =
+ StorageApiWritePayload.of(
+ hasUnknown(i) ? originalFor(i) : serializedFor(i).toByteArray(),
+ hasUnknown(i) ? unknownFor(i) : null,
+ failsafeFor(i))
+ .withTimestamp(timestampFor(i));
+ if (hasHash(i)) {
+ payload = payload.withSchemaHash(hashFor(i));
+ }
+ return StoragePayloadWithDeadline.of(payload, deadlineFor(i));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ /** Builds a packet of {@code numRows} rows, with the given indices marked schema-mismatched. */
+ private static AppendRowsPacket packetOf(int numRows, int... mismatchedIndices) {
+ BitSet mismatchedSet = new BitSet();
+ for (int i : mismatchedIndices) {
+ mismatchedSet.set(i);
+ }
+
+ ProtoRows.Builder rows = ProtoRows.newBuilder();
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeRows = Lists.newArrayList();
+ List<StoragePayloadWithDeadline> mismatchedRows = Lists.newArrayList();
+
+ for (int i = 0; i < numRows; i++) {
+ if (mismatchedSet.get(i)) {
+ mismatchedRows.add(payloadFor(i));
+ } else {
+ rows.addSerializedRows(serializedFor(i));
+ timestamps.add(timestampFor(i));
+ failsafeRows.add(failsafeFor(i));
+ }
+ }
+
+ return AppendRowsPacket.create(rows.build(), timestamps, failsafeRows, mismatchedRows);
+ }
+
+ private static List<String> serializedRowsOf(AppendRowsPacket packet) {
+ return packet.getProtoRows().getSerializedRowsList().stream()
+ .map(ByteString::toStringUtf8)
+ .collect(Collectors.toList());
+ }
+
+ @Test
+ public void testGetSchemaMatchedRowsOnly_preservesMatchedRowsAndClearsMismatches() {
+ AppendRowsPacket original = packetOf(7, 1, 3, 6);
+ AppendRowsPacket matched = original.getSchemaMatchedRowsOnly();
+
+ assertEquals(4, matched.getProtoRows().getSerializedRowsCount());
+ assertTrue(matched.getSchemaMismatchedRows().isEmpty());
+
+ int[] expectedIndices = {0, 2, 4, 5};
+ for (int newIndex = 0; newIndex < expectedIndices.length; newIndex++) {
+ int originalIndex = expectedIndices[newIndex];
+ assertEquals(
+ serializedFor(originalIndex), matched.getProtoRows().getSerializedRows(newIndex));
+ assertEquals(timestampFor(originalIndex), matched.getTimestamps().get(newIndex));
+ assertEquals(failsafeFor(originalIndex), matched.getFailsafeTableRows().get(newIndex));
+ }
+ }
+
+ @Test
+ public void testGetSchemaMismatchedRows_preservesExactOriginalPayloads() throws Exception {
+ AppendRowsPacket original = packetOf(7, 1, 3, 6);
+ List<StoragePayloadWithDeadline> payloads = original.getSchemaMismatchedRows();
+
+ assertEquals(3, payloads.size());
+ int[] originalIndices = {1, 3, 6};
+ for (int newIndex = 0; newIndex < originalIndices.length; newIndex++) {
+ int originalIndex = originalIndices[newIndex];
+ StoragePayloadWithDeadline withDeadline = payloads.get(newIndex);
+ StorageApiWritePayload payload = withDeadline.getStoragePayload();
+
+ assertEquals(deadlineFor(originalIndex), withDeadline.getDeadline());
+ assertEquals(timestampFor(originalIndex), payload.getTimestamp());
+ assertEquals(failsafeFor(originalIndex), payload.getFailsafeTableRow());
+
+ if (hasHash(originalIndex)) {
+ assertArrayEquals(hashFor(originalIndex), payload.getSchemaHash());
+ } else {
+ assertNull(payload.getSchemaHash());
+ }
+ if (hasUnknown(originalIndex)) {
+ assertEquals(unknownFor(originalIndex), payload.getUnknownFields());
+ assertArrayEquals(originalFor(originalIndex), payload.getPayload());
+ } else {
+ assertNull(payload.getUnknownFields());
+ assertArrayEquals(serializedFor(originalIndex).toByteArray(), payload.getPayload());
+ }
+ }
+ }
+
+ @Test
+ public void testPartitionIsLossless() {
+ int[][] patterns = {
+ {}, {0}, {6}, {1, 3, 6}, {0, 1, 2, 3, 4, 5, 6}, {0, 6}, {2, 3, 4},
+ };
+
+ for (int[] pattern : patterns) {
+ AppendRowsPacket original = packetOf(7, pattern);
+ List<String> matched = serializedRowsOf(original.getSchemaMatchedRowsOnly());
+ List<StoragePayloadWithDeadline> mismatched = original.getSchemaMismatchedRows();
+
+ assertEquals(
+ "row count must be conserved for pattern " + java.util.Arrays.toString(pattern),
+ 7,
+ matched.size() + mismatched.size());
+ }
+ }
+
+ @Test
+ public void testGetSchemaMatchedRowsOnly_noMismatchesReturnsSameInstance() {
+ AppendRowsPacket original = packetOf(4);
+ assertSame(original, original.getSchemaMatchedRowsOnly());
+ }
+
+ @Test
+ public void testPartitionOfEmptyPacket() {
+ AppendRowsPacket empty = packetOf(0);
+
+ assertTrue(empty.getSchemaMismatchedRows().isEmpty());
+ assertEquals(0, empty.getSchemaMatchedRowsOnly().getProtoRows().getSerializedRowsCount());
+ }
+
+ @Test
+ public void testGetSchemaMatchedRowsOnly_allMismatchedYieldsEmptyPacket() {
+ AppendRowsPacket matched = packetOf(3, 0, 1, 2).getSchemaMatchedRowsOnly();
+
+ assertEquals(0, matched.getProtoRows().getSerializedRowsCount());
+ assertTrue(matched.getSchemaMismatchedRows().isEmpty());
+ }
+
+ @Test
+ public void testFromStorageApiWritePayload_partitionsMatchedAndMismatchedRowsDirectly()
+ throws Exception {
+ TableSchema tableSchema = TableSchema.newBuilder().build();
+ DescriptorProtos.DescriptorProto descriptor =
+ DescriptorProtos.DescriptorProto.newBuilder().setName("test").build();
+ AppendClientInfo appendClientInfo = AppendClientInfo.of(tableSchema, descriptor, client -> {});
+ SchemaChangeDetectorHelper helper =
+ new SchemaChangeDetectorHelper(false, true, new TableReference(), true);
+
+ List<StoragePayloadWithDeadline> inputs = Lists.newArrayList();
+ for (int i = 0; i < 4; i++) {
+ StorageApiWritePayload payload =
+ StorageApiWritePayload.of(
+ hasUnknown(i) ? originalFor(i) : serializedFor(i).toByteArray(),
+ hasUnknown(i) ? unknownFor(i) : null,
+ failsafeFor(i))
+ .withTimestamp(timestampFor(i))
+ .withSchemaHash(appendClientInfo.getTableSchemaHash());
+ inputs.add(StoragePayloadWithDeadline.of(payload, deadlineFor(i)));
+ }
+
+ AppendRowsPacket packet =
+ AppendRowsPacket.fromStorageApiWritePayload(
+ Iterators.peekingIterator(inputs.iterator()),
+ Long.MAX_VALUE,
+ helper,
+ new Instant(5_000),
+ appendClientInfo,
+ e -> false);
+
+ // Rows 0 and 3 have unknown fields (hasUnknown(0) and hasUnknown(3)) so they are mismatched;
+ // rows 1 and 2 have no unknown fields and match the empty schema.
+ assertEquals(2, packet.getProtoRows().getSerializedRowsCount());
+ assertEquals(serializedFor(1), packet.getProtoRows().getSerializedRows(0));
+ assertEquals(serializedFor(2), packet.getProtoRows().getSerializedRows(1));
+
+ assertEquals(2, packet.getSchemaMismatchedRows().size());
+ assertEquals(inputs.get(0), packet.getSchemaMismatchedRows().get(0));
+ assertEquals(inputs.get(3), packet.getSchemaMismatchedRows().get(1));
+ }
+}
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 de63120c93c..4391a66e27e 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 7310e2f850e..3fbc4f5aed5 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.io.gcp.testing.FakeJobService;
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;
@@ -2364,12 +2365,27 @@ public class BigQueryIOWriteTest implements Serializable {
@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
+ public void testUpdateTableSchemaConsistentCreateIfNeeded() throws Exception {
+ updateTableSchemaTest(true, true, true);
}
@Test
@@ -2459,13 +2475,23 @@ public class BigQueryIOWriteTest implements Serializable {
}
}
- public void updateTableSchemaTest(boolean useSet) throws Exception {
+ public void updateTableSchemaTest(boolean useSet, boolean withConsistentSchemaUpdate)
+ throws Exception {
+ updateTableSchemaTest(useSet, withConsistentSchemaUpdate, false);
+ }
+
+ public void updateTableSchemaTest(
+ boolean useSet, boolean withConsistentSchemaUpdate, boolean createIfNeeded) 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;
@@ -2480,19 +2506,19 @@ public class BigQueryIOWriteTest implements Serializable {
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")));
- fakeDatasetService.createTable(new Table().setTableReference(tableRef).setSchema(tableSchema));
+ if (!createIfNeeded) {
+ fakeDatasetService.createTable(
+ new Table().setTableReference(tableRef).setSchema(tableSchema));
+ }
LongFunction<TableRow> getRowSet =
(LongFunction<TableRow> & Serializable)
@@ -2539,6 +2565,8 @@ public class BigQueryIOWriteTest implements Serializable {
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<TableRow> tableRows =
p.apply(testStream.advanceWatermarkToInfinity())
@@ -2551,25 +2579,38 @@ public class BigQueryIOWriteTest implements Serializable {
Duration.standardSeconds(5), tableSchemaUpdated, fakeDatasetService)))
.setCoder(TableRowJsonCoder.of());
- tableRows.apply(
+ BigQueryIO.Write<TableRow> write =
BigQueryIO.writeTableRows()
.to(tableRef)
.withMethod(method)
- .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+ .withCreateDisposition(
+ createIfNeeded
+ ? BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED
+ : BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
.ignoreUnknownValues()
- .withAutoSchemaUpdate(true)
.withTestServices(fakeBqServices)
- .withoutValidation());
+ .withoutValidation();
+ if (createIfNeeded) {
+ write = write.withSchema(tableSchema);
+ }
+ if (withConsistentSchemaUpdate) {
+ write = write.withAutoSchemaUpdateConsistent(true, Duration.standardHours(1));
+ } else {
+ write = write.withAutoSchemaUpdate(true);
+ }
+ tableRows.apply(write);
p.run();
Iterable<TableRow> expectedDroppedValues =
- LongStream.range(0, 6)
+ (withConsistentSchemaUpdate ? LongStream.empty() : LongStream.range(0, 6))
.mapToObj(getRowSet)
.map(tr -> filterUnknownValues(tr, tableSchema.getFields()))
.collect(Collectors.toList());
Iterable<TableRow> expectedFullValues =
- LongStream.range(6, 10).mapToObj(getRowSet).collect(Collectors.toList());
+ (withConsistentSchemaUpdate ? LongStream.range(0, 10) : LongStream.range(6, 10))
+ .mapToObj(getRowSet)
+ .collect(Collectors.toList());
assertThat(
fakeDatasetService.getAllRows(
tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()),
@@ -2587,6 +2628,7 @@ public class BigQueryIOWriteTest implements Serializable {
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/SchemaChangeDetectorHelperBufferingTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperBufferingTest.java
new file mode 100644
index 00000000000..7f11f05ac4a
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperBufferingTest.java
@@ -0,0 +1,431 @@
+/*
+ * 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.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.google.api.services.bigquery.model.TableRow;
+import java.util.List;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.ReadableState;
+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.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+
+/**
+ * Tests for {@link SchemaChangeDetectorHelper#bufferMismatchedRows}, the state machine that decides
+ * whether a schema-mismatched row is parked in state for another retry or given up on and routed to
+ * the failed-rows collection.
+ *
+ * <p>This method owns three pieces of persistent state at once -- the buffer itself, the retry
+ * timer, and a watermark hold -- and they have to stay consistent with each other. In particular a
+ * timer or a hold must never be left behind when nothing was actually buffered, since there would
+ * then be no future callback that clears them.
+ */
+@RunWith(JUnit4.class)
+public class SchemaChangeDetectorHelperBufferingTest {
+
+ private static final Instant NOW = Instant.parse("2024-01-01T00:00:00Z");
+ private static final Duration RETRY_PERIOD = Duration.standardMinutes(1);
+ private static final Instant ELEMENT_TIMESTAMP = Instant.parse("2024-01-01T00:00:00Z");
+
+ private FakeBagState<StoragePayloadWithDeadline> bufferedBag;
+ private FakeValueState<Long> currentTimerValue;
+ private FakeValueState<Long> minPendingTimestamp;
+ private FakeTimer retryTimer;
+ private TableDestination tableDestination;
+ private Counter counter;
+ private AppendClientInfo appendClientInfo;
+
+ @SuppressWarnings("unchecked")
+ private DoFn.OutputReceiver<BigQueryStorageApiInsertError> failedRowsReceiver =
+ mock(DoFn.OutputReceiver.class);
+
+ @Before
+ @SuppressWarnings("unchecked")
+ public void setUp() {
+ bufferedBag = new FakeBagState<>();
+ currentTimerValue = new FakeValueState<>();
+ minPendingTimestamp = new FakeValueState<>();
+ retryTimer = new FakeTimer(NOW);
+ tableDestination = new TableDestination("project-id:dataset-id.table", null);
+ counter = Metrics.counter(SchemaChangeDetectorHelperBufferingTest.class, "failedRows");
+ appendClientInfo = mock(AppendClientInfo.class);
+ failedRowsReceiver = mock(DoFn.OutputReceiver.class);
+ }
+
+ /** A row that is still worth retrying, i.e. whose deadline is in the future. */
+ private static StoragePayloadWithDeadline liveRow(String name, Instant deadline)
+ throws Exception {
+ return row(name, deadline, null);
+ }
+
+ private static StoragePayloadWithDeadline row(
+ String name, Instant deadline, @Nullable Instant timestamp) throws Exception {
+ StorageApiWritePayload payload =
+ StorageApiWritePayload.of(
+ name.getBytes(java.nio.charset.StandardCharsets.UTF_8),
+ null,
+ new TableRow().set("name", name));
+ if (timestamp != null) {
+ payload = payload.withTimestamp(timestamp);
+ }
+ return StoragePayloadWithDeadline.of(payload, deadline);
+ }
+
+ private void buffer(
+ Iterable<StoragePayloadWithDeadline> rows, @Nullable AppendClientInfo clientInfo)
+ throws Exception {
+ SchemaChangeDetectorHelper.bufferMismatchedRows(
+ rows,
+ bufferedBag,
+ retryTimer,
+ currentTimerValue,
+ minPendingTimestamp,
+ tableDestination,
+ failedRowsReceiver,
+ clientInfo,
+ counter,
+ RETRY_PERIOD,
+ ELEMENT_TIMESTAMP);
+ }
+
+ private List<BigQueryStorageApiInsertError> capturedFailedRows() {
+ ArgumentCaptor<BigQueryStorageApiInsertError> captor =
+ ArgumentCaptor.forClass(BigQueryStorageApiInsertError.class);
+ verify(failedRowsReceiver, org.mockito.Mockito.atLeast(0))
+ .outputWithTimestamp(captor.capture(), any(Instant.class));
+ return captor.getAllValues();
+ }
+
+ private List<String> bufferedNames() throws Exception {
+ List<String> names = Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : bufferedBag.read()) {
+ names.add((String) row.getStoragePayload().getFailsafeTableRow().get("name"));
+ }
+ return names;
+ }
+
+ /**
+ * The {@code processElement} path passes a null client, meaning "no deadline enforcement here".
+ * Every row must be buffered even if its deadline has already lapsed, because that path has no
+ * way to convert a row to a failed row yet.
+ */
+ @Test
+ public void testNullAppendClientInfoBuffersEveryRowRegardlessOfDeadline() throws Exception {
+ StoragePayloadWithDeadline expired = liveRow("expired", NOW.minus(Duration.standardHours(1)));
+ StoragePayloadWithDeadline live = liveRow("live", NOW.plus(Duration.standardHours(1)));
+
+ buffer(ImmutableList.of(expired, live), null);
+
+ assertEquals(ImmutableList.of("expired", "live"), bufferedNames());
+ verify(failedRowsReceiver, never()).outputWithTimestamp(any(), any());
+ }
+
+ /** With a client available, rows past their deadline are given up on rather than re-buffered. */
+ @Test
+ public void testRowsPastDeadlineAreDeadLettered() throws Exception {
+ StoragePayloadWithDeadline expired = liveRow("expired", NOW.minus(Duration.standardHours(1)));
+ StoragePayloadWithDeadline live = liveRow("live", NOW.plus(Duration.standardHours(1)));
+
+ buffer(ImmutableList.of(expired, live), appendClientInfo);
+
+ assertEquals(ImmutableList.of("live"), bufferedNames());
+
+ List<BigQueryStorageApiInsertError> failed = capturedFailedRows();
+ assertEquals(1, failed.size());
+ assertEquals("expired", failed.get(0).getRow().get("name"));
+ assertEquals(
+ "Timed out waiting for table schema update (Mismatched schema)",
+ failed.get(0).getErrorMessage());
+ assertEquals(tableDestination.getTableReference(), failed.get(0).getTable());
+
+ StoragePayloadWithDeadline expiredWithUnknown =
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(
+ new byte[0], new TableRow().set("extra", "val"), new TableRow().set("name", "u")),
+ NOW.minus(Duration.standardHours(1)));
+ buffer(ImmutableList.of(expiredWithUnknown), appendClientInfo);
+ failed = capturedFailedRows();
+ assertEquals(2, failed.size());
+ assertEquals(
+ "Timed out waiting for table schema update. Unknown fields: [extra]",
+ failed.get(1).getErrorMessage());
+ }
+
+ /**
+ * A row is retried while its deadline is strictly in the future; a deadline exactly equal to the
+ * current time is already spent. This is the boundary that decides whether a row gets one more
+ * attempt or is dropped, so it is pinned explicitly.
+ */
+ @Test
+ public void testDeadlineExactlyAtCurrentTimeIsTreatedAsExpired() throws Exception {
+ buffer(ImmutableList.of(liveRow("boundary", NOW)), appendClientInfo);
+
+ assertTrue(bufferedNames().isEmpty());
+ assertEquals(1, capturedFailedRows().size());
+ }
+
+ /**
+ * Regression guard: when every row is dead-lettered nothing is left in the buffer, so no retry
+ * timer and no watermark hold may be left behind. A timer set here would fire against an empty
+ * bag, and a hold would pin the output watermark with no data behind it.
+ */
+ @Test
+ public void testNoTimerOrHoldWhenNothingIsBuffered() throws Exception {
+ buffer(
+ ImmutableList.of(
+ liveRow("a", NOW.minus(Duration.standardHours(1))),
+ liveRow("b", NOW.minus(Duration.standardHours(2)))),
+ appendClientInfo);
+
+ assertTrue(bufferedNames().isEmpty());
+ assertNull("no retry timer may be scheduled", retryTimer.setTime);
+ assertNull("no watermark hold may be taken", retryTimer.outputTimestamp);
+ assertNull("timer bookkeeping must stay untouched", currentTimerValue.read());
+ assertNull("watermark bookkeeping must stay untouched", minPendingTimestamp.read());
+ }
+
+ /** With nothing previously scheduled, the retry lands one retry period out. */
+ @Test
+ public void testTimerScheduledAtRetryPeriodWhenNonePending() throws Exception {
+ buffer(ImmutableList.of(liveRow("a", NOW.plus(Duration.standardHours(1)))), appendClientInfo);
+
+ assertEquals(NOW.plus(RETRY_PERIOD), retryTimer.setTime);
+ assertEquals(Long.valueOf(NOW.plus(RETRY_PERIOD).getMillis()), currentTimerValue.read());
+ }
+
+ /**
+ * An already-scheduled retry wins. Rescheduling on every arrival would let a steady trickle of
+ * mismatched rows push the retry out indefinitely, so the earlier rows would never be retried.
+ */
+ @Test
+ public void testExistingTimerIsNotPushedOutByLaterRows() throws Exception {
+ Instant alreadyScheduled = NOW.plus(Duration.standardSeconds(5));
+ currentTimerValue.write(alreadyScheduled.getMillis());
+
+ buffer(ImmutableList.of(liveRow("a", NOW.plus(Duration.standardHours(1)))), appendClientInfo);
+
+ assertEquals(alreadyScheduled, retryTimer.setTime);
+ assertEquals(Long.valueOf(alreadyScheduled.getMillis()), currentTimerValue.read());
+ }
+
+ /** The watermark hold has to cover the earliest buffered row, not the most recent one. */
+ @Test
+ public void testHoldTracksEarliestBufferedRow() throws Exception {
+ Instant early = ELEMENT_TIMESTAMP.minus(Duration.standardHours(2));
+ Instant late = ELEMENT_TIMESTAMP.minus(Duration.standardMinutes(1));
+ Instant deadline = NOW.plus(Duration.standardHours(1));
+
+ buffer(
+ ImmutableList.of(row("late", deadline, late), row("early", deadline, early)),
+ appendClientInfo);
+
+ assertEquals(early, retryTimer.outputTimestamp);
+ assertEquals(Long.valueOf(early.getMillis()), minPendingTimestamp.read());
+ }
+
+ /** A hold already taken for an earlier row must not be released by a later one. */
+ @Test
+ public void testHoldIsNotReleasedByLaterRow() throws Exception {
+ Instant alreadyHeld = ELEMENT_TIMESTAMP.minus(Duration.standardHours(5));
+ minPendingTimestamp.write(alreadyHeld.getMillis());
+
+ buffer(
+ ImmutableList.of(
+ row(
+ "later",
+ NOW.plus(Duration.standardHours(1)),
+ ELEMENT_TIMESTAMP.minus(Duration.standardMinutes(1)))),
+ appendClientInfo);
+
+ assertEquals(alreadyHeld, retryTimer.outputTimestamp);
+ assertEquals(Long.valueOf(alreadyHeld.getMillis()), minPendingTimestamp.read());
+ }
+
+ /** Rows carrying no timestamp of their own fall back to the enclosing element's timestamp. */
+ @Test
+ public void testFallsBackToElementTimestampWhenRowHasNone() throws Exception {
+ buffer(ImmutableList.of(liveRow("a", NOW.plus(Duration.standardHours(1)))), appendClientInfo);
+
+ assertEquals(ELEMENT_TIMESTAMP, retryTimer.outputTimestamp);
+ assertEquals(Long.valueOf(ELEMENT_TIMESTAMP.getMillis()), minPendingTimestamp.read());
+ }
+
+ /** Buffering appends to whatever is already parked rather than replacing it. */
+ @Test
+ public void testBufferingAccumulatesAcrossCalls() throws Exception {
+ Instant deadline = NOW.plus(Duration.standardHours(1));
+
+ buffer(ImmutableList.of(liveRow("first", deadline)), appendClientInfo);
+ buffer(ImmutableList.of(liveRow("second", deadline)), appendClientInfo);
+
+ assertEquals(ImmutableList.of("first", "second"), bufferedNames());
+ }
+
+ @Test
+ public void testDeadLetteredRowsAreCountedAndTimestamped() throws Exception {
+ Instant rowTimestamp = ELEMENT_TIMESTAMP.minus(Duration.standardMinutes(7));
+ buffer(
+ ImmutableList.of(row("expired", NOW.minus(Duration.standardHours(1)), rowTimestamp)),
+ appendClientInfo);
+
+ verify(failedRowsReceiver, times(1))
+ .outputWithTimestamp(any(), org.mockito.ArgumentMatchers.eq(rowTimestamp));
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Minimal in-memory state and timer doubles. These deliberately record what was written rather
+ // than emulating a runner, so that assertions can distinguish "never set" from "set to a value".
+ // ---------------------------------------------------------------------------------------------
+
+ private static class FakeBagState<T> implements BagState<T> {
+ private final List<T> contents = Lists.newArrayList();
+
+ @Override
+ public void add(T value) {
+ contents.add(value);
+ }
+
+ @Override
+ public ReadableState<Boolean> isEmpty() {
+ return new ReadableState<Boolean>() {
+ @Override
+ public Boolean read() {
+ return contents.isEmpty();
+ }
+
+ @Override
+ public ReadableState<Boolean> readLater() {
+ return this;
+ }
+ };
+ }
+
+ @Override
+ public Iterable<T> read() {
+ return ImmutableList.copyOf(contents);
+ }
+
+ @Override
+ public BagState<T> readLater() {
+ return this;
+ }
+
+ @Override
+ public void clear() {
+ contents.clear();
+ }
+ }
+
+ private static class FakeValueState<T> implements ValueState<T> {
+ private @Nullable T value = null;
+
+ @Override
+ public void write(T input) {
+ this.value = input;
+ }
+
+ @Override
+ public @Nullable T read() {
+ return value;
+ }
+
+ @Override
+ public ValueState<T> readLater() {
+ return this;
+ }
+
+ @Override
+ public void clear() {
+ value = null;
+ }
+ }
+
+ private static class FakeTimer implements Timer {
+ private final Instant currentRelativeTime;
+
+ private @Nullable Instant setTime = null;
+ private @Nullable Instant outputTimestamp = null;
+
+ FakeTimer(Instant currentRelativeTime) {
+ this.currentRelativeTime = currentRelativeTime;
+ }
+
+ @Override
+ public void set(Instant absoluteTime) {
+ this.setTime = absoluteTime;
+ }
+
+ @Override
+ public void setRelative() {
+ this.setTime = currentRelativeTime;
+ }
+
+ @Override
+ public void clear() {
+ this.setTime = null;
+ }
+
+ @Override
+ public Timer offset(Duration offset) {
+ return this;
+ }
+
+ @Override
+ public Timer align(Duration period) {
+ return this;
+ }
+
+ @Override
+ public Timer withOutputTimestamp(Instant outputTime) {
+ this.outputTimestamp = outputTime;
+ return this;
+ }
+
+ @Override
+ public Timer withNoOutputTimestamp() {
+ this.outputTimestamp = null;
+ return this;
+ }
+
+ @Override
+ public Instant getCurrentRelativeTime() {
+ return currentRelativeTime;
+ }
+ }
+}
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 00000000000..526259e16e6
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java
@@ -0,0 +1,307 @@
+/*
+ * 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<TableSchema> 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<TableSchema> 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<TableSchema> 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<TableSchema> 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<TableSchema> 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<BigQueryStorageApiInsertError> 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, ByteString.EMPTY, () -> 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, ByteString.EMPTY, () -> 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, ByteString.EMPTY, () -> 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, ByteString.EMPTY, () -> 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, ByteString.copyFrom(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 00000000000..bdfd696b431
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java
@@ -0,0 +1,242 @@
+/*
+ * 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 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 org.apache.beam.sdk.values.TimestampedValue;
+import org.joda.time.Duration;
+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 -> {});
+ }
+
+ static final Instant DEADLINE = Instant.now().plus(Duration.standardDays(1));
+
+ @Test
+ public void testBatchingBySplitSize() throws Exception {
+ List<StoragePayloadWithDeadline> 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, DEADLINE));
+ }
+
+ List<TimestampedValue<BigQueryStorageApiInsertError>> 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,
+ Instant.now(),
+ appendClientInfo,
+ schemaChangeDetectorHelper);
+
+ Iterator<AppendRowsPacket> 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<StoragePayloadWithDeadline> payloads = new ArrayList<>();
+ // Payload of 10 bytes, 100 bytes, 10 bytes
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[10], null, null), DEADLINE));
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[100], null, null), DEADLINE));
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[10], null, null), DEADLINE));
+
+ List<TimestampedValue<BigQueryStorageApiInsertError>> 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,
+ Instant.now(),
+ appendClientInfo,
+ schemaChangeDetectorHelper);
+
+ Iterator<AppendRowsPacket> 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<StoragePayloadWithDeadline> payloads = new ArrayList<>();
+ AppendClientInfo appendClientInfo = createAppendClientInfo();
+ 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), DEADLINE));
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null).withSchemaHash(hash2),
+ DEADLINE));
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1), DEADLINE));
+
+ List<TimestampedValue<BigQueryStorageApiInsertError>> failedRows = new ArrayList<>();
+ SchemaChangeDetectorHelper schemaChangeDetectorHelper =
+ new SchemaChangeDetectorHelper(false, true, new TableReference(), true);
+
+ SplittingIterable iterable =
+ new SplittingIterable(
+ payloads,
+ 100,
+ failedRows::add,
+ Instant.now(),
+ appendClientInfo,
+ schemaChangeDetectorHelper);
+
+ Iterator<AppendRowsPacket> it = iterable.iterator();
+ assertTrue(it.hasNext());
+ AppendRowsPacket batch = it.next();
+
+ assertEquals(2, batch.getProtoRows().getSerializedRowsCount());
+ assertEquals(1, batch.getSchemaMismatchedRows().size());
+
+ // Check mismatched rows
+ StoragePayloadWithDeadline mismatched = batch.getSchemaMismatchedRows().get(0);
+ assertArrayEquals(new byte[0], mismatched.getStoragePayload().getPayload());
+ assertEquals(unknownFieldsRow, mismatched.getStoragePayload().getUnknownFields());
+ assertArrayEquals(hash2, mismatched.getStoragePayload().getSchemaHash());
+
+ // 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());
+ assertTrue(matched.getSchemaMismatchedRows().isEmpty());
+ }
+
+ @Test
+ public void testFailedRows() throws Exception {
+ List<StoragePayloadWithDeadline> payloads = new ArrayList<>();
+ payloads.add(
+ StoragePayloadWithDeadline.of(
+ StorageApiWritePayload.of(new byte[0], null, null), DEADLINE));
+ // 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), DEADLINE));
+
+ List<TimestampedValue<BigQueryStorageApiInsertError>> 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,
+ Instant.now(),
+ appendClientInfo,
+ schemaChangeDetectorHelper);
+
+ Iterator<AppendRowsPacket> 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/StorageApiSchemaMismatchDrainTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSchemaMismatchDrainTest.java
new file mode 100644
index 00000000000..8ff950ebfb0
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSchemaMismatchDrainTest.java
@@ -0,0 +1,641 @@
+/*
+ * 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.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.services.bigquery.model.Table;
+import com.google.api.services.bigquery.model.TableFieldSchema;
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import com.google.api.services.bigquery.model.TableSchema;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.Method;
+import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
+import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
+import org.apache.beam.sdk.io.gcp.testing.FakeJobService;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.StreamingOptions;
+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.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestStream;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
+import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter;
+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.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.rules.TestRule;
+import org.junit.rules.Timeout;
+import org.junit.runner.Description;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+import org.junit.runners.model.Statement;
+
+/**
+ * Tests the drain path taken when a consistent schema update never arrives.
+ *
+ * <p>When {@code withAutoSchemaUpdateConsistent} is enabled, rows whose schema does not match the
+ * destination table are buffered in state and retried rather than dropped. If the table schema is
+ * never widened, those buffered rows have to be released when the window expires -- otherwise a
+ * draining pipeline would either hang or silently discard them. That release happens in the
+ * {@code @OnWindowExpiration} handlers of {@link BufferMismatchedRows} (at-least-once) and {@link
+ * StorageApiWritesShardedRecords} (exactly-once), which are the two paths exercised here.
+ *
+ * <p>These tests never update the table schema, so every input row must end up in the failed-rows
+ * collection exactly once.
+ */
+@RunWith(Parameterized.class)
+@SuppressWarnings({"nullness"})
+public class StorageApiSchemaMismatchDrainTest implements Serializable {
+
+ private static final String DRAIN_MESSAGE =
+ "Timed out waiting for table schema update in OnWindowExpiration";
+
+ private static final int NUM_ROWS = 5;
+
+ @Parameters(name = "useAtLeastOnce={0}")
+ public static Iterable<Object[]> data() {
+ return ImmutableList.of(new Object[] {false}, new Object[] {true});
+ }
+
+ @Parameter(0)
+ public boolean useAtLeastOnce;
+
+ private transient PipelineOptions options;
+ private transient TemporaryFolder testFolder = new TemporaryFolder();
+ private transient TestPipeline p;
+
+ @Rule public transient Timeout globalTimeout = Timeout.seconds(300);
+
+ @Rule
+ public final transient TestRule folderThenPipeline =
+ new TestRule() {
+ @Override
+ public Statement apply(final Statement base, final Description description) {
+ Statement withPipeline =
+ new Statement() {
+ @Override
+ public void evaluate() throws Throwable {
+ options = TestPipeline.testingPipelineOptions();
+ BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
+ bqOptions.setProject("project-id");
+ bqOptions.setTempLocation(testFolder.getRoot().getAbsolutePath());
+ bqOptions.setUseStorageWriteApi(true);
+ if (useAtLeastOnce) {
+ bqOptions.setUseStorageWriteApiAtLeastOnce(true);
+ }
+ bqOptions.setStorageWriteApiTriggeringFrequencySec(1);
+ options.as(StreamingOptions.class).setStreaming(true);
+ p = TestPipeline.fromOptions(options);
+ p.apply(base, description).evaluate();
+ }
+ };
+ return testFolder.apply(withPipeline, description);
+ }
+ };
+
+ private FakeDatasetService fakeDatasetService = new FakeDatasetService();
+ private FakeJobService fakeJobService = new FakeJobService();
+ private FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices()
+ .withDatasetService(fakeDatasetService)
+ .withJobService(fakeJobService);
+
+ @Before
+ public void setUp() throws ExecutionException, IOException, InterruptedException {
+ FakeDatasetService.setUp();
+ BigQueryIO.clearStaticCaches();
+ fakeDatasetService.createDataset("project-id", "dataset-id", "", "", null);
+ }
+
+ /** A row whose {@code extra} field is absent from the destination table schema. */
+ private static TableRow rowWithUnknownField(long i) {
+ return new TableRow()
+ .set("number", Long.toString(i))
+ .set("name", "name" + i)
+ .set("extra", "extra" + i);
+ }
+
+ private TableReference createNarrowTable() throws IOException, InterruptedException {
+ TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset-id.table");
+ TableSchema narrowSchema =
+ new TableSchema()
+ .setFields(
+ ImmutableList.of(
+ new TableFieldSchema().setName("number").setType("INTEGER"),
+ new TableFieldSchema().setName("name").setType("STRING")));
+ fakeDatasetService.createTable(new Table().setTableReference(tableRef).setSchema(narrowSchema));
+ return tableRef;
+ }
+
+ /**
+ * Builds a pipeline whose destination table schema is never widened, so the rows stay mismatched
+ * until the window expires.
+ *
+ * @param strictTimeout the per-row retry deadline given to {@code withAutoSchemaUpdateConsistent}
+ */
+ private PCollection<BigQueryStorageApiInsertError> buildPipeline(
+ TableReference tableRef, Duration strictTimeout) {
+ p.enableAbandonedNodeEnforcement(false);
+
+ BigQueryOptions bqOptions = p.getOptions().as(BigQueryOptions.class);
+ // Keep GroupIntoBatches from buffering rows away from the schema-mismatch machinery.
+ bqOptions.setStorageApiAppendThresholdBytes(1);
+ bqOptions.setNumStorageWriteApiStreams(1);
+ bqOptions.setStorageApiMismatchLocalRetryTimeMilliSec(0);
+ // Retry quickly, and bound the drain so the test does not sit in the backoff loop. The drain
+ // budget is shared across keys within a window, so this bounds the whole drain, not each key.
+ bqOptions.setStorageApiMismatchRetryTimeMilliSec(500);
+ bqOptions.setStorageApiMismatchDrainRetryTimeMilliSec(1000);
+
+ TestStream.Builder<Long> testStream =
+ TestStream.create(VarLongCoder.of()).advanceWatermarkTo(new Instant(0));
+ for (long i = 0; i < NUM_ROWS; i++) {
+ testStream = testStream.addElements(i);
+ }
+ // Let the retry timer fire at least once before the window expires.
+ testStream = testStream.advanceProcessingTime(Duration.standardSeconds(5));
+
+ BigQueryIO.Write<TableRow> write =
+ BigQueryIO.writeTableRows()
+ .to(tableRef)
+ .withMethod(
+ useAtLeastOnce ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API)
+ .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+ .ignoreUnknownValues()
+ .withAutoSchemaUpdateConsistent(true, strictTimeout)
+ .withTestServices(fakeBqServices)
+ .withoutValidation();
+
+ return p.apply(testStream.advanceWatermarkToInfinity())
+ .apply(
+ "getRow",
+ MapElements.into(TypeDescriptor.of(TableRow.class))
+ .via(
+ (SerializableFunction<Long, TableRow>)
+ StorageApiSchemaMismatchDrainTest::rowWithUnknownField))
+ .setCoder(TableRowJsonCoder.of())
+ .apply(write)
+ .getFailedStorageApiInserts();
+ }
+
+ private static PCollection<String> failedRowNames(
+ PCollection<BigQueryStorageApiInsertError> failedRows) {
+ return failedRows.apply(
+ "names",
+ MapElements.into(TypeDescriptors.strings())
+ .via(
+ (SerializableFunction<BigQueryStorageApiInsertError, String>)
+ e -> (String) e.getRow().get("name")));
+ }
+
+ private static List<String> expectedNames() {
+ return LongStream.range(0, NUM_ROWS).mapToObj(i -> "name" + i).collect(Collectors.toList());
+ }
+
+ /**
+ * Asserts that every expected row name reaches the dead-letter output exactly once.
+ *
+ * <p>{@code containsInAnyOrder} is an exact multiset match and so would already fail on a
+ * duplicate, but only with a whole-collection mismatch dump that does not say which row was
+ * doubled. The drain loop re-reads the same buffered rows on every round, and rows that expire
+ * mid-round are emitted from a different place than rows released at the drain deadline, so a
+ * double emission is the most likely way for it to regress. This asserts it directly and names
+ * the offending rows.
+ */
+ private static SerializableFunction<Iterable<String>, Void> assertEmittedExactlyOnce() {
+ return observed -> {
+ Set<String> seen = Sets.newHashSet();
+ List<String> duplicated = Lists.newArrayList();
+ for (String name : observed) {
+ if (!seen.add(name)) {
+ duplicated.add(name);
+ }
+ }
+ assertEquals(
+ "rows emitted more than once by the drain", ImmutableList.<String>of(), duplicated);
+ assertEquals("rows missing from the drain output", Sets.newHashSet(expectedNames()), seen);
+ return null;
+ };
+ }
+
+ /**
+ * With a retry deadline far in the future no row expires on its own, so the drain deadline is
+ * what releases them and every row must be dead-lettered with the drain message.
+ */
+ @Test
+ public void testWindowExpirationDrainsBufferedRowsWhenSchemaNeverUpdates() throws Exception {
+ TableReference tableRef = createNarrowTable();
+
+ PCollection<BigQueryStorageApiInsertError> failedRows =
+ buildPipeline(tableRef, Duration.standardHours(1));
+
+ PCollection<String> names = failedRowNames(failedRows);
+ PAssert.that("every buffered row is dead-lettered exactly once", names)
+ .containsInAnyOrder(expectedNames());
+ PAssert.that("the drain must not emit a row twice", names)
+ .satisfies(assertEmittedExactlyOnce());
+
+ PCollection<String> messages =
+ failedRows.apply(
+ "messages",
+ MapElements.into(TypeDescriptors.strings())
+ .via(
+ (SerializableFunction<BigQueryStorageApiInsertError, String>)
+ BigQueryStorageApiInsertError::getErrorMessage));
+ PAssert.that("all failures come from the window-expiration drain", messages)
+ .satisfies(
+ observed -> {
+ for (String message : observed) {
+ assertEquals(DRAIN_MESSAGE, message);
+ }
+ return null;
+ });
+
+ p.run();
+
+ // Nothing should have reached the table, since its schema never accepted these rows.
+ assertTrue(
+ fakeDatasetService
+ .getAllRows(tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId())
+ .isEmpty());
+ }
+
+ /**
+ * With a retry deadline that lapses while rows are still buffered, rows expire inside the drain
+ * loop rather than at its deadline. Whichever branch catches them, no row may be lost: an earlier
+ * version of the drain loop derived its pending set and its expired set from two different
+ * collections, so a row that expired inside {@code processElement} was filtered out of the
+ * pending set without ever being recorded as expired.
+ */
+ @Test
+ public void testExpiredRowsAreDeadLetteredAndNotDropped() throws Exception {
+ TableReference tableRef = createNarrowTable();
+
+ PCollection<BigQueryStorageApiInsertError> failedRows =
+ buildPipeline(tableRef, Duration.millis(1));
+
+ PCollection<String> names = failedRowNames(failedRows);
+ PAssert.that("no row may be lost when its retry deadline lapses", names)
+ .containsInAnyOrder(expectedNames());
+ PAssert.that("no row may be dead-lettered twice when its retry deadline lapses", names)
+ .satisfies(assertEmittedExactlyOnce());
+
+ p.run();
+ }
+
+ private static TableSchema widenedSchema() {
+ return new TableSchema()
+ .setFields(
+ ImmutableList.of(
+ new TableFieldSchema().setName("number").setType("INTEGER"),
+ new TableFieldSchema().setName("name").setType("STRING"),
+ new TableFieldSchema().setName("extra").setType("STRING")));
+ }
+
+ private static List<TableRow> expectedRowsWithExtra() {
+ return LongStream.range(0, NUM_ROWS)
+ .mapToObj(StorageApiSchemaMismatchDrainTest::rowWithUnknownField)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Verifies that rows buffered due to a schema mismatch are flushed with all unknown fields intact
+ * by the {@code @OnTimer} retry handler after the destination table schema is widened externally,
+ * even when no subsequent upstream elements arrive.
+ */
+ @Test
+ public void testBufferedRowsRecoverOnTimerWhenSchemaUpdatesWithNoSubsequentElements()
+ throws Exception {
+ TableReference tableRef = createNarrowTable();
+ p.enableAbandonedNodeEnforcement(false);
+
+ BigQueryOptions bqOptions = p.getOptions().as(BigQueryOptions.class);
+ bqOptions.setStorageApiAppendThresholdBytes(1);
+ bqOptions.setNumStorageWriteApiStreams(1);
+ bqOptions.setStorageApiMismatchLocalRetryTimeMilliSec(0);
+ bqOptions.setStorageApiMismatchRetryTimeMilliSec(500);
+
+ TestStream.Builder<Long> testStream =
+ TestStream.create(VarLongCoder.of()).advanceWatermarkTo(new Instant(0));
+ for (long i = 0; i < NUM_ROWS; i++) {
+ testStream = testStream.addElements(i);
+ }
+ // Advance past UpdateTableSchemaDoFn's 2s timer so the schema is widened in FakeDatasetService
+ // without emitting any new elements.
+ testStream = testStream.advanceProcessingTime(Duration.standardSeconds(5));
+ // Advance processing time again so the @OnTimer retry handler observes the widened schema and
+ // flushes the buffered rows before window expiration.
+ testStream = testStream.advanceProcessingTime(Duration.standardSeconds(10));
+
+ BigQueryIO.Write<TableRow> write =
+ BigQueryIO.writeTableRows()
+ .to(tableRef)
+ .withMethod(
+ useAtLeastOnce ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API)
+ .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+ .ignoreUnknownValues()
+ .withAutoSchemaUpdateConsistent(true, Duration.standardHours(1))
+ .withTestServices(fakeBqServices)
+ .withoutValidation();
+
+ PCollection<BigQueryStorageApiInsertError> failedRows =
+ p.apply(testStream.advanceWatermarkToInfinity())
+ .apply(
+ "getRow",
+ MapElements.into(TypeDescriptor.of(TableRow.class))
+ .via(
+ (SerializableFunction<Long, TableRow>)
+ StorageApiSchemaMismatchDrainTest::rowWithUnknownField))
+ .apply("addKey", WithKeys.of("project-id:dataset-id.table"))
+ .apply(
+ "updateSchemaOnTimer",
+ ParDo.of(
+ new UpdateTableSchemaDoFn(
+ Duration.standardSeconds(2), widenedSchema(), fakeDatasetService)))
+ .setCoder(TableRowJsonCoder.of())
+ .apply(write)
+ .getFailedStorageApiInserts();
+
+ PAssert.that("no rows should fail when schema updates before timeout", failedRows).empty();
+
+ p.run();
+
+ assertThat(
+ fakeDatasetService.getAllRows(
+ tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()),
+ containsInAnyOrder(expectedRowsWithExtra().toArray(new TableRow[0])));
+ }
+
+ /**
+ * Verifies that when {@code @OnTimer} does not fire before the window expires, the
+ * {@code @OnWindowExpiration} drain path refreshes the destination table schema and writes all
+ * buffered rows (including their newly added columns) to BigQuery rather than dropping or
+ * dead-lettering them.
+ */
+ @Test
+ public void testWindowExpirationDrainsAndWritesBufferedRowsWhenSchemaUpdatesBeforeDrain()
+ throws Exception {
+ TableReference tableRef = createNarrowTable();
+ p.enableAbandonedNodeEnforcement(false);
+
+ BigQueryOptions bqOptions = p.getOptions().as(BigQueryOptions.class);
+ bqOptions.setStorageApiAppendThresholdBytes(1);
+ bqOptions.setNumStorageWriteApiStreams(1);
+ bqOptions.setStorageApiMismatchLocalRetryTimeMilliSec(0);
+ // Schedule the retry timer far in the future (1 hour) so @OnTimer never fires during the test,
+ // forcing recovery to happen exclusively inside @OnWindowExpiration.
+ bqOptions.setStorageApiMismatchRetryTimeMilliSec(3_600_000);
+ bqOptions.setStorageApiMismatchDrainRetryTimeMilliSec(5_000);
+
+ TestStream.Builder<Long> testStream =
+ TestStream.create(VarLongCoder.of()).advanceWatermarkTo(new Instant(0));
+ for (long i = 0; i < NUM_ROWS; i++) {
+ testStream = testStream.addElements(i);
+ }
+ // Fire UpdateTableSchemaDoFn at t=2s (while retryTimer is still set for t=3600s), then
+ // immediately expire the window so @OnWindowExpiration must drain the buffered rows.
+ testStream = testStream.advanceProcessingTime(Duration.standardSeconds(3));
+
+ BigQueryIO.Write<TableRow> write =
+ BigQueryIO.writeTableRows()
+ .to(tableRef)
+ .withMethod(
+ useAtLeastOnce ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API)
+ .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+ .ignoreUnknownValues()
+ .withAutoSchemaUpdateConsistent(true, Duration.standardHours(1))
+ .withTestServices(fakeBqServices)
+ .withoutValidation();
+
+ PCollection<BigQueryStorageApiInsertError> failedRows =
+ p.apply(testStream.advanceWatermarkToInfinity())
+ .apply(
+ "getRow",
+ MapElements.into(TypeDescriptor.of(TableRow.class))
+ .via(
+ (SerializableFunction<Long, TableRow>)
+ StorageApiSchemaMismatchDrainTest::rowWithUnknownField))
+ .apply("addKey", WithKeys.of("project-id:dataset-id.table"))
+ .apply(
+ "updateSchemaOnTimer",
+ ParDo.of(
+ new UpdateTableSchemaDoFn(
+ Duration.standardSeconds(2), widenedSchema(), fakeDatasetService)))
+ .setCoder(TableRowJsonCoder.of())
+ .apply(write)
+ .getFailedStorageApiInserts();
+
+ PAssert.that("no rows should be dead-lettered when drain succeeds", failedRows).empty();
+
+ p.run();
+
+ assertThat(
+ fakeDatasetService.getAllRows(
+ tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()),
+ containsInAnyOrder(expectedRowsWithExtra().toArray(new TableRow[0])));
+ }
+
+ private static class UpdateTableSchemaDoFn extends DoFn<KV<String, TableRow>, TableRow> {
+ @TimerId("updateTimer")
+ private final TimerSpec updateTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ private final Duration timerOffset;
+ private final String updatedSchema;
+ private final FakeDatasetService fakeDatasetService;
+
+ UpdateTableSchemaDoFn(
+ Duration timerOffset, TableSchema updatedSchema, FakeDatasetService fakeDatasetService) {
+ this.timerOffset = timerOffset;
+ this.updatedSchema = BigQueryHelpers.toJsonString(updatedSchema);
+ this.fakeDatasetService = fakeDatasetService;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRow> element,
+ @TimerId("updateTimer") Timer updateTimer,
+ OutputReceiver<TableRow> o) {
+ updateTimer.offset(timerOffset).setRelative();
+ o.output(element.getValue());
+ }
+
+ @OnTimer("updateTimer")
+ public void onTimer(@Key String tableSpec) throws IOException {
+ fakeDatasetService.updateTableSchema(
+ BigQueryHelpers.parseTableSpec(tableSpec),
+ BigQueryHelpers.fromJsonString(updatedSchema, TableSchema.class));
+ }
+ }
+
+ private static class ConstantTableDynamicDestinations
+ extends DynamicDestinations<TableRow, String> {
+ private final String tableSpec;
+
+ ConstantTableDynamicDestinations(String tableSpec) {
+ this.tableSpec = tableSpec;
+ }
+
+ @Override
+ public String getDestination(org.apache.beam.sdk.values.ValueInSingleWindow<TableRow> element) {
+ return tableSpec;
+ }
+
+ @Override
+ public TableDestination getTable(String destination) {
+ return new TableDestination(BigQueryHelpers.parseTableSpec(destination), null);
+ }
+
+ @Override
+ public TableSchema getSchema(String destination) {
+ return null;
+ }
+ }
+
+ /**
+ * Verifies that {@link SchemaUpdateHoldingFn#onWindowExpiration} bounds its retry loop by {@link
+ * BigQueryOptions#getStorageApiMismatchDrainRetryTimeMilliSec()} and routes remaining
+ * unconvertible elements to the dead-letter output when the table schema is never widened.
+ */
+ @Test
+ public void testSchemaUpdateHoldingFnWindowExpirationDrainsToDeadLetterOnTimeout()
+ throws Exception {
+ createNarrowTable();
+ p.enableAbandonedNodeEnforcement(false);
+
+ BigQueryOptions bqOptions = p.getOptions().as(BigQueryOptions.class);
+ bqOptions.setStorageApiMismatchDrainRetryTimeMilliSec(500);
+
+ StorageApiDynamicDestinations<TableRow, String> dynamicDestinations =
+ new StorageApiDynamicDestinationsTableRow<>(
+ new ConstantTableDynamicDestinations("project-id:dataset-id.table"),
+ BigQueryIO.TableRowFormatFunction.fromSerializableFunction(tr -> tr),
+ null,
+ false,
+ BigQueryIO.Write.CreateDisposition.CREATE_NEVER,
+ false,
+ false,
+ true);
+
+ TupleTag<KV<String, StorageApiWritePayload>> successfulWritesTag =
+ new TupleTag<>("successfulWrites");
+ TupleTag<BigQueryStorageApiInsertError> failedWritesTag = new TupleTag<>("failedWrites");
+ TupleTag<KV<String, TableRow>> waitingTag = new TupleTag<>("waiting");
+ TupleTag<KV<String, com.google.cloud.bigquery.storage.v1.TableSchema>> newSchemasTag =
+ new TupleTag<>("newSchemas");
+ TupleTag<BadRecord> badRecordTag = new TupleTag<>("badRecords");
+
+ ConvertMessagesDoFn<String, TableRow> convertDoFn =
+ new ConvertMessagesDoFn<>(
+ dynamicDestinations,
+ fakeBqServices,
+ "StorageApiWrite",
+ failedWritesTag,
+ successfulWritesTag,
+ newSchemasTag,
+ waitingTag,
+ null,
+ BadRecordRouter.THROWING_ROUTER,
+ KvCoder.of(StringUtf8Coder.of(), TableRowJsonCoder.of()),
+ true);
+
+ TestStream.Builder<Long> testStream =
+ TestStream.create(VarLongCoder.of()).advanceWatermarkTo(new Instant(0));
+ for (long i = 0; i < NUM_ROWS; i++) {
+ testStream = testStream.addElements(i);
+ }
+
+ PCollectionTuple result =
+ p.apply(testStream.advanceWatermarkToInfinity())
+ .apply(
+ "toShardedRow",
+ MapElements.into(new TypeDescriptor<KV<ShardedKey<String>, TableRow>>() {})
+ .via(
+ (SerializableFunction<Long, KV<ShardedKey<String>, TableRow>>)
+ i ->
+ KV.of(
+ StorageApiConvertMessages.AssignShardFn.getShardedKey(
+ "project-id:dataset-id.table", 0, 1),
+ rowWithUnknownField(i))))
+ .setCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()),
+ NullableCoder.of(TableRowJsonCoder.of())))
+ .apply(
+ "holdAndDrain",
+ ParDo.of(new SchemaUpdateHoldingFn<>(TableRowJsonCoder.of(), convertDoFn))
+ .withOutputTags(
+ successfulWritesTag,
+ TupleTagList.of(ImmutableList.of(failedWritesTag, badRecordTag))));
+
+ result
+ .get(successfulWritesTag)
+ .setCoder(
+ KvCoder.of(
+ StringUtf8Coder.of(),
+ p.getSchemaRegistry().getSchemaCoder(StorageApiWritePayload.class)));
+ result.get(failedWritesTag).setCoder(BigQueryStorageApiInsertErrorCoder.of());
+ result.get(badRecordTag).setCoder(BadRecord.getCoder(p));
+
+ PCollection<String> failedNames = failedRowNames(result.get(failedWritesTag));
+ PAssert.that(
+ "SchemaUpdateHoldingFn must dead-letter all buffered rows on drain timeout",
+ failedNames)
+ .containsInAnyOrder(expectedNames());
+
+ p.run();
+ }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java
index c4aeffd25aa..18a94cf85b9 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateITBase.java
@@ -37,10 +37,14 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
@@ -49,6 +53,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition;
import org.apache.beam.sdk.io.gcp.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;
@@ -68,11 +73,16 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Immuta
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
+import org.junit.runner.Runner;
+import org.junit.runners.Parameterized;
+import org.junit.runners.ParentRunner;
+import org.junit.runners.model.RunnerScheduler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -88,6 +98,66 @@ abstract class StorageApiSinkSchemaUpdateITBase {
this.bigQueryDatasetId = bigQueryDatasetId;
}
+ /**
+ * Runs both the parameter sets and the {@code @Test} methods within each parameter set in
+ * parallel across a shared thread pool.
+ *
+ * <p>Every test writes to its own table (the table name is derived from the test method name and
+ * the parameter values) and the class holds no mutable static state, so all {@code @Test} methods
+ * and parameter sets are independent. Because {@link Parameterized}'s own scheduler only
+ * schedules the outer per-parameter {@link Runner}s (and each child {@link ParentRunner} defaults
+ * to a sequential scheduler for its {@code @Test} methods), we install a shared {@link
+ * RunnerScheduler} on each child {@link ParentRunner} as well.
+ */
+ public static class ParallelParameterized extends Parameterized {
+ private static final int MAX_PARALLEL_TESTS = 6;
+
+ public ParallelParameterized(Class<?> klass) throws Throwable {
+ super(klass);
+ final ExecutorService executor = Executors.newFixedThreadPool(MAX_PARALLEL_TESTS);
+ final RunnerScheduler methodScheduler =
+ new RunnerScheduler() {
+ @Override
+ public void schedule(Runnable childStatement) {
+ // execute() rather than submit(): the child statement already records its own
+ // failures with JUnit, and submit() would swallow them into an unchecked Future.
+ executor.execute(childStatement);
+ }
+
+ @Override
+ public void finished() {
+ // Shutdown and awaitTermination are handled by the outer suite scheduler once all
+ // parameter runners have scheduled their test methods.
+ }
+ };
+ for (Runner child : getChildren()) {
+ if (child instanceof ParentRunner) {
+ ((ParentRunner<?>) child).setScheduler(methodScheduler);
+ }
+ }
+ setScheduler(
+ new RunnerScheduler() {
+ @Override
+ public void schedule(Runnable childStatement) {
+ // Each parameter runner simply enqueues its @Test methods onto `executor` via
+ // `methodScheduler`, so run it inline without consuming a worker thread.
+ childStatement.run();
+ }
+
+ @Override
+ public void finished() {
+ executor.shutdown();
+ try {
+ executor.awaitTermination(2, TimeUnit.HOURS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ }
+ });
+ }
+ }
+
@Rule public TestName testName = new TestName();
private static final Logger LOG = LoggerFactory.getLogger(StorageApiSinkSchemaUpdateITBase.class);
@@ -125,8 +195,16 @@ abstract class StorageApiSinkSchemaUpdateITBase {
private static final int SCHEMA_PROPAGATION_TIMEOUT_MS = 60000;
// interval between checks
private static final int SCHEMA_PROPAGATION_CHECK_INTERVAL_MS = 5000;
- // wait for streams to recognize schema
- private static final int STREAM_RECOGNITION_DELAY_MS = 15000;
+ // Rather than pacing every element (which dominated the runtime of these tests), rows are
+ // emitted as fast as possible and wall-clock delays are inserted only where they are actually
+ // required. (a) Before the table schema is changed, so that the sink has already built its
+ // message converter / append client from the original schema.
+ private static final int PRE_SCHEMA_UPDATE_SETTLE_MS = 10000;
+ // (b) BigQuery only reports an updated table schema on an append *response*, so the sink must
+ // keep appending after the schema change for the new schema to ever reach it. The rows between
+ // the schema update and the first new-column row are spread across this window to provide that
+ // traffic. This is a wall-clock budget, independent of TOTAL_N.
+ private static final int SCHEMA_RECOGNITION_WINDOW_MS = 60000;
// trigger for updating the schema when the row counter reaches this value
private static final int SCHEMA_UPDATE_TRIGGER = 2;
// Long wait (in seconds) for Storage API streams to recognize the new schema.
@@ -194,8 +272,22 @@ abstract class StorageApiSinkSchemaUpdateITBase {
@SuppressWarnings("unused")
private final StateSpec<ValueState<Integer>> counter;
+ private final boolean waitForPropagation;
+ // Whether to sleep before updating the table schema so that the sink's message converter and
+ // append client are built from the original table schema before the schema change happens.
+ private final boolean settleBeforeSchemaUpdate;
+ // Whether to spread rows between the schema update and the first new-column row across
+ // SCHEMA_RECOGNITION_WINDOW_MS. Only needed by non-consistent autoSchemaUpdate tests where
+ // BigQuery only reports an updated table schema passively on append responses.
+ private final boolean spreadRowsOverRecognitionWindow;
+
public UpdateSchemaDoFn(
- String projectId, String datasetId, Map<String, TableSchema> newSchemas) {
+ String projectId,
+ String datasetId,
+ Map<String, TableSchema> newSchemas,
+ boolean waitForPropagation,
+ boolean settleBeforeSchemaUpdate,
+ boolean spreadRowsOverRecognitionWindow) {
this.projectId = projectId;
this.datasetId = datasetId;
Map<String, String> serializableSchemas = new HashMap<>();
@@ -205,6 +297,9 @@ abstract class StorageApiSinkSchemaUpdateITBase {
this.newSchemas = serializableSchemas;
this.bqClient = null;
this.counter = StateSpecs.value();
+ this.waitForPropagation = waitForPropagation;
+ this.settleBeforeSchemaUpdate = settleBeforeSchemaUpdate;
+ this.spreadRowsOverRecognitionWindow = spreadRowsOverRecognitionWindow;
}
@Setup
@@ -220,6 +315,11 @@ abstract class StorageApiSinkSchemaUpdateITBase {
// recognize it,
// ensuring that subsequent writers are created with the updated schema.
if (current == SCHEMA_UPDATE_TRIGGER) {
+ if (settleBeforeSchemaUpdate) {
+ // Give the rows emitted so far time to reach BigQuery, so that the sink's message
+ // converter and append client are built from the *original* table schema.
+ Thread.sleep(PRE_SCHEMA_UPDATE_SETTLE_MS);
+ }
for (Map.Entry<String, String> entry : newSchemas.entrySet()) {
bqClient.updateTableSchema(
projectId,
@@ -228,35 +328,48 @@ abstract class StorageApiSinkSchemaUpdateITBase {
BigQueryHelpers.fromJsonString(entry.getValue(), TableSchema.class));
}
- // check that schema update propagated fully
- long startTime = System.currentTimeMillis();
- long timeoutMillis = SCHEMA_PROPAGATION_TIMEOUT_MS;
- boolean schemaPropagated = false;
- while (System.currentTimeMillis() - startTime < timeoutMillis) {
- schemaPropagated = true;
- for (Map.Entry<String, String> entry : newSchemas.entrySet()) {
- TableSchema currentSchema =
- bqClient.getTableResource(projectId, datasetId, entry.getKey()).getSchema();
- TableSchema expectedSchema =
- BigQueryHelpers.fromJsonString(entry.getValue(), TableSchema.class);
- if (currentSchema.getFields().size() != expectedSchema.getFields().size()) {
- schemaPropagated = false;
+ if (waitForPropagation) {
+ // check that schema update propagated fully
+ long startTime = System.currentTimeMillis();
+ long timeoutMillis = SCHEMA_PROPAGATION_TIMEOUT_MS;
+ boolean schemaPropagated = false;
+ while (System.currentTimeMillis() - startTime < timeoutMillis) {
+ schemaPropagated = true;
+ for (Map.Entry<String, String> entry : newSchemas.entrySet()) {
+ TableSchema currentSchema =
+ bqClient.getTableResource(projectId, datasetId, entry.getKey()).getSchema();
+ TableSchema expectedSchema =
+ BigQueryHelpers.fromJsonString(entry.getValue(), TableSchema.class);
+ if (currentSchema.getFields().size() != expectedSchema.getFields().size()) {
+ schemaPropagated = false;
+ break;
+ }
+ }
+ if (schemaPropagated) {
break;
}
+ Thread.sleep(SCHEMA_PROPAGATION_CHECK_INTERVAL_MS);
}
- if (schemaPropagated) {
- break;
+ if (!schemaPropagated) {
+ LOG.warn("Schema update did not propagate fully within the timeout.");
+ } else {
+ LOG.info(
+ "Schema update propagated fully within the timeout - {}.",
+ System.currentTimeMillis() - startTime);
}
- Thread.sleep(SCHEMA_PROPAGATION_CHECK_INTERVAL_MS);
}
- if (!schemaPropagated) {
- LOG.warn("Schema update did not propagate fully within the timeout.");
- } else {
- LOG.info(
- "Schema update propagated fully within the timeout - {}.",
- System.currentTimeMillis() - startTime);
- // wait for streams to recognize the new schema
- Thread.sleep(STREAM_RECOGNITION_DELAY_MS);
+ }
+
+ // Spread the rows between the schema update and the first new-column row across the
+ // recognition window. These appends are what cause BigQuery to eventually return the
+ // updated schema on an append response; rows outside the window are emitted immediately.
+ if (spreadRowsOverRecognitionWindow) {
+ Object rowId = c.element().getValue().get("id");
+ if (rowId instanceof Number) {
+ long id = ((Number) rowId).longValue();
+ if (id > SCHEMA_UPDATE_TRIGGER && id < ORIGINAL_N) {
+ Thread.sleep(SCHEMA_RECOGNITION_WINDOW_MS / (ORIGINAL_N - SCHEMA_UPDATE_TRIGGER - 1));
+ }
}
}
@@ -278,9 +391,7 @@ abstract class StorageApiSinkSchemaUpdateITBase {
public TableRow apply(Long rowId) {
TableRow row = new TableRow();
row.set("id", rowId);
-
List<String> fields = rowId < ORIGINAL_N ? fieldNames : fieldNamesWithExtra;
-
for (String name : fields) {
String type = Iterables.get(Splitter.on('_').split(name), 0);
switch (type) {
@@ -345,15 +456,21 @@ abstract class StorageApiSinkSchemaUpdateITBase {
}
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
// append
p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(0);
+ p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchRetryTimeMilliSec(1000);
// 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);
@@ -361,7 +478,11 @@ abstract class StorageApiSinkSchemaUpdateITBase {
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<String> fieldNamesOrigin = new ArrayList<String>(Arrays.asList(FIELDS));
@@ -388,38 +509,46 @@ abstract class StorageApiSinkSchemaUpdateITBase {
Write<TableRow> write =
BigQueryIO.writeTableRows()
.to(tableSpec)
- .withAutoSchemaUpdate(useAutoSchemaUpdate)
.withMethod(method)
.withCreateDisposition(CreateDisposition.CREATE_NEVER)
.withWriteDisposition(WriteDisposition.WRITE_APPEND);
+ // These two are mutually exclusive: withAutoSchemaUpdateConsistent also sets the
+ // autoSchemaUpdate flag, so calling both would clobber withAutoSchemaUpdate().
+ if (consistentAutoUpdate) {
+ write = write.withAutoSchemaUpdateConsistent(true, Duration.standardMinutes(5));
+ } else {
+ write = write.withAutoSchemaUpdate(useAutoSchemaUpdate);
+ }
if (useInputSchema) {
write = write.withSchema(inputSchema);
}
if (useIgnoreUnknownValues) {
write = write.ignoreUnknownValues();
}
- // 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);
+ // Whether to wait before altering the table so the sink initializes with the original schema.
+ boolean settleBeforeSchemaUpdate =
+ changeTableSchema && (useAutoSchemaUpdate || consistentAutoUpdate || !useInputSchema);
+ // Passive (non-consistent) autoSchemaUpdate relies on BigQuery returning the updated schema on
+ // subsequent append responses, so it needs the post-update recognition window and longer
+ // triggering frequency.
+ boolean waitForPassiveSchemaRecognition =
+ changeTableSchema && useAutoSchemaUpdate && !consistentAutoUpdate;
if (method == Write.Method.STORAGE_WRITE_API) {
write =
write.withTriggeringFrequency(
- Duration.standardSeconds(waitLonger ? LONG_WAIT_SECONDS : 1));
+ waitForPassiveSchemaRecognition
+ ? Duration.standardSeconds(LONG_WAIT_SECONDS)
+ : Duration.millis(10));
}
- // set up and build pipeline
+ // set up and build pipeline.
+ // Rows are emitted as fast as possible; any wall-clock delay that the test needs is inserted
+ // by UpdateSchemaDoFn around the schema change itself.
Instant start = new Instant(0);
- Duration interval =
- waitLonger ? Duration.standardSeconds(LONG_WAIT_SECONDS) : Duration.millis(1);
- Duration stop =
- waitLonger
- ? Duration.standardSeconds((TOTAL_N - 1) * LONG_WAIT_SECONDS)
- : Duration.millis(TOTAL_N - 1);
+ Duration interval = Duration.millis(1);
+ Duration stop = Duration.millis(TOTAL_N - 1);
Function<Instant, Long> getIdFromInstant =
- waitLonger
- ? (Function<Instant, Long> & Serializable)
- (Instant instant) -> instant.getMillis() / (1000 * LONG_WAIT_SECONDS)
- : (Function<Instant, Long> & Serializable) (Instant instant) -> instant.getMillis();
+ (Function<Instant, Long> & Serializable) (Instant instant) -> instant.getMillis();
// Generates rows with original schema up for row IDs under ORIGINAL_N
// Then generates rows with updated schema for the rest
@@ -449,10 +578,18 @@ abstract class StorageApiSinkSchemaUpdateITBase {
"Update Schema",
ParDo.of(
new UpdateSchemaDoFn(
- PROJECT, bigQueryDatasetId, ImmutableMap.of(tableId, updatedSchema))));
+ PROJECT,
+ bigQueryDatasetId,
+ ImmutableMap.of(tableId, updatedSchema),
+ waitForPassiveSchemaRecognition,
+ // The *Consistent and !useInputSchema variants need the pre-update settle
+ // delay so that the sink's converter is built BEFORE the table schema
+ // changes, but they do NOT need the 60s post-update recognition window.
+ settleBeforeSchemaUpdate,
+ waitForPassiveSchemaRecognition)));
}
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 {
@@ -465,12 +602,15 @@ abstract class StorageApiSinkSchemaUpdateITBase {
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;
- checkRowCompleteness(tableSpec, expectedCount, checkNoDuplication);
- if (useIgnoreUnknownValues) {
- checkRowsWithUpdatedSchema(tableSpec, extraField, useAutoSchemaUpdate);
- }
+ int expectedCount = (useIgnoreUnknownValues || consistentAutoUpdate) ? TOTAL_N : ORIGINAL_N;
+ boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API);
+ verifyTable(
+ tableSpec,
+ expectedCount,
+ checkNoDuplication,
+ // The extra-field checks should only be performed when ignoreUnknownValues is set.
+ (useIgnoreUnknownValues || consistentAutoUpdate) ? extraField : null,
+ useAutoSchemaUpdate || consistentAutoUpdate);
}
private static class VerifyPCollectionSize
@@ -507,22 +647,30 @@ abstract class StorageApiSinkSchemaUpdateITBase {
}
}
- // Check the expected number of rows reached the table.
- // If using STORAGE_WRITE_API, check no duplication happened.
- private static void checkRowCompleteness(
- String tableSpec, int expectedCount, boolean checkNoDuplication)
+ // Fetches the table's rows once and runs every verification against that single snapshot:
+ // - the expected number of rows reached the table (and, if using STORAGE_WRITE_API, that no
+ // duplication happened), and
+ // - if extraField is non-null, that the extra field is present exactly on the rows that are
+ // expected to have it.
+ // These used to be two separate queries, but every BigqueryClient#queryUnflattened call creates
+ // and then deletes a temporary dataset, which costs several seconds.
+ private void verifyTable(
+ String tableSpec,
+ int expectedCount,
+ boolean checkNoDuplication,
+ @Nullable String extraField,
+ boolean useAutoSchemaUpdate)
throws IOException, InterruptedException {
- TableRow queryResponse =
- Iterables.getOnlyElement(
- BQ_CLIENT.queryUnflattened(
- String.format("SELECT COUNT(DISTINCT(id)), COUNT(id) FROM [%s]", tableSpec),
- PROJECT,
- true,
- false,
- bigQueryLocation));
-
- int distinctCount = Integer.parseInt((String) queryResponse.get("f0_"));
- int totalCount = Integer.parseInt((String) queryResponse.get("f1_"));
+ List<TableRow> actualRows =
+ BQ_CLIENT.queryUnflattened(
+ String.format("SELECT * FROM [%s]", tableSpec), PROJECT, true, false, bigQueryLocation);
+
+ Set<String> distinctIds = new HashSet<>();
+ for (TableRow row : actualRows) {
+ distinctIds.add((String) row.get("id"));
+ }
+ int distinctCount = distinctIds.size();
+ int totalCount = actualRows.size();
LOG.info("total distinct count = {}, total count = {}", distinctCount, totalCount);
@@ -530,17 +678,16 @@ abstract class StorageApiSinkSchemaUpdateITBase {
if (checkNoDuplication) {
assertEquals(distinctCount, totalCount);
}
+
+ if (extraField != null) {
+ checkRowsWithUpdatedSchema(actualRows, extraField, useAutoSchemaUpdate);
+ }
}
// Performs checks on the table's rows under different conditions.
// Note: these should only be performed when ignoreUnknownValues is set.
- public void checkRowsWithUpdatedSchema(
- String tableSpec, String extraField, boolean useAutoSchemaUpdate)
- throws IOException, InterruptedException {
- List<TableRow> actualRows =
- BQ_CLIENT.queryUnflattened(
- String.format("SELECT * FROM [%s]", tableSpec), PROJECT, true, false, bigQueryLocation);
-
+ private void checkRowsWithUpdatedSchema(
+ List<TableRow> actualRows, String extraField, boolean useAutoSchemaUpdate) {
for (TableRow row : actualRows) {
// Rows written to the table should not have the extra field if
// 1. The row has original schema
@@ -566,39 +713,53 @@ abstract class StorageApiSinkSchemaUpdateITBase {
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)
@@ -673,17 +834,10 @@ abstract class StorageApiSinkSchemaUpdateITBase {
Instant start = new Instant(0);
// 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.
- Duration interval =
- changeTableSchema ? Duration.standardSeconds(LONG_WAIT_SECONDS) : Duration.millis(1);
- Duration stop =
- changeTableSchema
- ? Duration.standardSeconds((numRows - 1) * LONG_WAIT_SECONDS)
- : Duration.millis(numRows - 1);
+ Duration interval = Duration.millis(1);
+ Duration stop = Duration.millis(numRows - 1);
Function<Instant, Long> getIdFromInstant =
- changeTableSchema
- ? (Function<Instant, Long> & Serializable)
- (Instant instant) -> instant.getMillis() / (1000 * LONG_WAIT_SECONDS)
- : (Function<Instant, Long> & Serializable) Instant::getMillis;
+ (Function<Instant, Long> & Serializable) Instant::getMillis;
// Generates rows with original schema up for row IDs under ORIGINAL_N
// Then generates rows with updated schema for the rest
@@ -714,7 +868,9 @@ abstract class StorageApiSinkSchemaUpdateITBase {
.apply("Add a dummy key", WithKeys.of(1))
.apply(
"Update Schema",
- ParDo.of(new UpdateSchemaDoFn(PROJECT, bigQueryDatasetId, updatedSchemas)));
+ ParDo.of(
+ new UpdateSchemaDoFn(
+ PROJECT, bigQueryDatasetId, updatedSchemas, true, true, true)));
}
WriteResult result = rows.apply("Stream to BigQuery", write);
@@ -731,8 +887,7 @@ abstract class StorageApiSinkSchemaUpdateITBase {
for (Map.Entry<String, Integer> expectedCount : expectedCounts.entrySet()) {
String dest = expectedCount.getKey();
- checkRowCompleteness(dest, expectedCount.getValue(), true);
- checkRowsWithUpdatedSchema(dest, extraFields.get(dest), true);
+ verifyTable(dest, expectedCount.getValue(), true, extraFields.get(dest), true);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java
index cd29d9fb65a..169fb4afc6e 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithInputSchemaIT.java
@@ -24,7 +24,10 @@ import org.junit.BeforeClass;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
-@RunWith(Parameterized.class)
+// Runs the two changeTableSchema cases in parallel. Splitting the suite into two concrete
+// classes additionally lets Gradle's maxParallelForks run the with/without-input-schema
+// halves in separate JVMs.
+@RunWith(StorageApiSinkSchemaUpdateITBase.ParallelParameterized.class)
public class StorageApiSinkSchemaUpdateWithInputSchemaIT extends StorageApiSinkSchemaUpdateITBase {
private static final String BIG_QUERY_DATASET_ID =
"storage_api_sink_schema_change_with_input_" + System.nanoTime();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java
index 9b02e0be396..d71d39a9c58 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateWithoutInputSchemaIT.java
@@ -24,7 +24,10 @@ import org.junit.BeforeClass;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
-@RunWith(Parameterized.class)
+// Runs the two changeTableSchema cases in parallel. Splitting the suite into two concrete
+// classes additionally lets Gradle's maxParallelForks run the with/without-input-schema
+// halves in separate JVMs.
+@RunWith(StorageApiSinkSchemaUpdateITBase.ParallelParameterized.class)
public class StorageApiSinkSchemaUpdateWithoutInputSchemaIT
extends StorageApiSinkSchemaUpdateITBase {
private static final String BIG_QUERY_DATASET_ID =
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 13724ab83e6..8e96f2a4a7f 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,13 +22,18 @@ import static org.junit.Assert.assertFalse;
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.ByteString;
import com.google.protobuf.DescriptorProtos;
import com.google.protobuf.Descriptors;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.UnknownFieldSet;
import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
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 +330,25 @@ public class UpgradeTableSchemaTest {
// 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(),
+ ByteString.copyFrom(payloadNoHash.getPayload()),
+ () -> hash1,
+ null));
// 2. Equal hash
- assertFalse(UpgradeTableSchema.isPayloadSchemaOutOfDate(payload, () -> hash1, null));
+ assertFalse(
+ UpgradeTableSchema.isPayloadSchemaOutOfDate(
+ payload.getSchemaHash(), ByteString.copyFrom(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(),
+ ByteString.copyFrom(payload.getPayload()),
+ () -> hash2,
+ () -> descriptor));
// 4. Different hash with unknown fields.
DynamicMessage unknownFieldSet =
@@ -348,7 +364,10 @@ public class UpgradeTableSchemaTest {
StorageApiWritePayload.of(unknownFieldSet.toByteArray(), null, null).withSchemaHash(hash1);
assertTrue(
UpgradeTableSchema.isPayloadSchemaOutOfDate(
- payloadWithUnknown, () -> hash2, () -> descriptor));
+ payloadWithUnknown.getSchemaHash(),
+ ByteString.copyFrom(payloadWithUnknown.getPayload()),
+ () -> hash2,
+ () -> descriptor));
// 5. Different hash with missing required fields.
DynamicMessage missingField = DynamicMessage.newBuilder(descriptor).buildPartial();
@@ -356,6 +375,189 @@ public class UpgradeTableSchemaTest {
StorageApiWritePayload.of(missingField.toByteArray(), null, null).withSchemaHash(hash1);
assertTrue(
UpgradeTableSchema.isPayloadSchemaOutOfDate(
- payloadMissingField, () -> hash2, () -> descriptor));
+ payloadMissingField.getSchemaHash(),
+ ByteString.copyFrom(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));
+ }
+
+ @Test
+ public void testMissingUnknownField_NestedAndRepeatedObjectWithTableCells() 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("single_struct")
+ .setNumber(1)
+ .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_MESSAGE)
+ .setTypeName("NestedMessage")
+ .build())
+ .addField(
+ DescriptorProtos.FieldDescriptorProto.newBuilder()
+ .setName("repeated_struct")
+ .setNumber(2)
+ .setLabel(DescriptorProtos.FieldDescriptorProto.Label.LABEL_REPEATED)
+ .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 nestedRowOk =
+ new TableRow().setF(Collections.singletonList(new TableCell().setV("value1")));
+ TableRow parentRow =
+ new TableRow()
+ .setF(
+ Arrays.asList(
+ new TableCell().setV(nestedRowOk),
+ new TableCell()
+ .setV(Collections.singletonList(new TableCell().setV(nestedRowOk)))));
+
+ assertFalse(UpgradeTableSchema.missingUnknownField(parentRow, () -> descriptor));
+
+ TableRow nestedRowWithUnknown =
+ new TableRow()
+ .setF(Arrays.asList(new TableCell().setV("value1"), new TableCell().setV("value2")));
+ TableRow parentWithSingleUnknown =
+ new TableRow()
+ .setF(
+ Arrays.asList(
+ new TableCell().setV(nestedRowWithUnknown),
+ new TableCell()
+ .setV(Collections.singletonList(new TableCell().setV(nestedRowOk)))));
+ assertTrue(UpgradeTableSchema.missingUnknownField(parentWithSingleUnknown, () -> descriptor));
+
+ TableRow parentWithRepeatedUnknown =
+ new TableRow()
+ .setF(
+ Arrays.asList(
+ new TableCell().setV(nestedRowOk),
+ new TableCell()
+ .setV(
+ Collections.singletonList(
+ new TableCell().setV(nestedRowWithUnknown)))));
+ assertTrue(UpgradeTableSchema.missingUnknownField(parentWithRepeatedUnknown, () -> descriptor));
}
}