refactor: change getHostSpecByStrategy and getHost methods to pass pr…

…operties
aws · Aug 24, 2023 · bffdcce · bffdcce
1 parent 6bfd40c
commit bffdcce
Show file tree

Hide file tree

Showing 21 changed files with 177 additions and 256 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPlugin.java
@@ -128,7 +128,7 @@ Connection forceConnect(
    * @throws UnsupportedOperationException if this {@link ConnectionPlugin} does not support the
    *                                       requested strategy
    */
-  HostSpec getHostSpecByStrategy(final HostRole role, final String strategy)
+  HostSpec getHostSpecByStrategy(final HostRole role, final String strategy, final Properties props)
       throws SQLException, UnsupportedOperationException;
 
   void initHostProvider(

diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java
@@ -381,7 +381,7 @@ public boolean acceptsStrategy(HostRole role, String strategy) throws SQLExcepti
    *                                       {@link ConnectionPlugin} instances do not support the
    *                                       requested strategy
    */
-  public HostSpec getHostSpecByStrategy(HostRole role, String strategy)
+  public HostSpec getHostSpecByStrategy(HostRole role, String strategy, Properties props)
       throws SQLException, UnsupportedOperationException {
     try {
       for (ConnectionPlugin plugin : this.plugins) {
@@ -392,7 +392,7 @@ public HostSpec getHostSpecByStrategy(HostRole role, String strategy)
 
         if (isSubscribed) {
           try {
-            final HostSpec host = plugin.getHostSpecByStrategy(role, strategy);
+            final HostSpec host = plugin.getHostSpecByStrategy(role, strategy, props);
             if (host != null) {
               return host;
             }

diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProvider.java
@@ -64,7 +64,7 @@ boolean acceptsUrl(
    * @throws UnsupportedOperationException if the strategy is unsupported by the provider
    */
   HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException, UnsupportedOperationException;
 
   /**

diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java
@@ -141,14 +141,14 @@ public boolean acceptsStrategy(HostRole role, String strategy) {
    * @throws UnsupportedOperationException if the available {@link ConnectionProvider} instances do
    *                                       not support the requested strategy
    */
-  public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy)
+  public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy, Properties props)
       throws SQLException, UnsupportedOperationException {
     HostSpec host = null;
     if (connProvider != null) {
       connProviderLock.readLock().lock();
       try {
         if (connProvider != null && connProvider.acceptsStrategy(role, strategy)) {
-          host = connProvider.getHostSpecByStrategy(hosts, role, strategy);
+          host = connProvider.getHostSpecByStrategy(hosts, role, strategy, props);
         }
       } catch (UnsupportedOperationException e) {
         // The custom provider does not support the provided strategy, ignore it and try with the default provider.
@@ -158,7 +158,7 @@ public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, Strin
     }
 
     if (host == null) {
-      host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy);
+      host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy, props);
     }
 
     return host;

diff --git a/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java
@@ -29,7 +29,6 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
-import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.PropertyUtils;
@@ -88,7 +87,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -97,7 +96,7 @@ public HostSpec getHostSpecByStrategy(
               new Object[] {strategy, DataSourceConnectionProvider.class}));
     }
 
-    return acceptedStrategies.get(strategy).getHost(hosts, role);
+    return acceptedStrategies.get(strategy).getHost(hosts, role, props);
   }
 
   /**
@@ -116,14 +115,6 @@ public Connection connect(
       final @NonNull HostSpec hostSpec,
       final @NonNull Properties props)
       throws SQLException {
-    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
-    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
-      final RoundRobinHostSelector roundRobinHostSelector =
-          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
-      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
-        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
-      }
-    }
 
     final Properties copy = PropertyUtils.copyProperties(props);
 

diff --git a/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java
@@ -27,7 +27,6 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
-import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.Messages;
@@ -83,7 +82,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -92,7 +91,7 @@ public HostSpec getHostSpecByStrategy(
               new Object[] {strategy, DriverConnectionProvider.class}));
     }
 
-    return acceptedStrategies.get(strategy).getHost(hosts, role);
+    return acceptedStrategies.get(strategy).getHost(hosts, role, props);
   }
 
   /**
@@ -111,14 +110,6 @@ public Connection connect(
       final @NonNull HostSpec hostSpec,
       final @NonNull Properties props)
       throws SQLException {
-    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
-    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
-      final RoundRobinHostSelector roundRobinHostSelector =
-          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
-      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
-        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
-      }
-    }
 
     final Properties copy = PropertyUtils.copyProperties(props);
     final ConnectInfo connectInfo =

diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java
@@ -33,7 +33,6 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.cleanup.CanReleaseResources;
 import software.amazon.jdbc.dialect.Dialect;
-import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.util.HikariCPSQLException;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.RdsUrlType;
@@ -47,12 +46,11 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
   private static final Logger LOGGER =
       Logger.getLogger(HikariPooledConnectionProvider.class.getName());
   private static final Map<String, HostSelector> acceptedStrategies =
-      Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
+      new HashMap<String, HostSelector>() {
         {
-          put(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, new LeastConnectionsHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
         }
-      });
+      };
   private static final RdsUtils rdsUtils = new RdsUtils();
   private static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
       new SlidingExpirationCache<>(
@@ -104,10 +102,8 @@ public HikariPooledConnectionProvider(
       HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
-    final LeastConnectionsHostSelector hostSelector =
-        (LeastConnectionsHostSelector) acceptedStrategies
-            .get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS);
-    hostSelector.setDatabasePools(databasePools);
+    final LeastConnectionsHostSelector hostSelector = new LeastConnectionsHostSelector(databasePools);
+    acceptedStrategies.put(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, hostSelector);
   }
 
   /**
@@ -144,10 +140,8 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     poolExpirationCheckNanos = poolExpirationNanos;
     databasePools.setCleanupIntervalNanos(poolCleanupNanos);
-    final LeastConnectionsHostSelector hostSelector =
-            (LeastConnectionsHostSelector) acceptedStrategies
-                .get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS);
-    hostSelector.setDatabasePools(databasePools);
+    final LeastConnectionsHostSelector hostSelector = new LeastConnectionsHostSelector(databasePools);
+    acceptedStrategies.put(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, hostSelector);
   }
 
   @Override
@@ -164,17 +158,17 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, Properties props)
       throws SQLException {
     final HostSpec selectedHost;
     switch (strategy) {
       case LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS:
         selectedHost =
-            acceptedStrategies.get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS).getHost(hosts, role);
+            acceptedStrategies.get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS).getHost(hosts, role, props);
         break;
 
       case RoundRobinHostSelector.STRATEGY_ROUND_ROBIN:
-        selectedHost = acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN).getHost(hosts, role);
+        selectedHost = acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN).getHost(hosts, role, props);
         break;
 
       default:
@@ -203,15 +197,6 @@ public Connection connect(
 
     ds.setPassword(props.getProperty(PropertyDefinition.PASSWORD.name));
 
-    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
-    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
-      final RoundRobinHostSelector roundRobinHostSelector =
-          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
-      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
-        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
-      }
-    }
-
     return ds.getConnection();
   }
 

diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/HostSelector.java
@@ -18,6 +18,7 @@
 
 import java.sql.SQLException;
 import java.util.List;
+import java.util.Properties;
 
 public interface HostSelector {
 
@@ -26,9 +27,10 @@ public interface HostSelector {
    *
    * @param hosts a list of available hosts to pick from
    * @param role  the desired host role - either a writer or a reader
+   * @param props a properties object containing any necessary parameters
    * @return a host matching the requested role
    * @throws SQLException if the host list does not contain any hosts matching the requested role or
    *                      an error occurs while selecting a host
    */
-  HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException;
+  HostSpec getHost(List<HostSpec> hosts, HostRole role, Properties props) throws SQLException;
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java
@@ -20,20 +20,22 @@
 import java.sql.SQLException;
 import java.util.List;
 import java.util.Map;
+import java.util.Properties;
 import java.util.stream.Collectors;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.SlidingExpirationCache;
 
 public class LeastConnectionsHostSelector implements HostSelector {
   public static final String STRATEGY_LEAST_CONNECTIONS = "leastConnections";
-  private SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
+  private final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
 
-  public void setDatabasePools(
+  public LeastConnectionsHostSelector(
       SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
     this.databasePools = databasePools;
   }
 
-  public HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException {
+  @Override
+  public HostSpec getHost(List<HostSpec> hosts, HostRole role, Properties props) throws SQLException {
     List<HostSpec> eligibleHosts = hosts.stream()
         .filter(hostSpec -> role.equals(hostSpec.getRole()))
         .sorted((hostSpec1, hostSpec2) ->

diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java
@@ -81,7 +81,7 @@ EnumSet<NodeChangeOptions> setCurrentConnection(
    *                                       {@link ConnectionPlugin} instances do not support the
    *                                       requested strategy
    */
-  HostSpec getHostSpecByStrategy(HostRole role, String strategy)
+  HostSpec getHostSpecByStrategy(HostRole role, String strategy, Properties props)
       throws SQLException, UnsupportedOperationException;
 
   /**

diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java
@@ -135,8 +135,8 @@ public boolean acceptsStrategy(HostRole role, String strategy) throws SQLExcepti
   }
 
   @Override
-  public HostSpec getHostSpecByStrategy(HostRole role, String strategy) throws SQLException {
-    return this.pluginManager.getHostSpecByStrategy(role, strategy);
+  public HostSpec getHostSpecByStrategy(HostRole role, String strategy, Properties props) throws SQLException {
+    return this.pluginManager.getHostSpecByStrategy(role, strategy, props);
   }
 
   @Override

diff --git a/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java
@@ -18,6 +18,7 @@
 
 import java.sql.SQLException;
 import java.util.List;
+import java.util.Properties;
 import java.util.Random;
 import java.util.stream.Collectors;
 import software.amazon.jdbc.util.Messages;
@@ -27,7 +28,7 @@ public class RandomHostSelector implements HostSelector {
   public static final String STRATEGY_RANDOM = "random";
 
   @Override
-  public HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException {
+  public HostSpec getHost(List<HostSpec> hosts, HostRole role, Properties props) throws SQLException {
     List<HostSpec> eligibleHosts = hosts.stream()
         .filter(hostSpec -> role.equals(hostSpec.getRole())).collect(Collectors.toList());
     if (eligibleHosts.size() == 0) {