From d3daa510d0398ecd7ec6ec0bc10efeb5b838b6bb Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Fri, 3 Jan 2025 13:54:39 -0800 Subject: [PATCH] consolidate cache cleanup in a single place --- .../java/software/amazon/jdbc/Driver.java | 43 +++++++ .../jdbc/HikariPooledConnectionProvider.java | 106 +++++------------- .../amazon/jdbc/HikariPoolsHolder.java | 46 ++++++++ .../jdbc/LeastConnectionsHostSelector.java | 16 ++- .../amazon/jdbc/PluginServiceImpl.java | 4 + .../amazon/jdbc/RoundRobinHostSelector.java | 3 +- .../MonitoringRdsHostListProvider.java | 10 +- .../plugin/AwsSecretsManagerCacheHolder.java | 36 ++++++ .../AwsSecretsManagerConnectionPlugin.java | 22 ++-- .../jdbc/plugin/efm2/MonitorServiceImpl.java | 2 +- .../FederatedAuthCacheHolder.java | 32 ++++++ .../federatedauth/FederatedAuthPlugin.java | 10 +- .../federatedauth/OktaAuthCacheHolder.java | 32 ++++++ .../plugin/federatedauth/OktaAuthPlugin.java | 10 +- .../jdbc/plugin/iam/IamAuthCacheHolder.java | 32 ++++++ .../plugin/iam/IamAuthConnectionPlugin.java | 12 +- .../limitless/LimitlessRouterServiceImpl.java | 5 + .../FastestResponseStrategyPlugin.java | 4 + .../HostResponseTimeServiceImpl.java | 11 ++ .../java/software/amazon/jdbc/util/Pair.java | 64 +++++++++++ .../software/amazon/jdbc/util/RdsUtils.java | 1 + .../jdbc/util/SlidingExpirationCache.java | 18 ++- .../container/TestDriverProvider.java | 17 +-- .../tests/AdvancedPerformanceTest.java | 2 +- .../container/tests/PerformanceTest.java | 6 +- .../HikariPooledConnectionProviderTest.java | 24 ++-- ...AwsSecretsManagerConnectionPluginTest.java | 38 +++---- .../FederatedAuthPluginTest.java | 8 +- .../federatedauth/OktaAuthPluginTest.java | 8 +- .../iam/IamAuthConnectionPluginTest.java | 18 +-- 30 files changed, 450 insertions(+), 190 deletions(-) create mode 100644 wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerCacheHolder.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthCacheHolder.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthCacheHolder.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthCacheHolder.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/util/Pair.java diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java index ffaf6abda..ed4f23ea6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/Driver.java +++ b/wrapper/src/main/java/software/amazon/jdbc/Driver.java @@ -33,6 +33,21 @@ import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import software.amazon.jdbc.authentication.AwsCredentialsManager; +import software.amazon.jdbc.dialect.DialectManager; +import software.amazon.jdbc.hostlistprovider.RdsHostListProvider; +import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider; +import software.amazon.jdbc.plugin.AwsSecretsManagerCacheHolder; +import software.amazon.jdbc.plugin.DataCacheConnectionPlugin; +import software.amazon.jdbc.plugin.OpenedConnectionTracker; +import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; +import software.amazon.jdbc.plugin.customendpoint.CustomEndpointPlugin; +import software.amazon.jdbc.plugin.efm.MonitorThreadContainer; +import software.amazon.jdbc.plugin.federatedauth.FederatedAuthCacheHolder; +import software.amazon.jdbc.plugin.federatedauth.OktaAuthCacheHolder; +import software.amazon.jdbc.plugin.iam.IamAuthCacheHolder; +import software.amazon.jdbc.plugin.limitless.LimitlessRouterServiceImpl; +import software.amazon.jdbc.plugin.strategy.fastestresponse.FastestResponseStrategyPlugin; +import software.amazon.jdbc.plugin.strategy.fastestresponse.HostResponseTimeServiceImpl; import software.amazon.jdbc.profile.ConfigurationProfile; import software.amazon.jdbc.profile.DriverConfigurationProfiles; import software.amazon.jdbc.states.ResetSessionStateOnCloseCallable; @@ -277,4 +292,32 @@ public static void setPrepareHostFunc(final Function func) { public static void resetPrepareHostFunc() { RdsUtils.resetPrepareHostFunc(); } + + public static void clearCaches() { + RdsUtils.clearCache(); + RdsHostListProvider.clearAll(); + PluginServiceImpl.clearCache(); + DialectManager.resetEndpointCache(); + MonitoringRdsHostListProvider.clearCache(); + CustomEndpointMonitorImpl.clearCache(); + OpenedConnectionTracker.clearCache(); + AwsSecretsManagerCacheHolder.clearCache(); + DataCacheConnectionPlugin.clearCache(); + FederatedAuthCacheHolder.clearCache(); + OktaAuthCacheHolder.clearCache(); + IamAuthCacheHolder.clearCache(); + LimitlessRouterServiceImpl.clearCache(); + RoundRobinHostSelector.clearCache(); + FastestResponseStrategyPlugin.clearCache(); + } + + public static void releaseResources() { + software.amazon.jdbc.plugin.efm2.MonitorServiceImpl.closeAllMonitors(); + MonitorThreadContainer.releaseInstance(); + ConnectionProviderManager.releaseResources(); + CustomEndpointPlugin.closeMonitors(); + HikariPoolsHolder.closeAllPools(); + HostResponseTimeServiceImpl.closeAllMonitors(); + clearCaches(); + } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java index b0106ca3e..b43d34dda 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java +++ b/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java @@ -39,6 +39,7 @@ import software.amazon.jdbc.targetdriverdialect.ConnectInfo; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.Messages; +import software.amazon.jdbc.util.Pair; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.RdsUtils; @@ -61,17 +62,22 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider, }); protected static final RdsUtils rdsUtils = new RdsUtils(); - protected static SlidingExpirationCache databasePools = - new SlidingExpirationCache<>( - (hikariDataSource) -> hikariDataSource.getHikariPoolMXBean().getActiveConnections() == 0, - HikariDataSource::close - ); protected static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30); protected final HikariPoolConfigurator poolConfigurator; protected final HikariPoolMapping poolMapping; protected final AcceptsUrlFunc acceptsUrlFunc; protected final LeastConnectionsHostSelector leastConnectionsHostSelector; + static { + HikariPoolsHolder.databasePools.setShouldDisposeFunc( + (hikariDataSource) -> { + if (hikariDataSource instanceof HikariDataSource) { + return ((HikariDataSource) hikariDataSource).getHikariPoolMXBean().getActiveConnections() == 0; + } + return true; + }); + } + /** * {@link HikariPooledConnectionProvider} constructor. This class can be passed to * {@link ConnectionProviderManager#setConnectionProvider} to enable internal connection pools for @@ -114,7 +120,7 @@ public HikariPooledConnectionProvider( this.poolConfigurator = hikariPoolConfigurator; this.poolMapping = mapping; this.acceptsUrlFunc = null; - this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools); + this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools); } /** @@ -151,8 +157,8 @@ public HikariPooledConnectionProvider( this.poolMapping = mapping; this.acceptsUrlFunc = null; poolExpirationCheckNanos = poolExpirationNanos; - databasePools.setCleanupIntervalNanos(poolCleanupNanos); - this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools); + HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos); + this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools); } /** @@ -193,8 +199,8 @@ public HikariPooledConnectionProvider( this.poolMapping = mapping; this.acceptsUrlFunc = acceptsUrlFunc; poolExpirationCheckNanos = poolExpirationNanos; - databasePools.setCleanupIntervalNanos(poolCleanupNanos); - this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools); + HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos); + this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools); } @@ -272,8 +278,8 @@ public Connection connect( final HostSpec finalHostSpec = connectionHostSpec; dialect.prepareConnectProperties(copy, protocol, finalHostSpec); - final HikariDataSource ds = databasePools.computeIfAbsent( - new PoolKey(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)), + final HikariDataSource ds = (HikariDataSource) HikariPoolsHolder.databasePools.computeIfAbsent( + Pair.create(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)), (lambdaPoolKey) -> createHikariDataSource(protocol, finalHostSpec, copy, targetDriverDialect), poolExpirationCheckNanos ); @@ -297,22 +303,7 @@ protected String getPoolKey(HostSpec hostSpec, Properties props) { @Override public void releaseResources() { - databasePools.getEntries().forEach((poolKey, pool) -> { - if (!pool.isClosed()) { - pool.close(); - } - }); - databasePools.clear(); - } - - // For testing purposes - public static void clearCache() { - databasePools.getEntries().forEach((poolKey, pool) -> { - if (!pool.isClosed()) { - pool.close(); - } - }); - databasePools.clear(); + HikariPoolsHolder.closeAllPools(); } /** @@ -378,7 +369,7 @@ protected void configurePool( * @return the number of active connection pools */ public int getHostCount() { - return databasePools.size(); + return HikariPoolsHolder.databasePools.size(); } /** @@ -388,8 +379,8 @@ public int getHostCount() { */ public Set getHosts() { return Collections.unmodifiableSet( - databasePools.getEntries().keySet().stream() - .map(poolKey -> poolKey.url) + HikariPoolsHolder.databasePools.getEntries().keySet().stream() + .map(poolKey -> (String) poolKey.getValue1()) .collect(Collectors.toSet())); } @@ -398,8 +389,8 @@ public Set getHosts() { * * @return a set containing every key associated with an active connection pool */ - public Set getKeys() { - return databasePools.getEntries().keySet(); + public Set getKeys() { + return HikariPoolsHolder.databasePools.getEntries().keySet(); } @Override @@ -413,7 +404,7 @@ public String getTargetName() { public void logConnections() { LOGGER.finest(() -> { final StringBuilder builder = new StringBuilder(); - databasePools.getEntries().forEach((key, dataSource) -> { + HikariPoolsHolder.databasePools.getEntries().forEach((key, dataSource) -> { builder.append("\t[ "); builder.append(key).append(":"); builder.append("\n\t {"); @@ -436,51 +427,8 @@ HikariDataSource createHikariDataSource( return new HikariDataSource(config); } - public static class PoolKey { - private final @NonNull String url; - private final @NonNull String extraKey; - - public PoolKey(final @NonNull String url, final @NonNull String extraKey) { - this.url = url; - this.extraKey = extraKey; - } - - public String getUrl() { - return this.url; - } - - @Override - public int hashCode() { - final int prime = 31; - int result = 1; - result = prime * result + ((url == null) ? 0 : url.hashCode()) + ((extraKey == null) ? 0 : extraKey.hashCode()); - return result; - } - - @Override - public boolean equals(final Object obj) { - if (this == obj) { - return true; - } - if (obj == null) { - return false; - } - if (getClass() != obj.getClass()) { - return false; - } - final PoolKey other = (PoolKey) obj; - return this.url.equals(other.url) && this.extraKey.equals(other.extraKey); - } - - @Override - public String toString() { - return "PoolKey [url=" + url + ", extraKey=" + extraKey + "]"; - } - - } - // For testing purposes only - void setDatabasePools(SlidingExpirationCache connectionPools) { - databasePools = connectionPools; + void setDatabasePools(SlidingExpirationCache connectionPools) { + HikariPoolsHolder.databasePools = connectionPools; } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java new file mode 100644 index 000000000..7bec68a49 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java @@ -0,0 +1,46 @@ +/* + * 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; + +import software.amazon.jdbc.util.Pair; +import software.amazon.jdbc.util.SlidingExpirationCache; + +public class HikariPoolsHolder { + static SlidingExpirationCache databasePools = + new SlidingExpirationCache<>( + null, + (hikariDataSource) -> { + try { + hikariDataSource.close(); + } catch (Exception ex) { + // ignore + } + } + ); + + public static void closeAllPools() { + databasePools.getEntries().forEach((poolKey, pool) -> { + try { + pool.close(); + } catch (Exception ex) { + // ignore + } + }); + databasePools.clear(); + + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java index ea9d14cfa..2ada66281 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java +++ b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java @@ -26,14 +26,15 @@ import org.checkerframework.checker.nullness.qual.Nullable; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.Messages; +import software.amazon.jdbc.util.Pair; import software.amazon.jdbc.util.SlidingExpirationCache; public class LeastConnectionsHostSelector implements HostSelector { public static final String STRATEGY_LEAST_CONNECTIONS = "leastConnections"; - private final SlidingExpirationCache databasePools; + private final SlidingExpirationCache databasePools; public LeastConnectionsHostSelector( - SlidingExpirationCache databasePools) { + SlidingExpirationCache databasePools) { this.databasePools = databasePools; } @@ -58,15 +59,18 @@ public HostSpec getHost( private int getNumConnections( final HostSpec hostSpec, - final SlidingExpirationCache databasePools) { + final SlidingExpirationCache databasePools) { int numConnections = 0; final String url = hostSpec.getUrl(); - for (final Map.Entry entry : + for (final Map.Entry entry : databasePools.getEntries().entrySet()) { - if (!url.equals(entry.getKey().getUrl())) { + if (!url.equals(entry.getKey().getValue1())) { continue; } - numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections(); + if (!(entry.getValue() instanceof HikariDataSource)) { + continue; + } + numConnections += ((HikariDataSource) entry.getValue()).getHikariPoolMXBean().getActiveConnections(); } return numConnections; } diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java index 326d020ce..a2a888a20 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java @@ -754,4 +754,8 @@ public T getPlugin(final Class pluginClazz) { } return null; } + + public static void clearCache() { + hostAvailabilityExpiringCache.clear(); + } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java index d2a9ff31e..3ea1a4779 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java +++ b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java @@ -251,8 +251,7 @@ private void updateCachedHostWeightPairsPropertiesForRoundRobinClusterInfo( } } - // For testing purposes only - public void clearCache() { + public static void clearCache() { roundRobinCache.clear(); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java index 52bffb6a2..0dc9a46e9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java @@ -93,6 +93,12 @@ public void clear() { } public static void clearCache() { + topologyCache.clear(); + primaryClusterIdCache.clear(); + suggestedPrimaryClusterIdCache.clear(); + } + + public static void closeAllMonitors() { monitors.getEntries().values().forEach(monitor -> { try { monitor.close(); @@ -101,9 +107,7 @@ public static void clearCache() { } }); monitors.clear(); - topologyCache.clear(); - primaryClusterIdCache.clear(); - suggestedPrimaryClusterIdCache.clear(); + clearCache(); } @Override diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerCacheHolder.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerCacheHolder.java new file mode 100644 index 000000000..f86c19c04 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerCacheHolder.java @@ -0,0 +1,36 @@ +/* + * 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.plugin; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import software.amazon.jdbc.plugin.AwsSecretsManagerConnectionPlugin.Secret; +import software.amazon.jdbc.util.Pair; + +/* The main plugin code AwsSecretsManagerConnectionPlugin depends on AWS SDK. + In order to avoid unnecessary dependencies, the plugin cache has been extracted into this + AwsSecretsManagerCacheHolder class. This cache holder class doesn't depend on AWS SDK and + can be safely cleared if needed. + */ +public class AwsSecretsManagerCacheHolder { + static final Map, Secret> secretsCache + = new ConcurrentHashMap<>(); + + public static void clearCache() { + secretsCache.clear(); + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java index 4032588b7..763650ba6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPlugin.java @@ -26,10 +26,8 @@ import java.sql.SQLException; import java.util.Collections; import java.util.HashSet; -import java.util.Map; import java.util.Properties; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; import java.util.function.BiFunction; import java.util.function.Function; import java.util.logging.Level; @@ -42,7 +40,6 @@ import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest; import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse; import software.amazon.awssdk.services.secretsmanager.model.SecretsManagerException; -import software.amazon.awssdk.utils.Pair; import software.amazon.jdbc.AwsWrapperProperty; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.JdbcCallable; @@ -50,6 +47,7 @@ import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.util.Messages; +import software.amazon.jdbc.util.Pair; import software.amazon.jdbc.util.RegionUtils; import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.telemetry.TelemetryContext; @@ -81,12 +79,10 @@ public class AwsSecretsManagerConnectionPlugin extends AbstractConnectionPlugin "The endpoint of the secret to retrieve."); protected static final RegionUtils regionUtils = new RegionUtils(); - protected static final Map, Secret> secretsCache = new ConcurrentHashMap<>(); - private static final Pattern SECRETS_ARN_PATTERN = Pattern.compile("^arn:aws:secretsmanager:(?[^:\\n]*):[^:\\n]*:([^:/\\n]*[:/])?(.*)$"); - final Pair secretKey; + final Pair secretKey; private final BiFunction secretsManagerClientFunc; private final Function getSecretValueRequestFunc; @@ -173,7 +169,7 @@ public AwsSecretsManagerConnectionPlugin(final PluginService pluginService, fina new Object[] {REGION_PROPERTY.name})); } - this.secretKey = Pair.of(secretId, region); + this.secretKey = Pair.create(secretId, region.id()); this.secretsManagerClientFunc = secretsManagerClientFunc; this.getSecretValueRequestFunc = getSecretValueRequestFunc; @@ -253,14 +249,14 @@ private boolean updateSecret(final HostSpec hostSpec, final boolean forceReFetch try { boolean fetched = false; - this.secret = secretsCache.get(this.secretKey); + this.secret = AwsSecretsManagerCacheHolder.secretsCache.get(this.secretKey); if (secret == null || forceReFetch) { try { this.secret = fetchLatestCredentials(hostSpec); if (this.secret != null) { fetched = true; - secretsCache.put(this.secretKey, this.secret); + AwsSecretsManagerCacheHolder.secretsCache.put(this.secretKey, this.secret); } } catch (final SecretsManagerException | JsonProcessingException exception) { LOGGER.log( @@ -322,8 +318,8 @@ Secret fetchLatestCredentials(final HostSpec hostSpec) throws SecretsManagerException, JsonProcessingException { final SecretsManagerClient client = secretsManagerClientFunc.apply( hostSpec, - this.secretKey.right()); - final GetSecretValueRequest request = getSecretValueRequestFunc.apply(this.secretKey.left()); + Region.of(this.secretKey.getValue2())); + final GetSecretValueRequest request = getSecretValueRequestFunc.apply(this.secretKey.getValue1()); final GetSecretValueResponse valueResponse; try { @@ -349,6 +345,10 @@ private void applySecretToProperties(final Properties properties) { } } + public static void clearCache() { + AwsSecretsManagerCacheHolder.clearCache(); + } + @JsonIgnoreProperties(ignoreUnknown = true) static class Secret { @JsonProperty("username") diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/MonitorServiceImpl.java index 8e380eb22..f0f10dd0a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/MonitorServiceImpl.java @@ -94,7 +94,7 @@ public MonitorServiceImpl(final @NonNull PluginService pluginService) { this.monitorInitializer = monitorInitializer; } - public static void clearCache() { + public static void closeAllMonitors() { monitors.getEntries().values().forEach(monitor -> { try { monitor.close(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthCacheHolder.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthCacheHolder.java new file mode 100644 index 000000000..0ab3ebac0 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthCacheHolder.java @@ -0,0 +1,32 @@ +/* + * 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.plugin.federatedauth; + +import java.util.concurrent.ConcurrentHashMap; +import software.amazon.jdbc.plugin.TokenInfo; + +/* The main plugin code FederatedAuthPlugin depends on AWS SDK. In order to avoid unnecessary dependencies, + the plugin cache has been extracted into this FederatedAuthCacheHolder class. This cache holder class doesn't depend + on AWS SDK and can be safely cleared if needed. + */ +public class FederatedAuthCacheHolder { + static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); + + public static void clearCache() { + tokenCache.clear(); + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java index 8e9b4fb1c..c9a6ba946 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPlugin.java @@ -49,7 +49,6 @@ public class FederatedAuthPlugin extends AbstractConnectionPlugin { - static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); private final CredentialsProviderFactory credentialsProviderFactory; private static final int DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60 - 30; private static final int DEFAULT_HTTP_TIMEOUT_MILLIS = 60000; @@ -144,7 +143,8 @@ public FederatedAuthPlugin(final PluginService pluginService, this.pluginService = pluginService; this.credentialsProviderFactory = credentialsProviderFactory; this.telemetryFactory = pluginService.getTelemetryFactory(); - this.cacheSizeGauge = telemetryFactory.createGauge("federatedAuth.tokenCache.size", () -> (long) tokenCache.size()); + this.cacheSizeGauge = telemetryFactory.createGauge("federatedAuth.tokenCache.size", + () -> (long) FederatedAuthCacheHolder.tokenCache.size()); this.fetchTokenCounter = telemetryFactory.createCounter("federatedAuth.fetchToken.count"); this.rdsUtils = rdsUtils; this.samlUtils = new SamlUtils(this.rdsUtils); @@ -197,7 +197,7 @@ private Connection connectInternal(final HostSpec hostSpec, final Properties pro port, region); - final TokenInfo tokenInfo = tokenCache.get(cacheKey); + final TokenInfo tokenInfo = FederatedAuthCacheHolder.tokenCache.get(cacheKey); final boolean isCachedToken = tokenInfo != null && !tokenInfo.isExpired(); @@ -256,12 +256,12 @@ private void updateAuthenticationToken( "AuthenticationToken.useCachedToken", new Object[] {token})); PropertyDefinition.PASSWORD.set(props, token); - tokenCache.put( + FederatedAuthCacheHolder.tokenCache.put( cacheKey, new TokenInfo(token, tokenExpiry)); } public static void clearCache() { - tokenCache.clear(); + FederatedAuthCacheHolder.clearCache(); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthCacheHolder.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthCacheHolder.java new file mode 100644 index 000000000..0b38936d8 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthCacheHolder.java @@ -0,0 +1,32 @@ +/* + * 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.plugin.federatedauth; + +import java.util.concurrent.ConcurrentHashMap; +import software.amazon.jdbc.plugin.TokenInfo; + +/* The main plugin code OktaAuthPlugin depends on AWS SDK. In order to avoid unnecessary dependencies, + the plugin cache has been extracted into this OktaAuthCacheHolder class. This cache holder class doesn't depend + on AWS SDK and can be safely cleared if needed. + */ +public class OktaAuthCacheHolder { + static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); + + public static void clearCache() { + tokenCache.clear(); + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java index e0cac79cd..e22888ee8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPlugin.java @@ -47,7 +47,6 @@ public class OktaAuthPlugin extends AbstractConnectionPlugin { - static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); private final CredentialsProviderFactory credentialsProviderFactory; private static final int DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60 - 30; private static final int DEFAULT_HTTP_TIMEOUT_MILLIS = 60000; @@ -118,7 +117,8 @@ public OktaAuthPlugin(PluginService pluginService, CredentialsProviderFactory cr this.samlUtils = new SamlUtils(this.rdsUtils); this.iamTokenUtility = tokenUtils; this.telemetryFactory = pluginService.getTelemetryFactory(); - this.cacheSizeGauge = telemetryFactory.createGauge("oktaAuth.tokenCache.size", () -> (long) tokenCache.size()); + this.cacheSizeGauge = telemetryFactory.createGauge("oktaAuth.tokenCache.size", + () -> (long) OktaAuthCacheHolder.tokenCache.size()); this.fetchTokenCounter = telemetryFactory.createCounter("oktaAuth.fetchToken.count"); } @@ -169,7 +169,7 @@ private Connection connectInternal(final HostSpec hostSpec, final Properties pro port, region); - final TokenInfo tokenInfo = tokenCache.get(cacheKey); + final TokenInfo tokenInfo = OktaAuthCacheHolder.tokenCache.get(cacheKey); final boolean isCachedToken = tokenInfo != null && !tokenInfo.isExpired(); @@ -228,12 +228,12 @@ private void updateAuthenticationToken( "AuthenticationToken.useCachedToken", new Object[] {token})); PropertyDefinition.PASSWORD.set(props, token); - tokenCache.put( + OktaAuthCacheHolder.tokenCache.put( cacheKey, new TokenInfo(token, tokenExpiry)); } public static void clearCache() { - tokenCache.clear(); + OktaAuthCacheHolder.clearCache(); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthCacheHolder.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthCacheHolder.java new file mode 100644 index 000000000..8cadb297b --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthCacheHolder.java @@ -0,0 +1,32 @@ +/* + * 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.plugin.iam; + +import java.util.concurrent.ConcurrentHashMap; +import software.amazon.jdbc.plugin.TokenInfo; + +/* The main plugin code IamAuthConnectionPlugin depends on AWS SDK. In order to avoid unnecessary dependencies, + the plugin cache has been extracted into this IamAuthCacheHolder class. This cache holder class doesn't depend + on AWS SDK and can be safely cleared if needed. + */ +public class IamAuthCacheHolder { + static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); + + public static void clearCache() { + tokenCache.clear(); + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java index 29613da32..bc1786b44 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPlugin.java @@ -55,7 +55,6 @@ public class IamAuthConnectionPlugin extends AbstractConnectionPlugin { add("forceConnect"); } }); - static final ConcurrentHashMap tokenCache = new ConcurrentHashMap<>(); private static final int DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60 - 30; public static final AwsWrapperProperty IAM_HOST = new AwsWrapperProperty( @@ -96,7 +95,8 @@ public IamAuthConnectionPlugin(final @NonNull PluginService pluginService) { this.iamTokenUtility = utility; this.pluginService = pluginService; this.telemetryFactory = pluginService.getTelemetryFactory(); - this.cacheSizeGauge = telemetryFactory.createGauge("iam.tokenCache.size", () -> (long) tokenCache.size()); + this.cacheSizeGauge = telemetryFactory.createGauge("iam.tokenCache.size", + () -> (long) IamAuthCacheHolder.tokenCache.size()); this.fetchTokenCounter = telemetryFactory.createCounter("iam.fetchToken.count"); } @@ -142,7 +142,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro host, port, region); - final TokenInfo tokenInfo = tokenCache.get(cacheKey); + final TokenInfo tokenInfo = IamAuthCacheHolder.tokenCache.get(cacheKey); final boolean isCachedToken = tokenInfo != null && !tokenInfo.isExpired(); if (isCachedToken) { @@ -167,7 +167,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro "AuthenticationToken.generatedNewToken", new Object[] {token})); PropertyDefinition.PASSWORD.set(props, token); - tokenCache.put( + IamAuthCacheHolder.tokenCache.put( cacheKey, new TokenInfo(token, tokenExpiry)); } @@ -203,7 +203,7 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Pro "AuthenticationToken.generatedNewToken", new Object[] {token})); PropertyDefinition.PASSWORD.set(props, token); - tokenCache.put( + IamAuthCacheHolder.tokenCache.put( cacheKey, new TokenInfo(token, tokenExpiry)); @@ -230,6 +230,6 @@ public Connection forceConnect( } public static void clearCache() { - tokenCache.clear(); + IamAuthCacheHolder.clearCache(); } } 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 993afe950..7502587d4 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 @@ -326,4 +326,9 @@ public void startMonitoring(final @NonNull HostSpec hostSpec, throw new RuntimeException(e); } } + + public static void clearCache() { + forceGetLimitlessRoutersLockMap.clear(); + limitlessRouterCache.clear(); + } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/FastestResponseStrategyPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/FastestResponseStrategyPlugin.java index d12a8df35..ddcb2a3c9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/FastestResponseStrategyPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/FastestResponseStrategyPlugin.java @@ -196,6 +196,10 @@ public void notifyNodeListChanged(final Map> this.hostResponseTimeService.setHosts(this.hosts); } + public static void clearCache() { + cachedFastestResponseHostByRole.clear(); + } + private static class ResponseTimeTuple { public HostSpec hostSpec; public int responseTime; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java index 782662838..1632be0b9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java @@ -109,4 +109,15 @@ public void setHosts(final @NonNull List hosts) { } }); } + + public static void closeAllMonitors() { + monitoringNodes.getEntries().values().forEach(monitor -> { + try { + monitor.close(); + } catch (Exception ex) { + // ignore + } + }); + monitoringNodes.clear(); + } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/Pair.java b/wrapper/src/main/java/software/amazon/jdbc/util/Pair.java new file mode 100644 index 000000000..300a16766 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/util/Pair.java @@ -0,0 +1,64 @@ +/* + * 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; + +public final class Pair { + + private final T1 value1; + private final T2 value2; + + private Pair(T1 value1, T2 value2) { + this.value1 = value1; + this.value2 = value2; + } + + public T1 getValue1() { + return this.value1; + } + + public T2 getValue2() { + return this.value2; + } + + @Override + public boolean equals(Object obj) { + if (!(obj instanceof Pair)) { + return false; + } + + Pair other = (Pair) obj; + + return (this.value1 == null && other.value1 == null || other.value1 != null && other.value1.equals(this.value1)) + && (this.value2 == null && other.value2 == null || other.value2 != null && other.value2.equals(value2)); + } + + @Override + public int hashCode() { + return getClass().hashCode() + + (this.value1 != null ? this.value1.hashCode() : 0) + + (this.value2 != null ? this.value2.hashCode() : 0); + } + + @Override + public String toString() { + return String.format("Pair(value1=%s, value2=%s)", this.value1, this.value2); + } + + public static Pair create(T1 value1, T2 value2) { + return new Pair<>(value1, value2); + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java index 415903845..90221cf19 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java @@ -423,6 +423,7 @@ private String getDnsGroup(final String host) { public static void clearCache() { cachedPatterns.clear(); + cachedDnsPatterns.clear(); } public static void setPrepareHostFunc(final Function func) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/SlidingExpirationCache.java b/wrapper/src/main/java/software/amazon/jdbc/util/SlidingExpirationCache.java index d384cac41..015780948 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/SlidingExpirationCache.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/SlidingExpirationCache.java @@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; public class SlidingExpirationCache { @@ -30,7 +31,7 @@ public class SlidingExpirationCache { protected final Map cache = new ConcurrentHashMap<>(); protected long cleanupIntervalNanos = TimeUnit.MINUTES.toNanos(10); protected final AtomicLong cleanupTimeNanos = new AtomicLong(System.nanoTime() + cleanupIntervalNanos); - protected final ShouldDisposeFunc shouldDisposeFunc; + protected final AtomicReference> shouldDisposeFunc = new AtomicReference<>(null); protected final ItemDisposalFunc itemDisposalFunc; /** @@ -38,7 +39,7 @@ public class SlidingExpirationCache { * as not-expired and renews its expiration time. */ public SlidingExpirationCache() { - this.shouldDisposeFunc = null; + this.shouldDisposeFunc.set(null); this.itemDisposalFunc = null; } @@ -56,7 +57,7 @@ public SlidingExpirationCache() { public SlidingExpirationCache( final ShouldDisposeFunc shouldDisposeFunc, final ItemDisposalFunc itemDisposalFunc) { - this.shouldDisposeFunc = shouldDisposeFunc; + this.shouldDisposeFunc.set(shouldDisposeFunc); this.itemDisposalFunc = itemDisposalFunc; } @@ -64,11 +65,15 @@ public SlidingExpirationCache( final ShouldDisposeFunc shouldDisposeFunc, final ItemDisposalFunc itemDisposalFunc, final long cleanupIntervalNanos) { - this.shouldDisposeFunc = shouldDisposeFunc; + this.shouldDisposeFunc.set(shouldDisposeFunc); this.itemDisposalFunc = itemDisposalFunc; this.cleanupIntervalNanos = cleanupIntervalNanos; } + public void setShouldDisposeFunc(final ShouldDisposeFunc shouldDisposeFunc) { + this.shouldDisposeFunc.set(shouldDisposeFunc); + } + /** * In addition to performing the logic defined by {@link Map#computeIfAbsent}, cleans up expired * entries if we have hit cleanup time. If an expired entry is requested and we have not hit @@ -257,8 +262,9 @@ public CacheItem(final V item, final long expirationTimeNano) { * false. */ boolean shouldCleanup() { - if (shouldDisposeFunc != null) { - return System.nanoTime() > expirationTimeNano && shouldDisposeFunc.shouldDispose(this.item); + final ShouldDisposeFunc tempShouldDisposeFunc = shouldDisposeFunc.get(); + if (tempShouldDisposeFunc != null) { + return System.nanoTime() > expirationTimeNano && tempShouldDisposeFunc.shouldDispose(this.item); } return System.nanoTime() > expirationTimeNano; } diff --git a/wrapper/src/test/java/integration/container/TestDriverProvider.java b/wrapper/src/test/java/integration/container/TestDriverProvider.java index d7200c1a2..474ae4842 100644 --- a/wrapper/src/test/java/integration/container/TestDriverProvider.java +++ b/wrapper/src/test/java/integration/container/TestDriverProvider.java @@ -55,7 +55,8 @@ import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider; import org.junit.platform.commons.util.AnnotationUtils; import software.amazon.jdbc.ConnectionProviderManager; -import software.amazon.jdbc.HikariPooledConnectionProvider; +import software.amazon.jdbc.Driver; +import software.amazon.jdbc.HikariPoolsHolder; import software.amazon.jdbc.dialect.DialectManager; import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider; import software.amazon.jdbc.plugin.OpenedConnectionTracker; @@ -224,22 +225,10 @@ private static void registerDrivers(final TestDriver testDriver) throws SQLExcep } private static void clearCaches() { - RdsUtils.clearCache(); - TestAuroraHostListProvider.clearCache(); - TestPluginServiceImpl.clearHostAvailabilityCache(); - DialectManager.resetEndpointCache(); + Driver.releaseResources(); TargetDriverDialectManager.resetCustomDialect(); - MonitorThreadContainer.releaseInstance(); - MonitorServiceImpl.clearCache(); - ConnectionProviderManager.releaseResources(); ConnectionProviderManager.resetProvider(); ConnectionProviderManager.resetConnectionInitFunc(); - software.amazon.jdbc.plugin.efm2.MonitorServiceImpl.clearCache(); - HikariPooledConnectionProvider.clearCache(); - MonitoringRdsHostListProvider.clearCache(); - CustomEndpointPlugin.closeMonitors(); - CustomEndpointMonitorImpl.clearCache(); - OpenedConnectionTracker.clearCache(); } private static void checkClusterHealth(final boolean makeSureFirstInstanceWriter) diff --git a/wrapper/src/test/java/integration/container/tests/AdvancedPerformanceTest.java b/wrapper/src/test/java/integration/container/tests/AdvancedPerformanceTest.java index 8f46bcfdc..c91843320 100644 --- a/wrapper/src/test/java/integration/container/tests/AdvancedPerformanceTest.java +++ b/wrapper/src/test/java/integration/container/tests/AdvancedPerformanceTest.java @@ -688,7 +688,7 @@ private void ensureClusterHealthy() throws InterruptedException { TestAuroraHostListProvider.clearCache(); TestPluginServiceImpl.clearHostAvailabilityCache(); MonitorThreadContainer.releaseInstance(); - MonitorServiceImpl.clearCache(); + MonitorServiceImpl.closeAllMonitors(); } private static Stream generateParams() { diff --git a/wrapper/src/test/java/integration/container/tests/PerformanceTest.java b/wrapper/src/test/java/integration/container/tests/PerformanceTest.java index 2d534dfde..8df6c1669 100644 --- a/wrapper/src/test/java/integration/container/tests/PerformanceTest.java +++ b/wrapper/src/test/java/integration/container/tests/PerformanceTest.java @@ -146,7 +146,7 @@ public void test_FailureDetectionTime_EnhancedMonitoringEnabled(final String efm OpenedConnectionTracker.clearCache(); MonitorThreadContainer.releaseInstance(); - MonitorServiceImpl.clearCache(); + MonitorServiceImpl.closeAllMonitors(); AuroraHostListProvider.clearAll(); MonitoringRdsHostListProvider.clearCache(); @@ -228,7 +228,7 @@ public void test_FailureDetectionTime_FailoverAndEnhancedMonitoringEnabled(final OpenedConnectionTracker.clearCache(); MonitorThreadContainer.releaseInstance(); - MonitorServiceImpl.clearCache(); + MonitorServiceImpl.closeAllMonitors(); AuroraHostListProvider.clearAll(); MonitoringRdsHostListProvider.clearCache(); @@ -316,7 +316,7 @@ private void test_FailoverTime_SocketTimeout(final String plugins) throws IOExce OpenedConnectionTracker.clearCache(); MonitorThreadContainer.releaseInstance(); - MonitorServiceImpl.clearCache(); + MonitorServiceImpl.closeAllMonitors(); AuroraHostListProvider.clearAll(); MonitoringRdsHostListProvider.clearCache(); diff --git a/wrapper/src/test/java/software/amazon/jdbc/HikariPooledConnectionProviderTest.java b/wrapper/src/test/java/software/amazon/jdbc/HikariPooledConnectionProviderTest.java index 087c03a79..1b7e4a431 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/HikariPooledConnectionProviderTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/HikariPooledConnectionProviderTest.java @@ -45,11 +45,11 @@ import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.HikariPooledConnectionProvider.PoolKey; import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.targetdriverdialect.ConnectInfo; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.Pair; import software.amazon.jdbc.util.SlidingExpirationCache; class HikariPooledConnectionProviderTest { @@ -129,8 +129,8 @@ void tearDown() throws Exception { void testConnectWithDefaultMapping() throws SQLException { when(mockHostSpec.getUrl()).thenReturn("url"); final Set expectedUrls = new HashSet<>(Collections.singletonList("url")); - final Set expectedKeys = new HashSet<>( - Collections.singletonList(new PoolKey("url", user1))); + final Set expectedKeys = new HashSet<>( + Collections.singletonList(Pair.create("url", user1))); provider = spy(new HikariPooledConnectionProvider((hostSpec, properties) -> mockConfig)); @@ -146,7 +146,7 @@ void testConnectWithDefaultMapping() throws SQLException { assertEquals(1, provider.getHostCount()); final Set hosts = provider.getHosts(); assertEquals(expectedUrls, hosts); - final Set keys = provider.getKeys(); + final Set keys = provider.getKeys(); assertEquals(expectedKeys, keys); } } @@ -154,8 +154,8 @@ void testConnectWithDefaultMapping() throws SQLException { @Test void testConnectWithCustomMapping() throws SQLException { when(mockHostSpec.getUrl()).thenReturn("url"); - final Set expectedKeys = new HashSet<>( - Collections.singletonList(new PoolKey("url", "url+someUniqueKey"))); + final Set expectedKeys = new HashSet<>( + Collections.singletonList(Pair.create("url", "url+someUniqueKey"))); provider = spy(new HikariPooledConnectionProvider( (hostSpec, properties) -> mockConfig, @@ -169,7 +169,7 @@ void testConnectWithCustomMapping() throws SQLException { try (Connection conn = provider.connect(protocol, mockDialect, mockTargetDriverDialect, mockHostSpec, props)) { assertEquals(mockConnection, conn); assertEquals(1, provider.getHostCount()); - final Set keys = provider.getKeys(); + final Set keys = provider.getKeys(); assertEquals(expectedKeys, keys); } } @@ -208,13 +208,13 @@ public void testLeastConnectionsStrategy() throws SQLException { assertEquals(readerUrl1Connection, selectedHost.getHost()); } - private SlidingExpirationCache getTestPoolMap() { - SlidingExpirationCache map = new SlidingExpirationCache<>(); - map.computeIfAbsent(new PoolKey(readerHost2Connection.getUrl(), user1), + private SlidingExpirationCache getTestPoolMap() { + SlidingExpirationCache map = new SlidingExpirationCache<>(); + map.computeIfAbsent(Pair.create(readerHost2Connection.getUrl(), user1), (key) -> dsWith1Connection, TimeUnit.MINUTES.toNanos(10)); - map.computeIfAbsent(new PoolKey(readerHost2Connection.getUrl(), user2), + map.computeIfAbsent(Pair.create(readerHost2Connection.getUrl(), user2), (key) -> dsWith1Connection, TimeUnit.MINUTES.toNanos(10)); - map.computeIfAbsent(new PoolKey(readerHost1Connection.getUrl(), user1), + map.computeIfAbsent(Pair.create(readerHost1Connection.getUrl(), user1), (key) -> dsWith1Connection, TimeUnit.MINUTES.toNanos(10)); return map; } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPluginTest.java index 041778e61..647a41f37 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/AwsSecretsManagerConnectionPluginTest.java @@ -52,7 +52,6 @@ import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest; import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse; import software.amazon.awssdk.services.secretsmanager.model.SecretsManagerException; -import software.amazon.awssdk.utils.Pair; import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.HostSpecBuilder; @@ -71,6 +70,7 @@ import software.amazon.jdbc.states.SessionStateService; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.Messages; +import software.amazon.jdbc.util.Pair; import software.amazon.jdbc.util.telemetry.GaugeCallable; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryCounter; @@ -92,7 +92,7 @@ public class AwsSecretsManagerConnectionPluginTest { private static final String TEST_SQL_ERROR = "SQL exception error message"; private static final String UNHANDLED_ERROR_CODE = "HY000"; private static final int TEST_PORT = 5432; - private static final Pair SECRET_CACHE_KEY = Pair.of(TEST_SECRET_ID, Region.of(TEST_REGION)); + private static final Pair SECRET_CACHE_KEY = Pair.create(TEST_SECRET_ID, TEST_REGION); private static final AwsSecretsManagerConnectionPlugin.Secret TEST_SECRET = new AwsSecretsManagerConnectionPlugin.Secret("testUser", "testPassword"); private static final HostSpec TEST_HOSTSPEC = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) @@ -155,7 +155,7 @@ public void init() throws SQLException { @AfterEach void cleanUp() throws Exception { closeable.close(); - AwsSecretsManagerConnectionPlugin.secretsCache.clear(); + AwsSecretsManagerCacheHolder.clearCache(); TEST_PROPS.clear(); } @@ -165,11 +165,11 @@ void cleanUp() throws Exception { @Test public void testConnectWithCachedSecrets() throws SQLException { // Add initial cached secret to be used for a connection. - AwsSecretsManagerConnectionPlugin.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); + AwsSecretsManagerCacheHolder.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); this.plugin.connect(TEST_PG_PROTOCOL, TEST_HOSTSPEC, TEST_PROPS, true, this.connectFunc); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(this.mockSecretsManagerClient, never()).getSecretValue(this.mockGetValueRequest); verify(this.connectFunc).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); @@ -187,7 +187,7 @@ public void testConnectWithNewSecrets() throws SQLException { this.plugin.connect(TEST_PG_PROTOCOL, TEST_HOSTSPEC, TEST_PROPS, true, this.connectFunc); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(this.mockSecretsManagerClient).getSecretValue(this.mockGetValueRequest); verify(connectFunc).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); @@ -210,7 +210,7 @@ public void testMissingRequiredParameters(final Properties properties) { */ @Test public void testFailedInitialConnectionWithUnhandledError() throws SQLException { - AwsSecretsManagerConnectionPlugin.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); + AwsSecretsManagerCacheHolder.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); final SQLException failedFirstConnectionGenericException = new SQLException(TEST_SQL_ERROR, UNHANDLED_ERROR_CODE); doThrow(failedFirstConnectionGenericException).when(connectFunc).call(); @@ -257,7 +257,7 @@ public void testConnectWithNewSecretsAfterTryingWithCachedSecrets( // Fail the initial connection attempt with cached secret. // Second attempt should be successful. - AwsSecretsManagerConnectionPlugin.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); + AwsSecretsManagerCacheHolder.secretsCache.put(SECRET_CACHE_KEY, TEST_SECRET); final SQLException failedFirstConnectionAccessException = new SQLException(TEST_SQL_ERROR, accessError); doThrow(failedFirstConnectionAccessException).when(connectFunc).call(); @@ -275,7 +275,7 @@ public void testConnectWithNewSecretsAfterTryingWithCachedSecrets( true, this.connectFunc)); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(this.mockSecretsManagerClient).getSecretValue(this.mockGetValueRequest); verify(connectFunc, times(2)).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); @@ -305,7 +305,7 @@ public void testFailedToReadSecrets() throws SQLException { readSecretsFailedException.getMessage(), Messages.get( "AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials")); - assertEquals(0, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(0, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(this.mockSecretsManagerClient).getSecretValue(this.mockGetValueRequest); verify(this.connectFunc, never()).call(); } @@ -332,7 +332,7 @@ public void testFailedToGetSecrets() throws SQLException { getSecretsFailedException.getMessage(), Messages.get( "AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials")); - assertEquals(0, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(0, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(this.mockSecretsManagerClient).getSecretValue(this.mockGetValueRequest); verify(this.connectFunc, never()).call(); } @@ -374,7 +374,7 @@ public void testFailedInitialConnectionWithWrappedGenericError(final String acce true, this.connectFunc)); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(connectFunc).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); assertEquals(TEST_PASSWORD, TEST_PROPS.get(PropertyDefinition.PASSWORD.name)); @@ -415,7 +415,7 @@ public void testConnectWithWrappedMySQLException() throws SQLException { true, this.connectFunc)); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(connectFunc).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); assertEquals(TEST_PASSWORD, TEST_PROPS.get(PropertyDefinition.PASSWORD.name)); @@ -456,7 +456,7 @@ public void testConnectWithWrappedPostgreSQLException() throws SQLException { true, this.connectFunc)); - assertEquals(1, AwsSecretsManagerConnectionPlugin.secretsCache.size()); + assertEquals(1, AwsSecretsManagerCacheHolder.secretsCache.size()); verify(connectFunc).call(); assertEquals(TEST_USERNAME, TEST_PROPS.get(PropertyDefinition.USER.name)); assertEquals(TEST_PASSWORD, TEST_PROPS.get(PropertyDefinition.PASSWORD.name)); @@ -476,8 +476,8 @@ public void testConnectViaARN(final String arn, final Region expectedRegionParse (host, r) -> mockSecretsManagerClient, (id) -> mockGetValueRequest)); - final Pair secret = this.plugin.secretKey; - assertEquals(expectedRegionParsedFromARN, secret.right()); + final Pair secret = this.plugin.secretKey; + assertEquals(expectedRegionParsedFromARN, Region.of(secret.getValue2())); } @ParameterizedTest @@ -496,10 +496,10 @@ public void testConnectionWithRegionParameterAndARN(final String arn, final Regi (host, r) -> mockSecretsManagerClient, (id) -> mockGetValueRequest)); - final Pair secret = this.plugin.secretKey; + final Pair secret = this.plugin.secretKey; // The region specified in `secretsManagerRegion` should override the region parsed from ARN. - assertNotEquals(regionParsedFromARN, secret.right()); - assertEquals(expectedRegion, secret.right()); + assertNotEquals(regionParsedFromARN, Region.of(secret.getValue2())); + assertEquals(expectedRegion, Region.of(secret.getValue2())); } private static Stream provideExceptionCodeForDifferentDrivers() { diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPluginTest.java index 34d3a5212..7b72dae6e 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/FederatedAuthPluginTest.java @@ -117,7 +117,7 @@ void testCachedToken() throws SQLException { new FederatedAuthPlugin(mockPluginService, mockCredentialsProviderFactory); String key = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + DEFAULT_PORT + ":iamUser"; - FederatedAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + FederatedAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); plugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); @@ -134,7 +134,7 @@ void testExpiredCachedToken() throws SQLException { String someExpiredToken = "someExpiredToken"; TokenInfo expiredTokenInfo = new TokenInfo( someExpiredToken, Instant.now().minusMillis(300000)); - FederatedAuthPlugin.tokenCache.put(key, expiredTokenInfo); + FederatedAuthCacheHolder.tokenCache.put(key, expiredTokenInfo); spyPlugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); verify(mockIamTokenUtils).generateAuthenticationToken(mockAwsCredentialsProvider, @@ -173,7 +173,7 @@ void testSpecifiedIamHostPortRegion() throws SQLException { props.setProperty(FederatedAuthPlugin.IAM_REGION.name, expectedRegion.toString()); final String key = "us-west-2:pg.testdb.us-west-2.rds.amazonaws.com:" + expectedPort + ":iamUser"; - FederatedAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + FederatedAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); FederatedAuthPlugin plugin = new FederatedAuthPlugin(mockPluginService, mockCredentialsProviderFactory, mockRdsUtils, mockIamTokenUtils); @@ -195,7 +195,7 @@ void testIdpCredentialsFallback() throws SQLException { new FederatedAuthPlugin(mockPluginService, mockCredentialsProviderFactory, mockRdsUtils, mockIamTokenUtils); String key = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + DEFAULT_PORT + ":iamUser"; - FederatedAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + FederatedAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); plugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPluginTest.java index 6b529552f..910e06fe1 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/federatedauth/OktaAuthPluginTest.java @@ -112,7 +112,7 @@ void testCachedToken() throws SQLException { new OktaAuthPlugin(mockPluginService, mockCredentialsProviderFactory, mockRdsUtils, mockIamTokenUtils); String key = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + DEFAULT_PORT + ":iamUser"; - OktaAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + OktaAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); plugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); @@ -129,7 +129,7 @@ void testExpiredCachedToken() throws SQLException { final String someExpiredToken = "someExpiredToken"; final TokenInfo expiredTokenInfo = new TokenInfo( someExpiredToken, Instant.now().minusMillis(300000)); - OktaAuthPlugin.tokenCache.put(key, expiredTokenInfo); + OktaAuthCacheHolder.tokenCache.put(key, expiredTokenInfo); spyPlugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); verify(mockIamTokenUtils).generateAuthenticationToken(mockAwsCredentialsProvider, @@ -168,7 +168,7 @@ void testSpecifiedIamHostPortRegion() throws SQLException { props.setProperty(OktaAuthPlugin.IAM_REGION.name, expectedRegion.toString()); final String key = "us-west-2:pg.testdb.us-west-2.rds.amazonaws.com:" + expectedPort + ":iamUser"; - OktaAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + OktaAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); OktaAuthPlugin plugin = new OktaAuthPlugin(mockPluginService, mockCredentialsProviderFactory, mockRdsUtils, mockIamTokenUtils); @@ -190,7 +190,7 @@ void testIdpCredentialsFallback() throws SQLException { new OktaAuthPlugin(mockPluginService, mockCredentialsProviderFactory, mockRdsUtils, mockIamTokenUtils); final String key = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + DEFAULT_PORT + ":iamUser"; - OktaAuthPlugin.tokenCache.put(key, TEST_TOKEN_INFO); + OktaAuthCacheHolder.tokenCache.put(key, TEST_TOKEN_INFO); plugin.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda); diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java index a1b5fdab9..92a0c0ff6 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java @@ -136,7 +136,7 @@ public static void registerDrivers() throws SQLException { @Test public void testPostgresConnectValidTokenInCache() throws SQLException { - IamAuthConnectionPlugin.tokenCache.put(PG_CACHE_KEY, + IamAuthCacheHolder.tokenCache.put(PG_CACHE_KEY, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); when(mockDialect.getDefaultPort()).thenReturn(DEFAULT_PG_PORT); @@ -148,7 +148,7 @@ public void testPostgresConnectValidTokenInCache() throws SQLException { public void testMySqlConnectValidTokenInCache() throws SQLException { props.setProperty(PropertyDefinition.USER.name, "mysqlUser"); props.setProperty(PropertyDefinition.PASSWORD.name, "mysqlPassword"); - IamAuthConnectionPlugin.tokenCache.put(MYSQL_CACHE_KEY, + IamAuthCacheHolder.tokenCache.put(MYSQL_CACHE_KEY, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); when(mockDialect.getDefaultPort()).thenReturn(DEFAULT_MYSQL_PORT); @@ -163,7 +163,7 @@ public void testPostgresConnectWithInvalidPortFallbacksToHostPort() throws SQLEx final String cacheKeyWithNewPort = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + PG_HOST_SPEC_WITH_PORT.getPort() + ":postgresqlUser"; - IamAuthConnectionPlugin.tokenCache.put(cacheKeyWithNewPort, + IamAuthCacheHolder.tokenCache.put(cacheKeyWithNewPort, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); testTokenSetInProps(PG_DRIVER_PROTOCOL, PG_HOST_SPEC_WITH_PORT); @@ -178,7 +178,7 @@ public void testPostgresConnectWithInvalidPortAndNoHostPortFallbacksToHostPort() final String cacheKeyWithNewPort = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + DEFAULT_PG_PORT + ":postgresqlUser"; - IamAuthConnectionPlugin.tokenCache.put(cacheKeyWithNewPort, + IamAuthCacheHolder.tokenCache.put(cacheKeyWithNewPort, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); testTokenSetInProps(PG_DRIVER_PROTOCOL, PG_HOST_SPEC); @@ -186,7 +186,7 @@ public void testPostgresConnectWithInvalidPortAndNoHostPortFallbacksToHostPort() @Test public void testConnectExpiredTokenInCache() throws SQLException { - IamAuthConnectionPlugin.tokenCache.put(PG_CACHE_KEY, + IamAuthCacheHolder.tokenCache.put(PG_CACHE_KEY, new TokenInfo(TEST_TOKEN, Instant.now().minusMillis(300000))); when(mockDialect.getDefaultPort()).thenReturn(DEFAULT_PG_PORT); @@ -204,7 +204,7 @@ public void testConnectEmptyCache() throws SQLException { @Test public void testConnectWithSpecifiedPort() throws SQLException { final String cacheKeyWithNewPort = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:1234:" + "postgresqlUser"; - IamAuthConnectionPlugin.tokenCache.put(cacheKeyWithNewPort, + IamAuthCacheHolder.tokenCache.put(cacheKeyWithNewPort, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); testTokenSetInProps(PG_DRIVER_PROTOCOL, PG_HOST_SPEC_WITH_PORT); @@ -216,7 +216,7 @@ public void testConnectWithSpecifiedIamDefaultPort() throws SQLException { props.setProperty(IamAuthConnectionPlugin.IAM_DEFAULT_PORT.name, iamDefaultPort); final String cacheKeyWithNewPort = "us-east-2:pg.testdb.us-east-2.rds.amazonaws.com:" + iamDefaultPort + ":postgresqlUser"; - IamAuthConnectionPlugin.tokenCache.put(cacheKeyWithNewPort, + IamAuthCacheHolder.tokenCache.put(cacheKeyWithNewPort, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); testTokenSetInProps(PG_DRIVER_PROTOCOL, PG_HOST_SPEC_WITH_PORT); @@ -227,7 +227,7 @@ public void testConnectWithSpecifiedRegion() throws SQLException { final String cacheKeyWithNewRegion = "us-west-1:pg.testdb.us-west-1.rds.amazonaws.com:" + DEFAULT_PG_PORT + ":" + "postgresqlUser"; props.setProperty(IamAuthConnectionPlugin.IAM_REGION.name, "us-west-1"); - IamAuthConnectionPlugin.tokenCache.put(cacheKeyWithNewRegion, + IamAuthCacheHolder.tokenCache.put(cacheKeyWithNewRegion, new TokenInfo(TEST_TOKEN, Instant.now().plusMillis(300000))); when(mockDialect.getDefaultPort()).thenReturn(DEFAULT_PG_PORT); @@ -294,6 +294,6 @@ private void testGenerateToken( verify(mockLambda, times(1)).call(); assertEquals(GENERATED_TOKEN, PropertyDefinition.PASSWORD.getString(props)); - assertEquals(GENERATED_TOKEN, IamAuthConnectionPlugin.tokenCache.get(PG_CACHE_KEY).getToken()); + assertEquals(GENERATED_TOKEN, IamAuthCacheHolder.tokenCache.get(PG_CACHE_KEY).getToken()); } }