switch limitless monitor to use query helper

Author: aaronchung-bitquill

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

@@ -130,7 +130,7 @@ private Connection connectInternal(
       final boolean isInitialConnection,
       final JdbcCallable<Connection, SQLException> connectFunc) throws SQLException {
     final Dialect dialect = this.pluginService.getDialect();
-    if (AuroraLimitlessDialect.class.isAssignableFrom(dialect.getClass())) {
+    if (dialect instanceof AuroraLimitlessDialect) {
       return connectInternalWithDialect(driverProtocol, hostSpec, props, isInitialConnection, connectFunc);
     } else {
       return connectInternalWithoutDialect(driverProtocol, hostSpec, props, isInitialConnection, connectFunc);
@@ -204,7 +204,7 @@ private Connection connectInternalWithoutDialect(
     final Connection conn = connectFunc.call();
 
     final Dialect dialect = this.pluginService.getDialect();
-    if (AuroraLimitlessDialect.class.isAssignableFrom(dialect.getClass())) {
+    if (dialect instanceof AuroraLimitlessDialect) {
       throw new SQLException(""); // TODO: add message
     }
 

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

@@ -75,7 +75,9 @@ public List<HostSpec> queryForLimitlessRouters(final Connection conn, final int
     }
   }
 
-  private List<HostSpec> mapResultSetToHostSpecList(final ResultSet resultSet, final int hostPortToMap) throws SQLException {
+  private List<HostSpec> mapResultSetToHostSpecList(
+      final ResultSet resultSet,
+      final int hostPortToMap) throws SQLException {
 
     List<HostSpec> hosts = new ArrayList<>();
     while (resultSet.next()) {

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

@@ -65,6 +65,7 @@ public class LimitlessRouterMonitor implements AutoCloseable, Runnable {
   private final AtomicReference<List<HostSpec>> limitlessRouters;
   private final @NonNull Properties props;
   private final @NonNull PluginService pluginService;
+  protected final LimitlessQueryHelper queryHelper;
   private final TelemetryFactory telemetryFactory;
   private Connection monitoringConn = null;
   final Executor networkTimeoutExecutor = new SynchronousExecutor();
@@ -89,7 +90,6 @@ public LimitlessRouterMonitor(
     this.limitlessRouterEndpointQuery = getLimitlessRouterEndpointQuery();
     this.limitlessRouters = new AtomicReference<>(limitlessRouters);
     this.props = PropertyUtils.copyProperties(props);
-
     props.stringPropertyNames().stream()
         .filter(p -> p.startsWith(MONITORING_PROPERTY_PREFIX))
         .forEach(
@@ -103,7 +103,7 @@ public LimitlessRouterMonitor(
 
     this.intervalMs = intervalMs;
     this.telemetryFactory = this.pluginService.getTelemetryFactory();
-
+    this.queryHelper = new LimitlessQueryHelper(this.pluginService);
     this.threadPool.submit(this);
     this.threadPool.shutdown(); // No more task are accepted by pool.
   }
@@ -153,7 +153,8 @@ public void run() {
         if (this.monitoringConn == null || this.monitoringConn.isClosed()) {
           continue;
         }
-        List<HostSpec> newLimitlessRouters = queryForLimitlessRouters(this.monitoringConn);
+        List<HostSpec> newLimitlessRouters = queryHelper.queryForLimitlessRouters(this.monitoringConn,
+            this.hostSpec.getPort());
         this.limitlessRouters.set(Collections.unmodifiableList(newLimitlessRouters));
         RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(this.props, newLimitlessRouters);
         LOGGER.finest(Utils.logTopology(limitlessRouters.get(), "[limitlessRouterMonitor]"));
@@ -190,7 +191,8 @@ public List<HostSpec> forceGetLimitlessRouters() throws SQLException {
       if (this.monitoringConn == null || this.monitoringConn.isClosed()) {
         throw new SQLException(Messages.get("LimitlessRouterMonitor.forceGetLimitlessRoutersFailed"));
       }
-      List<HostSpec> newLimitlessRouters = queryForLimitlessRouters(this.monitoringConn);
+      List<HostSpec> newLimitlessRouters = queryHelper.queryForLimitlessRouters(this.monitoringConn,
+          this.hostSpec.getPort());
       this.limitlessRouters.set(Collections.unmodifiableList(newLimitlessRouters));
       RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(this.props, newLimitlessRouters);
       LOGGER.finest(Utils.logTopology(limitlessRouters.get(), "[limitlessRouterMonitor]"));

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

@@ -18,13 +18,11 @@
 
 import java.sql.Connection;
 import java.sql.SQLException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.Supplier;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.AwsWrapperProperty;
 import software.amazon.jdbc.HostSpec;
@@ -70,7 +68,9 @@ public LimitlessRouterServiceImpl(final @NonNull PluginService pluginService) {
     this(pluginService, LimitlessRouterMonitor::new);
   }
 
-  public LimitlessRouterServiceImpl(final @NonNull PluginService pluginService, final LimitlessRouterMonitorInitializer limitlessRouterMonitorInitializer) {
+  public LimitlessRouterServiceImpl(
+      final @NonNull PluginService pluginService,
+      final LimitlessRouterMonitorInitializer limitlessRouterMonitorInitializer) {
     this.pluginService = pluginService;
     this.limitlessRouterMonitorInitializer = limitlessRouterMonitorInitializer;
     this.queryHelper = new LimitlessQueryHelper(pluginService);
@@ -95,7 +95,8 @@ public List<HostSpec> forceGetLimitlessRoutersWithConn(
 
     forceGetLimitlessRoutersLock.lock();
     try {
-      final List<HostSpec> limitlessRouters = limitlessRouterCache.get(this.pluginService.getHostListProvider().getClusterId(), cacheExpirationNano);
+      final List<HostSpec> limitlessRouters =
+          limitlessRouterCache.get(this.pluginService.getHostListProvider().getClusterId(), cacheExpirationNano);
       if (!Utils.isNullOrEmpty(limitlessRouters)) {
         return limitlessRouters;
       }