Graph size multi shard support - changes in JdbcIOWrapper

.gitignore

@@ -29,3 +29,6 @@
 
 # Derby log
 **/derby.log*
+
+# build log
+**/build.log*

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

@@ -137,8 +137,7 @@ public ImmutableMap<String, ImmutableMap<String, SourceColumnType>> discoverTabl
     return result.build();
   }
 
-  @VisibleForTesting
-  protected static void convertException(Exception e) {
+  public static void convertException(Exception e) {
     Throwable cause = e;
     while (cause != null) {
       if (cause instanceof SchemaDiscoveryException) {
@@ -226,6 +225,10 @@ protected static <T> List<List<T>> partitionWork(List<T> items, int batchSize) {
     return Lists.partition(items, batchSize);
   }
 
+  public static long getParallelism() {
+    return THREAD_POOL_SIZE;
+  }
+
   private <T> T doRetries(SchemaDiscoveryOperation<T> operation) throws SchemaDiscoveryException {
 
     BackOff backoff = this.fluentBackoff.backoff();

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

@@ -110,14 +110,14 @@ private static PCollection<SourceTableReference> getSourceTableReference(
                 "CreateZeroCounts." + groupName,
                 Create.of(
                     tableReferences.stream()
-                        .map(ref -> KV.of(ref.sourceTableName(), 0L))
+                        .map(ref -> KV.of(ref.sourceTableSchemaUUID(), 0L))
                         .collect(Collectors.toList())))
             .setCoder(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
 
     // 2. Actual counts from the reader.
     PCollection<KV<String, Long>> actualCounts =
         groupRows
-            .apply("ExtractTableNames." + groupName, ParDo.of(new ExtractTableNameFn()))
+            .apply("ExtractTableNames." + groupName, ParDo.of(new ExtractTableIdFn()))
             .apply("CountPerTable." + groupName, Count.perElement());
 
     // 3. Aggregate and Emit Completions.

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/transform/ExtractTableNameFn.java → v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/transform/ExtractTableIdFn.java

@@ -23,13 +23,9 @@
  * DoFn to extract the table name from a {@link SourceRow} and delimit it to match the format used
  * in {@link SourceTableReference}.
  */
-class ExtractTableNameFn extends DoFn<SourceRow, String> {
+class ExtractTableIdFn extends DoFn<SourceRow, String> {
   @ProcessElement
   public void processElement(@Element SourceRow row, OutputReceiver<String> out) {
-    out.output(delimitIdentifier(row.tableName()));
-  }
-
-  private String delimitIdentifier(String identifier) {
-    return "\"" + identifier.replaceAll("\"", "\"\"") + "\"";
+    out.output(row.tableSchemaUUID());
   }
 }

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

@@ -32,7 +32,7 @@ class GroupCompletionDoFn extends DoFn<KV<String, Long>, SourceTableReference> {
   public GroupCompletionDoFn(ImmutableList<SourceTableReference> tableReferences) {
     this.tableReferencesMap =
         tableReferences.stream()
-            .collect(Collectors.toMap(SourceTableReference::sourceTableName, ref -> ref));
+            .collect(Collectors.toMap(SourceTableReference::sourceTableSchemaUUID, ref -> ref));
   }
 
   @ProcessElement

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/schema/SchemaDiscoveryImplTest.java

@@ -226,6 +226,7 @@ public void testSchemaDiscoveryImpl() throws RetriableSchemaDiscoveryException {
         .isEqualTo(ImmutableMap.of("testTable", ImmutableMap.of()));
     verify(mockRetriableSchemaDiscovery, times(expectedCallsCount))
         .discoverTableSchema(any(), any(), any());
+    assertThat(SchemaDiscoveryImpl.getParallelism()).isEqualTo(4L);
   }
 
   @Test

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/transform/ExtractTableNameFnTest.java → v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/transform/ExtractTableIdFnTest.java

@@ -31,9 +31,9 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Test class for {@link ExtractTableNameFn}. */
+/** Test class for {@link ExtractTableIdFn}. */
 @RunWith(JUnit4.class)
-public class ExtractTableNameFnTest implements Serializable {
+public class ExtractTableIdFnTest implements Serializable {
   @Rule public final transient TestPipeline pipeline = TestPipeline.create();
 
   @Test
@@ -48,9 +48,9 @@ public void testExtractTableNameFn_simpleName() {
             .build();
 
     PCollection<String> output =
-        pipeline.apply(Create.of(row)).apply(ParDo.of(new ExtractTableNameFn()));
+        pipeline.apply(Create.of(row)).apply(ParDo.of(new ExtractTableIdFn()));
 
-    PAssert.that(output).containsInAnyOrder("\"table1\"");
+    PAssert.that(output).containsInAnyOrder(tableSchema.tableSchemaUUID());
     pipeline.run();
   }
 
@@ -66,10 +66,10 @@ public void testExtractTableNameFn_withQuotes() {
             .build();
 
     PCollection<String> output =
-        pipeline.apply(Create.of(row)).apply(ParDo.of(new ExtractTableNameFn()));
+        pipeline.apply(Create.of(row)).apply(ParDo.of(new ExtractTableIdFn()));
 
-    // Delimit logic: table"with"quotes -> "table""with""quotes"
-    PAssert.that(output).containsInAnyOrder("\"table\"\"with\"\"quotes\"");
+    // tableSchemaUUID is unique and stable
+    PAssert.that(output).containsInAnyOrder(tableSchema.tableSchemaUUID());
     pipeline.run();
   }
 }

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/transform/GroupCompletionDoFnTest.java

@@ -57,7 +57,7 @@ public void testGroupCompletionDoFn_matchingTable() {
     GroupCompletionDoFn fn = new GroupCompletionDoFn(tableReferences);
 
     PCollection<SourceTableReference> output =
-        pipeline.apply(Create.of(KV.of("table1", 100L))).apply(ParDo.of(fn));
+        pipeline.apply(Create.of(KV.of(ref1.sourceTableSchemaUUID(), 100L))).apply(ParDo.of(fn));
 
     PAssert.that(output).containsInAnyOrder(ref1.toBuilder().setRecordCount(100L).build());
     pipeline.run();

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/PipelineControllerTest.java

@@ -394,7 +394,9 @@ public void shardedDbConfigContainerMySqlTest() {
     sourceDbToSpannerOptions.setUsername(testUser);
     sourceDbToSpannerOptions.setPassword(testPassword);
     sourceDbToSpannerOptions.setTables("table1,table2");
-    mockedStaticJdbcIoWrapper.when(() -> JdbcIoWrapper.of(any())).thenReturn(mockJdbcIoWrapper);
+    mockedStaticJdbcIoWrapper
+        .when(() -> JdbcIoWrapper.of((JdbcIOWrapperConfig) any()))
+        .thenReturn(mockJdbcIoWrapper);
 
     Shard shard =
         new Shard("shard1", "localhost", "3306", "user", "password", null, null, null, null);