Backup shard config bulk

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/cloudsql/CloudSqlShardOrchestrator.java

@@ -420,38 +420,28 @@ 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<String, List<String>> entry : requestedShardMap.entrySet()) {
       String instanceName = entry.getKey();
       String ip = instanceIpMap.get(instanceName);
       List<String> 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 =

v2/sourcedb-to-spanner/README.md

@@ -62,7 +62,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.
@@ -86,7 +86,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://<source_ip>:3306/<mysql_db_name>",username=<mysql user>,password=<mysql pass>,instanceId="<spanner instanceid>",databaseId="<spanner_database_id>",projectId="$PROJECT",outputDirectory=gs://<gcs-dir> \
+        --parameters sourceConfigURL="gs://<bucket-name>/source-config.json",username=<mysql user>,password=<mysql pass>,instanceId="<spanner instanceid>",databaseId="<spanner_database_id>",projectId="$PROJECT",outputDirectory=gs://<gcs-dir> \
         --additional-experiments=disable_runner_v2
 ```
 #### Replaying DLQ entries.

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 = "<instanceId>"
     databaseId = "<databaseId>"
     projectId = "<projectId>"
-    sourceConfigURL = "jdbc:mysql://some-host:3306/sampledb"
+    sourceConfigURL = "gs://your-bucket/source-config.json"
     username = "<username>"
     password = "<password>"
     outputDirectory = "gs://your-bucket/dir"

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.

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/options/OptionsToConfigBuilder.java

@@ -24,6 +24,7 @@
 import com.google.cloud.teleport.v2.reader.io.jdbc.iowrapper.config.SQLDialect;
 import com.google.cloud.teleport.v2.reader.io.schema.SourceSchemaReference;
 import com.google.cloud.teleport.v2.source.mysql.reader.io.jdbc.iowrapper.config.defaults.MySqlConfigDefaults;
+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;
@@ -63,15 +64,11 @@ public static String extractWorkerZone(PipelineOptions options) {
 
   public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults(
       SourceDbToSpannerOptions options,
+      Shard shard,
       List<String> 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();
@@ -88,14 +85,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,
@@ -112,7 +108,6 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfigWithDefaults(
   public static JdbcIOWrapperConfig getJdbcIOWrapperConfig(
       SQLDialect sqlDialect,
       List<String> tables,
-      String sourceDbURL,
       String host,
       String connectionProperties,
       int port,
@@ -157,13 +152,12 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig(
       builder = builder.setMaxConnections(maxConnections);
     }
 
+    String sourceDbURL = "";
     switch (sqlDialect) {
       case MYSQL:
-        if (sourceDbURL == null) {
-          sourceDbURL = "jdbc:mysql://" + host + ":" + port + "/" + dbName;
-          if (StringUtils.isNotBlank(connectionProperties)) {
-            sourceDbURL = sourceDbURL + "?" + connectionProperties;
-          }
+        sourceDbURL = "jdbc:mysql://" + host + ":" + port + "/" + dbName;
+        if (StringUtils.isNotBlank(connectionProperties)) {
+          sourceDbURL = sourceDbURL + "?" + connectionProperties;
         }
         for (Entry<String, String> entry :
             MySqlConfigDefaults.DEFAULT_MYSQL_URL_PROPERTIES.entrySet()) {
@@ -172,14 +166,14 @@ public static JdbcIOWrapperConfig getJdbcIOWrapperConfig(
         sourceDbURL = mysqlSetCursorModeIfNeeded(sqlDialect, sourceDbURL, fetchSize);
         break;
       case POSTGRESQL:
-        if (sourceDbURL == null) {
-          sourceDbURL = "jdbc:postgresql://" + host + ":" + port + "/" + dbName;
-        }
+        sourceDbURL = "jdbc:postgresql://" + host + ":" + port + "/" + dbName;
         sourceDbURL = sourceDbURL + "?currentSchema=" + sourceSchemaReference.jdbc().namespace();
         if (StringUtils.isNotBlank(connectionProperties)) {
           sourceDbURL = sourceDbURL + "&" + connectionProperties;
         }
         break;
+      default:
+        throw new IllegalArgumentException("Unsupported SQL Dialect: " + sqlDialect);
     }
 
     builder.setSourceDbURL(sourceDbURL);

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();
 
@@ -252,17 +248,6 @@ public interface SourceDbToSpannerOptions extends CommonTemplateOptions {
 
   void setTransformationCustomParameters(String value);
 
-  @TemplateParameter.Text(
-      order = 20,
-      optional = true,
-      description = "Namespace",
-      helpText =
-          "Namespace to exported. For PostgreSQL, if no namespace is provided, 'public' will be used")
-  @Default.String("")
-  String getNamespace();
-
-  void setNamespace(String value);
-
   @TemplateParameter.Text(
       order = 21,
       optional = true,

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/cassandra/reader/io/cassandra/iowrapper/CassandraIOWrapperFactory.java

@@ -15,23 +15,27 @@
  */
 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;
 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.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.
@@ -44,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();
 
@@ -54,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,
@@ -71,7 +74,7 @@ private static CassandraIOWrapperFactory create(
           default -> CassandraDialect.OSS;
         };
     return new AutoValue_CassandraIOWrapperFactory(
-        gcsConfigPath,
+        optionsMap,
         numPartions,
         cassandraDialect,
         astraDBToken,
@@ -80,35 +83,50 @@ private static CassandraIOWrapperFactory create(
         astraDBRegion);
   }
 
-  public static CassandraIOWrapperFactory fromPipelineOptions(SourceDbToSpannerOptions options) {
-    String gcsPath = options.getSourceConfigURL();
-    // Implementation Details. the pipeline options are strings.
+  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");
-    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;
+      astraDBToken = GuardedStringValueProvider.create(astraConfig.getAstraToken());
+      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(),
-        GuardedStringValueProvider.create(options.getAstraDBToken()),
-        options.getAstraDBDatabaseId(),
-        options.getAstraDBKeySpace(),
-        options.getAstraDBRegion());
+        astraDBToken,
+        astraDBDatabaseId,
+        astraDBKeyspace,
+        astraDBRegion);
   }
 
   /** Create an {@link IoWrapper} instance for a list of SourceTables. */
   @Override
   public IoWrapper getIOWrapper(List<String> sourceTables, OnSignal<?> waitOnSignal) {
     /** TODO(vardhanvthigle@) incorporate waitOnSignal */
     return new CassandraIoWrapper(
-        gcsConfigPath(),
+        optionsMap(),
         sourceTables,
         numPartitions(),
         cassandraDialect(),