diff --git a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json index 0942404b0ebf..540e8267eaf7 100644 --- a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "revision": 3, + "revision": 5, "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner" } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java index 7a628f244408..c1ee4f8db138 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadSpannerSchema.java +++ b/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,25 @@ 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. + * + *

Used by SchemaTransforms during expansion (including cross-language expansion services that + * do not ship DirectRunner). + */ + public static SpannerSchema getSpannerSchema( + SpannerConfig config, Dialect dialect, Set 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 allowedTableNames) { + Set 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 +105,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 +116,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 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 +208,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: diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangestreamsReadSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangestreamsReadSchemaTransformProvider.java index 26d6e757cb81..8f74f1176724 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/SpannerChangestreamsReadSchemaTransformProvider.java +++ b/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 TABLE_SCHEMAS = new HashMap<>(); - private static Schema getTableSchema(SpannerChangestreamsReadConfiguration config) { - Pipeline miniPipeline = Pipeline.create(); - PCollectionView 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() { - @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()); } diff --git a/sdks/python/apache_beam/yaml/extended_tests/databases/spanner_cdc.yaml b/sdks/python/apache_beam/yaml/extended_tests/databases/spanner_cdc.yaml new file mode 100644 index 000000000000..b6bb227ba7c0 --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/databases/spanner_cdc.yaml @@ -0,0 +1,64 @@ +# +# 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. +# + +fixtures: + - name: SPANNER_CDC + type: "apache_beam.yaml.integration_tests.temp_spanner_change_stream" + config: + project: "apache-beam-testing" + instance: "beam-test" + num_rows: 2 + +pipelines: + - name: read_spanner_cdc_bounded + pipeline: + type: composite + transforms: + - type: ReadFromSpannerCDC + name: ReadCdc + config: + project: "{SPANNER_CDC[PROJECT]}" + instance: "{SPANNER_CDC[INSTANCE]}" + database: "{SPANNER_CDC[DATABASE]}" + table: "{SPANNER_CDC[TABLE]}" + change_stream: "{SPANNER_CDC[CHANGE_STREAM]}" + start_at: "{SPANNER_CDC[START_AT]}" + end_at: "{SPANNER_CDC[END_AT]}" + - type: MapToFields + name: SelectOperation + input: ReadCdc + config: + language: python + fields: + operation: operation + - type: AssertEqual + input: SelectOperation + config: + elements: + - {operation: INSERT} + - {operation: INSERT} + - {operation: UPDATE} + - {operation: UPDATE} + - {operation: DELETE} + - {operation: DELETE} + - type: AssertEqual + input: ReadCdc.errors + config: + elements: [] + options: + project: "apache-beam-testing" + streaming: true diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 32794f4588d2..3e0f51e635dc 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -165,6 +165,104 @@ def temp_spanner_table(project, prefix='temp_spanner_db_'): spanner_client._delete_database() +def _spanner_commit_timestamp_str(commit_ts): + """Format a Spanner commit timestamp for SchemaTransform config fields.""" + text = commit_ts.isoformat() + if text.endswith('+00:00'): + return text[:-6] + 'Z' + return text + + +@contextlib.contextmanager +def temp_spanner_change_stream( + project='apache-beam-testing', instance='beam-test', num_rows=2): + """Context manager for a Spanner table with a change stream and seeded mutations. + + Creates a temporary database, table, and change stream, then inserts, updates, + and deletes ``num_rows`` rows. Yields identifiers and inclusive ``start_at`` / + ``end_at`` timestamps suitable for a bounded ``ReadFromSpannerCDC`` YAML IT. + + Args: + project (str): GCP project id. + instance (str): Spanner instance id. + num_rows (int): Number of rows to insert/update/delete. + + Yields: + dict: Keys ``PROJECT``, ``INSTANCE``, ``DATABASE``, ``TABLE``, + ``CHANGE_STREAM``, ``START_AT``, and ``END_AT``. + """ + from google.cloud import spanner + + client = spanner.Client(project=project) + spanner_instance = client.instance(instance) + database_id = f'temp-{uuid.uuid4().hex[:10]}' + # Unquoted GSQL identifiers must start with a letter. + table = f'T{uuid.uuid4().hex[:10]}' + change_stream = f'C{uuid.uuid4().hex[:10]}' + + database = spanner_instance.database( + database_id, + ddl_statements=[ + f'''CREATE TABLE {table} ( + SingerId INT64 NOT NULL, + FirstName STRING(1024), + LastName STRING(1024) + ) PRIMARY KEY (SingerId)''' + ]) + logging.info( + 'Creating Spanner CDC test database %s table %s', database_id, table) + database.create().result(timeout=600) + + logging.info('Creating change stream %s for table %s', change_stream, table) + database.update_ddl([f'CREATE CHANGE STREAM {change_stream} FOR {table}' + ]).result(timeout=600) + + def _commit_mutation(build_batch_fn): + # database.batch() returns BatchCheckout; mutations live on the Batch + # entered via the context manager. Commit timestamp is batch.committed. + with database.batch() as batch: + build_batch_fn(batch) + return batch.committed + + start_at = None + for singer_id in range(1, num_rows + 1): + commit_ts = _commit_mutation( + lambda batch, singer_id=singer_id: batch.insert( + table=table, columns=('SingerId', 'FirstName', 'LastName'), values=[ + (singer_id, f'First {singer_id}', f'Last {singer_id}'), ])) + if start_at is None: + start_at = commit_ts + + for singer_id in range(1, num_rows + 1): + _commit_mutation( + lambda batch, singer_id=singer_id: batch.update( + table=table, columns=('SingerId', 'FirstName', 'LastName'), values=[ + ( + singer_id, + f'Updated First {singer_id}', + f'Updated Last {singer_id}', ), ])) + + end_at = None + for singer_id in range(1, num_rows + 1): + end_at = _commit_mutation( + lambda batch, singer_id=singer_id: batch.delete( + table=table, keyset=spanner.KeySet(keys=[[singer_id]]))) + + try: + yield { + 'PROJECT': project, + 'INSTANCE': instance, + 'DATABASE': database_id, + 'TABLE': table, + 'CHANGE_STREAM': change_stream, + 'START_AT': _spanner_commit_timestamp_str(start_at), + 'END_AT': _spanner_commit_timestamp_str(end_at), + } + finally: + logging.info('Deleting Spanner CDC test database: %s', database_id) + database.drop() + + @contextlib.contextmanager def temp_bigquery_table(project, prefix='yaml_bq_it_'): """Context manager to create and clean up a temporary BigQuery dataset. diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index 253573c8069a..c6cdc0abf647 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/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' + # TFRecord - type: renaming transforms: