From b60077d5807f75d155d89f81fd5718b16d916e8b Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 25 May 2026 12:32:21 +0530 Subject: [PATCH 01/16] feat(Spanner): integrate SourceConfigParser to centralize shard configuration loading for SourceDbToSpanner pipelines --- .../v2/options/OptionsToConfigBuilder.java | 25 ++++------ .../v2/templates/PipelineController.java | 24 ++++++++++ .../v2/templates/SourceDbToSpanner.java | 6 +++ .../options/OptionsToConfigBuilderTest.java | 46 ++++++------------- 4 files changed, 54 insertions(+), 47 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java index 2271aeafc9..2d7df88499 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java @@ -21,6 +21,7 @@ import com.google.cloud.teleport.v2.reader.io.schema.SourceSchemaReference; import com.google.cloud.teleport.v2.source.SourceConnectorFactory; import com.google.cloud.teleport.v2.source.jdbc.AbstractJdbcSrcToSpSourceConnector; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import com.google.cloud.teleport.v2.spanner.migrations.utils.DataflowWorkerMachineTypeUtils; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -58,15 +59,11 @@ public static String extractWorkerZone(PipelineOptions options) { public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( SourceDbToSpannerOptions options, + Shard shard, List tables, String shardId, Wait.OnSignal waitOn) { SQLDialect sqlDialect = SQLDialect.valueOf(options.getSourceDbDialect()); - String sourceDbURL = options.getSourceConfigURL(); - String dbName = extractDbFromURL(sourceDbURL); - String username = options.getUsername(); - String password = options.getPassword(); - String namespace = options.getNamespace(); String jdbcDriverClassName = options.getJdbcDriverClassName(); String jdbcDriverJars = options.getJdbcDriverJars(); @@ -83,14 +80,13 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( return getJdbcIOWrapperConfig( sqlDialect, tables, - sourceDbURL, - null, - null, - 0, - username, - password, - dbName, - namespace, + shard.getHost(), + shard.getConnectionProperties(), + Integer.parseInt(shard.getPort()), + shard.getUserName(), + shard.getPassword(), + shard.getDbName(), + shard.getNamespace(), shardId, jdbcDriverClassName, jdbcDriverJars, @@ -107,7 +103,6 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( SQLDialect sqlDialect, List tables, - String sourceDbURL, String host, String connectionProperties, int port, @@ -154,7 +149,7 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( builder = builder.setMaxConnections(maxConnections); } - sourceDbURL = + String sourceDbURL = connector.getJdbcUrl( sourceDbURL, host, port, dbName, connectionProperties, namespace, fetchSize); diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java index f1327c0e0f..3492654bc2 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java @@ -24,8 +24,16 @@ import com.google.cloud.teleport.v2.spanner.migrations.schema.SchemaFileOverridesBasedMapper; import com.google.cloud.teleport.v2.spanner.migrations.schema.SchemaStringOverridesBasedMapper; import com.google.cloud.teleport.v2.spanner.migrations.schema.SessionBasedMapper; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConfigParser; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.spanner.SpannerSchema; +import com.google.cloud.teleport.v2.spanner.migrations.utils.ISecretManagerAccessor; +import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -215,4 +223,20 @@ static ISchemaMapper getSchemaMapper(SourceDbToSpannerOptions options, Ddl ddl) } return schemaMapper; } + + public static SourceConnectionConfig getSourceConnectionConfig( + String sourceType, String sourceShardsFilePath) { + ISecretManagerAccessor secretManagerAccessor = new SecretManagerAccessorImpl(); + SourceConfigParser sourceConfigParser = new SourceConfigParser(secretManagerAccessor); + SourceConnectionConfig sourceConnectionConfig; + try { + // Parse the source shards configuration file to respective + // SourceConnectionConfig. + LOG.info("Parsing source shards configuration file: {}", sourceShardsFilePath); + return sourceConfigParser.parseConfiguration(sourceType, sourceShardsFilePath); + } catch (Exception e) { + LOG.error("Error parsing source config", e); + throw new RuntimeException("Error parsing source config", e); + } + } } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index 8b047f74f3..701ea7fc6c 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -22,8 +22,11 @@ import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.source.ISrcToSpSourceConnector; import com.google.cloud.teleport.v2.source.SourceConnectorFactory; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.utils.DataflowWorkerMachineTypeUtils; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; @@ -121,6 +124,9 @@ static PipelineResult run(SourceDbToSpannerOptions options) { DataflowWorkerMachineTypeUtils.validateMachineSpecs(workerMachineType, 4); SpannerConfig spannerConfig = createSpannerConfig(options); + SourceConnectionConfig sourceConnectionConfig = + PipelineController.getSourceConnectionConfig( + options.getSourceDbDialect(), options.getSourceConfigURL()); // Decide type and source of migration ISrcToSpSourceConnector connector = SourceConnectorFactory.getSourceConnectorByDialect(options); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java index ab30ffc6af..b74404cf03 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java @@ -20,6 +20,7 @@ import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.JdbcIOWrapperConfig; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.List; @@ -44,37 +45,38 @@ public class OptionsToConfigBuilderTest { public void testConfigWithMySqlDefaultsFromOptions() { final String testDriverClassName = "org.apache.derby.jdbc.EmbeddedDriver"; final String testUrl = "jdbc:mysql://localhost:3306/testDB"; - final String testUser = "user"; - final String testPassword = "password"; SourceDbToSpannerOptions sourceDbToSpannerOptions = PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); sourceDbToSpannerOptions.setSourceDbDialect(SQLDialect.MYSQL.name()); - sourceDbToSpannerOptions.setSourceConfigURL(testUrl); sourceDbToSpannerOptions.setJdbcDriverClassName(testDriverClassName); sourceDbToSpannerOptions.setMaxConnections(150); sourceDbToSpannerOptions.setNumPartitions(4000); - sourceDbToSpannerOptions.setUsername(testUser); - sourceDbToSpannerOptions.setPassword(testPassword); sourceDbToSpannerOptions.setTables("table1,table2"); + Shard shard = new Shard("", "localhost", "3306", "user", "password", "testDB", "", "", ""); PCollection dummyPCollection = pipeline.apply(Create.of(1)); pipeline.run(); JdbcIOWrapperConfig config = OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - sourceDbToSpannerOptions, List.of("table1", "table2"), null, Wait.on(dummyPCollection)); + sourceDbToSpannerOptions, + shard, + List.of("table1", "table2"), + null, + Wait.on(dummyPCollection)); assertThat(config.jdbcDriverClassName()).isEqualTo(testDriverClassName); assertThat(config.sourceDbURL()) .isEqualTo( testUrl + "?allowMultiQueries=true&autoReconnect=true&maxReconnects=10&useCursorFetch=true"); assertThat(config.tables()).containsExactlyElementsIn(new String[] {"table1", "table2"}); - assertThat(config.dbAuth().getUserName().get()).isEqualTo(testUser); - assertThat(config.dbAuth().getPassword().get()).isEqualTo(testPassword); + assertThat(config.dbAuth().getUserName().get()).isEqualTo(shard.getUserName()); + assertThat(config.dbAuth().getPassword().get()).isEqualTo(shard.getPassword()); assertThat(config.waitOn()).isNotNull(); assertThat(config.maxFetchSize()).isNull(); sourceDbToSpannerOptions.setFetchSize(42); assertThat( OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( sourceDbToSpannerOptions, + shard, List.of("table1", "table2"), null, Wait.on(dummyPCollection)) @@ -90,7 +92,6 @@ public void testConfigWithMySqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.MYSQL, List.of("table1", "table2"), - null, "myhost", "testParam=testValue", 3306, @@ -114,7 +115,6 @@ public void testConfigWithMySqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.MYSQL, List.of("table1", "table2"), - null, "myhost", null, 3306, @@ -151,18 +151,17 @@ public void testConfigWithPostgreSQLDefaultsFromOptions() { SourceDbToSpannerOptions sourceDbToSpannerOptions = PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); sourceDbToSpannerOptions.setSourceDbDialect(SQLDialect.POSTGRESQL.name()); - sourceDbToSpannerOptions.setSourceConfigURL(testUrl); sourceDbToSpannerOptions.setJdbcDriverClassName(testDriverClassName); sourceDbToSpannerOptions.setMaxConnections(150); sourceDbToSpannerOptions.setNumPartitions(4000); - sourceDbToSpannerOptions.setUsername(testUser); - sourceDbToSpannerOptions.setPassword(testPassword); sourceDbToSpannerOptions.setTables("table1,table2,table3"); + Shard shard = new Shard("", "localhost", "5432", "user", "password", "testDB", "", "", ""); PCollection dummyPCollection = pipeline.apply(Create.of(1)); pipeline.run(); JdbcIOWrapperConfig config = OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( sourceDbToSpannerOptions, + shard, List.of("table1", "table2", "table3"), null, Wait.on(dummyPCollection)); @@ -183,7 +182,6 @@ public void testConfigWithPostgreSqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - null, "myhost", "testParam=testValue", 5432, @@ -206,7 +204,6 @@ public void testConfigWithPostgreSqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - null, "myhost", "", 5432, @@ -241,7 +238,6 @@ public void testConfigWithPostgreSqlUrlWithNamespace() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - null, "myhost", "", 5432, @@ -264,20 +260,6 @@ public void testConfigWithPostgreSqlUrlWithNamespace() { .isEqualTo("jdbc:postgresql://myhost:5432/mydb?currentSchema=mynamespace"); } - @Test - public void testURIParsingException() { - final String testUrl = "jd#bc://localhost"; - SourceDbToSpannerOptions sourceDbToSpannerOptions = - PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); - sourceDbToSpannerOptions.setSourceDbDialect(SQLDialect.MYSQL.name()); - sourceDbToSpannerOptions.setSourceConfigURL(testUrl); - assertThrows( - RuntimeException.class, - () -> - OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - sourceDbToSpannerOptions, new ArrayList<>(), null, null)); - } - @Test public void testaddParamToJdbcUrl() throws URISyntaxException { // No Parameters initially. @@ -340,13 +322,13 @@ public void testExtractWorkerZoneException() { public void testFetchSizeMinusOneBehavesLikeNull() { SourceDbToSpannerOptions options = PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); options.setSourceDbDialect(SQLDialect.MYSQL.name()); - options.setSourceConfigURL("jdbc:mysql://localhost:3306/testDB"); options.setJdbcDriverClassName("com.mysql.jdbc.Driver"); options.setFetchSize(-1); // Should be normalized to null + Shard shard = new Shard("", "localhost", "5432", "user", "password", "testDB", "", "", ""); JdbcIOWrapperConfig config = OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - options, List.of("table1"), null, null); + options, shard, List.of("table1"), null, null); assertThat(config.maxFetchSize()).isNull(); } From 21872c2a123c2182a07eed1ca82c8f6e25269977 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 22 Jun 2026 07:11:07 +0000 Subject: [PATCH 02/16] test(Spanner): Integration test fix. --- .../cloudsql/CloudSqlShardOrchestrator.java | 30 +++---- .../v2/options/SourceDbToSpannerOptions.java | 3 +- .../v2/templates/PipelineController.java | 1 - .../v2/templates/SourceDbToSpannerITBase.java | 71 +++++++++++++++- .../SourceDbToSpannerFTBase.java | 80 +++++++++---------- .../loadtesting/SourceDbToSpannerLTBase.java | 43 ++++++++-- 6 files changed, 153 insertions(+), 75 deletions(-) diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/cloudsql/CloudSqlShardOrchestrator.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/cloudsql/CloudSqlShardOrchestrator.java index 6e1704f104..519d32ea96 100644 --- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/cloudsql/CloudSqlShardOrchestrator.java +++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/cloudsql/CloudSqlShardOrchestrator.java @@ -420,9 +420,7 @@ protected void createLogicalDatabases() { protected String generateAndUploadConfig(String artifactName) { LOG.info("Generating and uploading shard configuration..."); JSONObject config = new JSONObject(); - config.put("configType", "dataflow"); - JSONObject shardConfigBulk = new JSONObject(); - JSONArray dataShards = new JSONArray(); + JSONArray shardConfigs = new JSONArray(); int shardIdx = 0; for (Map.Entry> entry : requestedShardMap.entrySet()) { @@ -430,28 +428,20 @@ protected String generateAndUploadConfig(String artifactName) { String ip = instanceIpMap.get(instanceName); List dbNames = entry.getValue(); - JSONObject dataShard = new JSONObject(); - dataShard.put("dataShardId", instanceName); - dataShard.put("host", ip); - dataShard.put("port", port); - dataShard.put("user", username); - dataShard.put("password", password); - - JSONArray databases = new JSONArray(); for (String dbName : dbNames) { - JSONObject db = new JSONObject(); - db.put("dbName", dbName); - db.put("databaseId", String.format("%s%02d%s", "shard_", shardIdx, dbName)); - db.put("refDataShardId", instanceName); - databases.put(db); + JSONObject shardConfig = new JSONObject(); + shardConfig.put("logicalShardId", String.format("%s%02d_%s", "shard_", shardIdx, dbName)); + shardConfig.put("host", ip); + shardConfig.put("port", port); + shardConfig.put("user", username); + shardConfig.put("password", password); + shardConfig.put("dbName", dbName); + shardConfigs.put(shardConfig); } shardIdx++; - dataShard.put("databases", databases); - dataShards.put(dataShard); } - shardConfigBulk.put("dataShards", dataShards); - config.put("shardConfigurationBulk", shardConfigBulk); + config.put("shardConfigs", shardConfigs); String configContent = config.toString(); GcsArtifact artifact = diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java index 0b8a70e749..c66ae233f4 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java @@ -256,8 +256,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { order = 20, optional = true, description = "Namespace", - helpText = - "Namespace to be exported. For PostgreSQL, if no namespace is provided, 'public' will be used. Note: Custom non-public namespaces are currently unsupported for PostgreSQL and will cause the job to fail immediately.") + helpText = "This field is no longer used.") @Default.String("") String getNamespace(); diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java index 3492654bc2..7944d4b2c4 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java @@ -32,7 +32,6 @@ import com.google.cloud.teleport.v2.spanner.migrations.utils.ISecretManagerAccessor; import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl; import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import java.util.HashMap; import java.util.List; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java index b0af7bc064..87c0593709 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java @@ -19,8 +19,11 @@ import com.google.cloud.spanner.Dialect; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.transformation.CustomTransformation; import com.google.common.io.Resources; +import com.google.gson.Gson; import java.io.BufferedWriter; import java.io.FileWriter; import java.io.IOException; @@ -235,7 +238,9 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob( } }; if (sourceResourceManager instanceof JDBCResourceManager) { - params.putAll(getJdbcParameters((JDBCResourceManager) sourceResourceManager)); + params.putAll( + getJdbcParameters( + (JDBCResourceManager) sourceResourceManager, gcsPathPrefix, jobParameters)); } else if (sourceResourceManager instanceof CassandraResourceManager) { params.putAll( getCassandraParameters((CassandraResourceManager) sourceResourceManager, gcsPathPrefix)); @@ -262,6 +267,9 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob( // overridden parameters if (jobParameters != null) { for (Map.Entry entry : jobParameters.entrySet()) { + if ("namespace".equals(entry.getKey())) { + continue; + } params.put(entry.getKey(), entry.getValue()); } } @@ -282,13 +290,70 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob( return jobInfo; } - private Map getJdbcParameters(JDBCResourceManager jdbcResourceManager) { + protected String createAndUploadShardConfigToGcs( + String gcsPathPrefix, + JDBCResourceManager jdbcResourceManager, + Map jobParameters) + throws IOException { + Shard shard = new Shard(); + shard.setLogicalShardId("Shard1"); + shard.setUser(jdbcResourceManager.getUsername()); + shard.setPassword(jdbcResourceManager.getPassword()); + if (jdbcResourceManager instanceof PostgresResourceManager pgRm) { + shard.setHost(pgRm.getHost()); + shard.setPort(String.valueOf(pgRm.getPort())); + shard.setDbName(pgRm.getDatabaseName()); + } else if (jdbcResourceManager instanceof MySQLResourceManager mySqlRm) { + shard.setHost(mySqlRm.getHost()); + shard.setPort(String.valueOf(mySqlRm.getPort())); + shard.setDbName(mySqlRm.getDatabaseName()); + } else if (jdbcResourceManager + instanceof org.apache.beam.it.gcp.cloudsql.CloudSqlResourceManager cloudRm) { + shard.setHost(cloudRm.getHost()); + shard.setPort(String.valueOf(cloudRm.getPort())); + shard.setDbName(cloudRm.getDatabaseName()); + } else { + throw new IllegalArgumentException( + "Unsupported JDBC resource manager type: " + jdbcResourceManager.getClass().getName()); + } + + if (jobParameters != null && jobParameters.containsKey("namespace")) { + shard.setNamespace(jobParameters.get("namespace")); + } + + JdbcShardConfig jdbcShardConfig = new JdbcShardConfig(); + jdbcShardConfig.setShardConfigs(List.of(shard)); + String shardFileContents = new Gson().toJson(jdbcShardConfig); + LOG.info("Shard file contents: {}", shardFileContents); + + String configBasePath = (gcsPathPrefix == null) ? "null" : gcsPathPrefix; + if (configBasePath.endsWith("/")) { + configBasePath = configBasePath.substring(0, configBasePath.length() - 1); + } + String configGcsPath = getGcsPath(configBasePath + "/shard.json"); + + gcsClient.createArtifact(configBasePath + "/shard.json", shardFileContents); + + return configGcsPath; + } + + private Map getJdbcParameters( + JDBCResourceManager jdbcResourceManager, + String gcsPathPrefix, + Map jobParameters) { Map params = new HashMap<>() { { put("sourceDbDialect", sqlDialectFrom(jdbcResourceManager)); - put("sourceConfigURL", jdbcResourceManager.getUri()); + try { + put( + "sourceConfigURL", + createAndUploadShardConfigToGcs( + gcsPathPrefix, jdbcResourceManager, jobParameters)); + } catch (IOException e) { + throw new RuntimeException(e); + } put("username", jdbcResourceManager.getUsername()); put("password", jdbcResourceManager.getPassword()); put("jdbcDriverClassName", driverClassNameFrom(jdbcResourceManager)); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java index 10040578be..670757cdf2 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java @@ -19,7 +19,10 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.cloud.teleport.v2.spanner.migrations.constants.Constants; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.transformation.CustomTransformation; +import com.google.gson.Gson; import com.google.pubsub.v1.SubscriptionName; import com.google.pubsub.v1.TopicName; import java.io.BufferedReader; @@ -45,8 +48,6 @@ import org.apache.beam.it.gcp.spanner.SpannerResourceManager; import org.apache.beam.it.gcp.storage.GcsResourceManager; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; -import org.json.JSONArray; -import org.json.JSONObject; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -116,6 +117,21 @@ protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( CloudSqlResourceManager cloudSqlResourceManager, CustomTransformation customTransformation) throws IOException { + + Shard shard = new Shard(); + shard.setLogicalShardId("Shard1"); + shard.setHost(cloudSqlResourceManager.getHost()); + shard.setPort(String.valueOf(cloudSqlResourceManager.getPort())); + shard.setDbName(cloudSqlResourceManager.getDatabaseName()); + shard.setUser(cloudSqlResourceManager.getUsername()); + shard.setPassword(cloudSqlResourceManager.getPassword()); + + JdbcShardConfig jdbcShardConfig = new JdbcShardConfig(); + jdbcShardConfig.setShardConfigs(List.of(shard)); + String shardFileContents = new Gson().toJson(jdbcShardConfig); + LOG.info("Shard file contents: {}", shardFileContents); + gcsResourceManager.createArtifact("input/shard.json", shardFileContents); + // launch dataflow template FlexTemplateDataflowJobResourceManager.Builder flexTemplateBuilder = FlexTemplateDataflowJobResourceManager.builder(jobName) @@ -125,9 +141,7 @@ protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( .addParameter("databaseId", spannerResourceManager.getDatabaseId()) .addParameter("projectId", PROJECT) .addParameter("outputDirectory", getGcsPath("output", gcsResourceManager)) - .addParameter("sourceConfigURL", cloudSqlResourceManager.getUri()) - .addParameter("username", cloudSqlResourceManager.getUsername()) - .addParameter("password", cloudSqlResourceManager.getPassword()) + .addParameter("sourceConfigURL", getGcsPath("input/shard.json", gcsResourceManager)) .addParameter("jdbcDriverClassName", "com.mysql.jdbc.Driver") .addParameter("workerMachineType", "n2-standard-4") .addEnvironmentVariable( @@ -189,51 +203,29 @@ protected PipelineLauncher.LaunchInfo launchShardedBulkDataflowJob( protected void createAndUploadBulkShardConfigToGcs( ArrayList dataShardsList, GcsResourceManager gcsResourceManager) { - JSONObject bulkConfig = new JSONObject(); - bulkConfig.put("configType", "dataflow"); - - JSONObject shardConfigBulk = new JSONObject(); - - JSONObject schemaSourceJson = new JSONObject(); - schemaSourceJson.put("dataShardId", ""); - schemaSourceJson.put("host", ""); - schemaSourceJson.put("user", ""); - schemaSourceJson.put("password", ""); - schemaSourceJson.put("port", ""); - schemaSourceJson.put("dbName", ""); - shardConfigBulk.put("schemaSource", schemaSourceJson); - - JSONArray dataShardsArray = new JSONArray(); + List shards = new ArrayList<>(); if (dataShardsList != null) { for (DataShard shardData : dataShardsList) { - JSONObject shardJson = new JSONObject(); - - shardJson.put("dataShardId", shardData.dataShardId); - shardJson.put("host", shardData.host); - shardJson.put("user", shardData.user); - shardJson.put("password", shardData.password); - shardJson.put("port", shardData.port); - shardJson.put("dbName", shardData.dbName); - shardJson.put("namespace", shardData.namespace); - shardJson.put("connectionProperties", shardData.connectionProperties); - - JSONArray databasesArray = new JSONArray(); - - for (Database dbData : shardData.databases) { - JSONObject dbJson = new JSONObject(); - dbJson.put("dbName", dbData.dbName); - dbJson.put("databaseId", dbData.databaseId); - dbJson.put("refDataShardId", dbData.refDataShardId); - databasesArray.put(dbJson); + Shard shard = new Shard(); + shard.setLogicalShardId(shardData.dataShardId); + shard.setHost(shardData.host); + shard.setUser(shardData.user); + shard.setPassword(shardData.password); + shard.setPort(shardData.port); + shard.setDbName(shardData.dbName); + if (shardData.namespace != null) { + shard.setNamespace(shardData.namespace); + } + if (shardData.connectionProperties != null) { + shard.setConnectionProperties(shardData.connectionProperties); } - shardJson.put("databases", databasesArray); - dataShardsArray.put(shardJson); + shards.add(shard); } } - shardConfigBulk.put("dataShards", dataShardsArray); - bulkConfig.put("shardConfigurationBulk", shardConfigBulk); - String shardFileContents = bulkConfig.toString(); + JdbcShardConfig jdbcShardConfig = new JdbcShardConfig(); + jdbcShardConfig.setShardConfigs(shards); + String shardFileContents = new Gson().toJson(jdbcShardConfig); LOG.info("Shard file contents: {}", shardFileContents); gcsResourceManager.createArtifact("input/shard-bulk.json", shardFileContents); } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java index ac77f80079..ba1388ecfc 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java @@ -19,11 +19,14 @@ import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.text.ParseException; import java.time.Duration; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -34,6 +37,7 @@ import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.it.conditions.ConditionCheck; import org.apache.beam.it.gcp.TemplateLoadTestBase; +import org.apache.beam.it.gcp.artifacts.utils.ArtifactUtils; import org.apache.beam.it.gcp.secretmanager.SecretManagerResourceManager; import org.apache.beam.it.gcp.spanner.SpannerResourceManager; import org.apache.beam.it.gcp.spanner.conditions.SpannerRowsCheck; @@ -175,11 +179,40 @@ protected String getOutputDirectory() { } protected Map getJdbcParameters(StaticJDBCResource jdbcResource) { - return getJdbcParameters( - jdbcResource.getconnectionURL(), - jdbcResource.username(), - jdbcResource.password(), - driverClassName()); + try { + return getJdbcParameters( + createAndUploadShardConfigToGcs(jdbcResource), + jdbcResource.username(), + jdbcResource.password(), + driverClassName()); + } catch (IOException e) { + throw new RuntimeException("Failed to create and upload shard config", e); + } + } + + protected String createAndUploadShardConfigToGcs(StaticJDBCResource jdbcResource) + throws IOException { + Shard shard = new Shard(); + shard.setLogicalShardId("Shard1"); + shard.setUser(jdbcResource.username()); + shard.setPassword(jdbcResource.password()); + shard.setHost(jdbcResource.hostname()); + shard.setPort(String.valueOf(jdbcResource.port())); + shard.setDbName(jdbcResource.database()); + + JdbcShardConfig jdbcShardConfig = new JdbcShardConfig(); + jdbcShardConfig.setShardConfigs(Collections.singletonList(shard)); + String shardFileContents = new Gson().toJson(jdbcShardConfig); + gcsResourceManager.createArtifact("input/shard.json", shardFileContents); + return getGcsPath("input/shard.json", gcsResourceManager); + } + + protected String getGcsPath(String artifactId, GcsResourceManager gcsResourceManager) { + return ArtifactUtils.getFullGcsPath( + gcsResourceManager.getBucket(), + getClass().getSimpleName(), + gcsResourceManager.runId(), + artifactId); } protected Map getJdbcParameters( From ac7b26d34006b0b2d02d0f9b5aaa2780dd6ab968 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Tue, 30 Jun 2026 13:42:55 +0000 Subject: [PATCH 03/16] feat(Spanner): refactor CassandraIOWrapperFactory to accept SourceConnectionConfig instead of pipeline options for improved configuration handling. --- .../iowrapper/CassandraIOWrapperFactory.java | 27 ++++++++++++++---- .../CassandraIOWrapperFactoryTest.java | 28 +++++++++++-------- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java index 1f843fdd10..bbc6f8db59 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java @@ -21,6 +21,8 @@ import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.IoWrapper; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper.CassandraDataSource.CassandraDialect; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.common.base.Preconditions; import java.util.List; import javax.annotation.Nullable; @@ -80,7 +82,8 @@ private static CassandraIOWrapperFactory create( astraDBRegion); } - public static CassandraIOWrapperFactory fromPipelineOptions(SourceDbToSpannerOptions options) { + public static CassandraIOWrapperFactory fromConfig( + SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { String gcsPath = options.getSourceConfigURL(); // Implementation Details. the pipeline options are strings. Preconditions.checkArgument( @@ -93,14 +96,28 @@ public static CassandraIOWrapperFactory fromPipelineOptions(SourceDbToSpannerOpt options.getSourceDbDialect().equals(SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT) || StringUtils.startsWith(gcsPath, "gs://"), "GCS path Expected in place of `" + gcsPath + "`."); + + GuardedStringValueProvider astraDBToken = GuardedStringValueProvider.create(""); + String astraDBDatabaseId = ""; + String astraDBKeyspace = ""; + String astraDBRegion = ""; + + if (sourceConnectionConfig instanceof AstraConnectionConfig) { + AstraConnectionConfig astraConfig = (AstraConnectionConfig) sourceConnectionConfig; + astraDBToken = GuardedStringValueProvider.create(astraConfig.getAstraToken()); + astraDBDatabaseId = astraConfig.getDatabaseId(); + astraDBKeyspace = astraConfig.getKeySpace(); + astraDBRegion = astraConfig.getAstraDbRegion(); + } + return CassandraIOWrapperFactory.create( options.getSourceConfigURL(), options.getNumPartitions(), options.getSourceDbDialect(), - GuardedStringValueProvider.create(options.getAstraDBToken()), - options.getAstraDBDatabaseId(), - options.getAstraDBKeySpace(), - options.getAstraDBRegion()); + astraDBToken, + astraDBDatabaseId, + astraDBKeyspace, + astraDBRegion); } /** Create an {@link IoWrapper} instance for a list of SourceTables. */ diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java index d173b32776..0f57a40a3b 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java @@ -36,6 +36,9 @@ import com.google.cloud.teleport.v2.reader.io.schema.SourceTableReference; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper.CassandraDataSource.CassandraDialect; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.schema.CassandraSchemaReference; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.net.InetSocketAddress; @@ -138,12 +141,9 @@ public void testCassandraIoWrapperFactoryOssBasic() { when(mockOptions.getSourceDbDialect()).thenReturn("CASSANDRA"); when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); when(mockOptions.getNumPartitions()).thenReturn(null); - when(mockOptions.getAstraDBToken()).thenReturn(""); - when(mockOptions.getAstraDBDatabaseId()).thenReturn(""); - when(mockOptions.getAstraDBRegion()).thenReturn(""); - when(mockOptions.getAstraDBKeySpace()).thenReturn(""); + CassandraConnectionConfig mockSourceConfig = mock(CassandraConnectionConfig.class); CassandraIOWrapperFactory cassandraIOWrapperFactory = - CassandraIOWrapperFactory.fromPipelineOptions(mockOptions); + CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); assertThat(cassandraIOWrapperFactory.gcsConfigPath()).isEqualTo(testConfigPath); assertThat(cassandraIOWrapperFactory.getIOWrapper(TABLES_TO_READ, null).discoverTableSchema()) .isEqualTo(ImmutableList.of(mockSourceSchema)); @@ -163,12 +163,15 @@ public void testCassandraIoWrapperFactoryAstraBasic() { when(mockOptions.getSourceDbDialect()).thenReturn("ASTRA_DB"); when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); when(mockOptions.getNumPartitions()).thenReturn(null); - when(mockOptions.getAstraDBToken()).thenReturn("AstraCS:testToken"); - when(mockOptions.getAstraDBDatabaseId()).thenReturn("testId"); - when(mockOptions.getAstraDBRegion()).thenReturn("testRegion"); - when(mockOptions.getAstraDBKeySpace()).thenReturn("testKeyspace"); + + AstraConnectionConfig mockSourceConfig = mock(AstraConnectionConfig.class); + when(mockSourceConfig.getAstraToken()).thenReturn("AstraCS:testToken"); + when(mockSourceConfig.getDatabaseId()).thenReturn("testId"); + when(mockSourceConfig.getAstraDbRegion()).thenReturn("testRegion"); + when(mockSourceConfig.getKeySpace()).thenReturn("testKeyspace"); + CassandraIOWrapperFactory cassandraIOWrapperFactory = - CassandraIOWrapperFactory.fromPipelineOptions(mockOptions); + CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); assertThat(cassandraIOWrapperFactory.gcsConfigPath()).isEqualTo(testConfigPath); assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.ASTRA); assertThat(cassandraIOWrapperFactory.astraDBKeyspace()).isEqualTo("testKeyspace"); @@ -185,11 +188,12 @@ public void testCassandraIoWrapperFactoryExceptions() { mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); when(mockOptions.getSourceDbDialect()).thenReturn("MYSQL").thenReturn("CASSANDRA"); when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); + SourceConnectionConfig mockConfig = mock(SourceConnectionConfig.class); assertThrows( IllegalArgumentException.class, - () -> CassandraIOWrapperFactory.fromPipelineOptions(mockOptions)); + () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); assertThrows( IllegalArgumentException.class, - () -> CassandraIOWrapperFactory.fromPipelineOptions(mockOptions)); + () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); } } From a9be7e4f088ffd1945f4a26de795b1c40c44a31f Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Wed, 1 Jul 2026 06:37:28 +0000 Subject: [PATCH 04/16] fix: resolve Astra token parsing issue and add integration test for AstraDbToSpanner --- .../iowrapper/AstraDbDataSource.java | 2 +- .../templates/AstraDbToSpannerSimpleIT.java | 228 ++++++++++++++++++ .../v2/templates/SourceDbToSpannerITBase.java | 8 +- 3 files changed, 235 insertions(+), 3 deletions(-) create mode 100644 v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java index 689d3ff54a..6812287f77 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java @@ -103,7 +103,7 @@ public Builder setAstraToken(String value) { astraToken = SecretManagerUtils.getSecret(value); } LOG.info("Astra Token is parsed"); - return this.setAstraToken(GuardedStringValueProvider.create(value)); + return this.setAstraToken(GuardedStringValueProvider.create(astraToken)); } public abstract Builder setKeySpace(String value); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java new file mode 100644 index 0000000000..2f5cefc08c --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java @@ -0,0 +1,228 @@ +/* + * 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.templates; + +import static com.google.common.truth.Truth.assertThat; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.dtsx.astra.sdk.db.AstraDBOpsClient; +import com.dtsx.astra.sdk.db.DbOpsClient; +import com.dtsx.astra.sdk.db.domain.Database; +import com.dtsx.astra.sdk.db.domain.DatabaseCreationRequest; +import com.dtsx.astra.sdk.db.domain.DatabaseStatusType; +import com.dtsx.astra.sdk.utils.ApiLocator; +import com.google.cloud.spanner.Struct; +import com.google.cloud.teleport.metadata.TemplateIntegrationTest; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; +import com.google.gson.Gson; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.Serializable; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.it.common.PipelineLauncher; +import org.apache.beam.it.common.PipelineOperator; +import org.apache.beam.it.common.TestProperties; +import org.apache.beam.it.common.utils.ResourceManagerUtils; +import org.apache.beam.it.gcp.spanner.SpannerResourceManager; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Integration test for {@link com.google.cloud.teleport.v2.templates.SourceDbToSpanner} from Astra + * DB. + */ +@RunWith(JUnit4.class) +@Category(TemplateIntegrationTest.class) +@TemplateIntegrationTest(SourceDbToSpanner.class) +public class AstraDbToSpannerSimpleIT extends SourceDbToSpannerITBase implements Serializable { + + private static final Logger LOGGER = LoggerFactory.getLogger(AstraDbToSpannerSimpleIT.class); + + private static final long NUM_ROWS = 50L; + private static final String ASTRA_DB = "dataflow_integration_tests"; + private static final String ASTRA_DB_REGION = TestProperties.region(); + private static final String ASTRA_KS = "beam"; + private static final String ASTRA_TBL = "scientist"; + + private static DbOpsClient dbClient; + private SpannerResourceManager spannerResourceManager; + + @Before + public void setup() throws Exception { + spannerResourceManager = + SpannerResourceManager.builder(testName, PROJECT, REGION).maybeUseStaticInstance().build(); + + // Create Spanner table + String spannerDdl = + String.format( + "CREATE TABLE %s (" + + " person_department STRING(MAX)," + + " person_id INT64," + + " person_name STRING(MAX)," + + ") PRIMARY KEY(person_department, person_id)", + ASTRA_TBL); + spannerResourceManager.executeDdlStatement(spannerDdl); + + // Setup Astra Db + createOrResumeAstraDatabase(); + // Setup Astra Data + createAndPopulateTables(); + LOGGER.info("Initialization Successful."); + } + + @Test + public void testAstraDbToSpanner() throws IOException { + // Generate shard.json + AstraConnectionConfig astraConfig = new AstraConnectionConfig(); + astraConfig.setAstraToken(dbClient.getToken()); + astraConfig.setDatabaseId(dbClient.getDatabaseId()); + astraConfig.setKeySpace(ASTRA_KS); + astraConfig.setAstraDbRegion(dbClient.get().getInfo().getRegion()); + + String configContents = new Gson().toJson(astraConfig); + artifactClient.createArtifact("input/shard.json", configContents); + String sourceConfigURL = getGcsPath("input/shard.json", artifactClient); + + Map jobParameters = new HashMap<>(); + jobParameters.put("sourceDbDialect", "ASTRA_DB"); + jobParameters.put("sourceConfigURL", sourceConfigURL); + jobParameters.put("outputDirectory", getGcsPath("output", artifactClient)); + jobParameters.put( + "ipConfiguration", "WORKER_IP_UNSPECIFIED"); // Require internet access for Astra API + + // Act + PipelineLauncher.LaunchInfo info = + launchDataflowJob(testName, null, null, null, spannerResourceManager, jobParameters, null); + LOGGER.debug("Pipeline is now running"); + + PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(info)); + + org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult(result).isLaunchFinished(); + LOGGER.debug("Destination Table has been populated."); + + // Optionally verify a row + List rows = + spannerResourceManager.readTableRecords( + ASTRA_TBL, List.of("person_department", "person_id", "person_name")); + assertThat(rows).isNotEmpty(); + } + + @After + public void tearDown() { + ResourceManagerUtils.cleanResources(spannerResourceManager); + } + + private static String test() { + return "AstraCS:" + HASH; + } + + @SuppressWarnings("BusyWait") + private void createOrResumeAstraDatabase() throws InterruptedException { + AstraDBOpsClient databasesClient = new AstraDBOpsClient(test()); + if (databasesClient.findByName(ASTRA_DB).findAny().isEmpty()) { + LOGGER.debug("Create a new Database {}", ASTRA_DB); + databasesClient.create( + DatabaseCreationRequest.builder() + .name(ASTRA_DB) + .keyspace(ASTRA_KS) + .cloudRegion(ASTRA_DB_REGION) + .build()); + } else { + LOGGER.debug("Database {} exists in source organization", ASTRA_DB); + } + dbClient = databasesClient.databaseByName(ASTRA_DB); + if (dbClient.get().getStatus() == DatabaseStatusType.HIBERNATED) { + resumeDb(dbClient.get()); + LOGGER.debug("Resuming as DB was Hibernated"); + } + while (dbClient.get().getStatus() != DatabaseStatusType.ACTIVE) { + Thread.sleep(5000); + LOGGER.debug("Waiting for DB to be ACTIVE...."); + } + } + + private void resumeDb(Database db) { + try { + HttpClient.newBuilder() + .version(HttpClient.Version.HTTP_2) + .connectTimeout(Duration.ofSeconds(20)) + .build() + .send( + HttpRequest.newBuilder() + .timeout(Duration.ofSeconds(20)) + .uri( + URI.create( + ApiLocator.getApiRestEndpoint(db.getId(), db.getInfo().getRegion()) + + "/v2/schemas/keyspace")) + .timeout(Duration.ofSeconds(20)) + .header("Content-Type", "application/json") + .header("X-Cassandra-Token", test()) + .GET() + .build(), + HttpResponse.BodyHandlers.ofString()); + } catch (Exception e) { + throw new IllegalStateException("Cannot resume database", e); + } + } + + private void createAndPopulateTables() { + try (CqlSession astraSession = + CqlSession.builder() + .withCloudSecureConnectBundle( + new ByteArrayInputStream(dbClient.downloadDefaultSecureConnectBundle())) + .withAuthCredentials("token", dbClient.getToken()) + .withKeyspace(ASTRA_KS) + .build()) { + astraSession.execute( + String.format( + "CREATE TABLE IF NOT EXISTS %s.%s(person_department text, person_id int, person_name text, PRIMARY KEY" + + "((person_department), person_id));", + ASTRA_KS, ASTRA_TBL)); + String[][] scientists = { + new String[] {"phys", "Einstein"}, + new String[] {"bio", "Darwin"}, + new String[] {"phys", "Copernicus"}, + new String[] {"bio", "Pasteur"}, + new String[] {"bio", "Curie"} + }; + for (int i = 0; i < NUM_ROWS; i++) { + int index = i % scientists.length; + String insertStr = + String.format( + "INSERT INTO %s.%s(person_department, person_id, person_name) values(" + + "'%s', %d, '%s');", + ASTRA_KS, ASTRA_TBL, scientists[index][0], i, scientists[index][1]); + astraSession.execute(insertStr); + } + } + } + + private static final String HASH = + "AIpXbGsYPQCXtrwExZvOktGw:3d5bae1547a667608f10ab2d2e89a90b936f8ff8a3e9111efe23fc818ef344fd"; +} diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java index 87c0593709..4bcbc59fe7 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java @@ -265,9 +265,13 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob( } // overridden parameters + String ipConfig = "WORKER_IP_PRIVATE"; if (jobParameters != null) { + if (jobParameters.containsKey("ipConfiguration")) { + ipConfig = jobParameters.get("ipConfiguration"); + } for (Map.Entry entry : jobParameters.entrySet()) { - if ("namespace".equals(entry.getKey())) { + if ("namespace".equals(entry.getKey()) || "ipConfiguration".equals(entry.getKey())) { continue; } params.put(entry.getKey(), entry.getValue()); @@ -282,7 +286,7 @@ protected PipelineLauncher.LaunchInfo launchDataflowJob( options.setParameters(params); options.addEnvironment("additionalExperiments", List.of("disable_runner_v2")); options.addEnvironment("numWorkers", 2); - options.addEnvironment("ipConfiguration", "WORKER_IP_PRIVATE"); + options.addEnvironment("ipConfiguration", ipConfig); // Run PipelineLauncher.LaunchInfo jobInfo = launchTemplate(options); assertThatPipeline(jobInfo).isRunning(); From b8fc6344793b73e5c8df3cda787e9c29123a4e78 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Wed, 1 Jul 2026 09:47:41 +0000 Subject: [PATCH 05/16] Nit fixes. --- .../teleport/v2/options/SourceDbToSpannerOptions.java | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java index c66ae233f4..bb11ca7a32 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java @@ -252,16 +252,6 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setTransformationCustomParameters(String value); - @TemplateParameter.Text( - order = 20, - optional = true, - description = "Namespace", - helpText = "This field is no longer used.") - @Default.String("") - String getNamespace(); - - void setNamespace(String value); - @TemplateParameter.Text( order = 21, optional = true, From 41af3453f15393c3e1be8564efd0a25fccb5c5ce Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Wed, 1 Jul 2026 10:42:13 +0000 Subject: [PATCH 06/16] Cassandra to use SourceConnectionConfig --- .../iowrapper/CassandraIOWrapperFactory.java | 31 +++---- .../iowrapper/CassandraIOWrapperHelper.java | 45 +++++------ .../iowrapper/CassandraIoWrapper.java | 5 +- .../CassandraIOWrapperFactoryTest.java | 18 ++--- .../CassandraIOWrapperHelperTest.java | 81 ++++++++----------- .../iowrapper/CassandraIoWrapperTest.java | 8 +- 6 files changed, 80 insertions(+), 108 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java index bbc6f8db59..78e93535fd 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java @@ -15,6 +15,7 @@ */ package com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper; +import com.datastax.oss.driver.api.core.config.OptionsMap; import com.google.auto.value.AutoValue; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.reader.IoWrapperFactory; @@ -22,18 +23,19 @@ import com.google.cloud.teleport.v2.reader.io.IoWrapper; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper.CassandraDataSource.CassandraDialect; import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.common.base.Preconditions; import java.util.List; import javax.annotation.Nullable; import org.apache.beam.sdk.transforms.Wait.OnSignal; -import org.apache.commons.lang3.StringUtils; @AutoValue public abstract class CassandraIOWrapperFactory implements IoWrapperFactory { - /** GCS Path for Cassandra Driver Config. */ - public abstract String gcsConfigPath(); + /** Options Map for Cassandra Driver Config. */ + @Nullable + public abstract OptionsMap optionsMap(); /** * Number of partitions to read from. Defaults to Null. @@ -46,7 +48,6 @@ public abstract class CassandraIOWrapperFactory implements IoWrapperFactory { /** Cassandra Dialect. */ public abstract CassandraDataSource.CassandraDialect cassandraDialect(); - /** Astra DB options. Empty for OSS dialect. */ /** Astra DB Token. * */ public abstract GuardedStringValueProvider astraDBToken(); @@ -56,11 +57,11 @@ public abstract class CassandraIOWrapperFactory implements IoWrapperFactory { /** Astra DB Keyspace. * */ public abstract String astraDBKeyspace(); - /** Astra DB Keyspace. * */ + /** Astra DB Region. * */ public abstract String astraDBRegion(); private static CassandraIOWrapperFactory create( - String gcsConfigPath, + OptionsMap optionsMap, Integer numPartions, String sourceDialect, GuardedStringValueProvider astraDBToken, @@ -73,7 +74,7 @@ private static CassandraIOWrapperFactory create( default -> CassandraDialect.OSS; }; return new AutoValue_CassandraIOWrapperFactory( - gcsConfigPath, + optionsMap, numPartions, cassandraDialect, astraDBToken, @@ -84,23 +85,18 @@ private static CassandraIOWrapperFactory create( public static CassandraIOWrapperFactory fromConfig( SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { - String gcsPath = options.getSourceConfigURL(); - // Implementation Details. the pipeline options are strings. Preconditions.checkArgument( options.getSourceDbDialect().equals(SourceDbToSpannerOptions.CASSANDRA_SOURCE_DIALECT) || options .getSourceDbDialect() .equals(SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT), "Unexpected Dialect " + options.getSourceDbDialect() + " for Cassandra Source"); - Preconditions.checkArgument( - options.getSourceDbDialect().equals(SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT) - || StringUtils.startsWith(gcsPath, "gs://"), - "GCS path Expected in place of `" + gcsPath + "`."); GuardedStringValueProvider astraDBToken = GuardedStringValueProvider.create(""); String astraDBDatabaseId = ""; String astraDBKeyspace = ""; String astraDBRegion = ""; + OptionsMap optionsMap = null; if (sourceConnectionConfig instanceof AstraConnectionConfig) { AstraConnectionConfig astraConfig = (AstraConnectionConfig) sourceConnectionConfig; @@ -108,10 +104,15 @@ public static CassandraIOWrapperFactory fromConfig( astraDBDatabaseId = astraConfig.getDatabaseId(); astraDBKeyspace = astraConfig.getKeySpace(); astraDBRegion = astraConfig.getAstraDbRegion(); + } else if (sourceConnectionConfig instanceof CassandraConnectionConfig) { + optionsMap = ((CassandraConnectionConfig) sourceConnectionConfig).getOptionsMap(); + } else { + throw new IllegalArgumentException( + "Unsupported source connection config type: " + sourceConnectionConfig); } return CassandraIOWrapperFactory.create( - options.getSourceConfigURL(), + optionsMap, options.getNumPartitions(), options.getSourceDbDialect(), astraDBToken, @@ -125,7 +126,7 @@ public static CassandraIOWrapperFactory fromConfig( public IoWrapper getIOWrapper(List sourceTables, OnSignal waitOnSignal) { /** TODO(vardhanvthigle@) incorporate waitOnSignal */ return new CassandraIoWrapper( - gcsConfigPath(), + optionsMap(), sourceTables, numPartitions(), cassandraDialect(), diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java index 277ef828b0..e91098eb1c 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java @@ -17,6 +17,7 @@ import static com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper.CassandraDefaults.DEFAULT_CASSANDRA_SCHEMA_DISCOVERY_BACKOFF; +import com.datastax.oss.driver.api.core.config.OptionsMap; import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.datasource.DataSource; import com.google.cloud.teleport.v2.reader.io.exception.SchemaDiscoveryException; @@ -34,7 +35,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import java.io.FileNotFoundException; import java.util.List; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.values.PBegin; @@ -48,7 +48,7 @@ class CassandraIOWrapperHelper { private static final Logger LOG = LoggerFactory.getLogger(CassandraIOWrapperHelper.class); static DataSource buildDataSource( - String gcsPath, + OptionsMap optionsMap, Integer numPartitions, CassandraDataSource.CassandraDialect cassandraDialect, GuardedStringValueProvider astraDBToken, @@ -56,30 +56,25 @@ static DataSource buildDataSource( String astraDBKeyspace, String astraDBRegion) { DataSource dataSource; - try { - dataSource = - switch (cassandraDialect) { - case ASTRA -> DataSource.ofCassandra( - // TODO: Astra: Build from Pipeline Options. - CassandraDataSource.ofAstra( - AstraDbDataSource.builder() - .setAstraToken(astraDBToken) - .setDatabaseId(astraDBDatabaseId) - .setKeySpace(astraDBKeyspace) - .setAstraDbRegion(astraDBRegion) - .build())); + dataSource = + switch (cassandraDialect) { + case ASTRA -> DataSource.ofCassandra( + // TODO: Astra: Build from Pipeline Options. + CassandraDataSource.ofAstra( + AstraDbDataSource.builder() + .setAstraToken(astraDBToken) + .setDatabaseId(astraDBDatabaseId) + .setKeySpace(astraDBKeyspace) + .setAstraDbRegion(astraDBRegion) + .build())); - default -> DataSource.ofCassandra( - CassandraDataSource.ofOss( - CassandraDataSourceOss.builder() - .setOptionsMapFromGcsFile(gcsPath) - .setNumPartitions(numPartitions) - .build())); - }; - } catch (FileNotFoundException e) { - LOG.error("Unable to find driver config file in {}. Cause ", gcsPath, e); - throw (new SchemaDiscoveryException(e)); - } + default -> DataSource.ofCassandra( + CassandraDataSource.ofOss( + CassandraDataSourceOss.builder() + .setOptionsMap(optionsMap) + .setNumPartitions(numPartitions) + .build())); + }; return dataSource; } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapper.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapper.java index 95ffa176d8..5881f56e69 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapper.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapper.java @@ -15,6 +15,7 @@ */ package com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper; +import com.datastax.oss.driver.api.core.config.OptionsMap; import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.IoWrapper; import com.google.cloud.teleport.v2.reader.io.datasource.DataSource; @@ -41,7 +42,7 @@ public final class CassandraIoWrapper implements IoWrapper { tableReaders; public CassandraIoWrapper( - String gcsPath, + OptionsMap optionsMap, List sourceTables, @Nullable Integer numPartitions, CassandraDialect cassandraDialect, @@ -51,7 +52,7 @@ public CassandraIoWrapper( String astraDBRegion) { DataSource dataSource = CassandraIOWrapperHelper.buildDataSource( - gcsPath, + optionsMap, numPartitions, cassandraDialect, astraDBToken, diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java index 0f57a40a3b..707aac6380 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java @@ -99,7 +99,7 @@ public void setup() { .when( () -> CassandraIOWrapperHelper.buildDataSource( - TEST_BUCKET_CASSANDRA_CONFIG_CONF, + null, null, CassandraDialect.OSS, GuardedStringValueProvider.create(""), @@ -135,16 +135,15 @@ public void cleanup() { @Test public void testCassandraIoWrapperFactoryOssBasic() { - String testConfigPath = TEST_BUCKET_CASSANDRA_CONFIG_CONF; SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); when(mockOptions.getSourceDbDialect()).thenReturn("CASSANDRA"); - when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); when(mockOptions.getNumPartitions()).thenReturn(null); CassandraConnectionConfig mockSourceConfig = mock(CassandraConnectionConfig.class); + when(mockSourceConfig.getOptionsMap()).thenReturn(null); CassandraIOWrapperFactory cassandraIOWrapperFactory = CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); - assertThat(cassandraIOWrapperFactory.gcsConfigPath()).isEqualTo(testConfigPath); + assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(null); assertThat(cassandraIOWrapperFactory.getIOWrapper(TABLES_TO_READ, null).discoverTableSchema()) .isEqualTo(ImmutableList.of(mockSourceSchema)); assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.OSS); @@ -157,11 +156,9 @@ public void testCassandraIoWrapperFactoryOssBasic() { @Test public void testCassandraIoWrapperFactoryAstraBasic() { - String testConfigPath = ""; SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); when(mockOptions.getSourceDbDialect()).thenReturn("ASTRA_DB"); - when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); when(mockOptions.getNumPartitions()).thenReturn(null); AstraConnectionConfig mockSourceConfig = mock(AstraConnectionConfig.class); @@ -172,7 +169,7 @@ public void testCassandraIoWrapperFactoryAstraBasic() { CassandraIOWrapperFactory cassandraIOWrapperFactory = CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); - assertThat(cassandraIOWrapperFactory.gcsConfigPath()).isEqualTo(testConfigPath); + assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(null); assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.ASTRA); assertThat(cassandraIOWrapperFactory.astraDBKeyspace()).isEqualTo("testKeyspace"); assertThat(cassandraIOWrapperFactory.astraDBRegion()).isEqualTo("testRegion"); @@ -183,17 +180,12 @@ public void testCassandraIoWrapperFactoryAstraBasic() { @Test public void testCassandraIoWrapperFactoryExceptions() { - String testConfigPath = "smt-test-bucket/test-conf.conf"; SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); - when(mockOptions.getSourceDbDialect()).thenReturn("MYSQL").thenReturn("CASSANDRA"); - when(mockOptions.getSourceConfigURL()).thenReturn(testConfigPath); + when(mockOptions.getSourceDbDialect()).thenReturn("MYSQL"); SourceConnectionConfig mockConfig = mock(SourceConnectionConfig.class); assertThrows( IllegalArgumentException.class, () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); - assertThrows( - IllegalArgumentException.class, - () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java index 714951ce30..1beefd8802 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java @@ -37,8 +37,6 @@ import static com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.testutils.BasicTestSchema.TEST_KEYSPACE; import static com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.testutils.BasicTestSchema.TEST_TABLES; import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; -import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockStatic; import com.datastax.oss.driver.api.core.config.OptionsMap; @@ -54,6 +52,7 @@ import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.schema.CassandraSchemaDiscovery; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.schema.CassandraSchemaReference; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.testutils.SharedEmbeddedCassandra; +import com.google.cloud.teleport.v2.spanner.migrations.utils.CassandraDriverConfigLoader; import com.google.cloud.teleport.v2.spanner.migrations.utils.JarFileReader; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -62,7 +61,6 @@ import java.net.URL; import java.util.List; import java.util.stream.Collectors; -import org.apache.beam.sdk.io.cassandra.CassandraIO; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.values.PBegin; @@ -96,58 +94,43 @@ public static void stopEmbeddedCassandra() throws Exception { } @Test - public void testBuildDataSource() { - + public void testBuildDataSource() throws Exception { String testGcsPath = "gs://smt-test-bucket/cassandraConfig.conf"; URL testUrl = Resources.getResource("CassandraUT/test-cassandra-config.conf"); - CassandraIO.Read mockCassandraIORead = mock(CassandraIO.Read.class); + OptionsMap optionsMap; try (MockedStatic mockFileReader = mockStatic(JarFileReader.class)) { - mockFileReader .when(() -> JarFileReader.saveFilesLocally(testGcsPath)) - .thenReturn(new URL[] {testUrl}) - .thenReturn(new URL[] {testUrl}) - /* Empty URL List to test FileNotFoundException handling. */ - .thenReturn(new URL[] {}); - - DataSource dataSource = - CassandraIOWrapperHelper.buildDataSource( - testGcsPath, - null, - CassandraDialect.OSS, - GuardedStringValueProvider.create(""), - "", - "", - ""); - assertThat(dataSource.cassandra().oss().loggedKeySpace()).isEqualTo("test-keyspace"); - assertThat(dataSource.cassandra().oss().localDataCenter()).isEqualTo("datacenter1"); - assertThat(dataSource.cassandra().oss().numPartitions()).isEqualTo(null); - assertThat( - CassandraIOWrapperHelper.buildDataSource( - testGcsPath, - 42, - CassandraDialect.OSS, - GuardedStringValueProvider.create(""), - "", - "", - "") - .cassandra() - .oss() - .numPartitions()) - .isEqualTo(42); - assertThrows( - SchemaDiscoveryException.class, - () -> - CassandraIOWrapperHelper.buildDataSource( - testGcsPath, - null, - CassandraDialect.OSS, - GuardedStringValueProvider.create(""), - "", - "", - "")); + .thenReturn(new URL[] {testUrl}); + optionsMap = CassandraDriverConfigLoader.getOptionsMapFromFile(testGcsPath); } + + DataSource dataSource = + CassandraIOWrapperHelper.buildDataSource( + optionsMap, + null, + CassandraDialect.OSS, + GuardedStringValueProvider.create(""), + "", + "", + ""); + assertThat(dataSource.cassandra().oss().loggedKeySpace()).isEqualTo("test-keyspace"); + assertThat(dataSource.cassandra().oss().localDataCenter()).isEqualTo("datacenter1"); + assertThat(dataSource.cassandra().oss().numPartitions()).isEqualTo(null); + assertThat( + CassandraIOWrapperHelper.buildDataSource( + optionsMap, + 42, + CassandraDialect.OSS, + GuardedStringValueProvider.create(""), + "", + "", + "") + .cassandra() + .oss() + .numPartitions()) + .isEqualTo(42); } @Test @@ -275,7 +258,7 @@ public void testAstra() { DataSource dataSource = CassandraIOWrapperHelper.buildDataSource( - "", + null, null, CassandraDialect.ASTRA, testAstraDbToken, diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java index 9b3b1eb7f9..aad79eadf3 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java @@ -95,7 +95,7 @@ public void testCassandraIoWrapperBasic() { .when( () -> CassandraIOWrapperHelper.buildDataSource( - testGcsPath, + null, null, CassandraDialect.OSS, GuardedStringValueProvider.create(""), @@ -107,7 +107,7 @@ public void testCassandraIoWrapperBasic() { .when( () -> CassandraIOWrapperHelper.buildDataSource( - "", + null, null, CassandraDialect.ASTRA, astraDataSource.cassandra().astra().astraToken(), @@ -151,7 +151,7 @@ public void testCassandraIoWrapperBasic() { CassandraIoWrapper cassandraIoWrapper = new CassandraIoWrapper( - testGcsPath, + null, tablesToRead, null, CassandraDialect.OSS, @@ -165,7 +165,7 @@ public void testCassandraIoWrapperBasic() { CassandraIoWrapper cassandraIoWrapperAstra = new CassandraIoWrapper( - "", + null, tablesToRead, null, CassandraDialect.ASTRA, From 035716887fce30651e88569b8494b90fe2bff874 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 2 Jul 2026 06:29:02 +0000 Subject: [PATCH 07/16] Extend unit test coverage --- .../v2/templates/PipelineControllerTest.java | 2 + .../v2/templates/SourceDbToSpannerTest.java | 47 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java index 4822e1b071..cb30e425fa 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java @@ -21,6 +21,7 @@ import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.when; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; @@ -64,6 +65,7 @@ import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.mockito.Mock; +import org.mockito.MockedConstruction; import org.mockito.MockedStatic; import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java index 7e971450bf..811ec77ddd 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java @@ -28,6 +28,9 @@ import com.google.cloud.teleport.v2.common.CommonTemplateJvmInitializer; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.joda.time.Duration; @@ -123,4 +126,48 @@ public void testValidateOptions_NonPostgresDialectSucceeds() { .thenReturn(SourceDbToSpannerOptions.MYSQL_SOURCE_DIALECT); SourceDbToSpanner.validateOptions(mockOptions); } + + @Test + public void testRun_ValidationFailures() { + SourceDbToSpannerOptions mockOptions = + PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); + mockOptions.setProjectId("testProject"); + mockOptions.setInstanceId("testInstance"); + mockOptions.setDatabaseId("testDatabaseId"); + mockOptions + .as(org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions.class) + .setWorkerMachineType("n2-standard-4"); + + try (MockedStatic mockedPipelineController = + mockStatic(PipelineController.class)) { + + // Test MYSQL dialect with wrong config type (AstraConnectionConfig) + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.MYSQL_SOURCE_DIALECT); + mockedPipelineController + .when(() -> PipelineController.getSourceConnectionConfig(any(), any())) + .thenReturn(mock(AstraConnectionConfig.class)); + assertThrows(IllegalArgumentException.class, () -> SourceDbToSpanner.run(mockOptions)); + + // Test POSTGRESQL dialect with wrong config type (CassandraConnectionConfig) + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.PG_SOURCE_DIALECT); + mockedPipelineController + .when(() -> PipelineController.getSourceConnectionConfig(any(), any())) + .thenReturn(mock(CassandraConnectionConfig.class)); + assertThrows(IllegalArgumentException.class, () -> SourceDbToSpanner.run(mockOptions)); + + // Test CASSANDRA dialect with wrong config type (JdbcShardConfig) + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.CASSANDRA_SOURCE_DIALECT); + mockedPipelineController + .when(() -> PipelineController.getSourceConnectionConfig(any(), any())) + .thenReturn(mock(JdbcShardConfig.class)); + assertThrows(IllegalArgumentException.class, () -> SourceDbToSpanner.run(mockOptions)); + + // Test ASTRADB dialect with wrong config type (JdbcShardConfig) + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT); + mockedPipelineController + .when(() -> PipelineController.getSourceConnectionConfig(any(), any())) + .thenReturn(mock(JdbcShardConfig.class)); + assertThrows(IllegalArgumentException.class, () -> SourceDbToSpanner.run(mockOptions)); + } + } } From 6effde963010d08514f2e1bce0a4cae421e190b2 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 2 Jul 2026 06:47:45 +0000 Subject: [PATCH 08/16] Astra DB Integration test update --- .../iowrapper/CassandraIOWrapperHelper.java | 1 - .../options/OptionsToConfigBuilderTest.java | 1 - .../CassandraIOWrapperHelperTest.java | 1 - ...rSimpleIT.java => AstraDbToSpannerIT.java} | 86 +++++++++++++------ .../loadtesting/SourceDbToSpannerLTBase.java | 1 + 5 files changed, 63 insertions(+), 27 deletions(-) rename v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/{AstraDbToSpannerSimpleIT.java => AstraDbToSpannerIT.java} (74%) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java index e91098eb1c..0313251184 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelper.java @@ -20,7 +20,6 @@ import com.datastax.oss.driver.api.core.config.OptionsMap; import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.datasource.DataSource; -import com.google.cloud.teleport.v2.reader.io.exception.SchemaDiscoveryException; import com.google.cloud.teleport.v2.reader.io.row.SourceRow; import com.google.cloud.teleport.v2.reader.io.schema.SchemaDiscovery; import com.google.cloud.teleport.v2.reader.io.schema.SchemaDiscoveryImpl; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java index b74404cf03..7c90d0ee35 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java @@ -22,7 +22,6 @@ import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import java.net.URISyntaxException; -import java.util.ArrayList; import java.util.List; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java index 1beefd8802..fbae2ea60c 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperHelperTest.java @@ -43,7 +43,6 @@ import com.datastax.oss.driver.api.core.config.TypedDriverOption; import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.datasource.DataSource; -import com.google.cloud.teleport.v2.reader.io.exception.SchemaDiscoveryException; import com.google.cloud.teleport.v2.reader.io.row.SourceRow; import com.google.cloud.teleport.v2.reader.io.schema.SourceSchema; import com.google.cloud.teleport.v2.reader.io.schema.SourceSchemaReference; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java similarity index 74% rename from v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java rename to v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java index 2f5cefc08c..54ea70753d 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerSimpleIT.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java @@ -36,6 +36,7 @@ import java.net.http.HttpRequest; import java.net.http.HttpResponse; import java.time.Duration; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -44,6 +45,7 @@ import org.apache.beam.it.common.TestProperties; import org.apache.beam.it.common.utils.ResourceManagerUtils; import org.apache.beam.it.gcp.spanner.SpannerResourceManager; +import org.apache.beam.it.gcp.spanner.matchers.SpannerAsserts; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -60,21 +62,32 @@ @RunWith(JUnit4.class) @Category(TemplateIntegrationTest.class) @TemplateIntegrationTest(SourceDbToSpanner.class) -public class AstraDbToSpannerSimpleIT extends SourceDbToSpannerITBase implements Serializable { +public class AstraDbToSpannerIT extends SourceDbToSpannerITBase implements Serializable { - private static final Logger LOGGER = LoggerFactory.getLogger(AstraDbToSpannerSimpleIT.class); + private static final Logger LOGGER = LoggerFactory.getLogger(AstraDbToSpannerIT.class); private static final long NUM_ROWS = 50L; private static final String ASTRA_DB = "dataflow_integration_tests"; private static final String ASTRA_DB_REGION = TestProperties.region(); private static final String ASTRA_KS = "beam"; - private static final String ASTRA_TBL = "scientist"; + + private static final String[][] SCIENTISTS = { + new String[] {"phys", "Einstein"}, + new String[] {"bio", "Darwin"}, + new String[] {"phys", "Copernicus"}, + new String[] {"bio", "Pasteur"}, + new String[] {"bio", "Curie"} + }; + + private String astraTable; private static DbOpsClient dbClient; private SpannerResourceManager spannerResourceManager; @Before public void setup() throws Exception { + astraTable = "scientist_" + testId.replaceAll("-", "_"); + spannerResourceManager = SpannerResourceManager.builder(testName, PROJECT, REGION).maybeUseStaticInstance().build(); @@ -86,7 +99,7 @@ public void setup() throws Exception { + " person_id INT64," + " person_name STRING(MAX)," + ") PRIMARY KEY(person_department, person_id)", - ASTRA_TBL); + astraTable); spannerResourceManager.executeDdlStatement(spannerDdl); // Setup Astra Db @@ -113,29 +126,61 @@ public void testAstraDbToSpanner() throws IOException { jobParameters.put("sourceDbDialect", "ASTRA_DB"); jobParameters.put("sourceConfigURL", sourceConfigURL); jobParameters.put("outputDirectory", getGcsPath("output", artifactClient)); + jobParameters.put("tables", astraTable); jobParameters.put( "ipConfiguration", "WORKER_IP_UNSPECIFIED"); // Require internet access for Astra API - // Act PipelineLauncher.LaunchInfo info = launchDataflowJob(testName, null, null, null, spannerResourceManager, jobParameters, null); - LOGGER.debug("Pipeline is now running"); + LOGGER.info("Pipeline is now running."); PipelineOperator.Result result = pipelineOperator().waitUntilDone(createConfig(info)); org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult(result).isLaunchFinished(); - LOGGER.debug("Destination Table has been populated."); + LOGGER.info("Destination Table has been populated."); - // Optionally verify a row + // Verify the row count matches what was inserted List rows = spannerResourceManager.readTableRecords( - ASTRA_TBL, List.of("person_department", "person_id", "person_name")); - assertThat(rows).isNotEmpty(); + astraTable, List.of("person_department", "person_id", "person_name")); + assertThat(rows).hasSize(Math.toIntExact(NUM_ROWS)); + + // Verify row data content using SpannerAsserts like other ITs + List> expectedData = getExpectedData(); + + SpannerAsserts.assertThatStructs(rows).hasRecordsUnorderedCaseInsensitiveColumns(expectedData); + } + + private List> getExpectedData() { + List> expectedData = new ArrayList<>(); + for (int i = 0; i < NUM_ROWS; i++) { + int index = i % SCIENTISTS.length; + Map expectedRow = new HashMap<>(); + expectedRow.put("person_department", SCIENTISTS[index][0]); + expectedRow.put("person_id", i); + expectedRow.put("person_name", SCIENTISTS[index][1]); + expectedData.add(expectedRow); + } + return expectedData; } @After public void tearDown() { ResourceManagerUtils.cleanResources(spannerResourceManager); + + if (dbClient != null) { + try (CqlSession astraSession = + CqlSession.builder() + .withCloudSecureConnectBundle( + new ByteArrayInputStream(dbClient.downloadDefaultSecureConnectBundle())) + .withAuthCredentials("token", dbClient.getToken()) + .withKeyspace(ASTRA_KS) + .build()) { + astraSession.execute(String.format("DROP TABLE IF EXISTS %s.%s;", ASTRA_KS, astraTable)); + } catch (Exception e) { + LOGGER.warn("Failed to drop Astra table", e); + } + } } private static String test() { @@ -146,7 +191,7 @@ private static String test() { private void createOrResumeAstraDatabase() throws InterruptedException { AstraDBOpsClient databasesClient = new AstraDBOpsClient(test()); if (databasesClient.findByName(ASTRA_DB).findAny().isEmpty()) { - LOGGER.debug("Create a new Database {}", ASTRA_DB); + LOGGER.info("Create a new Database {}", ASTRA_DB); databasesClient.create( DatabaseCreationRequest.builder() .name(ASTRA_DB) @@ -154,16 +199,16 @@ private void createOrResumeAstraDatabase() throws InterruptedException { .cloudRegion(ASTRA_DB_REGION) .build()); } else { - LOGGER.debug("Database {} exists in source organization", ASTRA_DB); + LOGGER.info("Database {} exists in source organization.", ASTRA_DB); } dbClient = databasesClient.databaseByName(ASTRA_DB); if (dbClient.get().getStatus() == DatabaseStatusType.HIBERNATED) { resumeDb(dbClient.get()); - LOGGER.debug("Resuming as DB was Hibernated"); + LOGGER.info("Resuming as DB was Hibernated."); } while (dbClient.get().getStatus() != DatabaseStatusType.ACTIVE) { Thread.sleep(5000); - LOGGER.debug("Waiting for DB to be ACTIVE...."); + LOGGER.info("Waiting for DB to be ACTIVE."); } } @@ -203,21 +248,14 @@ private void createAndPopulateTables() { String.format( "CREATE TABLE IF NOT EXISTS %s.%s(person_department text, person_id int, person_name text, PRIMARY KEY" + "((person_department), person_id));", - ASTRA_KS, ASTRA_TBL)); - String[][] scientists = { - new String[] {"phys", "Einstein"}, - new String[] {"bio", "Darwin"}, - new String[] {"phys", "Copernicus"}, - new String[] {"bio", "Pasteur"}, - new String[] {"bio", "Curie"} - }; + ASTRA_KS, astraTable)); for (int i = 0; i < NUM_ROWS; i++) { - int index = i % scientists.length; + int index = i % SCIENTISTS.length; String insertStr = String.format( "INSERT INTO %s.%s(person_department, person_id, person_name) values(" + "'%s', %d, '%s');", - ASTRA_KS, ASTRA_TBL, scientists[index][0], i, scientists[index][1]); + ASTRA_KS, astraTable, SCIENTISTS[index][0], i, SCIENTISTS[index][1]); astraSession.execute(insertStr); } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java index ba1388ecfc..314bc8b6bd 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java @@ -21,6 +21,7 @@ import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; +import com.google.gson.Gson; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.text.ParseException; From 7a7fddd70781cfd8be27387268fdbf02162a42e1 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 6 Jul 2026 10:04:04 +0000 Subject: [PATCH 09/16] feat: updated the documentation to reflect the config changes. --- v2/sourcedb-to-spanner/README.md | 4 ++-- .../README_Sourcedb_to_Spanner.md | 4 ++-- .../README_Sourcedb_to_Spanner_Flex.md | 2 +- .../v2/options/SourceDbToSpannerOptions.java | 18 +++++++----------- .../cassandra/iowrapper/AstraDbDataSource.java | 2 +- .../SourceConfig/astra-connection-config.json | 6 ++++++ .../SourceConfig/cassandra-driver-config.conf | 11 +++++++++++ .../SourceConfig/jdbc-shard-config.json | 14 ++++++++++++++ .../single-job-bulk-migration/variables.tf | 2 +- 9 files changed, 45 insertions(+), 18 deletions(-) create mode 100644 v2/sourcedb-to-spanner/src/test/resources/SourceConfig/astra-connection-config.json create mode 100644 v2/sourcedb-to-spanner/src/test/resources/SourceConfig/cassandra-driver-config.conf create mode 100644 v2/sourcedb-to-spanner/src/test/resources/SourceConfig/jdbc-shard-config.json diff --git a/v2/sourcedb-to-spanner/README.md b/v2/sourcedb-to-spanner/README.md index 1d911db78a..3744428b41 100644 --- a/v2/sourcedb-to-spanner/README.md +++ b/v2/sourcedb-to-spanner/README.md @@ -65,7 +65,7 @@ mvn test ### Executing Template #### Required Parameters -* **sourceConfigURL** (Configuration to connect to the source database): Can be the JDBC URL or the location of the sharding config. (Example: jdbc:mysql://10.10.10.10:3306/testdb or gs://test1/shard.conf). Refer to src/main/scripts/create_simple_shard_config.bash for steps to generate a shard configuration. +* **sourceConfigURL** (Source connection config file URL): The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. Refer to src/main/scripts/create_simple_shard_config.bash for steps to generate a shard configuration. * **username** (username of the source database): The username which can be used to connect to the source database. * **password** (username of the source database): The username which can be used to connect to the source database. * **instanceId** (Cloud Spanner Instance Id.): The destination Cloud Spanner instance. @@ -89,7 +89,7 @@ export JOB_NAME="${IMAGE_NAME}-`date +%Y%m%d-%H%M%S-%N`" gcloud dataflow flex-template run ${JOB_NAME} \ --project=${PROJECT} --region=us-central1 \ --template-file-gcs-location=${TEMPLATE_IMAGE_SPEC} \ - --parameters sourceConfigURL="jdbc:mysql://:3306/",username=,password=,instanceId="",databaseId="",projectId="$PROJECT",outputDirectory=gs:// \ + --parameters sourceConfigURL="gs:///source-config.json",username=,password=,instanceId="",databaseId="",projectId="$PROJECT",outputDirectory=gs:// \ --additional-experiments=disable_runner_v2 ``` #### Replaying DLQ entries. diff --git a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md index 445931ea4d..289e1d82ac 100644 --- a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md +++ b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md @@ -18,7 +18,7 @@ on [Metadata Annotations](https://github.com/GoogleCloudPlatform/DataflowTemplat ## Parameters #### Required Parameters -* **sourceConfigURL** (Configuration to connect to the source database): Can be the JDBC URL or the location of the sharding config. (Example: jdbc:mysql://10.10.10.10:3306/testdb or gs://test1/shard.conf) +* **sourceConfigURL** (Source connection config file URL.): The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. * **username** (username of the source database): The username which can be used to connect to the source database. * **password** (username of the source database): The username which can be used to connect to the source database. * **instanceId** (Cloud Spanner Instance Id.): The destination Cloud Spanner instance. @@ -519,7 +519,7 @@ resource "google_dataflow_flex_template_job" "sourcedb_to_spanner_flex" { instanceId = "" databaseId = "" projectId = "" - sourceConfigURL = "jdbc:mysql://some-host:3306/sampledb" + sourceConfigURL = "gs://your-bucket/source-config.json" username = "" password = "" outputDirectory = "gs://your-bucket/dir" diff --git a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md index 0d423e2ccb..4e8c8aa97e 100644 --- a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md +++ b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md @@ -30,7 +30,7 @@ on [Metadata Annotations](https://github.com/GoogleCloudPlatform/DataflowTemplat * **sourceDbDialect**: Possible values are `CASSANDRA`, `MYSQL` and `POSTGRESQL`. Defaults to: MYSQL. * **jdbcDriverJars**: The comma-separated list of driver JAR files. For example, `gs://your-bucket/driver_jar1.jar,gs://your-bucket/driver_jar2.jar`. Defaults to empty. * **jdbcDriverClassName**: The JDBC driver class name. For example, `com.mysql.jdbc.Driver`. Defaults to: com.mysql.jdbc.Driver. -* **sourceConfigURL**: The URL to connect to the source database host. This can be either: 1. A JDBC connection URL for a single source database, which must contain the host, port and source db name and can optionally contain properties like autoReconnect, maxReconnects etc. Format: `jdbc:{mysql|postgresql}://{host}:{port}/{dbName}?{parameters}`. For example,`jdbc:mysql://127.4.5.30:3306/my-db?autoReconnect=true&maxReconnects=10&unicode=true&characterEncoding=UTF-8`. 2. A Cloud Storage path to a shard config file for sharded migrations. For example, `gs://my-bucket/my-shard-config.yaml`. This parameter is required except for ASTRA_DB source. Defaults to empty. +* **sourceConfigURL**: The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. Defaults to empty. * **username**: The username to be used for the JDBC connection. Defaults to empty. * **password**: The password to be used for the JDBC connection. Defaults to empty. * **tables**: Tables to migrate from source. Defaults to empty. diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java index bb11ca7a32..61150db087 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java @@ -69,19 +69,15 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { @TemplateParameter.Text( order = 4, optional = true, - regexes = {"(^jdbc:mysql://.*|^jdbc:postgresql://.*|^gs://.*|^$)"}, + regexes = {"(^.+$)"}, groupName = "Source", - description = "Source database connection URL or shard config path.", + description = "Source connection config file URL.", helpText = - "The URL to connect to the source database host. This can be either:" - + " 1. A JDBC connection URL for a single source database, which" - + " must contain the host, port and source db name and can" - + " optionally contain properties like autoReconnect," - + " maxReconnects etc. Format: `jdbc:{mysql|postgresql}://{host}:{port}/{dbName}?{parameters}`." - + " For example,`jdbc:mysql://127.4.5.30:3306/my-db?autoReconnect=true&maxReconnects=10&unicode=true&characterEncoding=UTF-8`." - + " 2. A Cloud Storage path to a shard config file for sharded" - + " migrations. For example, `gs://my-bucket/my-shard-config.yaml`." - + " This parameter is required except for ASTRA_DB source.") + "The URL of the source connection config file. The file format is dependent on the source type." + + " For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json))." + + " For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json))." + + " For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf))." + + " This parameter is required.") @Default.String("") String getSourceConfigURL(); diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java index 6812287f77..689d3ff54a 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/AstraDbDataSource.java @@ -103,7 +103,7 @@ public Builder setAstraToken(String value) { astraToken = SecretManagerUtils.getSecret(value); } LOG.info("Astra Token is parsed"); - return this.setAstraToken(GuardedStringValueProvider.create(astraToken)); + return this.setAstraToken(GuardedStringValueProvider.create(value)); } public abstract Builder setKeySpace(String value); diff --git a/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/astra-connection-config.json b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/astra-connection-config.json new file mode 100644 index 0000000000..ea4ba3e789 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/astra-connection-config.json @@ -0,0 +1,6 @@ +{ + "databaseId": "test-database-id", + "astraToken": "AstraCS:test-token", + "keySpace": "test-keyspace", + "astraDbRegion": "us-east1" +} diff --git a/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/cassandra-driver-config.conf b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/cassandra-driver-config.conf new file mode 100644 index 0000000000..bfdf5a3af1 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/cassandra-driver-config.conf @@ -0,0 +1,11 @@ +datastax-java-driver { + basic.contact-points = ["127.0.0.1:9042", "127.0.0.1:9043"] + basic.session-keyspace = "test-keyspace" + basic.load-balancing-policy { + local-datacenter = "datacenter1" + } + advanced.auth-provider { + username = "testUserName" + password = "testPassword1234@" + } +} diff --git a/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/jdbc-shard-config.json b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/jdbc-shard-config.json new file mode 100644 index 0000000000..e429e593a0 --- /dev/null +++ b/v2/sourcedb-to-spanner/src/test/resources/SourceConfig/jdbc-shard-config.json @@ -0,0 +1,14 @@ +{ + "shardConfigs": [ + { + "logicalShardId": "shard1", + "host": "10.0.0.1", + "port": "3306", + "user": "db_user", + "password": "db_password", + "dbName": "test_db", + "namespace": "public", + "connectionProperties": "autoReconnect=true&maxReconnects=10" + } + ] +} diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf index 1e1231264f..8a0d6f4015 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf @@ -49,7 +49,7 @@ variable "jdbc_driver_class_name" { variable "source_config_url" { type = string - description = "JDBC connection url for the source database. Ex- jdbc:mysql://127.4.5.30:3306/my-db?autoReconnect=true&maxReconnects=10&unicode=true&characterEncoding=UTF-8" + description = "Source connection config file URL. The file format is dependent on the source type." } variable "username" { From 00def7aec5be551425ac529ce255c11cd573a7b2 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Tue, 7 Jul 2026 11:32:02 +0530 Subject: [PATCH 10/16] refactor: propagate SourceConnectionConfig through connector interfaces and update JDBC migration logic to utilize shard configurations --- .../v2/options/OptionsToConfigBuilder.java | 3 +- .../v2/options/SourceDbToSpannerOptions.java | 67 +------------------ .../v2/source/ISrcToSpSourceConnector.java | 6 +- .../CassandraSrcToSpSourceConnector.java | 9 ++- .../AbstractJdbcSrcToSpSourceConnector.java | 22 +++--- .../jdbc/ShardedJdbcDbConfigContainer.java | 5 +- .../SingleInstanceJdbcDbConfigContainer.java | 7 +- .../mysql/MySqlSrcToSpSourceConnector.java | 9 +-- .../PostgresSrcToSpSourceConnector.java | 5 +- .../v2/templates/PipelineController.java | 31 ++++----- .../v2/templates/SourceDbToSpanner.java | 4 +- .../CassandraIOWrapperFactoryTest.java | 20 ++++++ .../iowrapper/CassandraIoWrapperTest.java | 6 +- .../ShardedJdbcDbConfigContainerTest.java | 4 -- ...ngleInstanceJdbcDbConfigContainerTest.java | 13 +++- .../mysql/MySqlSourceConnectorTest.java | 22 ++---- .../v2/templates/PipelineControllerTest.java | 2 - 17 files changed, 94 insertions(+), 141 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java index 2d7df88499..2397d727a8 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java @@ -150,8 +150,7 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( } String sourceDbURL = - connector.getJdbcUrl( - sourceDbURL, host, port, dbName, connectionProperties, namespace, fetchSize); + connector.getJdbcUrl(host, port, dbName, connectionProperties, namespace, fetchSize); builder.setSourceDbURL(sourceDbURL); if (!StringUtils.isEmpty(shardId)) { diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java index 61150db087..64e7845cfa 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java @@ -68,8 +68,6 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { @TemplateParameter.Text( order = 4, - optional = true, - regexes = {"(^.+$)"}, groupName = "Source", description = "Source connection config file URL.", helpText = @@ -77,33 +75,13 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { + " For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json))." + " For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json))." + " For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf))." - + " This parameter is required.") + + " This parameter is required.", + example = "gs://your-bucket/source-config.json") @Default.String("") String getSourceConfigURL(); void setSourceConfigURL(String url); - @TemplateParameter.Text( - order = 5, - optional = true, - regexes = {"^.+$"}, - description = "JDBC connection username.", - helpText = "The username to be used for the JDBC connection.") - @Default.String("") - String getUsername(); // Make optional - - void setUsername(String username); - - @TemplateParameter.Password( - order = 6, - optional = true, - description = "JDBC connection password.", - helpText = "The password to be used for the JDBC connection.") - @Default.String("") - String getPassword(); // make optional - - void setPassword(String password); - @TemplateParameter.Text( order = 7, optional = true, @@ -350,47 +328,6 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setUniformizationStageCountHint(Long value); - @TemplateParameter.Text( - order = 28, - optional = true, - description = "Astra DB token", - helpText = - "AstraDB token, ignored for non-AstraDB dialects. This token is used to automatically download the securebundle by the tempalte.") - @Default.String("") - String getAstraDBToken(); - - void setAstraDBToken(String value); - - @TemplateParameter.Text( - order = 29, - optional = true, - description = "Astra DB databaseID", - helpText = "AstraDB databaseID, ignored for non-AstraDB dialects") - @Default.String("") - String getAstraDBDatabaseId(); - - void setAstraDBDatabaseId(String value); - - @TemplateParameter.Text( - order = 30, - optional = true, - description = "Astra DB keySpace", - helpText = "AstraDB keySpace, ignored for non-AstraDB dialects") - @Default.String("") - String getAstraDBKeySpace(); - - void setAstraDBKeySpace(String value); - - @TemplateParameter.Text( - order = 31, - optional = true, - description = "Astra DB Region", - helpText = "AstraDB region, ignored for non-AstraDB dialects") - @Default.String("") - String getAstraDBRegion(); - - void setAstraDBRegion(String value); - @TemplateParameter.Text( order = 32, optional = true, diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java index a37f66c1ab..b9d392cd6e 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java @@ -16,6 +16,7 @@ package com.google.cloud.teleport.v2.source; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; @@ -48,5 +49,8 @@ default String getDlqSourceType() { * @return The pipeline result. */ PipelineResult executeMigration( - SourceDbToSpannerOptions options, Pipeline pipeline, SpannerConfig spannerConfig); + SourceDbToSpannerOptions options, + SourceConnectionConfig sourceConnectionConfig, + Pipeline pipeline, + SpannerConfig spannerConfig); } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java index 2a4134546c..180a68d2c9 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java @@ -19,6 +19,7 @@ import com.google.cloud.teleport.v2.source.ISrcToSpSourceConnector; import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.iowrapper.CassandraIOWrapperFactory; import com.google.cloud.teleport.v2.spanner.migrations.constants.Constants; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.templates.DbConfigContainerDefaultImpl; import com.google.cloud.teleport.v2.templates.PipelineController; import org.apache.beam.sdk.Pipeline; @@ -39,11 +40,15 @@ public String getDlqSourceType() { @Override public PipelineResult executeMigration( - SourceDbToSpannerOptions options, Pipeline pipeline, SpannerConfig spannerConfig) { + SourceDbToSpannerOptions options, + SourceConnectionConfig sourceConnectionConfig, + Pipeline pipeline, + SpannerConfig spannerConfig) { return PipelineController.executeMigrationForDbConfigContainer( options, pipeline, spannerConfig, - new DbConfigContainerDefaultImpl(CassandraIOWrapperFactory.fromPipelineOptions(options))); + new DbConfigContainerDefaultImpl( + CassandraIOWrapperFactory.fromConfig(options, sourceConnectionConfig))); } } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java index 0885a27360..a4b9562307 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java @@ -23,8 +23,8 @@ import com.google.cloud.teleport.v2.reader.io.schema.typemapping.UnifiedTypeMapping; import com.google.cloud.teleport.v2.source.ISrcToSpSourceConnector; import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; -import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl; -import com.google.cloud.teleport.v2.spanner.migrations.utils.ShardFileReader; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.templates.DbConfigContainer; import com.google.cloud.teleport.v2.templates.PipelineController; import com.google.common.collect.ImmutableMap; @@ -33,6 +33,7 @@ import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; +import org.apache.parquet.Preconditions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,9 +45,16 @@ public abstract class AbstractJdbcSrcToSpSourceConnector implements ISrcToSpSour @Override public PipelineResult executeMigration( - SourceDbToSpannerOptions options, Pipeline pipeline, SpannerConfig spannerConfig) { + SourceDbToSpannerOptions options, + SourceConnectionConfig sourceConnectionConfig, + Pipeline pipeline, + SpannerConfig spannerConfig) { + Preconditions.checkArgument( + (sourceConnectionConfig instanceof JdbcShardConfig), + "Source config is not type of JdbcShardConfig."); DbConfigContainer dbConfigContainer; - if (options.getSourceConfigURL().startsWith("gs://")) { + List shards = ((JdbcShardConfig) sourceConnectionConfig).getShardConfigs(); + if (shards.size() > 1) { // TODO // Merge logical shards into 1 physical shard // Populate completion per shard @@ -54,16 +62,13 @@ public PipelineResult executeMigration( // Write to common DLQ ? SQLDialect sqlDialect = SQLDialect.valueOf(options.getSourceDbDialect()); - List shards = - new ShardFileReader(new SecretManagerAccessorImpl()) - .readForwardMigrationShardingConfig(options.getSourceConfigURL()); LOG.info( "running migration for {} shards: {}", shards.stream().count(), shards.stream().map(Shard::getHost).collect(Collectors.toList())); dbConfigContainer = new ShardedJdbcDbConfigContainer(shards, sqlDialect, options); } else { - dbConfigContainer = new SingleInstanceJdbcDbConfigContainer(options); + dbConfigContainer = new SingleInstanceJdbcDbConfigContainer(options, shards.get(0)); } return PipelineController.executeMigrationForDbConfigContainer( options, pipeline, spannerConfig, dbConfigContainer); @@ -81,7 +86,6 @@ public PipelineResult executeMigration( /** Gets the JDBC URL with source-specific properties added. */ public abstract String getJdbcUrl( - String jdbcUrl, String host, int port, String dbName, diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java index 5fb57bb7f1..493aa907c5 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java @@ -24,7 +24,6 @@ import com.google.common.collect.ImmutableList; import java.util.List; import java.util.Map; -import java.util.Optional; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.transforms.Wait; @@ -63,20 +62,18 @@ public JdbcIoWrapperConfigGroup getJdbcIoWrapperConfigGroup( // If a namespace is configured for a shard uses that, otherwise uses the namespace // configured in the options if there is one. - String namespace = Optional.ofNullable(shard.getNamespace()).orElse(options.getNamespace()); String dbName = entry.getKey(); JdbcIOWrapperConfig shardConfig = OptionsToConfigBuilder.getJdbcIOWrapperConfig( sqlDialect, sourceTables, - null, shard.getHost(), shard.getConnectionProperties(), Integer.parseInt(shard.getPort()), shard.getUserName(), shard.getPassword(), dbName, - namespace, + shard.getNamespace(), shardId, options.getJdbcDriverClassName(), options.getJdbcDriverJars(), diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java index b8280f21d0..8a22295c17 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java @@ -18,15 +18,18 @@ import com.google.cloud.teleport.v2.options.OptionsToConfigBuilder; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.JdbcIoWrapperConfigGroup; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import java.util.List; import org.apache.beam.sdk.transforms.Wait; /** Implementation of {@link JdbcDbConfigContainer} for single instance JDBC migration. */ public class SingleInstanceJdbcDbConfigContainer implements JdbcDbConfigContainer { private SourceDbToSpannerOptions options; + private Shard shard; - public SingleInstanceJdbcDbConfigContainer(SourceDbToSpannerOptions options) { + public SingleInstanceJdbcDbConfigContainer(SourceDbToSpannerOptions options, Shard shard) { this.options = options; + this.shard = shard; } @Override @@ -35,7 +38,7 @@ public JdbcIoWrapperConfigGroup getJdbcIoWrapperConfigGroup( return JdbcIoWrapperConfigGroup.builder() .addShardConfig( OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - options, sourceTables, null, waitOnSignal)) + options, shard, sourceTables, null, waitOnSignal)) .build(); } } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/mysql/MySqlSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/mysql/MySqlSrcToSpSourceConnector.java index 3d15290ed6..b14b8d86b3 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/mysql/MySqlSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/mysql/MySqlSrcToSpSourceConnector.java @@ -142,18 +142,15 @@ public SourceSchemaReference getSourceSchemaReference(String dbName, String name @Override public String getJdbcUrl( - String jdbcUrl, String host, int port, String dbName, String connectionProperties, String namespace, Integer fetchSize) { - if (jdbcUrl == null) { - jdbcUrl = "jdbc:mysql://" + host + ":" + port + "/" + dbName; - if (StringUtils.isNotBlank(connectionProperties)) { - jdbcUrl = jdbcUrl + "?" + connectionProperties; - } + String jdbcUrl = "jdbc:mysql://" + host + ":" + port + "/" + dbName; + if (StringUtils.isNotBlank(connectionProperties)) { + jdbcUrl = jdbcUrl + "?" + connectionProperties; } for (Entry entry : MySqlConfigDefaults.DEFAULT_MYSQL_URL_PROPERTIES.entrySet()) { diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/postgres/PostgresSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/postgres/PostgresSrcToSpSourceConnector.java index 01a890cea3..11b7007c4f 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/postgres/PostgresSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/postgres/PostgresSrcToSpSourceConnector.java @@ -181,16 +181,13 @@ public SourceSchemaReference getSourceSchemaReference(String dbName, String name @Override public String getJdbcUrl( - String jdbcUrl, String host, int port, String dbName, String connectionProperties, String namespace, Integer fetchSize) { - if (jdbcUrl == null) { - jdbcUrl = "jdbc:postgresql://" + host + ":" + port + "/" + dbName; - } + String jdbcUrl = "jdbc:postgresql://" + host + ":" + port + "/" + dbName; if (StringUtils.isBlank(namespace)) { namespace = "public"; } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java index 7944d4b2c4..fea05cd40b 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java @@ -24,15 +24,12 @@ import com.google.cloud.teleport.v2.spanner.migrations.schema.SchemaFileOverridesBasedMapper; import com.google.cloud.teleport.v2.spanner.migrations.schema.SchemaStringOverridesBasedMapper; import com.google.cloud.teleport.v2.spanner.migrations.schema.SessionBasedMapper; -import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; -import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConfigParser; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.spanner.SpannerSchema; import com.google.cloud.teleport.v2.spanner.migrations.utils.ISecretManagerAccessor; import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl; import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.ImmutableList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -223,19 +220,19 @@ static ISchemaMapper getSchemaMapper(SourceDbToSpannerOptions options, Ddl ddl) return schemaMapper; } - public static SourceConnectionConfig getSourceConnectionConfig( - String sourceType, String sourceShardsFilePath) { - ISecretManagerAccessor secretManagerAccessor = new SecretManagerAccessorImpl(); - SourceConfigParser sourceConfigParser = new SourceConfigParser(secretManagerAccessor); - SourceConnectionConfig sourceConnectionConfig; - try { - // Parse the source shards configuration file to respective - // SourceConnectionConfig. - LOG.info("Parsing source shards configuration file: {}", sourceShardsFilePath); - return sourceConfigParser.parseConfiguration(sourceType, sourceShardsFilePath); - } catch (Exception e) { - LOG.error("Error parsing source config", e); - throw new RuntimeException("Error parsing source config", e); - } + public static SourceConnectionConfig getSourceConnectionConfig( + String sourceType, String sourceShardsFilePath) { + ISecretManagerAccessor secretManagerAccessor = new SecretManagerAccessorImpl(); + SourceConfigParser sourceConfigParser = new SourceConfigParser(secretManagerAccessor); + SourceConnectionConfig sourceConnectionConfig; + try { + // Parse the source shards configuration file to respective + // SourceConnectionConfig. + LOG.info("Parsing source shards configuration file: {}", sourceShardsFilePath); + return sourceConfigParser.parseConfiguration(sourceType, sourceShardsFilePath); + } catch (Exception e) { + LOG.error("Error parsing source config", e); + throw new RuntimeException("Error parsing source config", e); } + } } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index 701ea7fc6c..304ed7e924 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -22,11 +22,9 @@ import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.source.ISrcToSpSourceConnector; import com.google.cloud.teleport.v2.source.SourceConnectorFactory; -import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.utils.DataflowWorkerMachineTypeUtils; import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Preconditions; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; @@ -130,7 +128,7 @@ static PipelineResult run(SourceDbToSpannerOptions options) { // Decide type and source of migration ISrcToSpSourceConnector connector = SourceConnectorFactory.getSourceConnectorByDialect(options); - return connector.executeMigration(options, pipeline, spannerConfig); + return connector.executeMigration(options, sourceConnectionConfig, pipeline, spannerConfig); } @VisibleForTesting diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java index 707aac6380..1b057941c8 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java @@ -188,4 +188,24 @@ public void testCassandraIoWrapperFactoryExceptions() { IllegalArgumentException.class, () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); } + + @Test + public void testCassandraIoWrapperFactoryOssWithOptionsMap() { + SourceDbToSpannerOptions mockOptions = + mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); + when(mockOptions.getSourceDbDialect()).thenReturn("CASSANDRA"); + when(mockOptions.getNumPartitions()).thenReturn(null); + CassandraConnectionConfig mockSourceConfig = mock(CassandraConnectionConfig.class); + OptionsMap mockOptionsMap = OptionsMap.driverDefaults(); + when(mockSourceConfig.getOptionsMap()).thenReturn(mockOptionsMap); + CassandraIOWrapperFactory cassandraIOWrapperFactory = + CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); + assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(mockOptionsMap); + assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.OSS); + assertThat(cassandraIOWrapperFactory.astraDBKeyspace()).isEqualTo(""); + assertThat(cassandraIOWrapperFactory.astraDBRegion()).isEqualTo(""); + assertThat(cassandraIOWrapperFactory.astraDBDatabaseId()).isEqualTo(""); + assertThat(cassandraIOWrapperFactory.astraDBToken()) + .isEqualTo(GuardedStringValueProvider.create("")); + } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java index aad79eadf3..7eb149b648 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIoWrapperTest.java @@ -91,11 +91,13 @@ public void testCassandraIoWrapperBasic() { ImmutableMap.of(ImmutableList.of(mockSourceTableReference), mockTableReader); try (MockedStatic mockCassandraIoWrapperHelper = mockStatic(CassandraIOWrapperHelper.class)) { + OptionsMap mockOptionsMap = Mockito.mock(OptionsMap.class); + mockCassandraIoWrapperHelper .when( () -> CassandraIOWrapperHelper.buildDataSource( - null, + mockOptionsMap, null, CassandraDialect.OSS, GuardedStringValueProvider.create(""), @@ -151,7 +153,7 @@ public void testCassandraIoWrapperBasic() { CassandraIoWrapper cassandraIoWrapper = new CassandraIoWrapper( - null, + mockOptionsMap, tablesToRead, null, CassandraDialect.OSS, diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainerTest.java index 78b0f3880a..1ff818a5ad 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainerTest.java @@ -76,8 +76,6 @@ public void shardedDbConfigContainerMySqlTest() { sourceDbToSpannerOptions.setJdbcDriverClassName(testDriverClassName); sourceDbToSpannerOptions.setMaxConnections(150); sourceDbToSpannerOptions.setNumPartitions(4000); - sourceDbToSpannerOptions.setUsername(testUser); - sourceDbToSpannerOptions.setPassword(testPassword); sourceDbToSpannerOptions.setTables("table1,table2"); mockedStaticJdbcIoWrapper .when(() -> JdbcIoWrapper.of(any(JdbcIoWrapperConfigGroup.class))) @@ -139,8 +137,6 @@ public void shardedDbConfigContainerPGTest() { sourceDbToSpannerOptions.setJdbcDriverClassName(testDriverClassName); sourceDbToSpannerOptions.setMaxConnections(150); sourceDbToSpannerOptions.setNumPartitions(4000); - sourceDbToSpannerOptions.setUsername(testUser); - sourceDbToSpannerOptions.setPassword(testPassword); sourceDbToSpannerOptions.setTables("table1,table2"); Shard shard = diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainerTest.java index 5b5a6e7d39..2ad69276a6 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainerTest.java @@ -21,6 +21,7 @@ import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.JdbcIOWrapperConfig; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.JdbcIoWrapperConfigGroup; import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import com.google.common.io.Resources; import java.nio.file.Paths; import java.util.List; @@ -55,8 +56,6 @@ public void testSingleInstanceJdbcDbConfigContainer() { sourceDbToSpannerOptions.setJdbcDriverClassName(testDriverClassName); sourceDbToSpannerOptions.setMaxConnections(150); sourceDbToSpannerOptions.setNumPartitions(4000); - sourceDbToSpannerOptions.setUsername(testUser); - sourceDbToSpannerOptions.setPassword(testPassword); String sessionFilePath = Paths.get(Resources.getResource("session-file-with-dropped-column.json").getPath()) .toString(); @@ -64,8 +63,16 @@ public void testSingleInstanceJdbcDbConfigContainer() { PCollection dummyPCollection = pipeline.apply(Create.of(1)); pipeline.run(); + + Shard shard = new Shard(); + shard.setHost("localhost"); + shard.setPort("3306"); + shard.setDbName("testDB"); + shard.setUser(testUser); + shard.setPassword(testPassword); + SingleInstanceJdbcDbConfigContainer dbConfigContainer = - new SingleInstanceJdbcDbConfigContainer(sourceDbToSpannerOptions); + new SingleInstanceJdbcDbConfigContainer(sourceDbToSpannerOptions, shard); JdbcIoWrapperConfigGroup configGroup = dbConfigContainer.getJdbcIoWrapperConfigGroup( List.of("table1", "table2"), Wait.on(dummyPCollection)); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/mysql/MySqlSourceConnectorTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/mysql/MySqlSourceConnectorTest.java index c0f7fc44d1..2272a39561 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/mysql/MySqlSourceConnectorTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/mysql/MySqlSourceConnectorTest.java @@ -35,7 +35,7 @@ public void testGetSourceType() { @Test public void testGetJdbcUrl_constructsUrl() { - String url = connector.getJdbcUrl(null, "localhost", 3306, "test_db", null, null, null); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", null, null, null); assertThat(url).startsWith("jdbc:mysql://localhost:3306/test_db"); // Should contain defaults assertThat(url).contains("allowMultiQueries=true"); @@ -47,8 +47,7 @@ public void testGetJdbcUrl_constructsUrl() { @Test public void testGetJdbcUrl_withConnectionProperties() { - String url = - connector.getJdbcUrl(null, "localhost", 3306, "test_db", "param1=value1", null, null); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", "param1=value1", null, null); assertThat(url).startsWith("jdbc:mysql://localhost:3306/test_db?param1=value1"); assertThat(url).contains("allowMultiQueries=true"); assertThat(url).contains("useCursorFetch=true"); @@ -56,31 +55,25 @@ public void testGetJdbcUrl_withConnectionProperties() { @Test public void testGetJdbcUrl_withFetchSizeNull_enablesCursorFetch() { - String url = - connector.getJdbcUrl( - "jdbc:mysql://localhost:3306/test_db", null, 0, null, null, null, null); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", null, null, null); assertThat(url).contains("useCursorFetch=true"); } @Test public void testGetJdbcUrl_withFetchSizePositive_enablesCursorFetch() { - String url = - connector.getJdbcUrl("jdbc:mysql://localhost:3306/test_db", null, 0, null, null, null, 42); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", null, null, 42); assertThat(url).contains("useCursorFetch=true"); } @Test public void testGetJdbcUrl_withFetchSizeZero_disablesCursorFetch() { - String url = - connector.getJdbcUrl("jdbc:mysql://localhost:3306/test_db", null, 0, null, null, null, 0); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", null, null, 0); assertThat(url).doesNotContain("useCursorFetch"); } @Test public void testGetJdbcUrl_preservesExistingParams() { - String url = - connector.getJdbcUrl( - "jdbc:mysql://localhost:3306/test_db?param=value", null, 0, null, null, null, null); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", "param=value", null, null); assertThat(url).startsWith("jdbc:mysql://localhost:3306/test_db?param=value"); assertThat(url).contains("allowMultiQueries=true"); assertThat(url).contains("useCursorFetch=true"); @@ -95,8 +88,7 @@ public void testGetJdbcUrl_withFetchSizeMinusOne_enablesCursorFetch() { // it would be treated as != 0, so it would enable cursor mode. // However, the main propagation test testFetchSizeMinusOneBehavesLikeNull // covers the normalization. - String url = - connector.getJdbcUrl("jdbc:mysql://localhost:3306/test_db", null, 0, null, null, null, -1); + String url = connector.getJdbcUrl("localhost", 3306, "test_db", null, null, -1); assertThat(url).contains("useCursorFetch=true"); } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java index cb30e425fa..4822e1b071 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java @@ -21,7 +21,6 @@ import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.when; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; @@ -65,7 +64,6 @@ import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.mockito.Mock; -import org.mockito.MockedConstruction; import org.mockito.MockedStatic; import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; From a6f89e659914854026589541e0a0aab5056248a8 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Tue, 7 Jul 2026 14:26:01 +0000 Subject: [PATCH 11/16] refactor: simplify Shard configuration handling by passing the Shard object directly to builder methods --- .../v2/options/OptionsToConfigBuilder.java | 43 ++++------ .../iowrapper/CassandraIOWrapperFactory.java | 33 ++------ .../jdbc/ShardedJdbcDbConfigContainer.java | 48 ++++------- .../SingleInstanceJdbcDbConfigContainer.java | 2 +- .../options/OptionsToConfigBuilderTest.java | 80 +++++++++---------- .../v2/templates/SourceDbToSpannerITBase.java | 2 - .../loadtesting/SourceDbToSpannerLTBase.java | 2 - 7 files changed, 76 insertions(+), 134 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java index 2397d727a8..bec7819607 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java @@ -58,11 +58,7 @@ public static String extractWorkerZone(PipelineOptions options) { } public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( - SourceDbToSpannerOptions options, - Shard shard, - List tables, - String shardId, - Wait.OnSignal waitOn) { + SourceDbToSpannerOptions options, Shard shard, List tables, Wait.OnSignal waitOn) { SQLDialect sqlDialect = SQLDialect.valueOf(options.getSourceDbDialect()); String jdbcDriverClassName = options.getJdbcDriverClassName(); @@ -80,14 +76,7 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( return getJdbcIOWrapperConfig( sqlDialect, tables, - shard.getHost(), - shard.getConnectionProperties(), - Integer.parseInt(shard.getPort()), - shard.getUserName(), - shard.getPassword(), - shard.getDbName(), - shard.getNamespace(), - shardId, + shard, jdbcDriverClassName, jdbcDriverJars, maxConnections, @@ -103,14 +92,7 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults( public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( SQLDialect sqlDialect, List tables, - String host, - String connectionProperties, - int port, - String username, - String password, - String dbName, - String namespace, - String shardId, + Shard shard, String jdbcDriverClassName, String jdbcDriverJars, long maxConnections, @@ -125,15 +107,16 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( SourceConnectorFactory.getSourceJdbcConnectorByDialect(sqlDialect); JdbcIOWrapperConfig.Builder builder = connector.getJdbcIOWrapperConfigBuilder(); SourceSchemaReference sourceSchemaReference = - connector.getSourceSchemaReference(dbName, namespace); + connector.getSourceSchemaReference(shard.getDbName(), shard.getNamespace()); builder = builder .setSourceSchemaReference(sourceSchemaReference) .setDbAuth( LocalCredentialsProvider.builder() .setUserName( - username) // TODO - support taking username and password from url as well - .setPassword(password) + shard.getUserName()) // TODO - support taking username and password from url + // as well + .setPassword(shard.getPassword()) .build()) .setJdbcDriverClassName(jdbcDriverClassName) .setJdbcDriverJars(jdbcDriverJars); @@ -150,11 +133,17 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( } String sourceDbURL = - connector.getJdbcUrl(host, port, dbName, connectionProperties, namespace, fetchSize); + connector.getJdbcUrl( + shard.getHost(), + Integer.parseInt(shard.getPort()), + shard.getDbName(), + shard.getConnectionProperties(), + shard.getNamespace(), + fetchSize); builder.setSourceDbURL(sourceDbURL); - if (!StringUtils.isEmpty(shardId)) { - builder.setShardID(shardId); + if (!StringUtils.isEmpty(shard.getLogicalShardId())) { + builder.setShardID(shard.getLogicalShardId()); } if (waitOn != null) { diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java index 78e93535fd..a8478ac184 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java @@ -60,29 +60,6 @@ public abstract class CassandraIOWrapperFactory implements IoWrapperFactory { /** Astra DB Region. * */ public abstract String astraDBRegion(); - private static CassandraIOWrapperFactory create( - OptionsMap optionsMap, - Integer numPartions, - String sourceDialect, - GuardedStringValueProvider astraDBToken, - String astraDBDatabaseId, - String astraDBKeyspace, - String astraDBRegion) { - CassandraDataSource.CassandraDialect cassandraDialect = - switch (sourceDialect) { - case SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT -> CassandraDialect.ASTRA; - default -> CassandraDialect.OSS; - }; - return new AutoValue_CassandraIOWrapperFactory( - optionsMap, - numPartions, - cassandraDialect, - astraDBToken, - astraDBDatabaseId, - astraDBKeyspace, - astraDBRegion); - } - public static CassandraIOWrapperFactory fromConfig( SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { Preconditions.checkArgument( @@ -110,11 +87,15 @@ public static CassandraIOWrapperFactory fromConfig( throw new IllegalArgumentException( "Unsupported source connection config type: " + sourceConnectionConfig); } - - return CassandraIOWrapperFactory.create( + CassandraDataSource.CassandraDialect cassandraDialect = + switch (options.getSourceDbDialect()) { + case SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT -> CassandraDialect.ASTRA; + default -> CassandraDialect.OSS; + }; + return new AutoValue_CassandraIOWrapperFactory( optionsMap, options.getNumPartitions(), - options.getSourceDbDialect(), + cassandraDialect, astraDBToken, astraDBDatabaseId, astraDBKeyspace, diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java index 493aa907c5..46b957d823 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/ShardedJdbcDbConfigContainer.java @@ -23,7 +23,6 @@ import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import com.google.common.collect.ImmutableList; import java.util.List; -import java.util.Map; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.transforms.Wait; @@ -56,37 +55,22 @@ public JdbcIoWrapperConfigGroup getJdbcIoWrapperConfigGroup( JdbcIoWrapperConfigGroup.builder().setSourceDbDialect(sqlDialect); for (Shard shard : shards) { // TODO Move towards clubbing all physical shards together in a single connection pool. - for (Map.Entry entry : shard.getDbNameToLogicalShardIdMap().entrySet()) { - // Read data from source - String shardId = entry.getValue(); - - // If a namespace is configured for a shard uses that, otherwise uses the namespace - // configured in the options if there is one. - String dbName = entry.getKey(); - JdbcIOWrapperConfig shardConfig = - OptionsToConfigBuilder.getJdbcIOWrapperConfig( - sqlDialect, - sourceTables, - shard.getHost(), - shard.getConnectionProperties(), - Integer.parseInt(shard.getPort()), - shard.getUserName(), - shard.getPassword(), - dbName, - shard.getNamespace(), - shardId, - options.getJdbcDriverClassName(), - options.getJdbcDriverJars(), - options.getMaxConnections(), - options.getNumPartitions(), - waitOnSignal, - options.getFetchSize(), - options.getUniformizationStageCountHint(), - options.getProjectId(), - workerZone, - options.as(DataflowPipelineWorkerPoolOptions.class).getWorkerMachineType()); - jdbcIoWrapperConfigGroupBuilder.addShardConfig(shardConfig); - } + JdbcIOWrapperConfig shardConfig = + OptionsToConfigBuilder.getJdbcIOWrapperConfig( + sqlDialect, + sourceTables, + shard, + options.getJdbcDriverClassName(), + options.getJdbcDriverJars(), + options.getMaxConnections(), + options.getNumPartitions(), + waitOnSignal, + options.getFetchSize(), + options.getUniformizationStageCountHint(), + options.getProjectId(), + workerZone, + options.as(DataflowPipelineWorkerPoolOptions.class).getWorkerMachineType()); + jdbcIoWrapperConfigGroupBuilder.addShardConfig(shardConfig); } return jdbcIoWrapperConfigGroupBuilder.build(); } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java index 8a22295c17..61dbe9ea26 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/SingleInstanceJdbcDbConfigContainer.java @@ -38,7 +38,7 @@ public JdbcIoWrapperConfigGroup getJdbcIoWrapperConfigGroup( return JdbcIoWrapperConfigGroup.builder() .addShardConfig( OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - options, shard, sourceTables, null, waitOnSignal)) + options, shard, sourceTables, waitOnSignal)) .build(); } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java index 7c90d0ee35..69c627e77a 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilderTest.java @@ -59,7 +59,6 @@ public void testConfigWithMySqlDefaultsFromOptions() { sourceDbToSpannerOptions, shard, List.of("table1", "table2"), - null, Wait.on(dummyPCollection)); assertThat(config.jdbcDriverClassName()).isEqualTo(testDriverClassName); assertThat(config.sourceDbURL()) @@ -77,7 +76,6 @@ public void testConfigWithMySqlDefaultsFromOptions() { sourceDbToSpannerOptions, shard, List.of("table1", "table2"), - null, Wait.on(dummyPCollection)) .maxFetchSize()) .isEqualTo(42); @@ -91,14 +89,17 @@ public void testConfigWithMySqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.MYSQL, List.of("table1", "table2"), - "myhost", - "testParam=testValue", - 3306, - "myuser", - "mypassword", - "mydb", - null, - null, + new Shard( + null, + "myhost", + "3306", + "myuser", + "mypassword", + "mydb", + null, + null, + "testParam=testValue", + ""), "com.mysql.jdbc.Driver", "mysql-jar", 10, @@ -114,14 +115,7 @@ public void testConfigWithMySqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.MYSQL, List.of("table1", "table2"), - "myhost", - null, - 3306, - "myuser", - "mypassword", - "mydb", - null, - null, + new Shard(null, "myhost", "3306", "myuser", "mypassword", "mydb", null, null, null, ""), "com.mysql.jdbc.Driver", "mysql-jar", 10, @@ -162,7 +156,6 @@ public void testConfigWithPostgreSQLDefaultsFromOptions() { sourceDbToSpannerOptions, shard, List.of("table1", "table2", "table3"), - null, Wait.on(dummyPCollection)); assertThat(config.jdbcDriverClassName()).isEqualTo(testDriverClassName); assertThat(config.sourceDbURL()).isEqualTo(testUrl + "?currentSchema=public"); @@ -181,14 +174,17 @@ public void testConfigWithPostgreSqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - "myhost", - "testParam=testValue", - 5432, - "myuser", - "mypassword", - "mydb", - null, - null, + new Shard( + null, + "myhost", + "5432", + "myuser", + "mypassword", + "mydb", + null, + null, + "testParam=testValue", + ""), "com.mysql.jdbc.Driver", "mysql-jar", 10, @@ -203,14 +199,7 @@ public void testConfigWithPostgreSqlUrlFromOptions() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - "myhost", - "", - 5432, - "myuser", - "mypassword", - "mydb", - null, - null, + new Shard(null, "myhost", "5432", "myuser", "mypassword", "mydb", null, null, "", ""), "com.mysql.jdbc.Driver", "mysql-jar", 10, @@ -237,14 +226,17 @@ public void testConfigWithPostgreSqlUrlWithNamespace() { OptionsToConfigBuilder.getJdbcIOWrapperConfig( SQLDialect.POSTGRESQL, List.of("table1", "table2"), - "myhost", - "", - 5432, - "myuser", - "mypassword", - "mydb", - "mynamespace", - null, + new Shard( + null, + "myhost", + "5432", + "myuser", + "mypassword", + "mydb", + "mynamespace", + null, + "", + ""), "com.mysql.jdbc.Driver", "mysql-jar", 10, @@ -327,7 +319,7 @@ public void testFetchSizeMinusOneBehavesLikeNull() { Shard shard = new Shard("", "localhost", "5432", "user", "password", "testDB", "", "", ""); JdbcIOWrapperConfig config = OptionsToConfigBuilder.getJdbcIOWrapperConfigWithDefaults( - options, shard, List.of("table1"), null, null); + options, shard, List.of("table1"), null); assertThat(config.maxFetchSize()).isNull(); } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java index 4bcbc59fe7..cca21cd3b5 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerITBase.java @@ -358,8 +358,6 @@ private Map getJdbcParameters( } catch (IOException e) { throw new RuntimeException(e); } - put("username", jdbcResourceManager.getUsername()); - put("password", jdbcResourceManager.getPassword()); put("jdbcDriverClassName", driverClassNameFrom(jdbcResourceManager)); } }; diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java index 314bc8b6bd..c2be7f667b 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/loadtesting/SourceDbToSpannerLTBase.java @@ -221,8 +221,6 @@ protected Map getJdbcParameters( Map params = new HashMap<>(); params.put("sourceDbDialect", dialect.name()); params.put("sourceConfigURL", connectionUrl); - params.put("username", username); - params.put("password", password); params.put("jdbcDriverClassName", driverClassName); return params; } From 2e4fb33265e6a54b2350d87ba14683a00c02bef9 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 13 Jul 2026 13:26:12 +0530 Subject: [PATCH 12/16] fix(Spanner): removed stale validations. --- .../v2/templates/SourceDbToSpanner.java | 19 ---------- .../postgres/PostgresSourceConnectorTest.java | 5 +-- .../v2/templates/MySQLSingleShardIT.java | 14 +++---- .../v2/templates/SourceDbToSpannerTest.java | 38 ------------------- 4 files changed, 7 insertions(+), 69 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index 304ed7e924..73fb8122ce 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -31,7 +31,6 @@ import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.ValueProvider; -import org.apache.commons.lang3.StringUtils; /** * A template that copies data from a relational database using JDBC to an existing Spanner @@ -90,22 +89,6 @@ protected static SourceDbToSpannerOptions getSourceDbToSpannerOptions(String[] a return options; } - /** - * Validates the provided pipeline options. - * - * @param options The execution parameters to the pipeline. - * @throws IllegalArgumentException if the provided options are invalid for the pipeline. - */ - @VisibleForTesting - static void validateOptions(SourceDbToSpannerOptions options) { - if (SourceDbToSpannerOptions.PG_SOURCE_DIALECT.equals(options.getSourceDbDialect()) - && StringUtils.isNotBlank(options.getNamespace()) - && !options.getNamespace().equals("public")) { - throw new IllegalArgumentException( - "Non-public namespaces are currently unsupported for PostgreSQL migrations."); - } - } - /** * Create the pipeline with the supplied options. * @@ -114,8 +97,6 @@ static void validateOptions(SourceDbToSpannerOptions options) { */ @VisibleForTesting static PipelineResult run(SourceDbToSpannerOptions options) { - // TODO - Validate if options are as expected - validateOptions(options); Pipeline pipeline = Pipeline.create(options); String workerMachineType = pipeline.getOptions().as(DataflowPipelineWorkerPoolOptions.class).getWorkerMachineType(); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/postgres/PostgresSourceConnectorTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/postgres/PostgresSourceConnectorTest.java index c0a1745626..53b2bbdb61 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/postgres/PostgresSourceConnectorTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/postgres/PostgresSourceConnectorTest.java @@ -51,14 +51,13 @@ public void testGetSourceSchemaReference_withCustomNamespace() { @Test public void testGetJdbcUrl_constructsUrl() { - String url = connector.getJdbcUrl(null, "localhost", 5432, "test_db", null, null, null); + String url = connector.getJdbcUrl("localhost", 5432, "test_db", null, null, null); assertThat(url).isEqualTo("jdbc:postgresql://localhost:5432/test_db?currentSchema=public"); } @Test public void testGetJdbcUrl_withConnectionProperties() { - String url = - connector.getJdbcUrl(null, "localhost", 5432, "test_db", "ssl=true", "myschema", null); + String url = connector.getJdbcUrl("localhost", 5432, "test_db", "ssl=true", "myschema", null); assertThat(url) .isEqualTo("jdbc:postgresql://localhost:5432/test_db?currentSchema=myschema&ssl=true"); } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MySQLSingleShardIT.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MySQLSingleShardIT.java index 0258c1bed4..532ebd18ea 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MySQLSingleShardIT.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MySQLSingleShardIT.java @@ -80,10 +80,6 @@ public void cleanUp() { ResourceManagerUtils.cleanResources(spannerResourceManager, mySQLResourceManager); } - /** - * TODO: This IT is currently not complete since shard id population is pending on reader. This - * test needs to be updated whenever reader support is added. - */ @Test public void singleShardWithIdPopulationTest() throws Exception { loadSQLFileResource(mySQLResourceManager, MYSQL_DUMP_FILE_RESOURCE); @@ -107,10 +103,10 @@ public void singleShardWithIdPopulationTest() throws Exception { private List> getExpectedData() { return List.of( - Map.of(PKID, 1, NAME, "Alice", STATUS, "active", SHARD_ID, "NULL"), - Map.of(PKID, 2, NAME, "Bob", STATUS, "inactive", SHARD_ID, "NULL"), - Map.of(PKID, 3, NAME, "Carol", STATUS, "pending", SHARD_ID, "NULL"), - Map.of(PKID, 4, NAME, "David", STATUS, "complete", SHARD_ID, "NULL"), - Map.of(PKID, 5, NAME, "Emily", STATUS, "error", SHARD_ID, "NULL")); + Map.of(PKID, 1, NAME, "Alice", STATUS, "active", SHARD_ID, "Shard1"), + Map.of(PKID, 2, NAME, "Bob", STATUS, "inactive", SHARD_ID, "Shard1"), + Map.of(PKID, 3, NAME, "Carol", STATUS, "pending", SHARD_ID, "Shard1"), + Map.of(PKID, 4, NAME, "David", STATUS, "complete", SHARD_ID, "Shard1"), + Map.of(PKID, 5, NAME, "Emily", STATUS, "error", SHARD_ID, "Shard1")); } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java index 811ec77ddd..f31e27a82d 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java @@ -89,44 +89,6 @@ public void testCreateSpannerConfig() { assertEquals(config.getMaxCommitDelay().get(), Duration.millis(42L)); } - @Test - public void testValidateOptions_PostgresThrowsExceptionForCustomNamespace() { - SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class); - when(mockOptions.getSourceDbDialect()).thenReturn(SourceDbToSpannerOptions.PG_SOURCE_DIALECT); - when(mockOptions.getNamespace()).thenReturn("sales"); - - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, () -> SourceDbToSpanner.validateOptions(mockOptions)); - - assertThat(exception.getMessage()) - .isEqualTo("Non-public namespaces are currently unsupported for PostgreSQL migrations."); - } - - @Test - public void testValidateOptions_PostgresSucceedsForEmptyNamespace() { - SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class); - when(mockOptions.getSourceDbDialect()).thenReturn(SourceDbToSpannerOptions.PG_SOURCE_DIALECT); - when(mockOptions.getNamespace()).thenReturn(""); - SourceDbToSpanner.validateOptions(mockOptions); - } - - @Test - public void testValidateOptions_PostgresSucceedsForPublicNamespace() { - SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class); - when(mockOptions.getSourceDbDialect()).thenReturn(SourceDbToSpannerOptions.PG_SOURCE_DIALECT); - when(mockOptions.getNamespace()).thenReturn("public"); - SourceDbToSpanner.validateOptions(mockOptions); - } - - @Test - public void testValidateOptions_NonPostgresDialectSucceeds() { - SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class); - when(mockOptions.getSourceDbDialect()) - .thenReturn(SourceDbToSpannerOptions.MYSQL_SOURCE_DIALECT); - SourceDbToSpanner.validateOptions(mockOptions); - } - @Test public void testRun_ValidationFailures() { SourceDbToSpannerOptions mockOptions = From a053d4ae69e03117d70d83c30b541408f64c1a3d Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 27 Jul 2026 14:09:15 +0530 Subject: [PATCH 13/16] PR fixes --- .../AbstractJdbcSrcToSpSourceConnector.java | 2 +- .../v2/templates/PipelineController.java | 1 - .../v2/templates/SourceDbToSpanner.java | 29 +++++++++++++++++++ .../SourceDbToSpannerFTBase.java | 1 - 4 files changed, 30 insertions(+), 3 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java index a4b9562307..eb559fce19 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/jdbc/AbstractJdbcSrcToSpSourceConnector.java @@ -27,13 +27,13 @@ import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.templates.DbConfigContainer; import com.google.cloud.teleport.v2.templates.PipelineController; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import java.util.List; import java.util.stream.Collectors; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; -import org.apache.parquet.Preconditions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java index fea05cd40b..be1435be29 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/PipelineController.java @@ -224,7 +224,6 @@ public static SourceConnectionConfig getSourceConnectionConfig( String sourceType, String sourceShardsFilePath) { ISecretManagerAccessor secretManagerAccessor = new SecretManagerAccessorImpl(); SourceConfigParser sourceConfigParser = new SourceConfigParser(secretManagerAccessor); - SourceConnectionConfig sourceConnectionConfig; try { // Parse the source shards configuration file to respective // SourceConnectionConfig. diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index 73fb8122ce..4cffaf10c5 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -22,15 +22,19 @@ import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.source.ISrcToSpSourceConnector; import com.google.cloud.teleport.v2.source.SourceConnectorFactory; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.utils.DataflowWorkerMachineTypeUtils; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.ValueProvider; +import org.apache.commons.lang3.StringUtils; /** * A template that copies data from a relational database using JDBC to an existing Spanner @@ -107,6 +111,8 @@ static PipelineResult run(SourceDbToSpannerOptions options) { PipelineController.getSourceConnectionConfig( options.getSourceDbDialect(), options.getSourceConfigURL()); + validateOptions(options, sourceConnectionConfig); + // Decide type and source of migration ISrcToSpSourceConnector connector = SourceConnectorFactory.getSourceConnectorByDialect(options); return connector.executeMigration(options, sourceConnectionConfig, pipeline, spannerConfig); @@ -126,4 +132,27 @@ static SpannerConfig createSpannerConfig(SourceDbToSpannerOptions options) { } return spannerConfig; } + + /** + * Validates the provided pipeline options. + * + * @param options The execution parameters to the pipeline. + * @throws IllegalArgumentException if the provided options are invalid for the pipeline. + */ + @VisibleForTesting + static void validateOptions( + SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { + if (SourceDbToSpannerOptions.PG_SOURCE_DIALECT.equals(options.getSourceDbDialect())) { + Preconditions.checkArgument( + (sourceConnectionConfig instanceof JdbcShardConfig), + "Postgresql dialect should have JDBC source config."); + for (Shard shard : ((JdbcShardConfig) sourceConnectionConfig).getShardConfigs()) { + if (StringUtils.isNotBlank(shard.getNamespace()) + && !shard.getNamespace().equals("public")) { + throw new IllegalArgumentException( + "Non-public namespaces are currently unsupported for PostgreSQL migrations."); + } + } + } + } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java index 670757cdf2..29fe78816e 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/failureinjectiontesting/SourceDbToSpannerFTBase.java @@ -129,7 +129,6 @@ protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( JdbcShardConfig jdbcShardConfig = new JdbcShardConfig(); jdbcShardConfig.setShardConfigs(List.of(shard)); String shardFileContents = new Gson().toJson(jdbcShardConfig); - LOG.info("Shard file contents: {}", shardFileContents); gcsResourceManager.createArtifact("input/shard.json", shardFileContents); // launch dataflow template From 2cb0f537f41e2f3617dc2a506c9d25a499e7db63 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Mon, 27 Jul 2026 09:31:13 +0000 Subject: [PATCH 14/16] Readme fix. --- v2/sourcedb-to-spanner/README.md | 4 +- .../README_Sourcedb_to_Spanner.md | 36 +++------------ .../README_Sourcedb_to_Spanner_Flex.md | 36 +-------------- .../v2/options/OptionsToConfigBuilder.java | 18 +------- .../v2/options/SourceDbToSpannerOptions.java | 46 +++++++++---------- .../v2/source/ISrcToSpSourceConnector.java | 1 + .../CassandraSrcToSpSourceConnector.java | 3 +- .../iowrapper/CassandraIOWrapperFactory.java | 19 ++------ .../v2/templates/SourceDbToSpanner.java | 1 + .../CassandraIOWrapperFactoryTest.java | 25 +++------- .../v2/templates/AstraDbToSpannerIT.java | 2 +- .../v2/templates/SourceDbToSpannerTest.java | 38 +++++++++++++++ .../Sourcedb_to_Spanner_Flex/dataflow_job.tf | 18 ++------ .../main.tf | 2 - .../terraform.tfvars | 4 +- .../terraform_simple.tfvars | 4 +- .../variables.tf | 10 +--- .../samples/single-job-bulk-migration/main.tf | 2 - .../terraform.tfvars | 4 +- .../terraform_simple.tfvars | 4 +- .../single-job-bulk-migration/variables.tf | 9 ---- 21 files changed, 93 insertions(+), 193 deletions(-) diff --git a/v2/sourcedb-to-spanner/README.md b/v2/sourcedb-to-spanner/README.md index 3744428b41..bd66ed30d3 100644 --- a/v2/sourcedb-to-spanner/README.md +++ b/v2/sourcedb-to-spanner/README.md @@ -66,8 +66,6 @@ mvn test #### Required Parameters * **sourceConfigURL** (Source connection config file URL): The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. Refer to src/main/scripts/create_simple_shard_config.bash for steps to generate a shard configuration. -* **username** (username of the source database): The username which can be used to connect to the source database. -* **password** (username of the source database): The username which can be used to connect to the source database. * **instanceId** (Cloud Spanner Instance Id.): The destination Cloud Spanner instance. * **databaseId** (Cloud Spanner Database Id.): The destination Cloud Spanner database. * **projectId** (Cloud Spanner Project Id.): This is the name of the Cloud Spanner project. @@ -89,7 +87,7 @@ export JOB_NAME="${IMAGE_NAME}-`date +%Y%m%d-%H%M%S-%N`" gcloud dataflow flex-template run ${JOB_NAME} \ --project=${PROJECT} --region=us-central1 \ --template-file-gcs-location=${TEMPLATE_IMAGE_SPEC} \ - --parameters sourceConfigURL="gs:///source-config.json",username=,password=,instanceId="",databaseId="",projectId="$PROJECT",outputDirectory=gs:// \ + --parameters sourceConfigURL="gs:///source-config.json",instanceId="",databaseId="",projectId="$PROJECT",outputDirectory=gs:// \ --additional-experiments=disable_runner_v2 ``` #### Replaying DLQ entries. diff --git a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md index 289e1d82ac..f2e50072be 100644 --- a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md +++ b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner.md @@ -18,9 +18,7 @@ on [Metadata Annotations](https://github.com/GoogleCloudPlatform/DataflowTemplat ## Parameters #### Required Parameters -* **sourceConfigURL** (Source connection config file URL.): The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. -* **username** (username of the source database): The username which can be used to connect to the source database. -* **password** (username of the source database): The username which can be used to connect to the source database. +* **sourceConfigURL** (Source connection config file URL): The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. * **instanceId** (Cloud Spanner Instance Id.): The destination Cloud Spanner instance. * **databaseId** (Cloud Spanner Database Id.): The destination Cloud Spanner database. * **projectId** (Cloud Spanner Project Id.): This is the name of the Cloud Spanner project. @@ -126,8 +124,6 @@ export OUTPUT_DIRECTORY= ### Optional export JDBC_DRIVER_JARS="" export JDBC_DRIVER_CLASS_NAME=com.mysql.jdbc.Driver -export USERNAME="" -export PASSWORD="" export TABLES="" export NUM_PARTITIONS=0 export SPANNER_HOST=https://batch-spanner.googleapis.com @@ -145,8 +141,6 @@ gcloud dataflow flex-template run "sourcedb-to-spanner-flex-job" \ --parameters "jdbcDriverJars=$JDBC_DRIVER_JARS" \ --parameters "jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME" \ --parameters "sourceConfigURL=$SOURCE_CONFIG_URL" \ - --parameters "username=$USERNAME" \ - --parameters "password=$PASSWORD" \ --parameters "tables=$TABLES" \ --parameters "numPartitions=$NUM_PARTITIONS" \ --parameters "instanceId=$INSTANCE_ID" \ @@ -186,8 +180,6 @@ export OUTPUT_DIRECTORY= ### Optional export JDBC_DRIVER_JARS="" export JDBC_DRIVER_CLASS_NAME=com.mysql.jdbc.Driver -export USERNAME="" -export PASSWORD="" export TABLES="" export NUM_PARTITIONS=0 export SPANNER_HOST=https://batch-spanner.googleapis.com @@ -204,7 +196,7 @@ mvn clean package -PtemplatesRun \ -Dregion="$REGION" \ -DjobName="sourcedb-to-spanner-flex-job" \ -DtemplateName="Sourcedb_to_Spanner_Flex" \ --Dparameters="jdbcDriverJars=$JDBC_DRIVER_JARS,jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME,sourceConfigURL=$SOURCE_CONFIG_URL,username=$USERNAME,password=$PASSWORD,tables=$TABLES,numPartitions=$NUM_PARTITIONS,instanceId=$INSTANCE_ID,databaseId=$DATABASE_ID,projectId=$PROJECT_ID,spannerHost=$SPANNER_HOST,maxConnections=$MAX_CONNECTIONS,sessionFilePath=$SESSION_FILE_PATH,outputDirectory=$OUTPUT_DIRECTORY,disabledAlgorithms=$DISABLED_ALGORITHMS,extraFilesToStage=$EXTRA_FILES_TO_STAGE,defaultLogLevel=$DEFAULT_LOG_LEVEL" \ +-Dparameters="jdbcDriverJars=$JDBC_DRIVER_JARS,jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME,sourceConfigURL=$SOURCE_CONFIG_URL,tables=$TABLES,numPartitions=$NUM_PARTITIONS,instanceId=$INSTANCE_ID,databaseId=$DATABASE_ID,projectId=$PROJECT_ID,spannerHost=$SPANNER_HOST,maxConnections=$MAX_CONNECTIONS,sessionFilePath=$SESSION_FILE_PATH,outputDirectory=$OUTPUT_DIRECTORY,disabledAlgorithms=$DISABLED_ALGORITHMS,extraFilesToStage=$EXTRA_FILES_TO_STAGE,defaultLogLevel=$DEFAULT_LOG_LEVEL" \ -f v2/sourcedb-to-spanner ``` @@ -371,13 +363,7 @@ For bulk data migration from AstraDB to spanner, here are a few prerequisites yo 2. Ensure that the VPC has network connectivity to your AstraDB instance. #### Prerequisite-2: AstraDB credentials and related details You will need the following Astra DB details: -1. AstraDB token. - 1. The AstraDB token can be generated from the database page. - 2. Please ensure that the token remains valid till the duration of the migration. Depending on the size of the database, the migration can take a few hours. -2. AstraDB Database ID -3. AstraDB Region - Leave it empty for default region. -4. AstraDB Keyspace - The keyspace you want to migrate to spanner. -Note that the template will automatically download the security bundle from the database. +1. AstraDB connection config JSON file uploaded to GCS. See [sample](src/test/resources/SourceConfig/astra-connection-config.json). #### Prerequisite-3: Active Astra DB database Please ensure that the AstraDB instance is active (not hibernated) through the migration. @@ -419,13 +405,8 @@ eport MACHINE_TYPE="" export INSTANCE_ID= export DATABASE_ID= export PROJECT_ID= -## Either the token directly (starting with `AstraCS`), or URL to gcp secret store. -ASTRA_DB_APPLICATION_TOKEN="AstraCS:" -## Astra DB database ID. -ASTRA_DB_ID="" -ASTRA_DB_KEYSPACE="" -## Astra DB region. Leave empty for default region. -ASTRA_DB_REGION="" +## Path to Astra connection config file in GCS. +export SOURCE_CONFIG_URL="gs:///astra-connection-config.json" #### Stores DLQ. export OUTPUT_DIRECTORY= @@ -471,11 +452,8 @@ gcloud dataflow flex-template run "sourcedb-to-spanner-flex-job" \ --template-file-gcs-location "$TEMPLATE_SPEC_GCSPATH" \ --additional-experiments="[\"disable_runner_v2\"]" \ --parameters "sourceDbDialect=ASTRA_DB" \ + --parameters "sourceConfigURL=$SOURCE_CONFIG_URL" \ --parameters "insertOnlyModeForSpannerMutations=$INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS" \ - --parameters "astraDBToken=${ASTRA_DB_APPLICATION_TOKEN}" \ - --parameters "astraDBRegion=${ASTRA_DB_REGION}" \ - --parameters "astraDBDatabaseId=${ASTRA_DB_ID}" \ - --parameters "astraDBKeySpace=${ASTRA_DB_KEYSPACE}" \ --parameters "instanceId=$INSTANCE_ID" \ --parameters "databaseId=$DATABASE_ID" \ --parameters "projectId=$PROJECT_ID" \ @@ -520,8 +498,6 @@ resource "google_dataflow_flex_template_job" "sourcedb_to_spanner_flex" { databaseId = "" projectId = "" sourceConfigURL = "gs://your-bucket/source-config.json" - username = "" - password = "" outputDirectory = "gs://your-bucket/dir" # jdbcDriverJars = "gs://your-bucket/driver_jar1.jar,gs://your-bucket/driver_jar2.jar" diff --git a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md index 4e8c8aa97e..d6b8732519 100644 --- a/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md +++ b/v2/sourcedb-to-spanner/README_Sourcedb_to_Spanner_Flex.md @@ -31,8 +31,6 @@ on [Metadata Annotations](https://github.com/GoogleCloudPlatform/DataflowTemplat * **jdbcDriverJars**: The comma-separated list of driver JAR files. For example, `gs://your-bucket/driver_jar1.jar,gs://your-bucket/driver_jar2.jar`. Defaults to empty. * **jdbcDriverClassName**: The JDBC driver class name. For example, `com.mysql.jdbc.Driver`. Defaults to: com.mysql.jdbc.Driver. * **sourceConfigURL**: The URL of the source connection config file. The file format is dependent on the source type. For Astra, it will point to an Astra connection config file ([sample](src/test/resources/SourceConfig/astra-connection-config.json)). For JDBC, it will point to a JDBC sharding config file ([sample](src/test/resources/SourceConfig/jdbc-shard-config.json)). For Cassandra, it will point to a Cassandra driver config file ([sample](src/test/resources/SourceConfig/cassandra-driver-config.conf)). This parameter is required. Defaults to empty. -* **username**: The username to be used for the JDBC connection. Defaults to empty. -* **password**: The password to be used for the JDBC connection. Defaults to empty. * **tables**: Tables to migrate from source. Defaults to empty. * **numPartitions**: The number of partitions. This, along with the lower and upper bound, form partitions strides for generated WHERE clause expressions used to split the partition column evenly. When the input is less than 1, the number is set to 1. Defaults to: 0. * **fetchSize**: The number of rows to fetch per page read for JDBC source. If not set, the default of JdbcIO of 50_000 rows gets used. If source dialect is Mysql, please see the note below. This ultimately translated to Statement.setFetchSize call at Jdbc layer. It should ONLY be used if the default value throws memory errors.Note for MySql Source: FetchSize is ignored by the Mysql connector unless, `useCursorFetch=true` is also part of the connection properties.In case, the fetchSize parameter is explicitly set, for MySql dialect, the pipeline will add `useCursorFetch=true` to the connection properties by default. @@ -50,10 +48,6 @@ on [Metadata Annotations](https://github.com/GoogleCloudPlatform/DataflowTemplat * **columnOverrides**: These are the column name overrides from source to spanner. They are written in thefollowing format: [{SourceTableName1.SourceColumnName1, SourceTableName1.SpannerColumnName1}, {SourceTableName2.SourceColumnName1, SourceTableName2.SpannerColumnName1}]Note that the SourceTableName should remain the same in both the source and spanner pair. To override table names, use tableOverrides.The example shows mapping SingerName to TalentName and AlbumName to RecordName in Singers and Albums table respectively. For example, `[{Singers.SingerName, Singers.TalentName}, {Albums.AlbumName, Albums.RecordName}]`. Defaults to empty. * **schemaOverridesFilePath**: A file which specifies the table and the column name overrides from source to spanner. Defaults to empty. * **uniformizationStageCountHint**: Hint for number of uniformization stages. Currently Applicable only for jdbc based sources like MySQL or PostgreSQL. Leave 0 or default to disable uniformization. Set to -1 for a log(numPartition) number of stages. If your source primary key space is uniformly distributed (for example an auto-incrementing key with sparse holes), it's based to leave it disabled. If your keyspace is not uniform, you might encounter a laggard VM in your dataflow run. In such a case, you can set it to -1 to enable uniformization. Manually setting it to values other than 0 or -1 would help you fine tune the tradeoff of the overhead added by uniformization stages and the performance improvement due to better distribution of work. -* **astraDBToken**: AstraDB token, ignored for non-AstraDB dialects. This token is used to automatically download the securebundle by the tempalte. Defaults to empty. -* **astraDBDatabaseId**: AstraDB databaseID, ignored for non-AstraDB dialects. Defaults to empty. -* **astraDBKeySpace**: AstraDB keySpace, ignored for non-AstraDB dialects. Defaults to empty. -* **astraDBRegion**: AstraDB region, ignored for non-AstraDB dialects. Defaults to empty. * **failureInjectionParameter**: Failure injection parameter. Only used for testing. Defaults to empty. * **maxCommitDelay**: Maximum commit delay time to optimize write throughput in Spanner. Reference https://cloud.google.com/spanner/docs/throughput-optimized-writes.Set -1 to let spanner choose the default. Set to a positive value to override for best suited tradeoff of throughput vs latency.Defaults to -1. * **gcsOutputDirectory**: This directory is used to write the AVRO files of the records read from source. For example, `gs://your-bucket/your-path`. Defaults to empty. @@ -161,8 +155,6 @@ export SOURCE_DB_DIALECT=MYSQL export JDBC_DRIVER_JARS="" export JDBC_DRIVER_CLASS_NAME=com.mysql.jdbc.Driver export SOURCE_CONFIG_URL="" -export USERNAME="" -export PASSWORD="" export TABLES="" export NUM_PARTITIONS=0 export FETCH_SIZE= @@ -172,7 +164,6 @@ export SESSION_FILE_PATH="" export TRANSFORMATION_JAR_PATH="" export TRANSFORMATION_CLASS_NAME="" export TRANSFORMATION_CUSTOM_PARAMETERS="" -export NAMESPACE="" export INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS=false export BATCH_SIZE_FOR_SPANNER_MUTATIONS=-1 export SPANNER_PRIORITY=HIGH @@ -180,10 +171,6 @@ export TABLE_OVERRIDES="" export COLUMN_OVERRIDES="" export SCHEMA_OVERRIDES_FILE_PATH="" export UNIFORMIZATION_STAGE_COUNT_HINT=0 -export ASTRA_DBTOKEN="" -export ASTRA_DBDATABASE_ID="" -export ASTRA_DBKEY_SPACE="" -export ASTRA_DBREGION="" export FAILURE_INJECTION_PARAMETER="" export MAX_COMMIT_DELAY=-1 export GCS_OUTPUT_DIRECTORY="" @@ -198,8 +185,6 @@ gcloud dataflow flex-template run "sourcedb-to-spanner-flex-job" \ --parameters "jdbcDriverJars=$JDBC_DRIVER_JARS" \ --parameters "jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME" \ --parameters "sourceConfigURL=$SOURCE_CONFIG_URL" \ - --parameters "username=$USERNAME" \ - --parameters "password=$PASSWORD" \ --parameters "tables=$TABLES" \ --parameters "numPartitions=$NUM_PARTITIONS" \ --parameters "fetchSize=$FETCH_SIZE" \ @@ -213,7 +198,6 @@ gcloud dataflow flex-template run "sourcedb-to-spanner-flex-job" \ --parameters "transformationJarPath=$TRANSFORMATION_JAR_PATH" \ --parameters "transformationClassName=$TRANSFORMATION_CLASS_NAME" \ --parameters "transformationCustomParameters=$TRANSFORMATION_CUSTOM_PARAMETERS" \ - --parameters "namespace=$NAMESPACE" \ --parameters "insertOnlyModeForSpannerMutations=$INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS" \ --parameters "batchSizeForSpannerMutations=$BATCH_SIZE_FOR_SPANNER_MUTATIONS" \ --parameters "spannerPriority=$SPANNER_PRIORITY" \ @@ -221,10 +205,6 @@ gcloud dataflow flex-template run "sourcedb-to-spanner-flex-job" \ --parameters "columnOverrides=$COLUMN_OVERRIDES" \ --parameters "schemaOverridesFilePath=$SCHEMA_OVERRIDES_FILE_PATH" \ --parameters "uniformizationStageCountHint=$UNIFORMIZATION_STAGE_COUNT_HINT" \ - --parameters "astraDBToken=$ASTRA_DBTOKEN" \ - --parameters "astraDBDatabaseId=$ASTRA_DBDATABASE_ID" \ - --parameters "astraDBKeySpace=$ASTRA_DBKEY_SPACE" \ - --parameters "astraDBRegion=$ASTRA_DBREGION" \ --parameters "failureInjectionParameter=$FAILURE_INJECTION_PARAMETER" \ --parameters "maxCommitDelay=$MAX_COMMIT_DELAY" \ --parameters "gcsOutputDirectory=$GCS_OUTPUT_DIRECTORY" \ @@ -258,8 +238,6 @@ export SOURCE_DB_DIALECT=MYSQL export JDBC_DRIVER_JARS="" export JDBC_DRIVER_CLASS_NAME=com.mysql.jdbc.Driver export SOURCE_CONFIG_URL="" -export USERNAME="" -export PASSWORD="" export TABLES="" export NUM_PARTITIONS=0 export FETCH_SIZE= @@ -269,7 +247,6 @@ export SESSION_FILE_PATH="" export TRANSFORMATION_JAR_PATH="" export TRANSFORMATION_CLASS_NAME="" export TRANSFORMATION_CUSTOM_PARAMETERS="" -export NAMESPACE="" export INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS=false export BATCH_SIZE_FOR_SPANNER_MUTATIONS=-1 export SPANNER_PRIORITY=HIGH @@ -277,10 +254,6 @@ export TABLE_OVERRIDES="" export COLUMN_OVERRIDES="" export SCHEMA_OVERRIDES_FILE_PATH="" export UNIFORMIZATION_STAGE_COUNT_HINT=0 -export ASTRA_DBTOKEN="" -export ASTRA_DBDATABASE_ID="" -export ASTRA_DBKEY_SPACE="" -export ASTRA_DBREGION="" export FAILURE_INJECTION_PARAMETER="" export MAX_COMMIT_DELAY=-1 export GCS_OUTPUT_DIRECTORY="" @@ -294,7 +267,7 @@ mvn clean package -PtemplatesRun \ -Dregion="$REGION" \ -DjobName="sourcedb-to-spanner-flex-job" \ -DtemplateName="Sourcedb_to_Spanner_Flex" \ --Dparameters="sourceDbDialect=$SOURCE_DB_DIALECT,jdbcDriverJars=$JDBC_DRIVER_JARS,jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME,sourceConfigURL=$SOURCE_CONFIG_URL,username=$USERNAME,password=$PASSWORD,tables=$TABLES,numPartitions=$NUM_PARTITIONS,fetchSize=$FETCH_SIZE,instanceId=$INSTANCE_ID,databaseId=$DATABASE_ID,projectId=$PROJECT_ID,spannerHost=$SPANNER_HOST,maxConnections=$MAX_CONNECTIONS,sessionFilePath=$SESSION_FILE_PATH,outputDirectory=$OUTPUT_DIRECTORY,transformationJarPath=$TRANSFORMATION_JAR_PATH,transformationClassName=$TRANSFORMATION_CLASS_NAME,transformationCustomParameters=$TRANSFORMATION_CUSTOM_PARAMETERS,namespace=$NAMESPACE,insertOnlyModeForSpannerMutations=$INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS,batchSizeForSpannerMutations=$BATCH_SIZE_FOR_SPANNER_MUTATIONS,spannerPriority=$SPANNER_PRIORITY,tableOverrides=$TABLE_OVERRIDES,columnOverrides=$COLUMN_OVERRIDES,schemaOverridesFilePath=$SCHEMA_OVERRIDES_FILE_PATH,uniformizationStageCountHint=$UNIFORMIZATION_STAGE_COUNT_HINT,astraDBToken=$ASTRA_DBTOKEN,astraDBDatabaseId=$ASTRA_DBDATABASE_ID,astraDBKeySpace=$ASTRA_DBKEY_SPACE,astraDBRegion=$ASTRA_DBREGION,failureInjectionParameter=$FAILURE_INJECTION_PARAMETER,maxCommitDelay=$MAX_COMMIT_DELAY,gcsOutputDirectory=$GCS_OUTPUT_DIRECTORY,disabledAlgorithms=$DISABLED_ALGORITHMS,extraFilesToStage=$EXTRA_FILES_TO_STAGE" \ +-Dparameters="sourceDbDialect=$SOURCE_DB_DIALECT,jdbcDriverJars=$JDBC_DRIVER_JARS,jdbcDriverClassName=$JDBC_DRIVER_CLASS_NAME,sourceConfigURL=$SOURCE_CONFIG_URL,tables=$TABLES,numPartitions=$NUM_PARTITIONS,fetchSize=$FETCH_SIZE,instanceId=$INSTANCE_ID,databaseId=$DATABASE_ID,projectId=$PROJECT_ID,spannerHost=$SPANNER_HOST,maxConnections=$MAX_CONNECTIONS,sessionFilePath=$SESSION_FILE_PATH,outputDirectory=$OUTPUT_DIRECTORY,transformationJarPath=$TRANSFORMATION_JAR_PATH,transformationClassName=$TRANSFORMATION_CLASS_NAME,transformationCustomParameters=$TRANSFORMATION_CUSTOM_PARAMETERS,insertOnlyModeForSpannerMutations=$INSERT_ONLY_MODE_FOR_SPANNER_MUTATIONS,batchSizeForSpannerMutations=$BATCH_SIZE_FOR_SPANNER_MUTATIONS,spannerPriority=$SPANNER_PRIORITY,tableOverrides=$TABLE_OVERRIDES,columnOverrides=$COLUMN_OVERRIDES,schemaOverridesFilePath=$SCHEMA_OVERRIDES_FILE_PATH,uniformizationStageCountHint=$UNIFORMIZATION_STAGE_COUNT_HINT,failureInjectionParameter=$FAILURE_INJECTION_PARAMETER,maxCommitDelay=$MAX_COMMIT_DELAY,gcsOutputDirectory=$GCS_OUTPUT_DIRECTORY,disabledAlgorithms=$DISABLED_ALGORITHMS,extraFilesToStage=$EXTRA_FILES_TO_STAGE" \ -f v2/sourcedb-to-spanner ``` @@ -347,8 +320,6 @@ resource "google_dataflow_flex_template_job" "sourcedb_to_spanner_flex" { # jdbcDriverJars = "" # jdbcDriverClassName = "com.mysql.jdbc.Driver" # sourceConfigURL = "" - # username = "" - # password = "" # tables = "" # numPartitions = "0" # fetchSize = "" @@ -358,7 +329,6 @@ resource "google_dataflow_flex_template_job" "sourcedb_to_spanner_flex" { # transformationJarPath = "" # transformationClassName = "" # transformationCustomParameters = "" - # namespace = "" # insertOnlyModeForSpannerMutations = "false" # batchSizeForSpannerMutations = "-1" # spannerPriority = "HIGH" @@ -366,10 +336,6 @@ resource "google_dataflow_flex_template_job" "sourcedb_to_spanner_flex" { # columnOverrides = "" # schemaOverridesFilePath = "" # uniformizationStageCountHint = "0" - # astraDBToken = "" - # astraDBDatabaseId = "" - # astraDBKeySpace = "" - # astraDBRegion = "" # failureInjectionParameter = "" # maxCommitDelay = "-1" # gcsOutputDirectory = "" diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java index bec7819607..62a932d349 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java @@ -27,8 +27,6 @@ import com.google.common.collect.ImmutableList; import com.google.re2j.Matcher; import com.google.re2j.Pattern; -import java.net.URI; -import java.net.URISyntaxException; import java.util.List; import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.options.PipelineOptions; @@ -113,9 +111,7 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig( .setSourceSchemaReference(sourceSchemaReference) .setDbAuth( LocalCredentialsProvider.builder() - .setUserName( - shard.getUserName()) // TODO - support taking username and password from url - // as well + .setUserName(shard.getUserName()) .setPassword(shard.getPassword()) .build()) .setJdbcDriverClassName(jdbcDriverClassName) @@ -223,17 +219,5 @@ public static String addParamToJdbcUrl(String jdbcUrl, String paramName, String } } - private static String extractDbFromURL(String sourceDbUrl) { - URI uri; - try { - // Strip off the prefix 'jdbc:' which the library cannot handle. - uri = new URI(sourceDbUrl.substring(5)); - } catch (URISyntaxException e) { - throw new RuntimeException(String.format("Unable to parse url: %s", sourceDbUrl), e); - } - // Remove '/' before returning. - return uri.getPath().substring(1); - } - private OptionsToConfigBuilder() {} } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java index 64e7845cfa..1f00faf510 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/SourceDbToSpannerOptions.java @@ -83,7 +83,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setSourceConfigURL(String url); @TemplateParameter.Text( - order = 7, + order = 5, optional = true, description = "colon-separated names of the tables in the source database.", helpText = "Tables to migrate from source.") @@ -94,7 +94,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { /* TODO(pipelineController) allow per table NumPartitions. */ @TemplateParameter.Integer( - order = 8, + order = 6, optional = true, description = "The number of partitions.", helpText = @@ -107,7 +107,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setNumPartitions(Integer value); @TemplateParameter.Integer( - order = 9, + order = 7, optional = true, description = "The number of rows to fetch per page read for JDBC source.", helpText = @@ -120,7 +120,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setFetchSize(Integer value); @TemplateParameter.Text( - order = 10, + order = 8, groupName = "Target", description = "Cloud Spanner Instance Id.", helpText = "The destination Cloud Spanner instance.") @@ -129,7 +129,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setInstanceId(String value); @TemplateParameter.Text( - order = 11, + order = 9, groupName = "Target", regexes = {"^[a-z]([a-z0-9_-]{0,28})[a-z0-9]$"}, description = "Cloud Spanner Database Id.", @@ -139,7 +139,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setDatabaseId(String value); @TemplateParameter.ProjectId( - order = 12, + order = 10, groupName = "Target", description = "Cloud Spanner Project Id.", helpText = "This is the name of the Cloud Spanner project.") @@ -148,7 +148,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setProjectId(String projectId); @TemplateParameter.Text( - order = 13, + order = 11, optional = true, description = "Cloud Spanner Endpoint to call", helpText = "The Cloud Spanner endpoint to call in the template.", @@ -159,7 +159,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setSpannerHost(String value); @TemplateParameter.Integer( - order = 14, + order = 12, optional = true, description = "Maximum number of connections to Source database per worker", helpText = @@ -171,7 +171,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setMaxConnections(Integer value); @TemplateParameter.GcsReadFile( - order = 15, + order = 13, optional = true, description = "Session File Path in Cloud Storage, to provide mapping information in the form of a session file", @@ -184,7 +184,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setSessionFilePath(String value); @TemplateParameter.GcsReadFile( - order = 16, + order = 14, description = "Output directory for failed/skipped/filtered events", helpText = "This directory is used to dump the failed/skipped/filtered records in a migration.") @@ -193,7 +193,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setOutputDirectory(String value); @TemplateParameter.GcsReadFile( - order = 17, + order = 15, optional = true, description = "Custom jar location in Cloud Storage", helpText = @@ -204,7 +204,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setTransformationJarPath(String value); @TemplateParameter.Text( - order = 18, + order = 16, optional = true, description = "Custom class name", helpText = @@ -216,7 +216,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setTransformationClassName(String value); @TemplateParameter.Text( - order = 19, + order = 17, optional = true, description = "Custom parameters for transformation", helpText = @@ -227,7 +227,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setTransformationCustomParameters(String value); @TemplateParameter.Text( - order = 21, + order = 18, optional = true, description = "Use Inserts instead of Upserts for spanner mutations.", helpText = @@ -238,7 +238,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setInsertOnlyModeForSpannerMutations(Boolean value); @TemplateParameter.Text( - order = 22, + order = 19, optional = true, description = "BatchSize for Spanner Mutation.", helpText = @@ -249,7 +249,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setBatchSizeForSpannerMutations(Long value); @TemplateParameter.Enum( - order = 23, + order = 20, enumOptions = { @TemplateParameter.TemplateEnumOption("LOW"), @TemplateParameter.TemplateEnumOption("MEDIUM"), @@ -266,7 +266,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setSpannerPriority(Options.RpcPriority value); @TemplateParameter.Text( - order = 24, + order = 21, optional = true, description = "Table name overrides from source to spanner", regexes = @@ -282,7 +282,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setTableOverrides(String value); @TemplateParameter.Text( - order = 25, + order = 22, optional = true, regexes = "^\\[([[:space:]]*\\{[[:space:]]*[[:graph:]]+\\.[[:graph:]]+[[:space:]]*,[[:space:]]*[[:graph:]]+\\.[[:graph:]]+[[:space:]]*\\}[[:space:]]*(,[[:space:]]*)*)*\\]$", @@ -299,7 +299,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setColumnOverrides(String value); @TemplateParameter.Text( - order = 26, + order = 23, optional = true, description = "File based overrides from source to spanner", helpText = @@ -310,7 +310,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setSchemaOverridesFilePath(String value); @TemplateParameter.Text( - order = 27, + order = 24, optional = true, description = "Hint for number of uniformization stages. Currently Applicable only for jdc based sources like MySql or PG. Leave 0 or default to disable uniformization. Set to -1 for a log(numPartition) number of stages.", @@ -329,7 +329,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setUniformizationStageCountHint(Long value); @TemplateParameter.Text( - order = 32, + order = 25, optional = true, description = "Failure injection parameter", helpText = "Failure injection parameter. Only used for testing.") @@ -339,7 +339,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setFailureInjectionParameter(String value); @TemplateParameter.Text( - order = 33, + order = 26, optional = true, description = "Maximum commit delay time (in milliseconds) to optimize write throughput in Spanner. Reference https://cloud.google.com/spanner/docs/throughput-optimized-writes", @@ -353,7 +353,7 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions { void setMaxCommitDelay(Long value); @TemplateParameter.GcsWriteFolder( - order = 34, + order = 27, optional = true, description = "GCS directory for AVRO files", helpText = "This directory is used to write the AVRO files of the records read from source.", diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java index b9d392cd6e..e6b7d882b4 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/ISrcToSpSourceConnector.java @@ -44,6 +44,7 @@ default String getDlqSourceType() { * Executes the migration pipeline for the source database. * * @param options Pipeline options. + * @param sourceConnectionConfig Parsed source connection config. * @param pipeline The Beam pipeline. * @param spannerConfig Spanner configuration. * @return The pipeline result. diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java index 180a68d2c9..8c68f3f131 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/CassandraSrcToSpSourceConnector.java @@ -49,6 +49,7 @@ public PipelineResult executeMigration( pipeline, spannerConfig, new DbConfigContainerDefaultImpl( - CassandraIOWrapperFactory.fromConfig(options, sourceConnectionConfig))); + CassandraIOWrapperFactory.fromConfig( + sourceConnectionConfig, options.getNumPartitions()))); } } diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java index a8478ac184..cd274824bf 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java @@ -17,7 +17,6 @@ import com.datastax.oss.driver.api.core.config.OptionsMap; import com.google.auto.value.AutoValue; -import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; import com.google.cloud.teleport.v2.reader.IoWrapperFactory; import com.google.cloud.teleport.v2.reader.auth.dbauth.GuardedStringValueProvider; import com.google.cloud.teleport.v2.reader.io.IoWrapper; @@ -25,7 +24,6 @@ import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; -import com.google.common.base.Preconditions; import java.util.List; import javax.annotation.Nullable; import org.apache.beam.sdk.transforms.Wait.OnSignal; @@ -61,13 +59,7 @@ public abstract class CassandraIOWrapperFactory implements IoWrapperFactory { public abstract String astraDBRegion(); public static CassandraIOWrapperFactory fromConfig( - SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { - Preconditions.checkArgument( - options.getSourceDbDialect().equals(SourceDbToSpannerOptions.CASSANDRA_SOURCE_DIALECT) - || options - .getSourceDbDialect() - .equals(SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT), - "Unexpected Dialect " + options.getSourceDbDialect() + " for Cassandra Source"); + SourceConnectionConfig sourceConnectionConfig, Integer numPartitions) { GuardedStringValueProvider astraDBToken = GuardedStringValueProvider.create(""); String astraDBDatabaseId = ""; @@ -75,26 +67,23 @@ public static CassandraIOWrapperFactory fromConfig( String astraDBRegion = ""; OptionsMap optionsMap = null; + CassandraDataSource.CassandraDialect cassandraDialect = CassandraDialect.OSS; if (sourceConnectionConfig instanceof AstraConnectionConfig) { AstraConnectionConfig astraConfig = (AstraConnectionConfig) sourceConnectionConfig; astraDBToken = GuardedStringValueProvider.create(astraConfig.getAstraToken()); astraDBDatabaseId = astraConfig.getDatabaseId(); astraDBKeyspace = astraConfig.getKeySpace(); astraDBRegion = astraConfig.getAstraDbRegion(); + cassandraDialect = CassandraDialect.ASTRA; } else if (sourceConnectionConfig instanceof CassandraConnectionConfig) { optionsMap = ((CassandraConnectionConfig) sourceConnectionConfig).getOptionsMap(); } else { throw new IllegalArgumentException( "Unsupported source connection config type: " + sourceConnectionConfig); } - CassandraDataSource.CassandraDialect cassandraDialect = - switch (options.getSourceDbDialect()) { - case SourceDbToSpannerOptions.ASTRA_DB_SOURCE_DIALECT -> CassandraDialect.ASTRA; - default -> CassandraDialect.OSS; - }; return new AutoValue_CassandraIOWrapperFactory( optionsMap, - options.getNumPartitions(), + numPartitions, cassandraDialect, astraDBToken, astraDBDatabaseId, diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index 4cffaf10c5..bc7e62506a 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -137,6 +137,7 @@ static SpannerConfig createSpannerConfig(SourceDbToSpannerOptions options) { * Validates the provided pipeline options. * * @param options The execution parameters to the pipeline. + * @param sourceConnectionConfig Parsed source connection config. * @throws IllegalArgumentException if the provided options are invalid for the pipeline. */ @VisibleForTesting diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java index 1b057941c8..1a1ccb3744 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java @@ -38,6 +38,7 @@ import com.google.cloud.teleport.v2.source.cassandra.reader.io.cassandra.schema.CassandraSchemaReference; import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; +import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.SourceConnectionConfig; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -135,14 +136,10 @@ public void cleanup() { @Test public void testCassandraIoWrapperFactoryOssBasic() { - SourceDbToSpannerOptions mockOptions = - mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); - when(mockOptions.getSourceDbDialect()).thenReturn("CASSANDRA"); - when(mockOptions.getNumPartitions()).thenReturn(null); CassandraConnectionConfig mockSourceConfig = mock(CassandraConnectionConfig.class); when(mockSourceConfig.getOptionsMap()).thenReturn(null); CassandraIOWrapperFactory cassandraIOWrapperFactory = - CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); + CassandraIOWrapperFactory.fromConfig(mockSourceConfig, null); assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(null); assertThat(cassandraIOWrapperFactory.getIOWrapper(TABLES_TO_READ, null).discoverTableSchema()) .isEqualTo(ImmutableList.of(mockSourceSchema)); @@ -156,11 +153,6 @@ public void testCassandraIoWrapperFactoryOssBasic() { @Test public void testCassandraIoWrapperFactoryAstraBasic() { - SourceDbToSpannerOptions mockOptions = - mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); - when(mockOptions.getSourceDbDialect()).thenReturn("ASTRA_DB"); - when(mockOptions.getNumPartitions()).thenReturn(null); - AstraConnectionConfig mockSourceConfig = mock(AstraConnectionConfig.class); when(mockSourceConfig.getAstraToken()).thenReturn("AstraCS:testToken"); when(mockSourceConfig.getDatabaseId()).thenReturn("testId"); @@ -168,7 +160,7 @@ public void testCassandraIoWrapperFactoryAstraBasic() { when(mockSourceConfig.getKeySpace()).thenReturn("testKeyspace"); CassandraIOWrapperFactory cassandraIOWrapperFactory = - CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); + CassandraIOWrapperFactory.fromConfig(mockSourceConfig, null); assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(null); assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.ASTRA); assertThat(cassandraIOWrapperFactory.astraDBKeyspace()).isEqualTo("testKeyspace"); @@ -183,23 +175,18 @@ public void testCassandraIoWrapperFactoryExceptions() { SourceDbToSpannerOptions mockOptions = mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); when(mockOptions.getSourceDbDialect()).thenReturn("MYSQL"); - SourceConnectionConfig mockConfig = mock(SourceConnectionConfig.class); + SourceConnectionConfig mockConfig = mock(JdbcShardConfig.class); assertThrows( - IllegalArgumentException.class, - () -> CassandraIOWrapperFactory.fromConfig(mockOptions, mockConfig)); + IllegalArgumentException.class, () -> CassandraIOWrapperFactory.fromConfig(mockConfig, 1)); } @Test public void testCassandraIoWrapperFactoryOssWithOptionsMap() { - SourceDbToSpannerOptions mockOptions = - mock(SourceDbToSpannerOptions.class, Mockito.withSettings().serializable()); - when(mockOptions.getSourceDbDialect()).thenReturn("CASSANDRA"); - when(mockOptions.getNumPartitions()).thenReturn(null); CassandraConnectionConfig mockSourceConfig = mock(CassandraConnectionConfig.class); OptionsMap mockOptionsMap = OptionsMap.driverDefaults(); when(mockSourceConfig.getOptionsMap()).thenReturn(mockOptionsMap); CassandraIOWrapperFactory cassandraIOWrapperFactory = - CassandraIOWrapperFactory.fromConfig(mockOptions, mockSourceConfig); + CassandraIOWrapperFactory.fromConfig(mockSourceConfig, 1); assertThat(cassandraIOWrapperFactory.optionsMap()).isEqualTo(mockOptionsMap); assertThat(cassandraIOWrapperFactory.cassandraDialect()).isEqualTo(CassandraDialect.OSS); assertThat(cassandraIOWrapperFactory.astraDBKeyspace()).isEqualTo(""); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java index 54ea70753d..533d9e6eb9 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/AstraDbToSpannerIT.java @@ -97,7 +97,7 @@ public void setup() throws Exception { "CREATE TABLE %s (" + " person_department STRING(MAX)," + " person_id INT64," - + " person_name STRING(MAX)," + + " person_name STRING(MAX)" + ") PRIMARY KEY(person_department, person_id)", astraTable); spannerResourceManager.executeDdlStatement(spannerDdl); diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java index f31e27a82d..2db1fa665b 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/SourceDbToSpannerTest.java @@ -28,9 +28,11 @@ import com.google.cloud.teleport.v2.common.CommonTemplateJvmInitializer; import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions; +import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard; import com.google.cloud.teleport.v2.spanner.migrations.source.config.AstraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.CassandraConnectionConfig; import com.google.cloud.teleport.v2.spanner.migrations.source.config.JdbcShardConfig; +import java.util.Arrays; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.joda.time.Duration; @@ -132,4 +134,40 @@ public void testRun_ValidationFailures() { assertThrows(IllegalArgumentException.class, () -> SourceDbToSpanner.run(mockOptions)); } } + + @Test + public void testValidateOptions() { + SourceDbToSpannerOptions mockOptions = + PipelineOptionsFactory.as(SourceDbToSpannerOptions.class); + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.PG_SOURCE_DIALECT); + + AstraConnectionConfig astraConfig = new AstraConnectionConfig(); + IllegalArgumentException exception1 = + assertThrows( + IllegalArgumentException.class, + () -> SourceDbToSpanner.validateOptions(mockOptions, astraConfig)); + assertEquals("Postgresql dialect should have JDBC source config.", exception1.getMessage()); + + JdbcShardConfig jdbcConfig = new JdbcShardConfig(); + Shard shard1 = new Shard(); + shard1.setNamespace("public"); + Shard shard2 = new Shard(); + shard2.setNamespace("custom"); + jdbcConfig.setShardConfigs(Arrays.asList(shard1, shard2)); + + IllegalArgumentException exception2 = + assertThrows( + IllegalArgumentException.class, + () -> SourceDbToSpanner.validateOptions(mockOptions, jdbcConfig)); + assertEquals( + "Non-public namespaces are currently unsupported for PostgreSQL migrations.", + exception2.getMessage()); + + jdbcConfig.setShardConfigs(Arrays.asList(shard1)); + SourceDbToSpanner.validateOptions(mockOptions, jdbcConfig); + + mockOptions.setSourceDbDialect(SourceDbToSpannerOptions.MYSQL_SOURCE_DIALECT); + jdbcConfig.setShardConfigs(Arrays.asList(shard1, shard2)); + SourceDbToSpanner.validateOptions(mockOptions, jdbcConfig); + } } diff --git a/v2/sourcedb-to-spanner/terraform/Sourcedb_to_Spanner_Flex/dataflow_job.tf b/v2/sourcedb-to-spanner/terraform/Sourcedb_to_Spanner_Flex/dataflow_job.tf index e134ab6b74..a258ab5a60 100644 --- a/v2/sourcedb-to-spanner/terraform/Sourcedb_to_Spanner_Flex/dataflow_job.tf +++ b/v2/sourcedb-to-spanner/terraform/Sourcedb_to_Spanner_Flex/dataflow_job.tf @@ -45,23 +45,13 @@ variable "jdbcDriverClassName" { default = null } -variable "sourceDbURL" { +variable "sourceConfigURL" { type = string - description = "The JDBC connection URL string. For example, `jdbc:mysql://127.4.5.30:3306/my-db?autoReconnect=true&maxReconnects=10&unicode=true&characterEncoding=UTF-8`." + description = "Source connection config file URL. The file format is dependent on the source type." } -variable "username" { - type = string - description = "The username to be used for the JDBC connection. Defaults to empty." - default = null -} -variable "password" { - type = string - description = "The password to be used for the JDBC connection. Defaults to empty." - default = null -} variable "tables" { type = string @@ -262,9 +252,7 @@ resource "google_dataflow_flex_template_job" "generated" { parameters = { jdbcDriverJars = var.jdbcDriverJars jdbcDriverClassName = var.jdbcDriverClassName - sourceDbURL = var.sourceDbURL - username = var.username - password = var.password + sourceConfigURL = var.sourceConfigURL tables = var.tables numPartitions = tostring(var.numPartitions) instanceId = var.instanceId diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/main.tf b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/main.tf index e809882c12..07b76cbc6d 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/main.tf +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/main.tf @@ -37,8 +37,6 @@ resource "google_dataflow_flex_template_job" "generated" { jdbcDriverClassName = var.jdbc_driver_class_name maxConnections = tostring(var.max_connections) sourceConfigURL = var.source_config_url - username = var.username - password = var.password numPartitions = tostring(var.num_partitions) instanceId = var.instance_id databaseId = var.database_id diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform.tfvars b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform.tfvars index dc13da2dd5..60322848c1 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform.tfvars +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform.tfvars @@ -6,10 +6,8 @@ working_directory_bucket = "bucket-name" # example "test-bucket" working_directory_prefix = "path/to/working/directory" # should not start or end with a '/' jdbc_driver_jars = "gs://path/to/driver/jars/postgresql-42.7.3.jar" jdbc_driver_class_name = "org.postgresql.Driver" -source_config_url = "jdbc:postgresql://127.4.5.30:5432/my-db" +source_config_url = "gs://your-bucket/jdbc-shard-config.json" source_db_dialect = "POSTGRESQL" -username = "postgres" -password = "abc" num_partitions = 4000 max_connections = 320 instance_id = "my-spanner-instance" diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform_simple.tfvars b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform_simple.tfvars index 5802092da1..153394d6a1 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform_simple.tfvars +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/terraform_simple.tfvars @@ -5,12 +5,10 @@ project = "project-name" region = "us-central1" # Or your desired region working_directory_bucket = "bucket-name" # example "test-bucket" working_directory_prefix = "path/to/working/directory" # should not start or end with a '/' -source_config_url = "jdbc:postgresql://127.4.5.30:5432/my-db" +source_config_url = "gs://your-bucket/jdbc-shard-config.json" source_db_dialect = "POSTGRESQL" -username = "postgres" jdbc_driver_jars = "gs://path/to/driver/jars/postgresql-42.7.3.jar" jdbc_driver_class_name = "org.postgresql.Driver" -password = "abc" instance_id = "my-spanner-instance" database_id = "my-spanner-database" spanner_project_id = "my-spanner-project" \ No newline at end of file diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/variables.tf b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/variables.tf index 1fba149fbe..9b999aa3c7 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/variables.tf +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration-postgres/variables.tf @@ -49,18 +49,10 @@ variable "jdbc_driver_class_name" { variable "source_config_url" { type = string - description = "JDBC connection url for the source database. Ex- jdbc:postgresql://127.4.5.30:5432/my-db" + description = "Source connection config file URL. The file format is dependent on the source type." } -variable "username" { - type = string - description = "Username to log in to the specified source database" -} -variable "password" { - type = string - description = "Password to log in to the specified source database" -} variable "num_partitions" { type = number diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/main.tf b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/main.tf index 7cf3f8e17d..bbe34427d7 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/main.tf +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/main.tf @@ -53,8 +53,6 @@ resource "google_dataflow_flex_template_job" "generated" { jdbcDriverClassName = var.jdbc_driver_class_name maxConnections = tostring(var.max_connections) sourceConfigURL = var.source_config_url - username = var.username - password = var.password numPartitions = tostring(var.num_partitions) instanceId = var.instance_id databaseId = var.database_id diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform.tfvars b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform.tfvars index 78141543c7..29b76c7a24 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform.tfvars +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform.tfvars @@ -6,9 +6,7 @@ working_directory_bucket = "bucket-name" # example "test-bucket" working_directory_prefix = "path/to/working/directory" # should not start or end with a '/' jdbc_driver_jars = "gs://path/to/driver/jars" jdbc_driver_class_name = "com.mysql.jdbc.driver" -source_config_url = "jdbc:mysql://127.4.5.30:3306/my-db" -username = "root" -password = "abc" +source_config_url = "gs://your-bucket/jdbc-shard-config.json" num_partitions = 4000 max_connections = 320 instance_id = "my-spanner-instance" diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform_simple.tfvars b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform_simple.tfvars index 45f17267e5..ef71072846 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform_simple.tfvars +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/terraform_simple.tfvars @@ -5,9 +5,7 @@ project = "project-name" region = "us-central1" # Or your desired region working_directory_bucket = "bucket-name" # example "test-bucket" working_directory_prefix = "path/to/working/directory" # should not start or end with a '/' -source_config_url = "jdbc:mysql://127.4.5.30:3306/my-db" -username = "root" -password = "abc" +source_config_url = "gs://your-bucket/jdbc-shard-config.json" instance_id = "my-spanner-instance" database_id = "my-spanner-database" spanner_project_id = "my-spanner-project" \ No newline at end of file diff --git a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf index 8a0d6f4015..e6ecc412f5 100644 --- a/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf +++ b/v2/sourcedb-to-spanner/terraform/samples/single-job-bulk-migration/variables.tf @@ -52,15 +52,6 @@ variable "source_config_url" { description = "Source connection config file URL. The file format is dependent on the source type." } -variable "username" { - type = string - description = "Username to log in to the specified source database" -} - -variable "password" { - type = string - description = "Password to log in to the specified source database" -} variable "num_partitions" { type = number From 77121fc6a6506a597c68b90b885f9bf10b0d5079 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Tue, 4 Aug 2026 16:32:38 +0530 Subject: [PATCH 15/16] minor refactoring. --- .../v2/templates/SourceDbToSpanner.java | 48 +++++++++---------- .../CassandraIOWrapperFactoryTest.java | 2 - 2 files changed, 24 insertions(+), 26 deletions(-) diff --git a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java index bc7e62506a..8866c7e887 100644 --- a/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java +++ b/v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/SourceDbToSpanner.java @@ -93,6 +93,30 @@ protected static SourceDbToSpannerOptions getSourceDbToSpannerOptions(String[] a return options; } + /** + * Validates the provided pipeline options. TODO: move this to source connector. + * + * @param options The execution parameters to the pipeline. + * @param sourceConnectionConfig Parsed source connection config. + * @throws IllegalArgumentException if the provided options are invalid for the pipeline. + */ + @VisibleForTesting + static void validateOptions( + SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { + if (SourceDbToSpannerOptions.PG_SOURCE_DIALECT.equals(options.getSourceDbDialect())) { + Preconditions.checkArgument( + (sourceConnectionConfig instanceof JdbcShardConfig), + "Postgresql dialect should have JDBC source config."); + for (Shard shard : ((JdbcShardConfig) sourceConnectionConfig).getShardConfigs()) { + if (StringUtils.isNotBlank(shard.getNamespace()) + && !shard.getNamespace().equals("public")) { + throw new IllegalArgumentException( + "Non-public namespaces are currently unsupported for PostgreSQL migrations."); + } + } + } + } + /** * Create the pipeline with the supplied options. * @@ -132,28 +156,4 @@ static SpannerConfig createSpannerConfig(SourceDbToSpannerOptions options) { } return spannerConfig; } - - /** - * Validates the provided pipeline options. - * - * @param options The execution parameters to the pipeline. - * @param sourceConnectionConfig Parsed source connection config. - * @throws IllegalArgumentException if the provided options are invalid for the pipeline. - */ - @VisibleForTesting - static void validateOptions( - SourceDbToSpannerOptions options, SourceConnectionConfig sourceConnectionConfig) { - if (SourceDbToSpannerOptions.PG_SOURCE_DIALECT.equals(options.getSourceDbDialect())) { - Preconditions.checkArgument( - (sourceConnectionConfig instanceof JdbcShardConfig), - "Postgresql dialect should have JDBC source config."); - for (Shard shard : ((JdbcShardConfig) sourceConnectionConfig).getShardConfigs()) { - if (StringUtils.isNotBlank(shard.getNamespace()) - && !shard.getNamespace().equals("public")) { - throw new IllegalArgumentException( - "Non-public namespaces are currently unsupported for PostgreSQL migrations."); - } - } - } - } } diff --git a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java index 1a1ccb3744..327367995b 100644 --- a/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java +++ b/v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactoryTest.java @@ -60,8 +60,6 @@ @RunWith(MockitoJUnitRunner.class) public class CassandraIOWrapperFactoryTest { private MockedStatic mockCassandraIoWrapperHelper; - private static final String TEST_BUCKET_CASSANDRA_CONFIG_CONF = - "gs://smt-test-bucket/cassandraConfig.conf"; private static final ImmutableList TABLES_TO_READ = ImmutableList.of(BASIC_TEST_TABLE, PRIMITIVE_TYPES_TABLE); @Mock SourceSchema mockSourceSchema; From 2d3b2a3ea199bec59a2b69e295238cd4361379c5 Mon Sep 17 00:00:00 2001 From: Pratick Chokhani Date: Thu, 6 Aug 2026 08:12:45 +0000 Subject: [PATCH 16/16] refactor: unify shard configuration format and update E2E tests to use standardized shard config GCS file --- ...dAndReverseMigrationShardedEndToEndIT.java | 3 +- .../endtoend/EndToEndTestingITBase.java | 119 ++++++++---------- .../v2/templates/EndToEndTestingITBase.java | 69 ++++------ ...hardedBulkMigrationAndValidationE2EIT.java | 2 +- 4 files changed, 80 insertions(+), 113 deletions(-) diff --git a/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/BulkForwardAndReverseMigrationShardedEndToEndIT.java b/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/BulkForwardAndReverseMigrationShardedEndToEndIT.java index 60dd24458e..72228d15af 100644 --- a/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/BulkForwardAndReverseMigrationShardedEndToEndIT.java +++ b/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/BulkForwardAndReverseMigrationShardedEndToEndIT.java @@ -161,8 +161,7 @@ public void setUp() throws IOException, InterruptedException { "", "", databases); - createAndUploadBulkShardConfigToGcs( - new ArrayList<>(List.of(dataShard)), gcsResourceManager); + createAndUploadShardConfigToGcs(List.of(dataShard), gcsResourceManager); // create pubsub manager pubsubResourceManager = setUpPubSubResourceManager(); diff --git a/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/EndToEndTestingITBase.java b/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/EndToEndTestingITBase.java index 9395d4cc94..982c15fb21 100644 --- a/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/EndToEndTestingITBase.java +++ b/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/endtoend/EndToEndTestingITBase.java @@ -236,55 +236,34 @@ protected void createAndUploadReverseMultiShardConfigToGcs( gcsResourceManager.createArtifact("input/shard.json", shardFileContents); } - protected void createAndUploadBulkShardConfigToGcs( - ArrayList dataShardsList, GcsResourceManager gcsResourceManager) { - JSONObject bulkConfig = new JSONObject(); - bulkConfig.put("configType", "dataflow"); - - JSONObject shardConfigBulk = new JSONObject(); - - JSONObject schemaSourceJson = new JSONObject(); - schemaSourceJson.put("dataShardId", ""); - schemaSourceJson.put("host", ""); - schemaSourceJson.put("user", ""); - schemaSourceJson.put("password", ""); - schemaSourceJson.put("port", ""); - schemaSourceJson.put("dbName", ""); - shardConfigBulk.put("schemaSource", schemaSourceJson); - - JSONArray dataShardsArray = new JSONArray(); + protected void createAndUploadShardConfigToGcs( + List dataShardsList, GcsResourceManager gcsResourceManager) { + JSONObject config = new JSONObject(); + JSONArray shardConfigs = new JSONArray(); + if (dataShardsList != null) { for (DataShard shardData : dataShardsList) { JSONObject shardJson = new JSONObject(); - - shardJson.put("dataShardId", shardData.dataShardId); + shardJson.put("logicalShardId", shardData.dataShardId); shardJson.put("host", shardData.host); shardJson.put("user", shardData.user); shardJson.put("password", shardData.password); shardJson.put("port", shardData.port); shardJson.put("dbName", shardData.dbName); - shardJson.put("namespace", shardData.namespace); - shardJson.put("connectionProperties", shardData.connectionProperties); - - JSONArray databasesArray = new JSONArray(); - - for (Database dbData : shardData.databases) { - JSONObject dbJson = new JSONObject(); - dbJson.put("dbName", dbData.dbName); - dbJson.put("databaseId", dbData.databaseId); - dbJson.put("refDataShardId", dbData.refDataShardId); - databasesArray.put(dbJson); + if (shardData.namespace != null) { + shardJson.put("namespace", shardData.namespace); + } + if (shardData.connectionProperties != null) { + shardJson.put("connectionProperties", shardData.connectionProperties); } - shardJson.put("databases", databasesArray); - dataShardsArray.put(shardJson); + shardConfigs.put(shardJson); } } - shardConfigBulk.put("dataShards", dataShardsArray); - bulkConfig.put("shardConfigurationBulk", shardConfigBulk); - String shardFileContents = bulkConfig.toString(); + config.put("shardConfigs", shardConfigs); + String shardFileContents = config.toString(); LOG.info("Shard file contents: {}", shardFileContents); - gcsResourceManager.createArtifact("input/shard-bulk.json", shardFileContents); + gcsResourceManager.createArtifact("input/shard-config.json", shardFileContents); } protected void createAndUploadShardContextFileToGcs( @@ -316,40 +295,44 @@ protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( Boolean multiSharded) throws IOException { // launch dataflow template - if (multiSharded) { - flexTemplateDataflowJobResourceManager = - FlexTemplateDataflowJobResourceManager.builder(jobName) - .withTemplateName("Sourcedb_to_Spanner_Flex") - .withTemplateModulePath("v2/sourcedb-to-spanner") - .addParameter("instanceId", spannerResourceManager.getInstanceId()) - .addParameter("databaseId", spannerResourceManager.getDatabaseId()) - .addParameter("projectId", PROJECT) - .addParameter("outputDirectory", "gs://" + artifactBucketName) - .addParameter("sessionFilePath", getGcsPath("input/session.json", gcsResourceManager)) - .addParameter( - "sourceConfigURL", getGcsPath("input/shard-bulk.json", gcsResourceManager)) - .addEnvironmentVariable( - "additionalExperiments", Collections.singletonList("disable_runner_v2")) - .build(); - } else { - flexTemplateDataflowJobResourceManager = - FlexTemplateDataflowJobResourceManager.builder(jobName) - .withTemplateName("Sourcedb_to_Spanner_Flex") - .withTemplateModulePath("v2/sourcedb-to-spanner") - .addParameter("instanceId", spannerResourceManager.getInstanceId()) - .addParameter("databaseId", spannerResourceManager.getDatabaseId()) - .addParameter("projectId", PROJECT) - .addParameter("outputDirectory", "gs://" + artifactBucketName) - .addParameter("sessionFilePath", getGcsPath("input/session.json", gcsResourceManager)) - .addParameter("sourceConfigURL", cloudSqlResourceManager.getUri()) - .addParameter("username", cloudSqlResourceManager.getUsername()) - .addParameter("password", cloudSqlResourceManager.getPassword()) - .addParameter("jdbcDriverClassName", "com.mysql.jdbc.Driver") - .addEnvironmentVariable( - "additionalExperiments", Collections.singletonList("disable_runner_v2")) - .build(); + if (!multiSharded) { + DataShard dataShard = + new DataShard( + "shard1", + cloudSqlResourceManager.getHost(), + cloudSqlResourceManager.getUsername(), + cloudSqlResourceManager.getPassword(), + String.valueOf(cloudSqlResourceManager.getPort()), + cloudSqlResourceManager.getDatabaseName(), + null, + "useSSL=false&allowPublicKeyRetrieval=true", + new ArrayList<>()); + + ArrayList shards = new ArrayList<>(); + shards.add(dataShard); + createAndUploadShardConfigToGcs(shards, gcsResourceManager); } + FlexTemplateDataflowJobResourceManager.Builder builder = + FlexTemplateDataflowJobResourceManager.builder(jobName) + .withTemplateName("Sourcedb_to_Spanner_Flex") + .withTemplateModulePath("v2/sourcedb-to-spanner") + .addParameter("instanceId", spannerResourceManager.getInstanceId()) + .addParameter("databaseId", spannerResourceManager.getDatabaseId()) + .addParameter("projectId", PROJECT) + .addParameter("outputDirectory", "gs://" + artifactBucketName) + .addParameter("sessionFilePath", getGcsPath("input/session.json", gcsResourceManager)) + .addParameter( + "sourceConfigURL", getGcsPath("input/shard-config.json", gcsResourceManager)) + .addEnvironmentVariable( + "additionalExperiments", Collections.singletonList("disable_runner_v2")); + + if (!multiSharded) { + builder.addParameter("jdbcDriverClassName", "com.mysql.jdbc.Driver"); + } + + flexTemplateDataflowJobResourceManager = builder.build(); + // Run PipelineLauncher.LaunchInfo jobInfo = flexTemplateDataflowJobResourceManager.launchJob(); assertThatPipeline(jobInfo).isRunning(); diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/EndToEndTestingITBase.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/EndToEndTestingITBase.java index ebcee61e7c..240be7a6b4 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/EndToEndTestingITBase.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/EndToEndTestingITBase.java @@ -49,28 +49,15 @@ public record DataShard( public record Database(String dbName, String databaseId, String refDataShardId) {} - protected void createAndUploadBulkShardConfigToGcs( + protected void createAndUploadShardConfigToGcs( List dataShardsList, GcsResourceManager gcsResourceManager) { - JSONObject bulkConfig = new JSONObject(); - bulkConfig.put("configType", "dataflow"); + JSONObject config = new JSONObject(); + JSONArray shardConfigs = new JSONArray(); - JSONObject shardConfigBulk = new JSONObject(); - - JSONObject schemaSourceJson = new JSONObject(); - schemaSourceJson.put("dataShardId", ""); - schemaSourceJson.put("host", ""); - schemaSourceJson.put("user", ""); - schemaSourceJson.put("password", ""); - schemaSourceJson.put("port", ""); - schemaSourceJson.put("dbName", ""); - shardConfigBulk.put("schemaSource", schemaSourceJson); - - JSONArray dataShardsArray = new JSONArray(); if (dataShardsList != null) { for (DataShard shardData : dataShardsList) { JSONObject shardJson = new JSONObject(); - - shardJson.put("dataShardId", shardData.dataShardId()); + shardJson.put("logicalShardId", shardData.dataShardId()); shardJson.put("host", shardData.host()); shardJson.put("user", shardData.user()); shardJson.put("password", shardData.password()); @@ -78,25 +65,13 @@ protected void createAndUploadBulkShardConfigToGcs( shardJson.put("dbName", shardData.dbName()); shardJson.put("namespace", shardData.namespace()); shardJson.put("connectionProperties", shardData.connectionProperties()); - - JSONArray databasesArray = new JSONArray(); - - for (Database dbData : shardData.databases()) { - JSONObject dbJson = new JSONObject(); - dbJson.put("dbName", dbData.dbName()); - dbJson.put("databaseId", dbData.databaseId()); - dbJson.put("refDataShardId", dbData.refDataShardId()); - databasesArray.put(dbJson); - } - shardJson.put("databases", databasesArray); - dataShardsArray.put(shardJson); + shardConfigs.put(shardJson); } } - shardConfigBulk.put("dataShards", dataShardsArray); - bulkConfig.put("shardConfigurationBulk", shardConfigBulk); - String shardFileContents = bulkConfig.toString(); - gcsResourceManager.createArtifact("input/shard-bulk.json", shardFileContents); + config.put("shardConfigs", shardConfigs); + String shardFileContents = config.toString(); + gcsResourceManager.createArtifact("input/shard-config.json", shardFileContents); } protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( @@ -128,15 +103,25 @@ protected PipelineLauncher.LaunchInfo launchBulkDataflowJob( builder.addParameter("sessionFilePath", getGcsPath("session.json", gcsResourceManager)); } - if (multiSharded) { - builder.addParameter( - "sourceConfigURL", getGcsPath("input/shard-bulk.json", gcsResourceManager)); - } else { - builder.addParameter( - "sourceConfigURL", - cloudSqlResourceManager.getUri() + "?useSSL=false&allowPublicKeyRetrieval=true"); - builder.addParameter("username", cloudSqlResourceManager.getUsername()); - builder.addParameter("password", cloudSqlResourceManager.getPassword()); + if (!multiSharded) { + DataShard dataShard = + new DataShard( + "shard1", + cloudSqlResourceManager.getHost(), + cloudSqlResourceManager.getUsername(), + cloudSqlResourceManager.getPassword(), + String.valueOf(cloudSqlResourceManager.getPort()), + cloudSqlResourceManager.getDatabaseName(), + null, + "useSSL=false&allowPublicKeyRetrieval=true", + Collections.emptyList()); + createAndUploadShardConfigToGcs(Collections.singletonList(dataShard), gcsResourceManager); + } + + builder.addParameter( + "sourceConfigURL", getGcsPath("input/shard-config.json", gcsResourceManager)); + + if (!multiSharded) { builder.addParameter("jdbcDriverClassName", "com.mysql.cj.jdbc.Driver"); } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/ShardedBulkMigrationAndValidationE2EIT.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/ShardedBulkMigrationAndValidationE2EIT.java index 8e0a563048..2e6f5a9b32 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/ShardedBulkMigrationAndValidationE2EIT.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/ShardedBulkMigrationAndValidationE2EIT.java @@ -182,7 +182,7 @@ public void shardedMigrationAndValidationE2E() throws Exception { mySQLResourceManager2.getDatabaseName(), LOGICAL_SHARD_2, LOGICAL_SHARD_2)))); - createAndUploadBulkShardConfigToGcs(dataShards, gcsClient); + createAndUploadShardConfigToGcs(dataShards, gcsClient); // 3. Launch Bulk Pipeline (SourceDbToSpanner) with multiSharded=true String gcsOutputDirectory = "gs://" + artifactBucketName + "/" + testId;