feat: enable setting ipType connection option for r2dbc drivers (#937)

Author: shubha-rajan

core/src/main/java/com/google/cloud/sql/core/CoreSocketFactory.java

@@ -78,7 +78,7 @@ public final class CoreSocketFactory {
 
   private static final Logger logger = Logger.getLogger(CoreSocketFactory.class.getName());
 
-  private static final String DEFAULT_IP_TYPES = "PUBLIC,PRIVATE";
+  public static final String DEFAULT_IP_TYPES = "PUBLIC,PRIVATE";
 
   // Test properties, not for end-user use. May be changed or removed without notice.
   private static final String API_ROOT_URL_PROPERTY = "_CLOUD_SQL_API_ROOT_URL";
@@ -152,25 +152,17 @@ public static synchronized CoreSocketFactory getInstance() throws IOException {
     return coreSocketFactory;
   }
 
-  private CloudSqlInstance getCloudSqlInstance(String instanceName, boolean enableIamAuth) {
-    return instances.computeIfAbsent(
-        instanceName,
-        k -> {
-          try {
-            return new CloudSqlInstance(k, adminApi, enableIamAuth, credentialFactory, executor,
-                localKeyPair);
-          } catch (IOException e) {
-            throw new RuntimeException(e);
-          }
-        });
+  @VisibleForTesting
+  CloudSqlInstance getCloudSqlInstance(String instanceName) {
+    return getCloudSqlInstance(instanceName, false);
   }
 
-  private CloudSqlInstance getCloudSqlInstance(String instanceName) {
+  private CloudSqlInstance getCloudSqlInstance(String instanceName, boolean enableIamAuth) {
     return instances.computeIfAbsent(
         instanceName,
         k -> {
           try {
-            return new CloudSqlInstance(k, adminApi, false, credentialFactory, executor,
+            return new CloudSqlInstance(k, adminApi, enableIamAuth, credentialFactory, executor,
                 localKeyPair);
           } catch (IOException e) {
             throw new RuntimeException(e);
@@ -273,18 +265,28 @@ public static SslData getSslData(String csqlInstanceName, boolean enableIamAuth)
     return getInstance().getCloudSqlInstance(csqlInstanceName, enableIamAuth).getSslData();
   }
 
+  /**
+   * Returns data that can be used to establish Cloud SQL SSL connection.
+   */
   public static SslData getSslData(String csqlInstanceName) throws IOException {
     return getSslData(csqlInstanceName, false);
   }
 
   /**
-   * Returns preferred ip address that can be used to establish Cloud SQL connection.
+   * Returns default ip address that can be used to establish Cloud SQL connection.
    */
   public static String getHostIp(String csqlInstanceName) throws IOException {
     return getInstance().getHostIp(csqlInstanceName, listIpTypes(DEFAULT_IP_TYPES));
   }
 
-  private String getHostIp(String instanceName, List<String> ipTypes) {
+  /**
+   * Returns preferred ip address that can be used to establish Cloud SQL connection.
+   */
+  public static String getHostIp(String csqlInstanceName, String ipTypes) throws IOException {
+    return getInstance().getHostIp(csqlInstanceName, listIpTypes(ipTypes));
+  }
+
+  private String getHostIp(String instanceName,  List<String> ipTypes) {
     CloudSqlInstance instance = getCloudSqlInstance(instanceName);
     return instance.getPreferredIp(ipTypes);
   }

pom.xml

@@ -272,10 +272,10 @@
                 Necessary to allow guava into the assembly jars
                 com.google.guava:guava
 
-                Global test dependencies unused in r2dbc core (no tests currently):
-                junit:junit,com.google.truth:truth
+                Global test dependencies unused in r2dbc core:
+                junit:junit,com.google.truth:truth,org.bouncycastle:bcpkix-jdk15on
                 -->
-                org.ow2.asm:asm-util,com.github.jnr:jnr-unixsocket,org.postgresql:postgresql,junit:junit,com.google.truth:truth,com.microsoft.sqlserver:mssql-jdbc,com.google.guava:guava
+                org.ow2.asm:asm-util,com.github.jnr:jnr-unixsocket,org.postgresql:postgresql,junit:junit,com.google.truth:truth,com.microsoft.sqlserver:mssql-jdbc,com.google.guava:guava,org.bouncycastle:bcpkix-jdk15on
               </ignoredDependencies>
             </configuration>
             <executions>

r2dbc/core/pom.xml

@@ -132,7 +132,48 @@
       <artifactId>netty-handler</artifactId>
       <version>4.1.84.Final</version>
     </dependency>
+    <dependency>
+      <groupId>junit</groupId>
+      <artifactId>junit</artifactId>
+      <version>4.13.2</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>4.11.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.bouncycastle</groupId>
+      <artifactId>bcpkix-jdk15on</artifactId>
+      <version>1.70</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.apis</groupId>
+      <artifactId>google-api-services-sqladmin</artifactId>
+      <version>v1beta4-rev20230111-2.0.0</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
+ <build>
+   <plugins>
+     <plugin>
+       <groupId>org.apache.maven.plugins</groupId>
+       <artifactId>maven-jar-plugin</artifactId>
+       <version>3.3.0</version>
+       <executions>
+         <execution>
+           <goals>
+             <goal>test-jar</goal>
+           </goals>
+         </execution>
+       </executions>
+     </plugin>
+   </plugins>
+
+ </build>
 
   <repositories>
     <repository>

r2dbc/core/src/main/java/com/google/cloud/sql/core/CloudSqlConnectionFactory.java

@@ -19,6 +19,7 @@
 import static io.r2dbc.spi.ConnectionFactoryOptions.HOST;
 import static io.r2dbc.spi.ConnectionFactoryOptions.PORT;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
 import io.r2dbc.spi.ConnectionFactoryMetadata;
@@ -36,15 +37,18 @@ public class CloudSqlConnectionFactory implements ConnectionFactory {
   private Function<ConnectionFactoryOptions, ConnectionFactory> connectionFactoryFactory;
   private ConnectionFactoryOptions.Builder builder;
   private String csqlHostName;
+  private String ipTypes;
 
   /**
    * Creates an instance of ConnectionFactory that pulls and sets host ip before delegating.
    */
   public CloudSqlConnectionFactory(
       Function<ConnectionFactoryOptions, ConnectionFactory> connectionFactoryFactory,
+      String ipTypes,
       Builder builder,
       String csqlHostName) {
     this.connectionFactoryFactory = connectionFactoryFactory;
+    this.ipTypes = ipTypes;
     this.builder = builder;
     this.csqlHostName = csqlHostName;
   }
@@ -67,9 +71,20 @@ public ConnectionFactoryMetadata getMetadata() {
     }
   }
 
-  private ConnectionFactory getConnectionFactory() throws IOException {
-    String hostIp = CoreSocketFactory.getHostIp(csqlHostName);
+  @VisibleForTesting
+  void setBuilderHostAndPort() throws IOException {
+    String hostIp = CoreSocketFactory.getHostIp(csqlHostName, ipTypes);
     builder.option(HOST, hostIp).option(PORT, CoreSocketFactory.getDefaultServerProxyPort());
+  }
+
+  @VisibleForTesting
+  ConnectionFactoryOptions.Builder getBuilder() throws IOException {
+    return builder;
+  }
+
+
+  private ConnectionFactory getConnectionFactory() throws IOException {
+    setBuilderHostAndPort();
     return connectionFactoryFactory.apply(builder.build());
   }
 }

r2dbc/core/src/main/java/com/google/cloud/sql/core/GcpConnectionFactoryProvider.java

@@ -37,6 +37,7 @@
 public abstract class GcpConnectionFactoryProvider implements ConnectionFactoryProvider {
 
   public static final Option<String> UNIX_SOCKET = Option.valueOf("UNIX_SOCKET");
+  public static final Option<String> IP_TYPES = Option.valueOf("IP_TYPES");
   public static final Option<Boolean> ENABLE_IAM_AUTH = Option.valueOf("ENABLE_IAM_AUTH");
 
   private static Function<SslContextBuilder, SslContextBuilder> createSslCustomizer(
@@ -72,6 +73,7 @@ private static Function<SslContextBuilder, SslContextBuilder> createSslCustomize
    */
   abstract ConnectionFactory tcpConnectionFactory(
       Builder optionBuilder,
+      String ipTypes,
       Function<SslContextBuilder, SslContextBuilder> customizer,
       String csqlHostName);
 
@@ -110,7 +112,12 @@ public ConnectionFactory create(ConnectionFactoryOptions connectionFactoryOption
   private ConnectionFactory createFactory(
       ConnectionFactoryOptions connectionFactoryOptions) throws IOException {
     String connectionName = (String) connectionFactoryOptions.getRequiredValue(HOST);
-    String socket = (String) connectionFactoryOptions.getValue(UNIX_SOCKET);
+
+    String ipTypes = CoreSocketFactory.DEFAULT_IP_TYPES;
+    Object ipTypesObj = connectionFactoryOptions.getValue(IP_TYPES);
+    if (ipTypesObj != null) {
+      ipTypes = (String) ipTypesObj;
+    }
 
     Object iamAuthObj = connectionFactoryOptions.getValue(ENABLE_IAM_AUTH);
     Boolean enableIamAuth = false;
@@ -122,15 +129,17 @@ private ConnectionFactory createFactory(
 
     Builder optionBuilder = createBuilder(connectionFactoryOptions);
 
-
     // Precompute SSL Data to trigger the initial refresh to happen immediately,
     // and ensure enableIAMAuth is set correctly.
     CoreSocketFactory.getSslData(connectionName, enableIamAuth);
 
+    String socket = (String) connectionFactoryOptions.getValue(UNIX_SOCKET);
     if (socket != null) {
       return socketConnectionFactory(optionBuilder, socket);
     }
-    return tcpConnectionFactory(optionBuilder, createSslCustomizer(connectionName, enableIamAuth),
+
+    return tcpConnectionFactory(optionBuilder, ipTypes,
+        createSslCustomizer(connectionName, enableIamAuth),
         connectionName);
   }