Skip to content

Commit

Permalink
fix: limitless various minor fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
aaronchung-bitquill committed Nov 13, 2024
1 parent 7016110 commit 7498d8a
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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());
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"));
}
Expand All @@ -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;
}
Expand Down Expand Up @@ -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();
}
Expand Down

0 comments on commit 7498d8a

Please sign in to comment.