Refactoring Read With Uniform Partitions for multi shard graph size

Author: VardhanThigle

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/transforms/CollationMapperDoFn.java

@@ -15,9 +15,8 @@
  */
 package com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.transforms;
 
-import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
-
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.DataSourceProvider;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.UniformSplitterDBAdapter;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.stringmapper.CollationMapper;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.stringmapper.CollationReference;
@@ -27,48 +26,44 @@
 import javax.annotation.Nullable;
 import javax.sql.DataSource;
 import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.values.KV;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /** Discover the Collation Mapping information for a given {@link CollationReference}. */
 public class CollationMapperDoFn
-    extends DoFn<CollationReference, KV<CollationReference, CollationMapper>>
+    extends DoFn<KV<String, CollationReference>, KV<CollationReference, CollationMapper>>
     implements Serializable {
 
   private static final Logger logger = LoggerFactory.getLogger(CollationMapper.class);
-  private final SerializableFunction<Void, DataSource> dataSourceProviderFn;
+  private final DataSourceProvider dataSourceProvider;
   private final UniformSplitterDBAdapter dbAdapter;
+  private transient DataSourceManager dataSourceManager;
 
   @JsonIgnore private transient @Nullable DataSource dataSource;
 
   public CollationMapperDoFn(
-      SerializableFunction<Void, DataSource> dataSourceProviderFn,
-      UniformSplitterDBAdapter dbAdapter) {
-    this.dataSourceProviderFn = dataSourceProviderFn;
+      DataSourceProvider dataSourceProvider, UniformSplitterDBAdapter dbAdapter) {
+    this.dataSourceProvider = dataSourceProvider;
     this.dbAdapter = dbAdapter;
     this.dataSource = null;
   }
 
-  @Setup
-  public void setup() throws Exception {
-    dataSource = dataSourceProviderFn.apply(null);
-  }
-
-  private Connection acquireConnection() throws SQLException {
-    return checkStateNotNull(this.dataSource).getConnection();
+  @StartBundle
+  public void startBundle() throws Exception {
+    this.dataSourceManager =
+        DataSourceManagerImpl.builder().setDataSourceProvider(dataSourceProvider).build();
   }
 
   @ProcessElement
   public void processElement(
-      @Element CollationReference input,
+      @Element KV<String, CollationReference> input,
       OutputReceiver<KV<CollationReference, CollationMapper>> out)
       throws SQLException {
-
-    try (Connection conn = acquireConnection()) {
-      CollationMapper mapper = CollationMapper.fromDB(conn, dbAdapter, input);
-      out.output(KV.of(input, mapper));
+    DataSource dataSource = dataSourceManager.getDatasource(input.getKey());
+    try (Connection conn = dataSource.getConnection()) {
+      CollationMapper mapper = CollationMapper.fromDB(conn, dbAdapter, input.getValue());
+      out.output(KV.of(input.getValue(), mapper));
     } catch (Exception e) {
       logger.error(
           "Exception: {} while generating collationMapper for dataSource: {}, collationReference: {}",
@@ -79,4 +74,21 @@ public void processElement(
       throw e;
     }
   }
+
+  @FinishBundle
+  public void finishBundle() throws Exception {
+    cleanupDataSource();
+  }
+
+  @Teardown
+  public void tearDown() throws Exception {
+    cleanupDataSource();
+  }
+
+  void cleanupDataSource() {
+    if (this.dataSourceManager != null) {
+      this.dataSourceManager.closeAll();
+      this.dataSourceManager = null;
+    }
+  }
 }

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/transforms/CollationMapperTransform.java

@@ -16,21 +16,21 @@
 package com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.transforms;
 
 import com.google.auto.value.AutoValue;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.DataSourceProvider;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.UniformSplitterDBAdapter;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.stringmapper.CollationMapper;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.stringmapper.CollationReference;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import java.io.Serializable;
 import java.util.Map;
-import java.util.stream.Collectors;
-import javax.sql.DataSource;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollectionView;
 
@@ -44,10 +44,10 @@ public abstract class CollationMapperTransform
     implements Serializable {
 
   /** List of {@link CollationReference} to discover the mapping for. */
-  public abstract ImmutableList<CollationReference> collationReferences();
+  public abstract ImmutableList<KV<String, CollationReference>> collationReferences();
 
   /** Provider for connection pool. */
-  public abstract SerializableFunction<Void, DataSource> dataSourceProviderFn();
+  public abstract DataSourceProvider dataSourceProvider();
 
   /** Provider to dialect specific Collation mapping query. */
   public abstract UniformSplitterDBAdapter dbAdapter();
@@ -73,12 +73,10 @@ public PCollectionView<Map<CollationReference, CollationMapper>> expand(PBegin i
             .apply("To Empty Map View", View.asMap());
       }
       return input
-          .apply(
-              "Create-Collation-References",
-              Create.of(collationReferences().stream().distinct().collect(Collectors.toList())))
+          .apply("Create-Collation-References", Create.of(collationReferences()))
           .apply(
               "Generate-Mappers",
-              ParDo.of(new CollationMapperDoFn(dataSourceProviderFn(), dbAdapter())))
+              ParDo.of(new CollationMapperDoFn(dataSourceProvider(), dbAdapter())))
           .setCoder(
               KvCoder.of(
                   input.getPipeline().getCoderRegistry().getCoder(CollationReference.class),
@@ -96,13 +94,40 @@ public static Builder builder() {
 
   @AutoValue.Builder
   public abstract static class Builder {
+    private ImmutableList<CollationReference> collationReferencesToDiscover;
+
+    public Builder setCollationReferencesToDiscover(ImmutableList<CollationReference> value) {
+      this.collationReferencesToDiscover = value;
+      return this;
+    }
+
+    abstract Builder setCollationReferences(ImmutableList<KV<String, CollationReference>> value);
 
-    public abstract Builder setCollationReferences(ImmutableList<CollationReference> value);
+    public abstract Builder setDataSourceProvider(DataSourceProvider value);
 
-    public abstract Builder setDataSourceProviderFn(SerializableFunction<Void, DataSource> value);
+    abstract DataSourceProvider dataSourceProvider();
 
     public abstract Builder setDbAdapter(UniformSplitterDBAdapter value);
 
-    public abstract CollationMapperTransform build();
+    public abstract CollationMapperTransform autoBuild();
+
+    public CollationMapperTransform build() {
+      ImmutableList<CollationReference> deDupedRefs =
+          collationReferencesToDiscover.stream()
+              .distinct()
+              .collect(ImmutableList.toImmutableList());
+      ImmutableList<String> ids =
+          ImmutableList.copyOf(this.dataSourceProvider().getDataSourceIds());
+      Preconditions.checkState(ids.size() > 0, "No DataSources Configured for collation detection");
+      ImmutableList.Builder<KV<String, CollationReference>> collationReferencesBuilder =
+          ImmutableList.builder();
+      // Round-robin collations across available shards.
+      // All shards of the same database should have the same collation mapping.
+      for (int i = 0; i < deDupedRefs.size(); i++) {
+        collationReferencesBuilder.add(KV.of(ids.get(i % ids.size()), deDupedRefs.get(i)));
+      }
+      this.setCollationReferences(collationReferencesBuilder.build());
+      return autoBuild();
+    }
   }
 }

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/transforms/MultiTableReadAll.java

@@ -20,6 +20,8 @@
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 
 import com.google.auto.value.AutoValue;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.DataSourceProvider;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.DataSourceProviderImpl;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableIdentifier;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.TableReadSpecification;
 import com.google.common.annotations.VisibleForTesting;
@@ -83,7 +85,7 @@ public abstract class MultiTableReadAll<ParameterT, OutputT>
   private static final boolean DEFAULT_DISABLE_AUTO_COMMIT = true;
 
   @Pure
-  protected abstract @Nullable SerializableFunction<Void, DataSource> getDataSourceProviderFn();
+  protected abstract @Nullable DataSourceProvider getDataSourceProvider();
 
   @Pure
   protected abstract @Nullable ValueProvider<QueryProvider> getQueryProvider();
@@ -121,8 +123,8 @@ public static Builder builder() {
 
   @AutoValue.Builder
   abstract static class Builder<ParameterT, OutputT> {
-    abstract Builder<ParameterT, OutputT> setDataSourceProviderFn(
-        SerializableFunction<Void, DataSource> dataSourceProviderFn);
+    abstract Builder<ParameterT, OutputT> setDataSourceProvider(
+        DataSourceProvider dataSourceProvider);
 
     abstract Builder<ParameterT, OutputT> setQueryProvider(ValueProvider<QueryProvider> query);
 
@@ -151,8 +153,8 @@ abstract Builder<ParameterT, OutputT> setTableIdentifierFn(
    * @return a new transform instance with the data source configured.
    */
   public MultiTableReadAll<ParameterT, OutputT> withDataSourceConfiguration(
-      DataSourceConfiguration config) {
-    return withDataSourceProviderFn(DataSourceProviderFromDataSourceConfiguration.of(config));
+      String id, DataSourceConfiguration config) {
+    return withDataSourceProviderFn(id, DataSourceProviderFromDataSourceConfiguration.of(config));
   }
 
   /**
@@ -162,13 +164,34 @@ public MultiTableReadAll<ParameterT, OutputT> withDataSourceConfiguration(
    * @return a new transform instance.
    */
   public MultiTableReadAll<ParameterT, OutputT> withDataSourceProviderFn(
-      SerializableFunction<Void, DataSource> dataSourceProviderFn) {
-    if (getDataSourceProviderFn() != null) {
+      String id, SerializableFunction<Void, DataSource> dataSourceProviderFn) {
+    if (getDataSourceProvider() != null) {
       throw new IllegalArgumentException(
           "A dataSourceConfiguration or dataSourceProviderFn has "
               + "already been provided, and does not need to be provided again.");
     }
-    return toBuilder().setDataSourceProviderFn(dataSourceProviderFn).build();
+    return withDataSourceProvider(
+        DataSourceProviderImpl.builder().addDataSource(id, dataSourceProviderFn).build());
+  }
+
+  /**
+   * Configures a provider function for the data source.
+   *
+   * @param dataSourceProvider the data source provider function.
+   * @return a new transform instance.
+   */
+  public MultiTableReadAll<ParameterT, OutputT> withDataSourceProvider(
+      DataSourceProvider dataSourceProvider) {
+    if (dataSourceProvider == null) {
+      throw new IllegalArgumentException(
+          "DataSource can not be null "
+              + "already been provided, and does not need to be provided again.");
+    }
+    if (getDataSourceProvider() != null) {
+      throw new IllegalArgumentException(
+          "A dataSource has " + "already been provided, and does not need to be provided again.");
+    }
+    return toBuilder().setDataSourceProvider(dataSourceProvider).build();
   }
 
   /**
@@ -322,7 +345,7 @@ public PCollection<OutputT> expand(PCollection<ParameterT> input) {
             .apply(
                 ParDo.of(
                     new MultiTableReadFn<>(
-                        checkStateNotNull(getDataSourceProviderFn()),
+                        checkStateNotNull(getDataSourceProvider()),
                         checkStateNotNull(getQueryProvider()),
                         checkStateNotNull(getParameterSetter()),
                         getTableReadSpecifications(),
@@ -371,8 +394,8 @@ public void populateDisplayData(DisplayData.Builder builder) {
     if (getCoder() != null) {
       builder.add(DisplayData.item("coder", getCoder().getClass().getName()));
     }
-    if (getDataSourceProviderFn() instanceof HasDisplayData) {
-      ((HasDisplayData) getDataSourceProviderFn()).populateDisplayData(builder);
+    if (getDataSourceProvider() instanceof HasDisplayData) {
+      ((HasDisplayData) getDataSourceProvider()).populateDisplayData(builder);
     }
   }