normalize Spanner CDC

.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 7
+  "revision": 8
 }

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

@@ -23,6 +23,7 @@
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Statement;
 import io.opentelemetry.api.OpenTelemetry;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import org.apache.beam.sdk.options.PipelineOptions;
@@ -77,22 +78,23 @@ public ReadSpannerSchema(
     this.allowedTableNames = allowedTableNames == null ? new HashSet<>() : allowedTableNames;
   }
 
-  @Setup
-  public void setup(PipelineOptions options) throws Exception {
-    OpenTelemetry otel = options.as(SdkHarnessOptions.class).getOpenTelemetry();
-    spannerAccessor = SpannerAccessor.getOrCreate(config, otel);
-  }
-
-  @Teardown
-  public void teardown() throws Exception {
-    spannerAccessor.close();
+  /**
+   * Reads Spanner schema information without running a Beam pipeline.
+   *
+   * <p>Used by SchemaTransforms during expansion (including cross-language expansion services that
+   * do not ship DirectRunner).
+   */
+  public static SpannerSchema getSpannerSchema(
+      SpannerConfig config, Dialect dialect, Set<String> allowedTableNames) {
+    try (SpannerAccessor spannerAccessor = SpannerAccessor.getOrCreate(config)) {
+      return getSpannerSchema(spannerAccessor.getDatabaseClient(), dialect, allowedTableNames);
+    }
   }
 
-  @ProcessElement
-  public void processElement(ProcessContext c) throws Exception {
-    Dialect dialect = c.sideInput(dialectView);
+  static SpannerSchema getSpannerSchema(
+      DatabaseClient databaseClient, Dialect dialect, Set<String> allowedTableNames) {
+    Set<String> allowed = allowedTableNames == null ? Collections.emptySet() : allowedTableNames;
     SpannerSchema.Builder builder = SpannerSchema.builder(dialect);
-    DatabaseClient databaseClient = spannerAccessor.getDatabaseClient();
     try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) {
       ResultSet resultSet = readTableInfo(tx, dialect);
 
@@ -101,9 +103,7 @@ public void processElement(ProcessContext c) throws Exception {
         String columnName = resultSet.getString(1);
         String type = resultSet.getString(2);
         long cellsMutated = resultSet.getLong(3);
-        if (allowedTableNames.size() > 0 && !allowedTableNames.contains(tableName)) {
-          // If we want to filter out table names, and the current table name is not part
-          // of the allowed names, we exclude it.
+        if (!isTableAllowed(allowed, tableName)) {
           continue;
         }
         builder.addColumn(tableName, columnName, type, cellsMutated);
@@ -114,14 +114,46 @@ public void processElement(ProcessContext c) throws Exception {
         String tableName = resultSet.getString(0);
         String columnName = resultSet.getString(1);
         String ordering = resultSet.getString(2);
-
+        if (!isTableAllowed(allowed, tableName)) {
+          continue;
+        }
         builder.addKeyPart(tableName, columnName, "DESC".equalsIgnoreCase(ordering));
       }
     }
-    c.output(builder.build());
+    return builder.build();
+  }
+
+  private static boolean isTableAllowed(Set<String> allowedTableNames, String tableName) {
+    if (allowedTableNames.isEmpty()) {
+      return true;
+    }
+    for (String allowed : allowedTableNames) {
+      if (allowed.equalsIgnoreCase(tableName)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  @Setup
+  public void setup(PipelineOptions options) throws Exception {
+    OpenTelemetry otel = options.as(SdkHarnessOptions.class).getOpenTelemetry();
+    spannerAccessor = SpannerAccessor.getOrCreate(config, otel);
+  }
+
+  @Teardown
+  public void teardown() throws Exception {
+    spannerAccessor.close();
+  }
+
+  @ProcessElement
+  public void processElement(ProcessContext c) throws Exception {
+    c.output(
+        getSpannerSchema(
+            spannerAccessor.getDatabaseClient(), c.sideInput(dialectView), allowedTableNames));
   }
 
-  private ResultSet readTableInfo(ReadOnlyTransaction tx, Dialect dialect) {
+  private static ResultSet readTableInfo(ReadOnlyTransaction tx, Dialect dialect) {
     // retrieve schema information for all tables, as well as aggregating the
     // number of indexes that cover each column. this will be used to estimate
     // the number of cells (table column plus indexes) mutated in an upsert operation
@@ -174,7 +206,7 @@ private ResultSet readTableInfo(ReadOnlyTransaction tx, Dialect dialect) {
     return tx.executeQuery(Statement.of(statement));
   }
 
-  private ResultSet readPrimaryKeyInfo(ReadOnlyTransaction tx, Dialect dialect) {
+  private static ResultSet readPrimaryKeyInfo(ReadOnlyTransaction tx, Dialect dialect) {
     String statement = "";
     switch (dialect) {
       case GOOGLE_STANDARD_SQL:

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangestreamsReadSchemaTransformProvider.java

@@ -27,7 +27,6 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -36,7 +35,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.gcp.spanner.ReadSpannerSchema;
 import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
 import org.apache.beam.sdk.io.gcp.spanner.SpannerIO;
@@ -52,14 +50,11 @@
 import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
 import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
-import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.DoFn.FinishBundle;
 import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.values.PCollectionRowTuple;
 import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
@@ -304,42 +299,17 @@ public void finish(FinishBundleContext c) {
     }
   }
 
-  private static final HashMap<String, SpannerSchema> TABLE_SCHEMAS = new HashMap<>();
-
   private static Schema getTableSchema(SpannerChangestreamsReadConfiguration config) {
-    Pipeline miniPipeline = Pipeline.create();
-    PCollectionView<Dialect> sqlDialectView =
-        miniPipeline
-            .apply("Create Dialect", Create.of(Dialect.GOOGLE_STANDARD_SQL))
-            .apply("Dialect to View", View.asSingleton());
-    miniPipeline
-        .apply(Create.of((Void) null))
-        .apply(
-            ParDo.of(
-                    new ReadSpannerSchema(
-                        SpannerConfig.create()
-                            .withDatabaseId(config.getDatabaseId())
-                            .withInstanceId(config.getInstanceId())
-                            .withProjectId(config.getProjectId()),
-                        sqlDialectView,
-                        Sets.newHashSet(config.getTable())))
-                .withSideInput("dialect", sqlDialectView))
-        .apply(
-            ParDo.of(
-                new DoFn<SpannerSchema, String>() {
-                  @ProcessElement
-                  public void process(@DoFn.Element SpannerSchema schema) {
-                    TABLE_SCHEMAS.put(config.getTable(), schema);
-                  }
-                }))
-        .setCoder(StringUtf8Coder.of());
-    miniPipeline.run().waitUntilFinish();
-    // Clean up the static map from the object.
-    SpannerSchema finalSchemaObj = TABLE_SCHEMAS.remove(config.getTable());
-    if (finalSchemaObj == null) {
-      throw new RuntimeException(
-          String.format("Could not get schema for configuration %s", config));
-    }
+    // Query information_schema directly. A nested Pipeline would require DirectRunner,
+    // which is not on the GCP expansion-service classpath used by cross-language YAML.
+    SpannerSchema finalSchemaObj =
+        ReadSpannerSchema.getSpannerSchema(
+            SpannerConfig.create()
+                .withDatabaseId(config.getDatabaseId())
+                .withInstanceId(config.getInstanceId())
+                .withProjectId(config.getProjectId()),
+            Dialect.GOOGLE_STANDARD_SQL,
+            Sets.newHashSet(config.getTable()));
     return spannerSchemaToBeamSchema(finalSchemaObj, config.getTable());
   }
 

sdks/python/apache_beam/yaml/standard_io.yaml

@@ -514,6 +514,27 @@
       config:
         gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar'
 
+# Spanner CDC
+- type: renaming
+  transforms:
+    'ReadFromSpannerCDC': 'ReadFromSpannerCDC'
+  config:
+    mappings:
+      'ReadFromSpannerCDC':
+        project: 'project_id'
+        instance: 'instance_id'
+        database: 'database_id'
+        table: 'table'
+        change_stream: 'change_stream_name'
+        start_at: 'start_at_timestamp'
+        end_at: 'end_at_timestamp'
+    underlying_provider:
+      type: beamJar
+      transforms:
+        'ReadFromSpannerCDC': 'beam:schematransform:org.apache.beam:spanner_cdc_read:v1'
+      config:
+        gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar'
+
 # Firestore
 - type: renaming
   transforms: