diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfiguration.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfiguration.java new file mode 100644 index 0000000000..cb580bf6a2 --- /dev/null +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfiguration.java @@ -0,0 +1,152 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.config; + +import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; +import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper; +import com.google.gson.Gson; +import java.io.InputStream; +import java.io.Serializable; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.HashSet; +import java.util.NoSuchElementException; +import java.util.Set; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.fs.ResourceId; +import org.apache.commons.io.IOUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Configuration class for table-based filtering in Data Validation pipeline. Encapsulates parsing, + * matching, and validation of source and Spanner tables. + */ +public class TableConfiguration implements Serializable { + + private static final Logger LOG = LoggerFactory.getLogger(TableConfiguration.class); + + private final Set configuredSourceTables; + + private TableConfiguration(Set configuredSourceTables) { + this.configuredSourceTables = Collections.unmodifiableSet(configuredSourceTables); + } + + /** Creates an empty configuration with no filters. Useful for testing. */ + public static TableConfiguration empty() { + return new TableConfiguration(new HashSet<>()); + } + + /** + * Parses and validates table configuration from pipeline options. + * + * @param options The pipeline options. + * @return A TableConfiguration instance containing the configured source tables. + */ + public static TableConfiguration parseFromOptions(GCSSpannerDVOptions options) { + String tablesConfig = options.getTables(); + String tableConfigurationFilePath = options.getTableConfigurationFilePath(); + boolean hasTablesConfig = tablesConfig != null && !tablesConfig.trim().isEmpty(); + boolean hasTableConfigFile = + tableConfigurationFilePath != null && !tableConfigurationFilePath.trim().isEmpty(); + + if (hasTablesConfig && hasTableConfigFile) { + throw new IllegalArgumentException( + "Both --tables and --tableConfigurationFilePath are provided. Please configure only one of these parameters at a time."); + } + + Set configuredTables = new HashSet<>(); + + if (hasTablesConfig) { + for (String table : tablesConfig.split(",")) { + String trimmed = table.trim(); + if (!trimmed.isEmpty()) { + configuredTables.add(trimmed); + } + } + } else if (hasTableConfigFile) { + try { + ResourceId resourceId = FileSystems.matchNewResource(tableConfigurationFilePath, false); + try (InputStream stream = Channels.newInputStream(FileSystems.open(resourceId))) { + String result = IOUtils.toString(stream, StandardCharsets.UTF_8); + Gson gson = new Gson(); + TableConfigurationFile fileConfig = gson.fromJson(result, TableConfigurationFile.class); + + if (fileConfig != null && fileConfig.getTableNames() != null) { + for (String table : fileConfig.getTableNames()) { + String trimmed = table.trim(); + if (!trimmed.isEmpty()) { + configuredTables.add(trimmed); + } + } + } + } + } catch (Exception e) { + throw new RuntimeException( + "Failed to read JSON tableConfigurationFilePath: " + tableConfigurationFilePath, e); + } + } + + TableConfiguration config = new TableConfiguration(configuredTables); + + return config; + } + + public boolean hasFilters() { + return configuredSourceTables != null && !configuredSourceTables.isEmpty(); + } + + public Set getSourceTables() { + return configuredSourceTables; + } + + /** + * Checks if a source table is allowed by the configuration. + * + * @param sourceTableName The source table name. + * @return true if allowed or no filters are configured, false otherwise. + */ + public boolean isSourceTableAllowed(String sourceTableName) { + if (!hasFilters()) { + return true; + } + return configuredSourceTables.contains(sourceTableName); + } + + /** + * Checks if a Spanner table is allowed by the configuration. Translates the Spanner table name to + * its source table counterpart using the schema mapper. + * + * @param spannerTableName The Spanner table name. + * @param schemaMapper The schema mapper to translate the table name. + * @return true if allowed or no filters are configured, false otherwise. + */ + public boolean isSpannerTableAllowed(String spannerTableName, ISchemaMapper schemaMapper) { + if (!hasFilters()) { + return true; + } + try { + String sourceTable = schemaMapper.getSourceTableName("", spannerTableName); + return configuredSourceTables.contains(sourceTable); + } catch (NoSuchElementException e) { + LOG.warn( + "Could not map Spanner table '{}' back to a source table. Skipping validation.", + spannerTableName); + return false; + } + } +} diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfigurationFile.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfigurationFile.java new file mode 100644 index 0000000000..60d17faaaf --- /dev/null +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableConfigurationFile.java @@ -0,0 +1,52 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.config; + +import java.io.Serializable; +import java.util.List; +import java.util.Map; + +/** POJO representing the table configuration JSON file. */ +public class TableConfigurationFile implements Serializable { + + private final List tableNames; + + /** + * Future Extensibility: Map of Source Table Name -> Table-specific configuration. + * + *

Note: The `tableNames` list remains the absolute source of truth for the exhaustive list of + * tables to be validated. This map is strictly for providing advanced configurations (e.g., + * column filtering, sampling) for a subset of those tables. Tables cannot be implicitly included + * for validation by solely appearing in this map; they MUST be explicitly listed in `tableNames`. + * + *

This is currently a placeholder and is not yet processed by the pipeline logic. + */ + private final Map optionalConfigurations; + + public TableConfigurationFile( + List tableNames, Map optionalConfigurations) { + this.tableNames = tableNames; + this.optionalConfigurations = optionalConfigurations; + } + + public List getTableNames() { + return tableNames; + } + + public Map getOptionalConfigurations() { + return optionalConfigurations; + } +} diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableLevelConfig.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableLevelConfig.java new file mode 100644 index 0000000000..92b7199b1c --- /dev/null +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/TableLevelConfig.java @@ -0,0 +1,34 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.config; + +import java.io.Serializable; + +/** + * Placeholder POJO representing future advanced configurations for a specific table. + * + *

This is intended to support features like column-level validation or deterministic sampling in + * the future. + */ +public class TableLevelConfig implements Serializable { + + // Intentionally left empty for now. + // + // Example future fields: + // private List columnsToValidate; + // private SamplingConfig sampling; + +} diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/package-info.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/package-info.java new file mode 100644 index 0000000000..6fef2920e6 --- /dev/null +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/config/package-info.java @@ -0,0 +1,18 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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. + */ + +/** Configuration classes for Data Validation pipeline. */ +package com.google.cloud.teleport.v2.config; diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFn.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFn.java index aabc70f9f0..d8a8278d1a 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFn.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFn.java @@ -15,35 +15,48 @@ */ package com.google.cloud.teleport.v2.dofn; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.spanner.ddl.Ddl; +import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper; import java.util.List; import org.apache.beam.sdk.io.gcp.spanner.ReadOperation; import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.values.PCollectionView; public class CreateSpannerReadOpsFn extends DoFn { private final PCollectionView ddlView; + private final SerializableFunction schemaMapperProvider; + private final TableConfiguration tableConfig; - public CreateSpannerReadOpsFn(PCollectionView ddlView) { + public CreateSpannerReadOpsFn( + PCollectionView ddlView, + SerializableFunction schemaMapperProvider, + TableConfiguration tableConfig) { this.ddlView = ddlView; + this.schemaMapperProvider = schemaMapperProvider; + this.tableConfig = tableConfig; } // TODO: @aasthabharill to check if there's a better way to generalize dialect specific changes @ProcessElement public void processElement(ProcessContext c) { Ddl ddl = c.sideInput(ddlView); + ISchemaMapper schemaMapper = schemaMapperProvider.apply(ddl); List tableNames = ddl.getTablesOrderedByReference(); - tableNames.forEach( - tableName -> { - String quote = ddl.dialect() == com.google.cloud.spanner.Dialect.POSTGRESQL ? "\"" : "`"; - // We encode the tableName in the query itself to push table information dynamically - // and avoid table level stages. - String query = - String.format( - "SELECT *, '%s' as __tableName__ FROM %s%s%s", - tableName, quote, tableName, quote); - c.output(ReadOperation.create().withQuery(query)); - }); + + for (String tableName : tableNames) { + if (tableConfig != null && !tableConfig.isSpannerTableAllowed(tableName, schemaMapper)) { + continue; + } + String quote = ddl.dialect() == com.google.cloud.spanner.Dialect.POSTGRESQL ? "\"" : "`"; + // We encode the tableName in the query itself to push table information dynamically + // and avoid table level stages. + String query = + String.format( + "SELECT *, '%s' as __tableName__ FROM %s%s%s", tableName, quote, tableName, quote); + c.output(ReadOperation.create().withQuery(query)); + } } } diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/options/GCSSpannerDVOptions.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/options/GCSSpannerDVOptions.java index f1c18d0528..4a803dec28 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/options/GCSSpannerDVOptions.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/options/GCSSpannerDVOptions.java @@ -215,4 +215,25 @@ public interface GCSSpannerDVOptions extends PipelineOptions { String getTransformationCustomParameters(); void setTransformationCustomParameters(String value); + + @TemplateParameter.Text( + order = 16, + optional = true, + description = "Comma-separated list of source tables to validate", + helpText = "A comma-separated list of source tables to include in the validation run.") + @Default.String("") + String getTables(); + + void setTables(String value); + + @TemplateParameter.GcsReadFile( + order = 17, + optional = true, + description = "GCS path to a file containing a list of source tables to validate", + helpText = + "A GCS file path containing a JSON list of source tables to validate. This must be a JSON file with the structure `{\"tableNames\": [\"table1\", \"table2\"]}`.") + @Default.String("") + String getTableConfigurationFilePath(); + + void setTableConfigurationFilePath(String value); } diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java index 68cb526ba8..dd3ce817e1 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java @@ -21,6 +21,7 @@ import com.google.cloud.teleport.metadata.Template; import com.google.cloud.teleport.metadata.TemplateCategory; import com.google.cloud.teleport.v2.common.UncaughtExceptionLogger; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.dto.ComparisonRecord; import com.google.cloud.teleport.v2.fn.SchemaMapperProviderFn; import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; @@ -76,6 +77,8 @@ public static void main(String[] args) { public static PipelineResult run(GCSSpannerDVOptions options) { Pipeline pipeline = Pipeline.create(options); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + SpannerConfig spannerConfig = createSpannerConfig(options); // Fetch Spanner DDL using Info schema @@ -107,13 +110,14 @@ public static PipelineResult run(GCSSpannerDVOptions options) { options.getGcsInputDirectory(), ddlView, schemaMapperProvider, - customTransformation)); + customTransformation, + tableConfig)); // Get Spanner records hashes PCollection spannerRecords = pipeline.apply( "ReadSpannerRecords", - new SpannerReaderTransform(spannerConfig, ddlView, schemaMapperProvider)); + new SpannerReaderTransform(spannerConfig, ddlView, schemaMapperProvider, tableConfig)); PCollectionTuple inputs = PCollectionTuple.of(SOURCE_TAG, sourceRecords).and(SPANNER_TAG, spannerRecords); diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransform.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransform.java index 886b65991f..03ae880e0d 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransform.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransform.java @@ -16,13 +16,19 @@ package com.google.cloud.teleport.v2.transforms; import com.google.cloud.teleport.v2.coders.GenericRecordCoder; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.dofn.SourceHashFn; import com.google.cloud.teleport.v2.dto.ComparisonRecord; import com.google.cloud.teleport.v2.fn.IdentityGenericRecordFn; import com.google.cloud.teleport.v2.spanner.ddl.Ddl; import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper; import com.google.cloud.teleport.v2.spanner.migrations.transformation.CustomTransformation; +import java.util.ArrayList; +import java.util.List; import org.apache.beam.sdk.extensions.avro.io.AvroIO; +import org.apache.beam.sdk.io.FileIO; +import org.apache.beam.sdk.io.fs.EmptyMatchTreatment; +import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; @@ -38,36 +44,56 @@ public class SourceReaderTransform private final PCollectionView ddlView; private final SerializableFunction schemaMapperProvider; private final CustomTransformation customTransformation; + private final TableConfiguration tableConfig; public SourceReaderTransform( String gcsInputDirectory, PCollectionView ddlView, SerializableFunction schemaMapperProvider, - CustomTransformation customTransformation) { + CustomTransformation customTransformation, + TableConfiguration tableConfig) { this.gcsInputDirectory = gcsInputDirectory; this.ddlView = ddlView; this.schemaMapperProvider = schemaMapperProvider; this.customTransformation = customTransformation; + this.tableConfig = tableConfig; } @Override public @NotNull PCollection expand(PBegin input) { return input + .apply("CreateFilePatterns", Create.of(getFilePatterns(gcsInputDirectory, tableConfig))) + .apply( + "MatchFilePatterns", + FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.ALLOW)) + .apply( + "ReadMatchedFiles", + FileIO.readMatches() + .withDirectoryTreatment(FileIO.ReadMatches.DirectoryTreatment.PROHIBIT)) .apply( "ReadSourceAvroRecords", - AvroIO.parseGenericRecords(new IdentityGenericRecordFn()) - .from(createAvroFilePattern(gcsInputDirectory)) - .withCoder(GenericRecordCoder.of()) - .withHintMatchesManyFiles()) + AvroIO.parseFilesGenericRecords(new IdentityGenericRecordFn()) + .withCoder(GenericRecordCoder.of())) .apply( "CalculateSourceRecordsHash", ParDo.of(new SourceHashFn(ddlView, schemaMapperProvider, customTransformation)) .withSideInputs(ddlView)); } - private static String createAvroFilePattern(String inputPath) { + static List getFilePatterns(String gcsInputDirectory, TableConfiguration tableConfig) { + List filePatterns = new ArrayList<>(); String cleanPath = - inputPath.endsWith("/") ? inputPath.substring(0, inputPath.length() - 1) : inputPath; - return cleanPath + "/**.avro"; + gcsInputDirectory.endsWith("/") + ? gcsInputDirectory.substring(0, gcsInputDirectory.length() - 1) + : gcsInputDirectory; + + if (tableConfig == null || !tableConfig.hasFilters()) { + filePatterns.add(cleanPath + "/**.avro"); + } else { + for (String table : tableConfig.getSourceTables()) { + filePatterns.add(cleanPath + "/" + table + "/**.avro"); + } + } + return filePatterns; } } diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransform.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransform.java index 7a69c087c3..7d35df3d33 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransform.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransform.java @@ -17,6 +17,7 @@ import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.dofn.CreateSpannerReadOpsFn; import com.google.cloud.teleport.v2.dofn.SpannerHashFn; import com.google.cloud.teleport.v2.dto.ComparisonRecord; @@ -43,21 +44,26 @@ public class SpannerReaderTransform private final PCollectionView ddlView; private final SerializableFunction schemaMapperProvider; + private final TableConfiguration tableConfig; public SpannerReaderTransform( SpannerConfig spannerConfig, PCollectionView ddlView, - SerializableFunction schemaMapperProvider) { + SerializableFunction schemaMapperProvider, + TableConfiguration tableConfig) { this.spannerConfig = spannerConfig; this.ddlView = ddlView; this.schemaMapperProvider = schemaMapperProvider; + this.tableConfig = tableConfig; } @Override public @NotNull PCollection expand(PBegin p) { return p.apply("Pulse", Create.of((Void) null)) .apply( - "CreateReadOps", ParDo.of(new CreateSpannerReadOpsFn(ddlView)).withSideInputs(ddlView)) + "CreateReadOps", + ParDo.of(new CreateSpannerReadOpsFn(ddlView, schemaMapperProvider, tableConfig)) + .withSideInputs(ddlView)) .apply("ReadSpannerRecords", readFromSpanner()) .apply( "CalculateSpannerRecordsHash", diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/config/TableConfigurationTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/config/TableConfigurationTest.java new file mode 100644 index 0000000000..6f4ecbaf01 --- /dev/null +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/config/TableConfigurationTest.java @@ -0,0 +1,178 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.config; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; +import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper; +import java.io.File; +import java.io.FileWriter; +import java.io.IOException; +import java.util.NoSuchElementException; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +public class TableConfigurationTest { + + @Rule public TemporaryFolder tempFolder = new TemporaryFolder(); + + private GCSSpannerDVOptions options; + private ISchemaMapper mockSchemaMapper; + + @Before + public void setUp() { + options = PipelineOptionsFactory.create().as(GCSSpannerDVOptions.class); + mockSchemaMapper = mock(ISchemaMapper.class); + } + + @Test + public void testEmptyConfig() { + TableConfiguration config = TableConfiguration.empty(); + assertFalse(config.hasFilters()); + assertTrue(config.getSourceTables().isEmpty()); + assertTrue(config.isSourceTableAllowed("any_table")); + assertTrue(config.isSpannerTableAllowed("any_table", mockSchemaMapper)); + } + + @Test + public void testParseFromOptionsWithTables() throws IOException { + options.setTables("table1, table2,table3 "); + + File inputDir = tempFolder.newFolder("input"); + options.setGcsInputDirectory(inputDir.getAbsolutePath()); + new File(inputDir, "table1").mkdirs(); + new File(inputDir, "table1/data.avro").createNewFile(); + new File(inputDir, "table2").mkdirs(); + new File(inputDir, "table2/data.avro").createNewFile(); + new File(inputDir, "table3").mkdirs(); + new File(inputDir, "table3/data.avro").createNewFile(); + + TableConfiguration config = TableConfiguration.parseFromOptions(options); + + assertTrue(config.hasFilters()); + assertEquals(3, config.getSourceTables().size()); + assertTrue(config.getSourceTables().contains("table1")); + assertTrue(config.getSourceTables().contains("table2")); + assertTrue(config.getSourceTables().contains("table3")); + assertFalse(config.getSourceTables().contains("table4")); + } + + @Test + public void testParseFromOptionsWithTableConfigFile() throws IOException { + File tableConfigFile = tempFolder.newFile("tables.json"); + try (FileWriter writer = new FileWriter(tableConfigFile)) { + writer.write("{\"tableNames\": [\"tableA\", \" tableB \", \"\", \"tableC\"]}"); + } + options.setTableConfigurationFilePath(tableConfigFile.getAbsolutePath()); + + File inputDir = tempFolder.newFolder("input"); + options.setGcsInputDirectory(inputDir.getAbsolutePath()); + new File(inputDir, "tableA").mkdirs(); + new File(inputDir, "tableA/data.avro").createNewFile(); + new File(inputDir, "tableB").mkdirs(); + new File(inputDir, "tableB/data.avro").createNewFile(); + new File(inputDir, "tableC").mkdirs(); + new File(inputDir, "tableC/data.avro").createNewFile(); + + TableConfiguration config = TableConfiguration.parseFromOptions(options); + + assertTrue(config.hasFilters()); + assertEquals(3, config.getSourceTables().size()); + assertTrue(config.getSourceTables().contains("tableA")); + assertTrue(config.getSourceTables().contains("tableB")); + assertTrue(config.getSourceTables().contains("tableC")); + } + + @Test + public void testParseFromOptionsThrowsWhenBothProvided() { + options.setTables("table1"); + options.setTableConfigurationFilePath("gs://dummy/tables.txt"); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> TableConfiguration.parseFromOptions(options)); + assertTrue( + thrown.getMessage().contains("Please configure only one of these parameters at a time.")); + } + + @Test + public void testParseFromOptionsNoGcsInputDirectory() { + options.setTables("table1,table2"); + options.setGcsInputDirectory(null); + + TableConfiguration config = TableConfiguration.parseFromOptions(options); + assertTrue(config.hasFilters()); + assertEquals(2, config.getSourceTables().size()); + } + + @Test + public void testIsSourceTableAllowed() { + options.setTables("table1,table2"); + options.setGcsInputDirectory(null); + TableConfiguration config = TableConfiguration.parseFromOptions(options); + + assertTrue(config.isSourceTableAllowed("table1")); + assertTrue(config.isSourceTableAllowed("table2")); + assertFalse(config.isSourceTableAllowed("table3")); + } + + @Test + public void testIsSpannerTableAllowed() { + options.setTables("source_table1,source_table2"); + options.setGcsInputDirectory(null); + TableConfiguration config = TableConfiguration.parseFromOptions(options); + + when(mockSchemaMapper.getSourceTableName("", "spanner_table1")).thenReturn("source_table1"); + when(mockSchemaMapper.getSourceTableName("", "spanner_table2")).thenReturn("source_table2"); + when(mockSchemaMapper.getSourceTableName("", "spanner_table3")).thenReturn("source_table3"); + + assertTrue(config.isSpannerTableAllowed("spanner_table1", mockSchemaMapper)); + assertTrue(config.isSpannerTableAllowed("spanner_table2", mockSchemaMapper)); + assertFalse(config.isSpannerTableAllowed("spanner_table3", mockSchemaMapper)); + } + + @Test + public void testIsSpannerTableAllowedThrowsNoSuchElementException() { + options.setTables("source_table1"); + options.setGcsInputDirectory(null); + TableConfiguration config = TableConfiguration.parseFromOptions(options); + + when(mockSchemaMapper.getSourceTableName(anyString(), anyString())) + .thenThrow(new NoSuchElementException("Table not found")); + + assertFalse(config.isSpannerTableAllowed("unknown_table", mockSchemaMapper)); + } + + @Test + public void testParseFromOptionsThrowsWhenTableConfigFileFailsToRead() { + options.setTableConfigurationFilePath( + tempFolder.getRoot().getAbsolutePath() + "/non_existent_file.json"); + + RuntimeException thrown = + assertThrows(RuntimeException.class, () -> TableConfiguration.parseFromOptions(options)); + assertTrue(thrown.getMessage().contains("Failed to read JSON tableConfigurationFilePath")); + } +} diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFnTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFnTest.java index 212422cc8e..105e3079a8 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFnTest.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/dofn/CreateSpannerReadOpsFnTest.java @@ -20,9 +20,14 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.cloud.teleport.v2.config.TableConfiguration; +import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; import com.google.cloud.teleport.v2.spanner.ddl.Ddl; +import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper; +import com.google.cloud.teleport.v2.spanner.migrations.schema.IdentityMapper; import com.google.common.collect.ImmutableList; import org.apache.beam.sdk.io.gcp.spanner.ReadOperation; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.values.PCollectionView; import org.junit.Test; @@ -48,7 +53,8 @@ public void testProcessElement() { when(context.sideInput(ddlView)).thenReturn(ddl); // Create DoFn - CreateSpannerReadOpsFn doFn = new CreateSpannerReadOpsFn(ddlView); + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, IdentityMapper::new, TableConfiguration.empty()); // Execute doFn.processElement(context); @@ -81,7 +87,8 @@ public void testProcessElementPostgres() { when(context.sideInput(ddlView)).thenReturn(ddl); // Create DoFn - CreateSpannerReadOpsFn doFn = new CreateSpannerReadOpsFn(ddlView); + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, IdentityMapper::new, TableConfiguration.empty()); // Execute doFn.processElement(context); @@ -100,4 +107,127 @@ public void testProcessElementPostgres() { ReadOperation.create() .withQuery("SELECT *, 'Table2' as __tableName__ FROM \"Table2\"")); } + + @Test + public void testProcessElementWithConfiguredSubset() { + // Spanner DDL contains TableA, TableB, TableC. The config specifies TableA, TableC. + PCollectionView ddlView = mock(PCollectionView.class); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + Ddl ddl = mock(Ddl.class); + + when(ddl.dialect()).thenReturn(com.google.cloud.spanner.Dialect.GOOGLE_STANDARD_SQL); + when(ddl.getTablesOrderedByReference()) + .thenReturn(ImmutableList.of("TableA", "TableB", "TableC")); + when(context.sideInput(ddlView)).thenReturn(ddl); + + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("TableA,TableC"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, IdentityMapper::new, tableConfig); + + doFn.processElement(context); + + ArgumentCaptor argument = ArgumentCaptor.forClass(ReadOperation.class); + verify(context, times(2)).output(argument.capture()); + + // Only TableA and TableC ReadOperations are generated. TableB is skipped. + verify(context) + .output( + ReadOperation.create().withQuery("SELECT *, 'TableA' as __tableName__ FROM `TableA`")); + verify(context) + .output( + ReadOperation.create().withQuery("SELECT *, 'TableC' as __tableName__ FROM `TableC`")); + } + + @Test + public void testProcessElementWithMissingSpannerTable() { + // Configured Table Missing in Spanner: DDL contains TableA, TableB. Config specifies TableA, + // TableC. + PCollectionView ddlView = mock(PCollectionView.class); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + Ddl ddl = mock(Ddl.class); + + when(ddl.dialect()).thenReturn(com.google.cloud.spanner.Dialect.GOOGLE_STANDARD_SQL); + when(ddl.getTablesOrderedByReference()).thenReturn(ImmutableList.of("TableA", "TableB")); + when(context.sideInput(ddlView)).thenReturn(ddl); + + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("TableA,TableC"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, IdentityMapper::new, tableConfig); + + doFn.processElement(context); + + ArgumentCaptor argument = ArgumentCaptor.forClass(ReadOperation.class); + verify(context, times(1)).output(argument.capture()); + + // Only TableA is queried. TableC is naturally skipped because it's not in the DDL. + verify(context) + .output( + ReadOperation.create().withQuery("SELECT *, 'TableA' as __tableName__ FROM `TableA`")); + } + + @Test + public void testProcessElementCompleteMismatch() { + // DDL contains TableA. Config specifies TableB. + PCollectionView ddlView = mock(PCollectionView.class); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + Ddl ddl = mock(Ddl.class); + + when(ddl.dialect()).thenReturn(com.google.cloud.spanner.Dialect.GOOGLE_STANDARD_SQL); + when(ddl.getTablesOrderedByReference()).thenReturn(ImmutableList.of("TableA")); + when(context.sideInput(ddlView)).thenReturn(ddl); + + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("TableB"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, IdentityMapper::new, tableConfig); + + doFn.processElement(context); + + // Completes successfully with zero ReadOperations output. + verify(context, org.mockito.Mockito.never()).output(org.mockito.ArgumentMatchers.any()); + } + + @Test + public void testProcessElementWithSchemaMapper() { + // Table Config specifies source_table which was renamed to spanner_table in Spanner. + // SchemaMapper should successfully map spanner_table to source_table. + PCollectionView ddlView = mock(PCollectionView.class); + DoFn.ProcessContext context = mock(DoFn.ProcessContext.class); + Ddl ddl = mock(Ddl.class); + + when(ddl.dialect()).thenReturn(com.google.cloud.spanner.Dialect.GOOGLE_STANDARD_SQL); + when(ddl.getTablesOrderedByReference()).thenReturn(ImmutableList.of("spanner_table")); + when(context.sideInput(ddlView)).thenReturn(ddl); + + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("source_table"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + ISchemaMapper mockMapper = mock(ISchemaMapper.class); + when(mockMapper.getSourceTableName( + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq("spanner_table"))) + .thenReturn("source_table"); + + CreateSpannerReadOpsFn doFn = + new CreateSpannerReadOpsFn(ddlView, (d) -> mockMapper, tableConfig); + + doFn.processElement(context); + + ArgumentCaptor argument = ArgumentCaptor.forClass(ReadOperation.class); + verify(context, times(1)).output(argument.capture()); + + verify(context) + .output( + ReadOperation.create() + .withQuery("SELECT *, 'spanner_table' as __tableName__ FROM `spanner_table`")); + } } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVCoreMatchingIT.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVCoreMatchingIT.java index b15bec03e9..48eb6c6f23 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVCoreMatchingIT.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVCoreMatchingIT.java @@ -383,4 +383,109 @@ public void validationTestWithDuplicateAvroRecords() throws Exception { new MismatchedRecordDto( null, null, "AccountRoles", "[role_id:100]", "MISSING_IN_DESTINATION"))); } + + @Test + public void validationTestWithConfiguredTables() throws Exception { + + Instant t1 = Instant.parse("2024-01-01T10:00:00Z"); + + // 1. Create Source Avro records for Users and AccountRoles + GenericRecord usersRecord = + new GCSSpannerDVAvroSetupHelper.RecordBuilder( + GCSSpannerDVAvroSetupHelper.TableDef.USERS, null) + .set("user_id", 1L) + .set("event_id", "E1") + .set("full_name", "Alice") + .set("age", 30) + .set("created_at", t1) + .build(); + + GenericRecord rolesRecord = + new GCSSpannerDVAvroSetupHelper.RecordBuilder( + GCSSpannerDVAvroSetupHelper.TableDef.ACCOUNT_ROLES, null) + .set("role_id", 1) + .set("role_name", "ADMIN") + .build(); + + String gcsInputDirectory = getGcsPath("input"); + uploadAvroFileToGcs( + "input/Users/users.avro", + GCSSpannerDVAvroSetupHelper.TableDef.USERS.schema, + Arrays.asList(usersRecord)); + uploadAvroFileToGcs( + "input/AccountRoles/roles.avro", + GCSSpannerDVAvroSetupHelper.TableDef.ACCOUNT_ROLES.schema, + Arrays.asList(rolesRecord)); + + // 2. Inject Spanner Records (Destination) + spannerResourceManager.write( + Arrays.asList( + // Users: Identical record + Mutation.newInsertOrUpdateBuilder("Users_ConfiguredTables") + .set("user_id") + .to(1L) + .set("event_id") + .to("E1") + .set("full_name") + .to("Alice") + .set("age") + .to(30L) + .set("created_at") + .to(com.google.cloud.Timestamp.parseTimestamp(t1.toString())) + .build(), + // AccountRoles: Mismatched record (role_name = ADMINSTRATOR instead of ADMIN) + Mutation.newInsertOrUpdateBuilder("AccountRoles") + .set("role_id") + .to(1L) + .set("role_name") + .to("ADMINSTRATOR") + .build())); + + // Wait for Spanner's 20-second exact staleness read bound in SpannerReaderTransform + Thread.sleep(20000); + + // 3. Launch Pipeline configured to ONLY validate 'Users_ConfiguredTables' + LaunchConfig.Builder options = LaunchConfig.builder(testName, specPath); + LaunchInfo jobInfo = + launchDataflowJob( + options, + testName, + PROJECT, + spannerResourceManager, + bigQueryResourceManager.getDatasetId(), + gcsInputDirectory, + null, + null, + "[{Users, Users_ConfiguredTables}]", // table overrides + null, + null, + java.util.Map.of("tables", "Users")); // Table mapping to validate only Users + + pipelineOperator().waitUntilDone(createConfig(jobInfo)); + + // 4. Assert BigQuery Validation Results + // Note: If table filtering wasn't working, the result would have been MISMATCHED + // due to the discrepancy in the AccountRoles table. Since it's filtered, we expect a MATCH. + GCSSpannerDVTestAsserts.assertValidationSummary( + bigQueryResourceManager, + Arrays.asList( + new ValidationSummaryDto( + /* status= */ "MATCH", + /* totalTablesValidated= */ 1L, // Only Users_ConfiguredTables is validated + /* totalRowsMatched= */ 1L, + /* totalRowsMismatched= */ 0L, + /* tablesWithMismatches= */ ""))); + + GCSSpannerDVTestAsserts.assertTableValidationStats( + bigQueryResourceManager, + Arrays.asList( + new TableValidationStatsDto( + /* schemaName= */ null, + /* tableName= */ "Users_ConfiguredTables", + /* status= */ "MATCH", + /* sourceRowCount= */ 1L, + /* destinationRowCount= */ 1L, + /* matchedRowCount= */ 1L, + /* mismatchRowCount= */ 0L))); + } } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVTest.java index 013bff2160..f4259cd4d4 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVTest.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVTest.java @@ -16,9 +16,12 @@ package com.google.cloud.teleport.v2.templates; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -26,18 +29,40 @@ @RunWith(JUnit4.class) public class GCSSpannerDVTest { + private GCSSpannerDVOptions options; + + @Before + public void setUp() { + options = PipelineOptionsFactory.create().as(GCSSpannerDVOptions.class); + // Set required options to bypass early validation (if any) + options.setGcsInputDirectory("gs://dummy/input"); + options.setProjectId("test-project"); + options.setInstanceId("test-instance"); + options.setDatabaseId("test-database"); + options.setBigQueryDataset("test_dataset"); + } + @Test public void testCreateSpannerConfig() { - String[] args = - new String[] { - "--projectId=test-project", - "--instanceId=test-instance", - "--databaseId=test-database", - "--bigQueryDataset=test-dataset", - "--gcsInputDirectory=test-directory" - }; - GCSSpannerDVOptions options = - PipelineOptionsFactory.fromArgs(args).withValidation().as(GCSSpannerDVOptions.class); assertNotNull(GCSSpannerDV.createSpannerConfig(options)); } + + @Test + public void testRunThrowsExceptionWhenBothTableConfigsProvided() { + options.setTables("table1,table2"); + options.setTableConfigurationFilePath("gs://dummy/tables.json"); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> GCSSpannerDV.run(options)); + + assertTrue( + thrown.getMessage().contains("Please configure only one of these parameters at a time")); + } + + @Test + public void testRunThrowsExceptionWhenTableConfigurationFileFailsToRead() { + options.setTableConfigurationFilePath("non_existent_file.json"); + RuntimeException thrown = assertThrows(RuntimeException.class, () -> GCSSpannerDV.run(options)); + assertTrue(thrown.getMessage().contains("Failed to read JSON tableConfigurationFilePath")); + } } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransformTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransformTest.java index b427f3d9bf..cef9f2ea2a 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransformTest.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SourceReaderTransformTest.java @@ -18,7 +18,9 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.dto.ComparisonRecord; +import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; import com.google.cloud.teleport.v2.spanner.ddl.Ddl; import com.google.cloud.teleport.v2.spanner.migrations.schema.IdentityMapper; import java.io.File; @@ -31,6 +33,7 @@ import org.apache.avro.generic.GenericDatumWriter; import org.apache.avro.generic.GenericRecord; import org.apache.avro.io.DatumWriter; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; @@ -80,7 +83,8 @@ public void testReadAndMapAvroRecords() throws IOException { // FileIO in beam support a variety of paths dynamically, such as GCS, S3 and TempFolder // This allows us to pass a tempFolder into the same transform that accepts a GCS path SourceReaderTransform transform = - new SourceReaderTransform(inputPath, ddlView, IdentityMapper::new, null); + new SourceReaderTransform( + inputPath, ddlView, IdentityMapper::new, null, TableConfiguration.empty()); PCollection output = pipeline.apply(transform); @@ -123,14 +127,15 @@ public void testReadWithNoMatchingFiles() { // 2. Run Pipeline with input path that has no avro files String inputPath = tempFolder.getRoot().getAbsolutePath(); SourceReaderTransform transform = - new SourceReaderTransform(inputPath, ddlView, IdentityMapper::new, null); - - pipeline.apply(transform); + new SourceReaderTransform( + inputPath, ddlView, IdentityMapper::new, null, TableConfiguration.empty()); + PCollection output = pipeline.apply(transform); // AvroIO throws a RuntimeException when no files are found matching the pattern - // if withHintMatchesManyFiles is used (which uses match() internally). - RuntimeException e = assertThrows(RuntimeException.class, () -> pipeline.run()); - assertTrue(e.getMessage().contains("No files matched spec")); + // AvroIO.parseAllGenericRecords does not throw when it matches 0 files, it emits 0 elements. + org.apache.beam.sdk.testing.PAssert.that(output).empty(); + + pipeline.run(); } @Test @@ -162,7 +167,8 @@ public void testInvalidTable() throws IOException { // 3. Run Pipeline String inputPath = tempFolder.getRoot().getAbsolutePath(); SourceReaderTransform transform = - new SourceReaderTransform(inputPath, ddlView, IdentityMapper::new, null); + new SourceReaderTransform( + inputPath, ddlView, IdentityMapper::new, null, TableConfiguration.empty()); pipeline.apply(transform); @@ -204,7 +210,8 @@ public void testReadRecursively() throws IOException { // 3. Run Pipeline String inputPath = tempFolder.getRoot().getAbsolutePath(); SourceReaderTransform transform = - new SourceReaderTransform(inputPath, ddlView, IdentityMapper::new, null); + new SourceReaderTransform( + inputPath, ddlView, IdentityMapper::new, null, TableConfiguration.empty()); PCollection output = pipeline.apply(transform); @@ -228,6 +235,78 @@ public void testReadRecursively() throws IOException { pipeline.run(); } + @Test + public void testReadWithTableConfigFiltersTables() throws IOException { + // 1. Setup Ddl + Ddl ddl = + Ddl.builder() + .createTable("AllowedTable") + .column("id") + .int64() + .notNull() + .endColumn() + .column("name") + .string() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .createTable("SkippedTable") + .column("id") + .int64() + .notNull() + .endColumn() + .column("name") + .string() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .build(); + + PCollectionView ddlView = + pipeline.apply("CreateDDL", Create.of(ddl)).apply(View.asSingleton()); + + // 2. Create Avro files for both tables in separate directories + File allowedDir = tempFolder.newFolder("AllowedTable"); + createAvroFile(new File(allowedDir, "data.avro"), "AllowedTable", "1"); + File skippedDir = tempFolder.newFolder("SkippedTable"); + createAvroFile(new File(skippedDir, "data.avro"), "SkippedTable", "2"); + + // 3. Configure TableConfiguration to only allow "AllowedTable" + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("AllowedTable"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + // 4. Run Pipeline + String inputPath = tempFolder.getRoot().getAbsolutePath(); + SourceReaderTransform transform = + new SourceReaderTransform(inputPath, ddlView, IdentityMapper::new, null, tableConfig); + + PCollection output = pipeline.apply(transform); + + // 5. Verify only AllowedTable was read + PAssert.that(output) + .satisfies( + records -> { + int count = 0; + for (ComparisonRecord rec : records) { + count++; + if (!rec.getTableName().equals("AllowedTable")) { + throw new AssertionError("Expected AllowedTable, got " + rec.getTableName()); + } + } + if (count != 1) { + throw new AssertionError("Expected exactly 1 record, got " + count); + } + return null; + }); + + pipeline.run(); + } + private void createAvroFile(File file, String tableName, String id) throws IOException { Schema payloadSchema = SchemaBuilder.record("Payload") @@ -262,4 +341,34 @@ private void createAvroFile(File file, String tableName, String id) throws IOExc dataFileWriter.append(record); } } + + @Test + public void testGetFilePatternsNullConfig() { + java.util.List patterns = + SourceReaderTransform.getFilePatterns("gs://my-bucket/dir", null); + org.junit.Assert.assertEquals(1, patterns.size()); + org.junit.Assert.assertEquals("gs://my-bucket/dir/**.avro", patterns.get(0)); + } + + @Test + public void testGetFilePatternsEmptyConfig() { + java.util.List patterns = + SourceReaderTransform.getFilePatterns("gs://my-bucket/dir/", TableConfiguration.empty()); + org.junit.Assert.assertEquals(1, patterns.size()); + // Also tests that trailing slash is handled correctly + org.junit.Assert.assertEquals("gs://my-bucket/dir/**.avro", patterns.get(0)); + } + + @Test + public void testGetFilePatternsWithTables() { + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("Table1,Table2"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + java.util.List patterns = + SourceReaderTransform.getFilePatterns("gs://my-bucket/dir", tableConfig); + org.junit.Assert.assertEquals(2, patterns.size()); + org.junit.Assert.assertTrue(patterns.contains("gs://my-bucket/dir/Table1/**.avro")); + org.junit.Assert.assertTrue(patterns.contains("gs://my-bucket/dir/Table2/**.avro")); + } } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransformTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransformTest.java index c2e7cdb5b3..40ceededf0 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransformTest.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/SpannerReaderTransformTest.java @@ -17,14 +17,18 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; import com.google.cloud.spanner.Struct; +import com.google.cloud.teleport.v2.config.TableConfiguration; import com.google.cloud.teleport.v2.dto.ComparisonRecord; +import com.google.cloud.teleport.v2.options.GCSSpannerDVOptions; import com.google.cloud.teleport.v2.spanner.ddl.Ddl; import com.google.cloud.teleport.v2.spanner.migrations.schema.IdentityMapper; import java.io.Serializable; import org.apache.beam.sdk.io.gcp.spanner.ReadOperation; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; @@ -88,7 +92,8 @@ public void testReadAndMapRecords() { // 3. Create Transform with overridden readFromSpanner SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); SpannerReaderTransform transform = - new SpannerReaderTransform(spannerConfig, ddlView, IdentityMapper::new) { + new SpannerReaderTransform( + spannerConfig, ddlView, IdentityMapper::new, TableConfiguration.empty()) { @Override protected PTransform, PCollection> readFromSpanner() { return new PTransform, PCollection>() { @@ -132,7 +137,8 @@ public void testReadWithEmptyDdl() { // 2. Create Transform with overridden readFromSpanner SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); SpannerReaderTransform transform = - new SpannerReaderTransform(spannerConfig, ddlView, IdentityMapper::new) { + new SpannerReaderTransform( + spannerConfig, ddlView, IdentityMapper::new, TableConfiguration.empty()) { @Override protected PTransform<@NotNull PCollection, @NotNull PCollection> readFromSpanner() { @@ -197,7 +203,8 @@ public void testReadWithNullFields() { // 3. Create Transform SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); SpannerReaderTransform transform = - new SpannerReaderTransform(spannerConfig, ddlView, IdentityMapper::new) { + new SpannerReaderTransform( + spannerConfig, ddlView, IdentityMapper::new, TableConfiguration.empty()) { @Override protected PTransform<@NotNull PCollection, @NotNull PCollection> readFromSpanner() { @@ -235,9 +242,176 @@ public void testOriginalReadFromSpanner() { SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); SpannerReaderTransform transform = - new SpannerReaderTransform(spannerConfig, ddlView, IdentityMapper::new); + new SpannerReaderTransform( + spannerConfig, ddlView, IdentityMapper::new, TableConfiguration.empty()); assertNotNull(transform.readFromSpanner()); pipeline.run(); } + + @Test + public void testReadWithTableConfigFiltersTables() { + // 1. Setup Ddl with two tables + Ddl ddl = + Ddl.builder() + .createTable("AllowedTable") + .column("id") + .int64() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .createTable("SkippedTable") + .column("id") + .int64() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .build(); + + PCollectionView ddlView = + pipeline.apply("CreateDDL", Create.of(ddl)).apply(View.asSingleton()); + + // 2. Setup TableConfiguration with only one table + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("AllowedTable"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + // 3. Create Transform with overridden readFromSpanner to intercept and assert ReadOperations + SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); + SpannerReaderTransform transform = + new SpannerReaderTransform(spannerConfig, ddlView, IdentityMapper::new, tableConfig) { + @Override + protected PTransform<@NotNull PCollection, @NotNull PCollection> + readFromSpanner() { + return new PTransform<>() { + @Override + public @NotNull PCollection expand( + @NotNull PCollection input) { + // Assert that the pipeline only generated a ReadOperation for "AllowedTable" + PAssert.that(input) + .satisfies( + ops -> { + int count = 0; + for (ReadOperation op : ops) { + count++; + assertTrue( + "Expected ReadOperation for AllowedTable but got: " + + op.getQuery().getSql(), + op.getQuery().getSql().contains("AllowedTable")); + } + assertEquals(1, count); + return null; + }); + + // Return an empty PCollection of Structs to safely complete the pipeline + return input + .getPipeline() + .apply( + "MockEmptyRead", + Create.empty(org.apache.beam.sdk.values.TypeDescriptor.of(Struct.class))); + } + }; + } + }; + + // 4. Run Pipeline (PAssert runs during pipeline execution) + pipeline.apply(transform); + pipeline.run(); + } + + @Test + public void testReadWithTableConfigAndSchemaMapperFiltersTables() { + // 1. Setup Ddl with two tables (using Spanner names) + Ddl ddl = + Ddl.builder() + .createTable("spanner_mapped_table") + .column("id") + .int64() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .createTable("skipped_table") + .column("id") + .int64() + .notNull() + .endColumn() + .primaryKey() + .asc("id") + .end() + .endTable() + .build(); + + PCollectionView ddlView = + pipeline.apply("CreateDDL", Create.of(ddl)).apply(View.asSingleton()); + + // 2. Setup TableConfiguration with the Source name + GCSSpannerDVOptions options = PipelineOptionsFactory.as(GCSSpannerDVOptions.class); + options.setTables("source_mapped_table"); + TableConfiguration tableConfig = TableConfiguration.parseFromOptions(options); + + // 3. Create a Serializable SchemaMapper stub to translate spanner_mapped_table -> + // source_mapped_table + IdentityMapper stubMapper = + new IdentityMapper(ddl) { + @Override + public String getSourceTableName(String namespace, String spannerTableName) { + if ("spanner_mapped_table".equals(spannerTableName)) { + return "source_mapped_table"; + } + return super.getSourceTableName(namespace, spannerTableName); + } + }; + + // 4. Create Transform with overridden readFromSpanner + SpannerConfig spannerConfig = SpannerConfig.create().withProjectId("test-project"); + SpannerReaderTransform transform = + new SpannerReaderTransform(spannerConfig, ddlView, (d) -> stubMapper, tableConfig) { + @Override + protected PTransform<@NotNull PCollection, @NotNull PCollection> + readFromSpanner() { + return new PTransform<>() { + @Override + public @NotNull PCollection expand( + @NotNull PCollection input) { + // Assert that the pipeline correctly translated the spanner name and generated one + // ReadOperation + PAssert.that(input) + .satisfies( + ops -> { + int count = 0; + for (ReadOperation op : ops) { + count++; + assertTrue( + "Expected ReadOperation for spanner_mapped_table but got: " + + op.getQuery().getSql(), + op.getQuery().getSql().contains("spanner_mapped_table")); + } + assertEquals(1, count); + return null; + }); + + // Return an empty PCollection of Structs + return input + .getPipeline() + .apply( + "MockEmptyRead2", + Create.empty(org.apache.beam.sdk.values.TypeDescriptor.of(Struct.class))); + } + }; + } + }; + + // 5. Run Pipeline + pipeline.apply(transform); + pipeline.run(); + } } diff --git a/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVCoreMatchingIT/spanner-schema.sql b/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVCoreMatchingIT/spanner-schema.sql index 815ce495c3..2a85c67632 100644 --- a/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVCoreMatchingIT/spanner-schema.sql +++ b/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVCoreMatchingIT/spanner-schema.sql @@ -10,3 +10,11 @@ CREATE TABLE AccountRoles ( role_id INT64 NOT NULL, role_name STRING(MAX) ) PRIMARY KEY (role_id); + +CREATE TABLE Users_ConfiguredTables ( + user_id INT64 NOT NULL, + event_id STRING(MAX) NOT NULL, + full_name STRING(MAX), + age INT64, + created_at TIMESTAMP +) PRIMARY KEY (user_id, event_id); \ No newline at end of file diff --git a/v2/gcs-spanner-dv/terraform/Avro_to_Spanner_Data_Validator/dataflow_job.tf b/v2/gcs-spanner-dv/terraform/Avro_to_Spanner_Data_Validator/dataflow_job.tf index 274fd3af7d..39da403d64 100644 --- a/v2/gcs-spanner-dv/terraform/Avro_to_Spanner_Data_Validator/dataflow_job.tf +++ b/v2/gcs-spanner-dv/terraform/Avro_to_Spanner_Data_Validator/dataflow_job.tf @@ -123,6 +123,20 @@ variable "transformationCustomParameters" { default = null } +variable "tables" { + type = string + description = "A comma-separated list of source tables to include in the validation run. Defaults to empty." + default = null +} + +variable "tableConfigurationFilePath" { + type = string + description = <