auto update schema

Author: reuvenlax

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java

@@ -33,6 +33,7 @@
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.util.Preconditions;
 
 /**
  * Container class used by {@link StorageApiWritesShardedRecords} and {@link
@@ -145,15 +146,17 @@ boolean hasSchemaChanged(TableSchema updatedTableSchema) {
   public ByteString encodeUnknownFields(TableRow unknown, boolean ignoreUnknownValues)
       throws TableRowToStorageApiProto.SchemaConversionException {
     Message msg =
-        TableRowToStorageApiProto.messageFromTableRow(
-            getSchemaInformation(),
-            getDescriptorIgnoreRequired(),
-            unknown,
-            ignoreUnknownValues,
-            true,
-            null,
-            null,
-            null);
+        Preconditions.checkArgumentNotNull(
+            TableRowToStorageApiProto.messageFromTableRow(
+                getSchemaInformation(),
+                getDescriptorIgnoreRequired(),
+                unknown,
+                ignoreUnknownValues,
+                true,
+                null,
+                null,
+                null,
+                TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT));
     return msg.toByteString();
   }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -4202,9 +4202,19 @@ private <DestinationT> WriteResult continueExpandTyped(
         }
         return input.apply(batchLoads);
       } else if (method == Method.STORAGE_WRITE_API || method == Method.STORAGE_API_AT_LEAST_ONCE) {
+        boolean useSchemaUpdate =
+            getSchemaUpdateOptions() != null && !getSchemaUpdateOptions().isEmpty();
+        if (useSchemaUpdate) {
+          checkArgument(
+              !getAutoSchemaUpdate(),
+              "Schema update options are not supported when using auto schema update");
+          checkArgument(!getIgnoreUnknownValues());
+        }
         BigQueryOptions bqOptions = input.getPipeline().getOptions().as(BigQueryOptions.class);
         StorageApiDynamicDestinations<T, DestinationT> storageApiDynamicDestinations;
         if (getUseBeamSchema()) {
+          checkArgument(
+              !useSchemaUpdate, "SchemaUpdateOptions are not supported when using Beam schemas");
           // This ensures that the Beam rows are directly translated into protos for Storage API
           // writes, with no
           // need to round trip through JSON TableRow objects.
@@ -4216,6 +4226,9 @@ private <DestinationT> WriteResult continueExpandTyped(
                   getFormatRecordOnFailureFunction(),
                   getRowMutationInformationFn() != null);
         } else if (getWriteProtosClass() != null && getDirectWriteProtos()) {
+          checkArgument(
+              !useSchemaUpdate, "SchemaUpdateOptions are not supported when writing protos");
+
           // We could support both of these by falling back to
           // StorageApiDynamicDestinationsTableRow. This
           // would defeat the optimization (we would be forced to create a new dynamic proto message
@@ -4233,13 +4246,18 @@ private <DestinationT> WriteResult continueExpandTyped(
               !getIgnoreUnknownValues(),
               "ignoreUnknownValues not supported when using writeProtos."
                   + " Try setting withDirectWriteProtos(false)");
+          checkArgument(!useSchemaUpdate);
+
           storageApiDynamicDestinations =
               (StorageApiDynamicDestinations<T, DestinationT>)
                   new StorageApiDynamicDestinationsProto(
                       dynamicDestinations,
                       getWriteProtosClass(),
                       getFormatRecordOnFailureFunction());
         } else if (getAvroRowWriterFactory() != null) {
+          checkArgument(
+              !useSchemaUpdate, "SchemaUpdateOptions are not supported when writing avros");
+
           // we can configure the avro to storage write api proto converter for this
           // assuming the format function returns an Avro GenericRecord
           // and there is a schema defined
@@ -4248,7 +4266,6 @@ private <DestinationT> WriteResult continueExpandTyped(
                   || getDynamicDestinations() != null
                   || getSchemaFromView() != null,
               "A schema must be provided for avro rows to be used with StorageWrite API.");
-
           RowWriterFactory.AvroRowWriterFactory<T, GenericRecord, DestinationT>
               recordWriterFactory =
                   (RowWriterFactory.AvroRowWriterFactory<T, GenericRecord, DestinationT>)
@@ -4275,7 +4292,10 @@ private <DestinationT> WriteResult continueExpandTyped(
                   getRowMutationInformationFn() != null,
                   getCreateDisposition(),
                   getIgnoreUnknownValues(),
-                  getAutoSchemaUpdate());
+                  getAutoSchemaUpdate(),
+                  getSchemaUpdateOptions() == null
+                      ? Collections.emptySet()
+                      : getSchemaUpdateOptions());
         }
 
         int numShards = getStorageApiNumStreams(bqOptions);
@@ -4287,6 +4307,7 @@ private <DestinationT> WriteResult continueExpandTyped(
         StorageApiLoads<DestinationT, T> storageApiLoads =
             new StorageApiLoads<>(
                 destinationCoder,
+                elementCoder,
                 storageApiDynamicDestinations,
                 getRowMutationInformationFn(),
                 getCreateDisposition(),

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java

@@ -245,4 +245,12 @@ public interface BigQueryOptions
   Boolean getGroupFilesFileLoad();
 
   void setGroupFilesFileLoad(Boolean value);
+
+  @Hidden
+  @Description(
+      "The number of parallelization to use for buffering elements when upgrading table schemas.")
+  @Default.Integer(50)
+  Integer getSchemaUpgradeBufferingShards();
+
+  void setSchemaUpgradeBufferingShards(Integer value);
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java

@@ -204,6 +204,10 @@ <T> long insertAll(
     /** Patch BigQuery {@link Table} description. */
     Table patchTableDescription(TableReference tableReference, @Nullable String tableDescription)
         throws IOException, InterruptedException;
+
+    Table patchTableSchema(
+        TableReference tableReference, com.google.api.services.bigquery.model.TableSchema newSchema)
+        throws IOException, InterruptedException;
   }
 
   /** An interface to get, create and flush Cloud BigQuery STORAGE API write streams. */

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -1459,6 +1459,31 @@ public Table patchTableDescription(
           ALWAYS_RETRY);
     }
 
+    @Override
+    public Table patchTableSchema(
+        TableReference tableReference, com.google.api.services.bigquery.model.TableSchema newSchema)
+        throws IOException, InterruptedException {
+      Table newTable = new Table();
+      newTable.setSchema(newSchema);
+
+      Tables.Patch request =
+          client
+              .tables()
+              .patch(
+                  tableReference.getProjectId(),
+                  tableReference.getDatasetId(),
+                  tableReference.getTableId(),
+                  newTable);
+      return executeWithRetries(
+          request,
+          String.format(
+              "Unable to patch table: %s, aborting after %d retries.",
+              tableReference, MAX_RPC_RETRIES),
+          Sleeper.DEFAULT,
+          createDefaultBackoff(),
+          DONT_RETRY_INVALID_ARG_OR_PRECONDITION);
+    }
+
     @Override
     public void close() throws Exception {
       // Nothing to close
@@ -1664,6 +1689,11 @@ public void close() throws Exception {
         return !errorExtractor.itemNotFound(input);
       };
 
+  static final SerializableFunction<IOException, Boolean> DONT_RETRY_INVALID_ARG_OR_PRECONDITION =
+      input -> {
+        ApiErrorExtractor errorExtractor = new ApiErrorExtractor();
+        return !errorExtractor.preconditionNotMet(input) && !errorExtractor.badRequest(input);
+      };
   static final SerializableFunction<IOException, Boolean> ALWAYS_RETRY = input -> true;
 
   @VisibleForTesting

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MergeSchemaCombineFn.java

@@ -0,0 +1,76 @@
+/*
+ * 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.cloud.bigquery.storage.v1.TableSchema;
+import java.util.Iterator;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.transforms.Combine;
+
+/** Merge schemas. */
+public class MergeSchemaCombineFn extends Combine.CombineFn<TableSchema, TableSchema, TableSchema> {
+  @Override
+  public TableSchema createAccumulator() {
+    return TableSchema.newBuilder().build();
+  }
+
+  @Override
+  public TableSchema addInput(TableSchema accumulator, TableSchema input) {
+    if (input.equals(accumulator)) {
+      return accumulator;
+    }
+    return UpgradeTableSchema.mergeSchemas(accumulator, input);
+  }
+
+  @Override
+  public Coder<TableSchema> getAccumulatorCoder(
+      CoderRegistry registry, Coder<TableSchema> inputCoder) {
+    return inputCoder;
+  }
+
+  @Override
+  public Coder<TableSchema> getDefaultOutputCoder(
+      CoderRegistry registry, Coder<TableSchema> inputCoder) throws CannotProvideCoderException {
+    return inputCoder;
+  }
+
+  @Override
+  public TableSchema mergeAccumulators(Iterable<TableSchema> accumulators) {
+    checkNotNull(accumulators, "accumulators must be non-null");
+
+    Iterator<TableSchema> iter = accumulators.iterator();
+    if (!iter.hasNext()) {
+      return createAccumulator();
+    }
+
+    TableSchema merged = iter.next();
+    while (iter.hasNext()) {
+      merged = addInput(merged, iter.next());
+    }
+    return merged;
+  }
+
+  @Override
+  public TableSchema extractOutput(TableSchema accumulator) {
+    return accumulator;
+  }
+}

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/PatchTableSchemaDoFn.java

@@ -0,0 +1,117 @@
+/*
+ * 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.cloud.bigquery.storage.v1.TableSchema;
+import com.google.cloud.hadoop.util.ApiErrorExtractor;
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * This DoFn is responsible for updating a BigQuery's table schema. The input is a TableSchema
+ * containing only the schema delta (new fields, relaxed fields). It outputs elements for all
+ * updated tables, which act as notifcations to the buffering stage that the elements can be
+ * retried.
+ */
+public class PatchTableSchemaDoFn<DestinationT extends @NonNull Object, ElementT>
+    extends DoFn<KV<DestinationT, TableSchema>, KV<DestinationT, ElementT>> {
+  private final BigQueryServices bqServices;
+  private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
+  private TwoLevelMessageConverterCache<DestinationT, ElementT> messageConverters;
+  private transient BigQueryServices.@Nullable DatasetService datasetServiceInternal = null;
+
+  PatchTableSchemaDoFn(
+      String operationName,
+      BigQueryServices bqServices,
+      StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations) {
+    this.messageConverters = new TwoLevelMessageConverterCache<>(operationName);
+    this.bqServices = bqServices;
+    this.dynamicDestinations = dynamicDestinations;
+  }
+
+  private BigQueryServices.DatasetService getDatasetService(PipelineOptions pipelineOptions)
+      throws IOException {
+    if (datasetServiceInternal == null) {
+      datasetServiceInternal =
+          bqServices.getDatasetService(pipelineOptions.as(BigQueryOptions.class));
+    }
+    return datasetServiceInternal;
+  }
+
+  @ProcessElement
+  public void processElement(
+      @Element KV<DestinationT, TableSchema> element,
+      OutputReceiver<KV<DestinationT, @Nullable ElementT>> o,
+      ProcessContext context,
+      PipelineOptions pipelineOptions)
+      throws Exception {
+    dynamicDestinations.setSideInputAccessorFromProcessContext(context);
+    DestinationT destination = element.getKey();
+    TableSchema tableSchemaDiff = element.getValue();
+
+    StorageApiDynamicDestinations.MessageConverter<ElementT> messageConverter =
+        messageConverters.get(destination, dynamicDestinations, getDatasetService(pipelineOptions));
+    messageConverter.updateSchemaFromTable();
+
+    while (true) {
+      TableSchema baseSchema = messageConverter.getTableSchema();
+      TableSchema updatedSchema = UpgradeTableSchema.mergeSchemas(baseSchema, tableSchemaDiff);
+      if (baseSchema.equals(updatedSchema)) {
+        return;
+      }
+      // Check first to see if the schema still needs updating.
+      BackOff backoff =
+          new ExponentialBackOff.Builder()
+              .setMaxElapsedTimeMillis((int) TimeUnit.MINUTES.toMillis(1))
+              .build();
+      boolean schemaOutOfDate = false;
+      do {
+        try {
+          getDatasetService(pipelineOptions)
+              .patchTableSchema(
+                  dynamicDestinations.getTable(destination).getTableReference(),
+                  TableRowToStorageApiProto.protoSchemaToTableSchema(updatedSchema));
+          // Indicate that we've patched this schema.
+          o.output(KV.of(destination, null));
+          return;
+        } catch (IOException e) {
+          ApiErrorExtractor errorExtractor = new ApiErrorExtractor();
+          if (errorExtractor.preconditionNotMet(e) || errorExtractor.badRequest(e)) {
+            schemaOutOfDate = true;
+            break;
+          }
+        }
+      } while (BackOffUtils.next(com.google.api.client.util.Sleeper.DEFAULT, backoff));
+      if (schemaOutOfDate) {
+        // This could be due to an out-of-date schema.
+        messageConverter.updateSchemaFromTable();
+      } else {
+        // We ran out of retries.
+        throw new RuntimeException("Failed to patch table schema");
+      }
+    }
+  }
+}