[Not-for-review] Graph Size Optimization Milestone-2 multi shard read.

it/google-cloud-platform/pom.xml

@@ -61,6 +61,11 @@
             <artifactId>google-api-services-dataflow</artifactId>
             <version>${dataflow-api.version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.google.apis</groupId>
+            <artifactId>google-api-services-sqladmin</artifactId>
+            <version>${sqladmin-api.version}</version>
+        </dependency>
 
         <!-- JDBC dependencies for CloudSQL -->
         <dependency>

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

@@ -62,6 +62,27 @@ protected void configurePort() {
       }
     }
 
+    @Override
+    public Builder maybeUseStaticInstance(String host, int port, String userName, String password) {
+      super.maybeUseStaticInstance(host, port, userName, password);
+      return this;
+    }
+
+    public Builder setProjectId(String projectId) {
+      this.projectId = projectId;
+      return this;
+    }
+
+    public Builder setRegion(String region) {
+      this.region = region;
+      return this;
+    }
+
+    public Builder setCredentials(com.google.auth.oauth2.GoogleCredentials credentials) {
+      this.credentials = credentials;
+      return this;
+    }
+
     @Override
     public @NonNull CloudMySQLResourceManager build() {
       return new CloudMySQLResourceManager(this);

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

@@ -364,6 +364,27 @@ protected void configurePort() {
       }
     }
 
+    @Override
+    public Builder maybeUseStaticInstance(String host, int port, String userName, String password) {
+      super.maybeUseStaticInstance(host, port, userName, password);
+      return this;
+    }
+
+    public Builder setProjectId(String projectId) {
+      this.projectId = projectId;
+      return this;
+    }
+
+    public Builder setRegion(String region) {
+      this.region = region;
+      return this;
+    }
+
+    public Builder setCredentials(com.google.auth.oauth2.GoogleCredentials credentials) {
+      this.credentials = credentials;
+      return this;
+    }
+
     @Override
     public @NonNull CloudPostgresResourceManager build() {
       return new CloudPostgresResourceManager(this);

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

@@ -19,6 +19,7 @@
 
 import static org.apache.beam.it.gcp.cloudsql.CloudSqlResourceManagerUtils.generateDatabaseName;
 
+import com.google.auth.oauth2.GoogleCredentials;
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.beam.it.jdbc.AbstractJDBCResourceManager;
@@ -151,6 +152,10 @@ public void cleanupAll() {
   public abstract static class Builder
       extends AbstractJDBCResourceManager.Builder<@NonNull CloudSqlContainer<?>> {
 
+    protected String projectId;
+    protected String region;
+    protected GoogleCredentials credentials;
+
     private String dbName;
     private boolean usingCustomDb;
 
@@ -174,6 +179,31 @@ public Builder maybeUseStaticInstance() {
       return this;
     }
 
+    public Builder maybeUseStaticInstance(String host, int port, String userName, String password) {
+      this.setHost(host);
+      this.setPort(port);
+      this.setUsername(userName);
+      this.setPassword(password);
+      this.useStaticContainer();
+
+      return this;
+    }
+
+    public Builder setProjectId(String projectId) {
+      this.projectId = projectId;
+      return this;
+    }
+
+    public Builder setRegion(String region) {
+      this.region = region;
+      return this;
+    }
+
+    public Builder setCredentials(GoogleCredentials credentials) {
+      this.credentials = credentials;
+      return this;
+    }
+
     protected String getDefaultUsername() {
       return DEFAULT_JDBC_USERNAME;
     }

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/cloudsql/CloudMySQLResourceManagerTest.java

@@ -54,4 +54,21 @@ public void setUp() {
   public void testGetJDBCPrefixReturnsCorrectValue() {
     assertThat(testManager.getJDBCPrefix()).isEqualTo("mysql");
   }
+
+  @Test
+  public void testBuilder() {
+    CloudMySQLResourceManager.Builder builder = CloudMySQLResourceManager.builder(TEST_ID);
+
+    com.google.auth.oauth2.GoogleCredentials credentials =
+        org.mockito.Mockito.mock(com.google.auth.oauth2.GoogleCredentials.class);
+    builder
+        .setProjectId("test-project")
+        .setRegion("test-region")
+        .setCredentials(credentials)
+        .maybeUseStaticInstance("1.1.1.1", 3306, "u", "p");
+
+    assertThat(builder.projectId).isEqualTo("test-project");
+    assertThat(builder.region).isEqualTo("test-region");
+    assertThat(builder.credentials).isEqualTo(credentials);
+  }
 }

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/cloudsql/CloudPostgresResourceManagerTest.java

@@ -155,6 +155,23 @@ public void testCreateLogicalReplicationRollbackOnError() throws SQLException {
     verify(mockConnection).rollback();
   }
 
+  @Test
+  public void testBuilder() {
+    CloudPostgresResourceManager.Builder builder = CloudPostgresResourceManager.builder(TEST_ID);
+
+    com.google.auth.oauth2.GoogleCredentials credentials =
+        org.mockito.Mockito.mock(com.google.auth.oauth2.GoogleCredentials.class);
+    builder
+        .setProjectId("test-project")
+        .setRegion("test-region")
+        .setCredentials(credentials)
+        .maybeUseStaticInstance("1.1.1.1", 5432, "u", "p");
+
+    assertThat(builder.projectId).isEqualTo("test-project");
+    assertThat(builder.region).isEqualTo("test-region");
+    assertThat(builder.credentials).isEqualTo(credentials);
+  }
+
   /** Helper mock implementation of {@link CloudPostgresResourceManager} for testing. */
   private static class MockCloudPostgresResourceManager extends CloudPostgresResourceManager {
 

it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/cloudsql/CloudSqlResourceManagerTest.java

@@ -48,10 +48,16 @@ private static class MockCloudSqlResourceManager extends CloudSqlResourceManager
     private final boolean initialized;
     private boolean createdDatabase;
     private String lastRunSqlCommand;
+    private final String projectId;
+    private final String region;
+    private final com.google.auth.oauth2.GoogleCredentials credentials;
 
     private MockCloudSqlResourceManager(Builder builder) {
       super(builder);
       this.initialized = true;
+      this.projectId = builder.projectId;
+      this.region = builder.region;
+      this.credentials = builder.credentials;
     }
 
     @Override
@@ -183,6 +189,59 @@ public void testCleanupAllRemovesAllTablesWhenDBNotCreated() {
     assertThat(testManager.createdTables).isEmpty();
   }
 
+  @Test
+  public void testMaybeUseStaticInstanceWithHost() {
+    CloudSqlResourceManager.Builder builder =
+        new CloudSqlResourceManager.Builder(TEST_ID) {
+          @Override
+          public @NonNull CloudSqlResourceManager build() {
+            return new MockCloudSqlResourceManager(this);
+          }
+
+          @Override
+          protected void configurePort() {
+            this.setPort(1234);
+          }
+        };
+
+    String customHost = "10.1.1.1";
+    builder.maybeUseStaticInstance(customHost, 1234, "testUser", "testPassword");
+    CloudSqlResourceManager manager = (CloudSqlResourceManager) builder.build();
+
+    assertThat(manager.getHost()).isEqualTo(customHost);
+  }
+
+  @Test
+  public void testBuilder() {
+    CloudSqlResourceManager.Builder builder =
+        new CloudSqlResourceManager.Builder(TEST_ID) {
+          @Override
+          public @NonNull CloudSqlResourceManager build() {
+            return new MockCloudSqlResourceManager(this);
+          }
+
+          @Override
+          protected void configurePort() {
+            this.setPort(1234);
+          }
+        };
+
+    com.google.auth.oauth2.GoogleCredentials credentials =
+        org.mockito.Mockito.mock(com.google.auth.oauth2.GoogleCredentials.class);
+    builder
+        .setProjectId("test-project")
+        .setRegion("test-region")
+        .setCredentials(credentials)
+        .setHost(HOST)
+        .setPort(Integer.parseInt(PORT));
+
+    MockCloudSqlResourceManager manager = (MockCloudSqlResourceManager) builder.build();
+
+    assertThat(manager.projectId).isEqualTo("test-project");
+    assertThat(manager.region).isEqualTo("test-region");
+    assertThat(manager.credentials).isEqualTo(credentials);
+  }
+
   /*
    * Currently only supports static Cloud SQL instance which means jdbc port uses system property.
    */

pom.xml

@@ -96,6 +96,7 @@
     <truth.version>1.4.5</truth.version>
     <netty.version>4.2.12.Final</netty.version>
     <zookeeper.version>3.9.5</zookeeper.version>
+    <sqladmin-api.version>v1beta4-rev20240115-2.0.0</sqladmin-api.version>
 
     <!-- Drop pinned version once maven-dependency-plugin gets past plexus-archiver 4.8.0 -->
     <plexus-archiver.version>4.8.0</plexus-archiver.version>

v2/sourcedb-to-spanner/pom.xml

@@ -103,6 +103,12 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.google.apis</groupId>
+            <artifactId>google-api-services-sqladmin</artifactId>
+            <version>${sqladmin-api.version}</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.beam</groupId>
             <artifactId>beam-it-jdbc</artifactId>

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

@@ -15,14 +15,14 @@
  */
 package com.google.cloud.teleport.v2.source.reader.io.exception;
 
-import com.google.cloud.teleport.v2.source.reader.io.jdbc.iowrapper.config.JdbcIOWrapperConfig;
+import com.google.cloud.teleport.v2.source.reader.io.jdbc.iowrapper.config.JdbcIoWrapperConfigGroup;
 
 /**
  * Exception thrown when a suitable indexed column that can act as the partition column is not
  * found.
  *
  * <p>Please refer to {@link
- * com.google.cloud.teleport.v2.source.reader.io.jdbc.iowrapper.JdbcIoWrapper#of(JdbcIOWrapperConfig)}
+ * com.google.cloud.teleport.v2.source.reader.io.jdbc.iowrapper.JdbcIoWrapper#of(JdbcIoWrapperConfigGroup)}
  * for details on the cases where this is thrown.
  */
 public class SuitableIndexNotFoundException extends SchemaDiscoveryException {