Seamlessly propagate schema changes made after pipelines starts running for Spanner Change Streams to BigQuery template

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigquery/BigQueryResourceManager.java

@@ -145,7 +145,7 @@ private synchronized Dataset getDatasetIfExists(String datasetId) throws Illegal
    * @return the table, if it exists.
    * @throws IllegalStateException if the given table name does not exist in the dataset.
    */
-  private synchronized Table getTableIfExists(String tableId) throws IllegalStateException {
+  public synchronized Table getTableIfExists(String tableId) throws IllegalStateException {
     checkHasDataset();
     Table table = dataset.get(tableId);
     if (table == null) {

v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/BigQueryDynamicDestinations.java

@@ -24,17 +24,13 @@
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Dialect;
-import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.BigQueryUtils;
-import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.SpannerChangeStreamsUtils;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.SpannerToBigQueryUtils;
 import com.google.cloud.teleport.v2.transforms.BigQueryConverters;
 import com.google.common.collect.ImmutableSet;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
-import java.util.Optional;
 import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinations;
 import org.apache.beam.sdk.io.gcp.bigquery.TableDestination;
 import org.apache.beam.sdk.io.gcp.spanner.SpannerAccessor;
@@ -53,39 +49,18 @@
 
   private static final Logger LOG = LoggerFactory.getLogger(BigQueryDynamicDestinations.class);
 
-  private final Map<String, TrackedSpannerTable> spannerTableByName;
   private final String bigQueryProject, bigQueryDataset, bigQueryTableTemplate;
   private final Boolean useStorageWriteApi;
   private final ImmutableSet<String> ignoreFields;
 
   public static BigQueryDynamicDestinations of(
       BigQueryDynamicDestinationsOptions bigQueryDynamicDestinationsOptions) {
     Dialect dialect = getDialect(bigQueryDynamicDestinationsOptions.getSpannerConfig());
-    try (SpannerAccessor spannerAccessor =
-        SpannerAccessor.getOrCreate(bigQueryDynamicDestinationsOptions.getSpannerConfig())) {
-      Map<String, TrackedSpannerTable> spannerTableByName =
-          new SpannerChangeStreamsUtils(
-                  spannerAccessor.getDatabaseClient(),
-                  bigQueryDynamicDestinationsOptions.getChangeStreamName(),
-                  dialect)
-              .getSpannerTableByName();
-      return new BigQueryDynamicDestinations(
-          bigQueryDynamicDestinationsOptions, spannerTableByName);
-    } catch (RuntimeException e) {
-      String errorMessage =
-          String.format(
-              "Caught exception when getting BigQueryDynamicDestinations, message: %s,"
-                  + " cause: %s",
-              Optional.ofNullable(e.getMessage()), e.getCause());
-      LOG.error(errorMessage);
-      throw new RuntimeException(errorMessage, e);
-    }
+    return new BigQueryDynamicDestinations(bigQueryDynamicDestinationsOptions);
   }
 
   private BigQueryDynamicDestinations(
-      BigQueryDynamicDestinationsOptions bigQueryDynamicDestinationsOptions,
-      Map<String, TrackedSpannerTable> spannerTableByName) {
-    this.spannerTableByName = spannerTableByName;
+      BigQueryDynamicDestinationsOptions bigQueryDynamicDestinationsOptions) {
     this.ignoreFields = bigQueryDynamicDestinationsOptions.getIgnoreFields();
     this.bigQueryProject = bigQueryDynamicDestinationsOptions.getBigQueryProject();
     this.bigQueryDataset = bigQueryDynamicDestinationsOptions.getBigQueryDataset();
@@ -118,10 +93,8 @@
   @Override
   public TableSchema getSchema(KV<TableId, TableRow> destination) {
     TableRow tableRow = destination.getValue();
-    String spannerTableName =
-        (String) tableRow.get(BigQueryUtils.BQ_CHANGELOG_FIELD_NAME_TABLE_NAME);
-    TrackedSpannerTable spannerTable = spannerTableByName.get(spannerTableName);
-    List<TableFieldSchema> fields = getFields(spannerTable);
+    // Get List<TableFieldSchema> for both user columns and metadata columns.
+    List<TableFieldSchema> fields = getFields(tableRow);
     List<TableFieldSchema> filteredFields = new ArrayList<>();
     for (TableFieldSchema field : fields) {
       if (!ignoreFields.contains(field.getName())) {
@@ -132,9 +105,12 @@
     return new TableSchema().setFields(filteredFields);
   }
 
-  private List<TableFieldSchema> getFields(TrackedSpannerTable spannerTable) {
+  // Returns List<TableFieldSchema> for user columns and metadata columns based on the parameter
+  // TableRow.
+  private List<TableFieldSchema> getFields(TableRow tableRow) {
+    // Add all user data fields (excluding metadata fields stored in metadataColumns).
     List<TableFieldSchema> fields =
-        SpannerToBigQueryUtils.spannerColumnsToBigQueryIOFields(spannerTable.getAllColumns());
+        SpannerToBigQueryUtils.tableRowColumnsToBigQueryIOFields(tableRow, this.useStorageWriteApi);
 
     // Add all metadata fields.
     String requiredMode = Field.Mode.REQUIRED.name();

v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/FailsafeModJsonToTableRowTransformer.java

@@ -38,6 +38,7 @@
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerColumn;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.BigQueryUtils;
+import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.SchemaUpdateUtils;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.SpannerChangeStreamsUtils;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.SpannerToBigQueryUtils;
 import com.google.cloud.teleport.v2.values.FailsafeElement;
@@ -132,6 +133,7 @@
       private transient CallContextConfigurator callContextConfigurator;
       private transient boolean seenException;
       private Boolean useStorageWriteApi;
+      private Dialect dialect;
 
       public FailsafeModJsonToTableRowFn(
           SpannerConfig spannerConfig,
@@ -146,6 +148,7 @@
         this.transformDeadLetterOut = transformDeadLetterOut;
         this.ignoreFields = ignoreFields;
         this.useStorageWriteApi = useStorageWriteApi;
+        this.dialect = getDialect(spannerConfig);
       }
 
       private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
@@ -169,7 +172,6 @@
       public void setUp() {
         seenException = false;
         try {
-          Dialect dialect = getDialect(spannerConfig);
           spannerAccessor = SpannerAccessor.getOrCreate(spannerConfig);
           spannerTableByName =
               new SpannerChangeStreamsUtils(
@@ -181,7 +183,7 @@
                   "Caught exception when setting up FailsafeModJsonToTableRowFn, message: %s,"
                       + " cause: %s",
                   Optional.ofNullable(e.getMessage()), e.getCause()));
-          seenException = true;
+          throw new RuntimeException(e);
         }
         setUpCallContextConfigurator();
       }
@@ -252,6 +254,19 @@
         }
         String spannerTableName = mod.getTableName();
         TrackedSpannerTable spannerTable;
+        com.google.cloud.Timestamp spannerCommitTimestamp =
+            com.google.cloud.Timestamp.ofTimeSecondsAndNanos(
+                mod.getCommitTimestampSeconds(), mod.getCommitTimestampNanos());
+
+        // Detect schema updates (newly added tables/columns) from mod and propagate changes into
+        // spannerTableByName which stores schema information by table name.
+        // Not able to get schema update from DELETE mods as they have empty newValuesJson.
+        if (mod.getModType() != ModType.DELETE) {
+          spannerTableByName =
+              SchemaUpdateUtils.updateStoredSchemaIfNeeded(
+                  spannerAccessor, spannerChangeStream, dialect, mod, spannerTableByName);
+        }
+
         try {
           spannerTable = checkStateNotNull(spannerTableByName.get(spannerTableName));
 
@@ -262,9 +277,6 @@
           LOG.error(errorMessage);
           throw new RuntimeException(errorMessage, e);
         }
-        com.google.cloud.Timestamp spannerCommitTimestamp =
-            com.google.cloud.Timestamp.ofTimeSecondsAndNanos(
-                mod.getCommitTimestampSeconds(), mod.getCommitTimestampNanos());
 
         // Set metadata fields of the tableRow.
         TableRow tableRow = new TableRow();
@@ -277,32 +289,17 @@
             useStorageWriteApi);
         JSONObject keysJsonObject = new JSONObject(mod.getKeysJson());
         // Set Spanner key columns of the tableRow.
-        for (TrackedSpannerColumn spannerColumn : spannerTable.getPkColumns()) {
-          String spannerColumnName = spannerColumn.getName();
-          if (keysJsonObject.has(spannerColumnName)) {
-            tableRow.set(spannerColumnName, keysJsonObject.get(spannerColumnName));
-          } else {
-            String errorMessage =
-                String.format(
-                    "Caught exception when setting key column of the tableRow: Cannot find value"
-                        + " for key column %s",
-                    spannerColumnName);
-            LOG.error(errorMessage);
-            throw new IllegalArgumentException(errorMessage);
-          }
-        }
-
-        // For "DELETE" mod, we only need to set the key columns.
-        if (mod.getModType() == ModType.DELETE) {
-          return tableRow;
-        }
+        SpannerToBigQueryUtils.addSpannerPkColumnsToTableRow(
+            keysJsonObject, spannerTable.getPkColumns(), tableRow);
 
         // Set non-key columns of the tableRow.
         SpannerToBigQueryUtils.addSpannerNonPkColumnsToTableRow(
-            mod.getNewValuesJson(), spannerTable.getNonPkColumns(), tableRow);
+            mod.getNewValuesJson(), spannerTable.getNonPkColumns(), tableRow, mod.getModType());
 
         // For "INSERT" mod, we can get all columns from mod.
-        if (mod.getModType() == ModType.INSERT) {
+        // For "DELETE" mod, we only set the key columns. For all non-key columns, we already
+        // populated "null".
+        if (mod.getModType() == ModType.INSERT || mod.getModType() == ModType.DELETE) {
           return tableRow;
         }
 
@@ -383,8 +380,7 @@
         return tableRow;
       }
 
-      // Do a Spanner read to retrieve full row. The schema change is currently not supported. so we
-      // assume the schema isn't changed while the pipeline is running,
+      // Do a Spanner read to retrieve full row. Schema can change while the pipeline is running.
       private void readSpannerRow(
           String spannerTableName,
           com.google.cloud.spanner.Key key,

v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/SpannerChangeStreamsToBigQuery.java

@@ -25,6 +25,7 @@
 import com.google.cloud.teleport.v2.common.UncaughtExceptionLogger;
 import com.google.cloud.teleport.v2.options.SpannerChangeStreamsToBigQueryOptions;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.Mod;
+import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.ModColumnType;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.BigQueryUtils;
 import com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils.OptionsUtils;
 import com.google.cloud.teleport.v2.transforms.DLQWriteTransform;
@@ -35,6 +36,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
@@ -209,7 +211,8 @@
         .equals(BigQueryUtils.BQ_CHANGELOG_FIELD_NAME_TABLE_NAME)) {
       throw new IllegalArgumentException(
           String.format(
-              "bigQueryChangelogTableNameTemplate cannot be set to '{%s}'. This value is reserved for the Cloud Spanner table name.",
+              "bigQueryChangelogTableNameTemplate cannot be set to '{%s}'. This value is reserved"
+                  + " for the Cloud Spanner table name.",
               BigQueryUtils.BQ_CHANGELOG_FIELD_NAME_TABLE_NAME));
     }
 
@@ -386,19 +389,39 @@
                 .setBigQueryTableTemplate(options.getBigQueryChangelogTableNameTemplate())
                 .setUseStorageWriteApi(options.getUseStorageWriteApi())
                 .build();
-    WriteResult writeResult =
-        tableRowTuple
-            .get(failsafeModJsonToTableRow.transformOut)
-            .apply(
-                "Write To BigQuery",
-                BigQueryIO.<TableRow>write()
-                    .to(BigQueryDynamicDestinations.of(bigQueryDynamicDestinationsOptions))
-                    .withFormatFunction(element -> removeIntermediateMetadataFields(element))
-                    .withFormatRecordOnFailureFunction(element -> element)
-                    .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
-                    .withWriteDisposition(Write.WriteDisposition.WRITE_APPEND)
-                    .withExtendedErrorInfo()
-                    .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors()));
+    WriteResult writeResult;
+    if (!options.getUseStorageWriteApi()) {
+      writeResult =
+          tableRowTuple
+              .get(failsafeModJsonToTableRow.transformOut)
+              .apply(
+                  "Write To BigQuery",
+                  BigQueryIO.<TableRow>write()
+                      .to(BigQueryDynamicDestinations.of(bigQueryDynamicDestinationsOptions))
+                      .withFormatFunction(element -> removeIntermediateMetadataFields(element))
+                      .withFormatRecordOnFailureFunction(element -> element)
+                      .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
+                      .withWriteDisposition(Write.WriteDisposition.WRITE_APPEND)
+                      .withExtendedErrorInfo()
+                      .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors()));
+    } else {
+      writeResult =
+          tableRowTuple
+              .get(failsafeModJsonToTableRow.transformOut)
+              .apply(
+                  "Write To BigQuery",
+                  BigQueryIO.<TableRow>write()
+                      .to(BigQueryDynamicDestinations.of(bigQueryDynamicDestinationsOptions))
+                      .withFormatFunction(element -> removeIntermediateMetadataFields(element))
+                      .withFormatRecordOnFailureFunction(element -> element)
+                      .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
+                      .withWriteDisposition(Write.WriteDisposition.WRITE_APPEND)
+                      .ignoreUnknownValues()
+                      .withAutoSchemaUpdate(true) // only supported when using STORAGE_WRITE_API or
+                      // STORAGE_API_AT_LEAST_ONCE.
+                      .withExtendedErrorInfo()
+                      .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors()));
+    }
 
     PCollection<String> transformDlqJson =
         tableRowTuple
@@ -471,6 +494,8 @@
     for (String rowKey : rowKeys) {
       if (metadataFields.contains(rowKey)) {
         cleanTableRow.remove(rowKey);
+      } else if (rowKeys.contains("_type_" + rowKey)) {
+        cleanTableRow.remove("_type_" + rowKey);
       }
     }
 
@@ -496,6 +521,7 @@
                 input.isLastRecordInTransactionInPartition(),
                 input.getRecordSequence(),
                 input.getTableName(),
+                input.getRowType().stream().map(ModColumnType::new).collect(Collectors.toList()),
                 input.getModType(),
                 input.getValueCaptureType(),
                 input.getNumberOfRecordsInTransaction(),