From 7498d8a35ef1dd85d57fa7a0601f41dbcf79d7c2 Mon Sep 17 00:00:00 2001
From: Aaron Chung <aaronch@bitquilltech.com>
Date: Wed, 6 Nov 2024 16:43:14 -0800
Subject: [PATCH] fix: limitless various minor fixes

---
 .../limitless/LimitlessConnectionPlugin.java  | 21 +++++++++++++++++--
 .../limitless/LimitlessRouterServiceImpl.java |  6 ++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java
index 0243cfdd6..13459805b 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java
@@ -22,7 +22,6 @@
 import java.util.HashSet;
 import java.util.Properties;
 import java.util.Set;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Supplier;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -35,6 +34,7 @@
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.plugin.AbstractConnectionPlugin;
 import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.PropertyUtils;
 
 public class LimitlessConnectionPlugin extends AbstractConnectionPlugin {
   private static final Logger LOGGER = Logger.getLogger(LimitlessConnectionPlugin.class.getName());
@@ -65,13 +65,13 @@ public class LimitlessConnectionPlugin extends AbstractConnectionPlugin {
   protected final Properties properties;
   private final Supplier<LimitlessRouterService> limitlessRouterServiceSupplier;
   private LimitlessRouterService limitlessRouterService;
-  private static final ReentrantLock lock = new ReentrantLock();
   private static final Set<String> subscribedMethods =
       Collections.unmodifiableSet(new HashSet<String>() {
         {
           add("connect");
         }
       });
+  private static final String INTERNAL_CONNECT_PROPERTY_NAME = "784dd5c2-a77b-4c9f-a0a9-b4ea37395e6c";
 
   static {
     PropertyDefinition.registerPluginProperties(LimitlessConnectionPlugin.class);
@@ -107,6 +107,23 @@ public Connection connect(
       final JdbcCallable<Connection, SQLException> connectFunc)
       throws SQLException {
 
+    if (props.containsKey(INTERNAL_CONNECT_PROPERTY_NAME)) {
+      return connectFunc.call();
+    }
+
+    final Properties copyProps = PropertyUtils.copyProperties(props);
+    copyProps.setProperty(INTERNAL_CONNECT_PROPERTY_NAME, "true");
+    return connectInternal(driverProtocol, hostSpec, copyProps, isInitialConnection, connectFunc);
+  }
+
+  public Connection connectInternal(
+      final String driverProtocol,
+      final HostSpec hostSpec,
+      final Properties props,
+      final boolean isInitialConnection,
+      final JdbcCallable<Connection, SQLException> connectFunc)
+      throws SQLException {
+
     Connection conn = null;
 
     final Dialect dialect = this.pluginService.getDialect();
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java
index f46850ce9..778134d0f 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java
@@ -194,6 +194,7 @@ private void retryConnectWithLeastLoadedRouters(
           } else {
             try {
               context.setConnection(context.getConnectFunc().call());
+              return;
             } catch (final SQLException e) {
               throw new SQLException(Messages.get("LimitlessRouterServiceImpl.noRoutersAvailable"));
             }
@@ -208,7 +209,7 @@ private void retryConnectWithLeastLoadedRouters(
             HostRole.WRITER, HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT);
         LOGGER.finest(Messages.get(
             "LimitlessRouterServiceImpl.selectedHostForRetry",
-            new Object[] {selectedHostSpec.getHost()}));
+            new Object[] {selectedHostSpec != null ? selectedHostSpec.getHost() : "null"}));
         if (selectedHostSpec == null) {
           continue;
         }
@@ -289,8 +290,9 @@ protected void synchronouslyGetLimitlessRouters(final LimitlessConnectionContext
         limitlessRouterCache.put(
             this.pluginService.getHostListProvider().getClusterId(),
             newLimitlessRouters, LimitlessRouterServiceImpl.MONITOR_DISPOSAL_TIME_MS.getLong(context.getProps()));
+      } else {
+        throw new SQLException(Messages.get("LimitlessRouterServiceImpl.fetchedEmptyRouterList"));
       }
-      throw new SQLException(Messages.get("LimitlessRouterServiceImpl.fetchedEmptyRouterList"));
     } finally {
       lock.unlock();
     }