enable WeightedRandomHostSelector, and Limitless Plugin to use it

Author: aaronchung-bitquill

wrapper/src/main/java/software/amazon/jdbc/C3P0PooledConnectionProvider.java

@@ -50,6 +50,7 @@ public class C3P0PooledConnectionProvider implements PooledConnectionProvider, C
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
   protected static final long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);

wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java

@@ -53,6 +53,7 @@ public class DataSourceConnectionProvider implements ConnectionProvider {
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
   private final @NonNull DataSource dataSource;

wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java

@@ -50,6 +50,7 @@ public class DriverConnectionProvider implements ConnectionProvider {
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
 

wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java

@@ -57,6 +57,7 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
 

wrapper/src/main/java/software/amazon/jdbc/WeightedRandomHostSelector.java

@@ -104,8 +104,8 @@ public HostSpec getHost(
         return entry.getKey();
       }
     }
-    // TODO: proper messaging
-    throw new SQLException(Messages.get("HostSelector.TODO", new Object[] {role}));
+
+    throw new SQLException(Messages.get("HostSelector.weightedRandomUnableToGetHost", new Object[] {role}));
   }
 
   private Map<String, Integer> getHostWeightPairMap(final String hostWeightMapString) throws SQLException {
@@ -124,7 +124,6 @@ private Map<String, Integer> getHostWeightPairMap(final String hostWeightMapStri
     for (final String hostWeightPair : hostWeightPairs) {
       final Matcher matcher = HOST_WEIGHT_PAIRS_PATTERN.matcher(hostWeightPair);
       if (!matcher.matches()) {
-        // TODO: add this message
         throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
       }
 

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java

@@ -34,7 +34,7 @@
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.PropertyDefinition;
-import software.amazon.jdbc.RoundRobinHostSelector;
+import software.amazon.jdbc.WeightedRandomHostSelector;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.util.FullServicesContainer;
 import software.amazon.jdbc.util.HostSelectorUtils;
@@ -127,15 +127,15 @@ public void establishConnection(final LimitlessConnectionContext context) throws
       return;
     }
 
-    HostSelectorUtils.setHostWeightPairsProperty(RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS,
+    HostSelectorUtils.setHostWeightPairsProperty(WeightedRandomHostSelector.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS,
         context.getProps(),
         context.getLimitlessRouters());
     HostSpec selectedHostSpec;
     try {
       selectedHostSpec = this.pluginService.getHostSpecByStrategy(
           context.getLimitlessRouters(),
           HostRole.WRITER,
-          RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
+          WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM);
       LOGGER.fine(Messages.get(
           "LimitlessRouterServiceImpl.selectedHost",
           new Object[] {selectedHostSpec != null ? selectedHostSpec.getHost() : "null"}));

wrapper/src/main/java/software/amazon/jdbc/util/HostSelectorUtils.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package software.amazon.jdbc.util;
 
 import java.util.List;

wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties

@@ -194,6 +194,7 @@ HostSelector.noHostsMatchingRole=No hosts were found matching the requested ''{0
 HostSelector.roundRobinInvalidHostWeightPairs=The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
 HostSelector.roundRobinInvalidDefaultWeight=The provided default weight value is not valid. Weight values must be an integer greater than or equal to the default weight value of 1.
 HostSelector.weightedRandomInvalidHostWeightPairs=The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
+HostSelector.weightedRandomUnableToGetHost=Weighted Random strategy was unable to select a host.
 
 IamAuthConnectionPlugin.unhandledException=Unhandled exception: ''{0}''
 IamAuthConnectionPlugin.connectException=Error occurred while opening a connection: ''{0}''

wrapper/src/test/java/software/amazon/jdbc/util/HostSelectorUtilsTests.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package software.amazon.jdbc.util;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -34,7 +50,9 @@ void testSetHostWeightPairsProperty() {
             .build()
     );
     final Properties properties = new Properties();
-    HostSelectorUtils.setHostWeightPairsProperty(RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS, properties, hosts);
+    HostSelectorUtils.setHostWeightPairsProperty(
+        RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS,
+        properties, hosts);
 
     final String actualPropertyValue = properties.getProperty(
         RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS.name);