Read cassandra schema but not do anything

Author: Deep1998

v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/SpannerToSourceDb.java

@@ -18,6 +18,7 @@
 import static com.google.cloud.teleport.v2.spanner.migrations.constants.Constants.CASSANDRA_SOURCE_TYPE;
 import static com.google.cloud.teleport.v2.spanner.migrations.constants.Constants.MYSQL_SOURCE_TYPE;
 
+import com.datastax.oss.driver.api.core.CqlSession;
 import com.google.cloud.Timestamp;
 import com.google.cloud.teleport.metadata.Template;
 import com.google.cloud.teleport.metadata.TemplateCategory;
@@ -41,9 +42,16 @@
 import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.SessionFileReader;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.ShardFileReader;
+import com.google.cloud.teleport.v2.spanner.sourceddl.CassandraInformationSchemaScanner;
+import com.google.cloud.teleport.v2.spanner.sourceddl.MySqlInformationSchemaScanner;
+import com.google.cloud.teleport.v2.spanner.sourceddl.SourceInformationSchemaScanner;
+import com.google.cloud.teleport.v2.spanner.sourceddl.SourceSchema;
 import com.google.cloud.teleport.v2.templates.SpannerToSourceDb.Options;
 import com.google.cloud.teleport.v2.templates.changestream.TrimmedShardedDataChangeRecord;
 import com.google.cloud.teleport.v2.templates.constants.Constants;
+import com.google.cloud.teleport.v2.templates.dbutils.processor.SourceProcessorFactory;
+import com.google.cloud.teleport.v2.templates.exceptions.ConnectionException;
+import com.google.cloud.teleport.v2.templates.exceptions.UnsupportedSourceException;
 import com.google.cloud.teleport.v2.templates.transforms.AssignShardIdFn;
 import com.google.cloud.teleport.v2.templates.transforms.ConvertChangeStreamErrorRecordToFailsafeElementFn;
 import com.google.cloud.teleport.v2.templates.transforms.ConvertDlqRecordToTrimmedShardedDataChangeRecordFn;
@@ -56,6 +64,7 @@
 import com.google.cloud.teleport.v2.transforms.DLQWriteTransform;
 import com.google.cloud.teleport.v2.values.FailsafeElement;
 import com.google.common.base.Strings;
+import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -551,6 +560,33 @@ public static PipelineResult run(Options options) {
       LOG.info("Cassandra config is: {}", shards.get(0));
       shardingMode = Constants.SHARDING_MODE_SINGLE_SHARD;
     }
+    SourceInformationSchemaScanner scanner = null;
+    try {
+      SourceProcessorFactory.initializeConnectionHelper(
+          options.getSourceType(), shards, connectionPoolSizePerWorker);
+      if (options.getSourceType().equals(MYSQL_SOURCE_TYPE)) {
+        Connection connection =
+            (Connection)
+                SourceProcessorFactory.getConnectionToShard(options.getSourceType(), shards.get(0));
+        scanner = new MySqlInformationSchemaScanner(connection, shards.get(0).getDbName());
+      } else {
+        CqlSession session =
+            (CqlSession)
+                SourceProcessorFactory.getConnectionToShard(options.getSourceType(), shards.get(0));
+        scanner =
+            new CassandraInformationSchemaScanner(
+                session, ((CassandraShard) shards.get(0)).getKeySpaceName());
+      }
+    } catch (UnsupportedSourceException e) {
+      LOG.error("Unsupported source type: {}", options.getSourceType());
+      throw new IllegalArgumentException(e);
+    } catch (ConnectionException e) {
+      LOG.error("Error getting connection to shard: {}", shards.get(0));
+      throw new IllegalArgumentException(e);
+    }
+    SourceSchema sourceSchema = scanner.scan();
+    LOG.info("Source schema: {}", sourceSchema);
+
     if (shards.size() == 1 && !options.getIsShardedMigration()) {
       shardingMode = Constants.SHARDING_MODE_SINGLE_SHARD;
       Shard shard = shards.get(0);
@@ -562,6 +598,7 @@ public static PipelineResult run(Options options) {
     }
 
     if (options.getSourceType().equals(CASSANDRA_SOURCE_TYPE)) {
+      // TODO: Remove this once session file support is removed.
       Map<String, SpannerTable> spannerTableMap =
           SpannerSchema.convertDDLTableToSpannerTable(ddl.allTables());
       Map<String, NameAndCols> spannerTableNameColsMap =

v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/processor/SourceProcessorFactory.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.teleport.v2.templates.dbutils.processor;
 
+import com.datastax.oss.driver.api.core.CqlSession;
 import com.google.cloud.teleport.v2.spanner.migrations.shard.CassandraShard;
 import com.google.cloud.teleport.v2.spanner.migrations.shard.Shard;
 import com.google.cloud.teleport.v2.templates.constants.Constants;
@@ -27,8 +28,10 @@
 import com.google.cloud.teleport.v2.templates.dbutils.dml.CassandraDMLGenerator;
 import com.google.cloud.teleport.v2.templates.dbutils.dml.IDMLGenerator;
 import com.google.cloud.teleport.v2.templates.dbutils.dml.MySQLDMLGenerator;
+import com.google.cloud.teleport.v2.templates.exceptions.ConnectionException;
 import com.google.cloud.teleport.v2.templates.exceptions.UnsupportedSourceException;
 import com.google.cloud.teleport.v2.templates.models.ConnectionHelperRequest;
+import java.sql.Connection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -137,7 +140,7 @@ private static IConnectionHelper getConnectionHelper(String source)
                     "Invalid source type for connection helper: " + source));
   }
 
-  private static void initializeConnectionHelper(
+  public static void initializeConnectionHelper(
       String source, List<Shard> shards, int maxConnections) throws UnsupportedSourceException {
     IConnectionHelper connectionHelper = getConnectionHelper(source);
     if (!connectionHelper.isConnectionPoolInitialized()) {
@@ -177,4 +180,33 @@ private static Map<String, IDao> createSourceDaoMap(String source, List<Shard> s
     }
     return sourceDaoMap;
   }
+
+  /**
+   * Returns a connection for the specified shard.
+   *
+   * @param source The source database type (e.g. "mysql", "cassandra")
+   * @param shard The shard to get a connection for
+   * @return The connection object for the specified shard
+   * @throws UnsupportedSourceException If the source type is not supported
+   * @throws ConnectionException If there is an error getting the connection
+   */
+  public static Object getConnectionToShard(String source, Shard shard)
+      throws UnsupportedSourceException, ConnectionException {
+    IConnectionHelper connectionHelper = getConnectionHelper(source);
+    Function<Shard, String> urlGenerator =
+        Optional.ofNullable(connectionUrl.get(source))
+            .orElseThrow(
+                () ->
+                    new UnsupportedSourceException(
+                        "Invalid source type for URL generation: " + source));
+    String sqlURL = urlGenerator.apply(shard);
+
+    if (source.equals(Constants.SOURCE_MYSQL)) {
+      return (Connection) connectionHelper.getConnection(sqlURL + "/" + shard.getUserName());
+    } else if (source.equals(Constants.SOURCE_CASSANDRA)) {
+      return (CqlSession) connectionHelper.getConnection(sqlURL);
+    } else {
+      throw new UnsupportedSourceException("Unsupported source type: " + source);
+    }
+  }
 }