PipelineController Changes for Cassandra (#2086)

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

@@ -20,15 +20,20 @@
 
 /** Interface used by the SourcedbToSpanner pipeline to accept user input. */
 public interface SourceDbToSpannerOptions extends CommonTemplateOptions {
+  String CASSANDRA_SOURCE_DIALECT = "CASSANDRA";
+  String MYSQL_SOURCE_DIALECT = "MYSQL";
+  String PG_SOURCE_DIALECT = "POSTGRESQL";
+
   @TemplateParameter.Enum(
       order = 1,
       optional = true,
       enumOptions = {
-        @TemplateParameter.TemplateEnumOption("MYSQL"),
-        @TemplateParameter.TemplateEnumOption("POSTGRESQL")
+        @TemplateParameter.TemplateEnumOption(CASSANDRA_SOURCE_DIALECT),
+        @TemplateParameter.TemplateEnumOption(MYSQL_SOURCE_DIALECT),
+        @TemplateParameter.TemplateEnumOption(PG_SOURCE_DIALECT)
       },
-      description = "SQL Dialect of the source database",
-      helpText = "Possible values are `MYSQL` and `POSTGRESQL`.")
+      description = "Dialect of the source database",
+      helpText = "Possible values are `CASSANDRA`, `MYSQL` and `POSTGRESQL`.")
   @Default.String("MYSQL")
   String getSourceDbDialect();
 

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/IoWrapperFactory.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 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.source.reader;
+
+import com.google.cloud.teleport.v2.source.reader.io.IoWrapper;
+import java.util.List;
+import org.apache.beam.sdk.transforms.Wait;
+
+/** Factory to provide the {@link IoWrapper}. */
+public interface IoWrapperFactory {
+
+  /**
+   * Create an {@link IoWrapper} instance for a list of SourceTables.
+   *
+   * @param sourceTables
+   * @param waitOnSignal
+   * @return
+   */
+  IoWrapper getIOWrapper(List<String> sourceTables, Wait.OnSignal<?> waitOnSignal);
+}

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2024 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.source.reader.io.cassandra.iowrapper;
+
+import com.google.auto.value.AutoValue;
+import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions;
+import com.google.cloud.teleport.v2.source.reader.IoWrapperFactory;
+import com.google.cloud.teleport.v2.source.reader.io.IoWrapper;
+import com.google.common.base.Preconditions;
+import java.util.List;
+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();
+
+  private static CassandraIOWrapperFactory create(String gcsConfigPath) {
+    return new AutoValue_CassandraIOWrapperFactory(gcsConfigPath);
+  }
+
+  public static CassandraIOWrapperFactory fromPipelineOptions(SourceDbToSpannerOptions options) {
+    String gcsPath = options.getSourceConfigURL();
+    // Implementation Details. the pipeline options are strings.
+    Preconditions.checkArgument(
+        options.getSourceDbDialect().equals(SourceDbToSpannerOptions.CASSANDRA_SOURCE_DIALECT),
+        "Unexpected Dialect " + options.getSourceDbDialect() + " for Cassandra Source");
+    Preconditions.checkArgument(
+        StringUtils.startsWith(gcsPath, "gs://"),
+        "GCS path Expected in place of `" + gcsPath + "`.");
+    return CassandraIOWrapperFactory.create(options.getSourceConfigURL());
+  }
+
+  /** Create an {@link IoWrapper} instance for a list of SourceTables. */
+  @Override
+  public IoWrapper getIOWrapper(List<String> sourceTables, OnSignal<?> waitOnSignal) {
+    /** TODO(vardhanvthigle@) */
+    return null;
+  }
+}

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

@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2024 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 com.google.cloud.teleport.v2.source.reader.io.IoWrapper;
+import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.transforms.Wait;
+
+/**
+ * Container for DB Configuration. Implementations should encapsulate the relevant source
+ * configuration to help with invoking sharded or non-sharded migration.
+ */
+public interface DbConfigContainer {
+
+  /**
+   * Get a Unique id for the physical data source. For Non-sharded migration, the id can be returned
+   * as null.
+   *
+   * @return Unique id.
+   */
+  @Nullable
+  String getShardId();
+
+  /**
+   * For the spanner tables that contain the shard id column, returns the source table to
+   * shardColumn. For non-Sharded Migration, return empty Map.
+   *
+   * @param schemaMapper
+   * @param spannerTables
+   * @return
+   */
+  Map<String, String> getSrcTableToShardIdColumnMap(
+      ISchemaMapper schemaMapper, List<String> spannerTables);
+
+  /**
+   * Create an {@link IoWrapper} instance for a list of SourceTables.
+   *
+   * @param sourceTables
+   * @param waitOnSignal
+   * @return
+   */
+  IoWrapper getIOWrapper(List<String> sourceTables, Wait.OnSignal<?> waitOnSignal);
+}

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2024 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 com.google.cloud.teleport.v2.source.reader.IoWrapperFactory;
+import com.google.cloud.teleport.v2.source.reader.io.IoWrapper;
+import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.transforms.Wait.OnSignal;
+import org.jetbrains.annotations.Nullable;
+
+/** Default Implementation for {@link DbConfigContainer} for non-Sharded sources. */
+public final class DbConfigContainerDefaultImpl implements DbConfigContainer {
+  private IoWrapperFactory ioWrapperFactory;
+
+  public DbConfigContainerDefaultImpl(IoWrapperFactory ioWrapperFactory) {
+    this.ioWrapperFactory = ioWrapperFactory;
+  }
+
+  /** Create an {@link IoWrapper} instance for a list of SourceTables. */
+  @Override
+  public IoWrapper getIOWrapper(List<String> sourceTables, OnSignal<?> waitOnSignal) {
+    return this.ioWrapperFactory.getIOWrapper(sourceTables, waitOnSignal);
+  }
+
+  @Nullable
+  @Override
+  public String getShardId() {
+    return null;
+  }
+
+  @Override
+  public Map<String, String> getSrcTableToShardIdColumnMap(
+      ISchemaMapper schemaMapper, List<String> spannerTables) {
+    return new HashMap<>();
+  }
+}