diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java index bf18efb9b..3a327b19b 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java @@ -20,20 +20,15 @@ import java.sql.SQLException; import java.util.List; import java.util.Properties; -import java.util.concurrent.atomic.AtomicReference; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import software.amazon.jdbc.cleanup.CanReleaseResources; public class ConnectionProviderManager { - private static AtomicReference customConnectionProvider = new AtomicReference<>(null); - private final ConnectionProvider defaultProvider; private final @Nullable ConnectionProvider effectiveConnProvider; - private static ConnectionInitFunc connectionInitFunc = null; - /** * {@link ConnectionProviderManager} constructor. * @@ -57,9 +52,11 @@ public ConnectionProviderManager( * {@link ConnectionProvider#acceptsUrl} for more info. * * @param connProvider the {@link ConnectionProvider} to use to establish new connections + * @deprecated Use software.amazon.jdbc.Driver instead */ + @Deprecated public static void setConnectionProvider(ConnectionProvider connProvider) { - customConnectionProvider.set(connProvider); + Driver.setCustomConnectionProvider(connProvider); } /** @@ -78,9 +75,9 @@ public static void setConnectionProvider(ConnectionProvider connProvider) { public ConnectionProvider getConnectionProvider( String driverProtocol, HostSpec host, Properties props) { - final ConnectionProvider tmpCustomConnectionProvider = customConnectionProvider.get(); - if (tmpCustomConnectionProvider != null && tmpCustomConnectionProvider.acceptsUrl(driverProtocol, host, props)) { - return tmpCustomConnectionProvider; + final ConnectionProvider customConnectionProvider = Driver.getCustomConnectionProvider(); + if (customConnectionProvider != null && customConnectionProvider.acceptsUrl(driverProtocol, host, props)) { + return customConnectionProvider; } if (this.effectiveConnProvider != null && this.effectiveConnProvider.acceptsUrl(driverProtocol, host, props)) { @@ -110,8 +107,8 @@ public ConnectionProvider getDefaultProvider() { * return false */ public boolean acceptsStrategy(HostRole role, String strategy) { - final ConnectionProvider tmpCustomConnectionProvider = customConnectionProvider.get(); - if (tmpCustomConnectionProvider != null && tmpCustomConnectionProvider.acceptsStrategy(role, strategy)) { + final ConnectionProvider customConnectionProvider = Driver.getCustomConnectionProvider(); + if (customConnectionProvider != null && customConnectionProvider.acceptsStrategy(role, strategy)) { return true; } @@ -143,10 +140,10 @@ public boolean acceptsStrategy(HostRole role, String strategy) { public HostSpec getHostSpecByStrategy(List hosts, HostRole role, String strategy, Properties props) throws SQLException, UnsupportedOperationException { HostSpec host = null; - final ConnectionProvider tmpCustomConnectionProvider = customConnectionProvider.get(); + final ConnectionProvider customConnectionProvider = Driver.getCustomConnectionProvider(); try { - if (tmpCustomConnectionProvider != null && tmpCustomConnectionProvider.acceptsStrategy(role, strategy)) { - host = tmpCustomConnectionProvider.getHostSpecByStrategy(hosts, role, strategy, props); + if (customConnectionProvider != null && customConnectionProvider.acceptsStrategy(role, strategy)) { + host = customConnectionProvider.getHostSpecByStrategy(hosts, role, strategy, props); } } catch (UnsupportedOperationException e) { // The custom provider does not support the provided strategy, ignore it and try with the other providers. @@ -170,27 +167,43 @@ public HostSpec getHostSpecByStrategy(List hosts, HostRole role, Strin * Clears the non-default {@link ConnectionProvider} if it has been set. The default * ConnectionProvider will be used if the non-default ConnectionProvider has not been set or has * been cleared. + * + * @deprecated Use software.amazon.jdbc.Driver instead */ + @Deprecated public static void resetProvider() { - customConnectionProvider.set(null); + Driver.resetCustomConnectionProvider(); } /** * Releases any resources held by the available {@link ConnectionProvider} instances. */ public static void releaseResources() { - final ConnectionProvider tmpCustomConnectionProvider = customConnectionProvider.get(); - if (tmpCustomConnectionProvider instanceof CanReleaseResources) { - ((CanReleaseResources) tmpCustomConnectionProvider).releaseResources(); + final ConnectionProvider customConnectionProvider = Driver.getCustomConnectionProvider(); + if (customConnectionProvider instanceof CanReleaseResources) { + ((CanReleaseResources) customConnectionProvider).releaseResources(); } } + /** + * Sets a custom connection initialization function. It'll be used + * for every brand-new database connection. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void setConnectionInitFunc(final @NonNull ConnectionInitFunc func) { - connectionInitFunc = func; + Driver.setConnectionInitFunc(func); } + /** + * Resets a custom connection initialization function. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void resetConnectionInitFunc() { - connectionInitFunc = null; + Driver.resetConnectionInitFunc(); } public void initConnection( @@ -199,12 +212,12 @@ public void initConnection( final @NonNull HostSpec hostSpec, final @NonNull Properties props) throws SQLException { - final ConnectionInitFunc copy = connectionInitFunc; - if (copy == null) { + final ConnectionInitFunc connectionInitFunc = Driver.getConnectionInitFunc(); + if (connectionInitFunc == null) { return; } - copy.initConnection(connection, protocol, hostSpec, props); + connectionInitFunc.initConnection(connection, protocol, hostSpec, props); } public interface ConnectionInitFunc { diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java index ed4f23ea6..5c44dae6d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/Driver.java +++ b/wrapper/src/main/java/software/amazon/jdbc/Driver.java @@ -25,6 +25,7 @@ import java.util.Collections; import java.util.List; import java.util.Properties; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import java.util.logging.ConsoleHandler; import java.util.logging.Handler; @@ -32,8 +33,11 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; +import software.amazon.jdbc.ConnectionProviderManager.ConnectionInitFunc; import software.amazon.jdbc.authentication.AwsCredentialsManager; +import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.dialect.DialectManager; +import software.amazon.jdbc.exceptions.ExceptionHandler; import software.amazon.jdbc.hostlistprovider.RdsHostListProvider; import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider; import software.amazon.jdbc.plugin.AwsSecretsManagerCacheHolder; @@ -73,8 +77,25 @@ public class Driver implements java.sql.Driver { private static final Logger LOGGER = Logger.getLogger("software.amazon.jdbc.Driver"); private static @Nullable Driver registeredDriver; - private static ResetSessionStateOnCloseCallable resetSessionStateOnCloseCallable = null; - private static TransferSessionStateOnSwitchCallable transferSessionStateOnSwitchCallable = null; + private static final AtomicReference resetSessionStateOnCloseCallable = + new AtomicReference<>(null); + private static final AtomicReference transferSessionStateOnSwitchCallable = + new AtomicReference<>(null); + + private static final AtomicReference customExceptionHandler = + new AtomicReference<>(null); + + private static final AtomicReference customDialect = + new AtomicReference<>(null); + + private static final AtomicReference customTargetDriverDialect = + new AtomicReference<>(null); + + private static final AtomicReference customConnectionProvider = + new AtomicReference<>(null); + + private static final AtomicReference connectionInitFunc = + new AtomicReference<>(null); static { try { @@ -262,27 +283,27 @@ public Logger getParentLogger() throws SQLFeatureNotSupportedException { } public static void setResetSessionStateOnCloseFunc(final @NonNull ResetSessionStateOnCloseCallable func) { - resetSessionStateOnCloseCallable = func; + resetSessionStateOnCloseCallable.set(func); } public static void resetResetSessionStateOnCloseFunc() { - resetSessionStateOnCloseCallable = null; + resetSessionStateOnCloseCallable.set(null); } public static ResetSessionStateOnCloseCallable getResetSessionStateOnCloseFunc() { - return resetSessionStateOnCloseCallable; + return resetSessionStateOnCloseCallable.get(); } public static void setTransferSessionStateOnSwitchFunc(final @NonNull TransferSessionStateOnSwitchCallable func) { - transferSessionStateOnSwitchCallable = func; + transferSessionStateOnSwitchCallable.set(func); } public static void resetTransferSessionStateOnSwitchFunc() { - transferSessionStateOnSwitchCallable = null; + transferSessionStateOnSwitchCallable.set(null); } public static TransferSessionStateOnSwitchCallable getTransferSessionStateOnSwitchFunc() { - return transferSessionStateOnSwitchCallable; + return transferSessionStateOnSwitchCallable.get(); } public static void setPrepareHostFunc(final Function func) { @@ -293,6 +314,79 @@ public static void resetPrepareHostFunc() { RdsUtils.resetPrepareHostFunc(); } + public static void setCustomExceptionHandler(final ExceptionHandler exceptionHandler) { + customExceptionHandler.set(exceptionHandler); + } + + public static ExceptionHandler getCustomExceptionHandler() { + return customExceptionHandler.get(); + } + + public static void resetCustomExceptionHandler() { + customExceptionHandler.set(null); + } + + public static void setCustomDialect(final @NonNull Dialect dialect) { + customDialect.set(dialect); + } + + public static Dialect getCustomDialect() { + return customDialect.get(); + } + + public static void resetCustomDialect() { + customDialect.set(null); + } + + public static void setCustomTargetDriverDialect(final @NonNull TargetDriverDialect targetDriverDialect) { + customTargetDriverDialect.set(targetDriverDialect); + } + + public static TargetDriverDialect getCustomTargetDriverDialect() { + return customTargetDriverDialect.get(); + } + + public static void resetCustomTargetDriverDialect() { + customTargetDriverDialect.set(null); + } + + /** + * Setter that can optionally be called to request a non-default {@link ConnectionProvider}. The + * requested ConnectionProvider will be used to establish future connections unless it does not + * support a requested URL, in which case the default ConnectionProvider will be used. See + * {@link ConnectionProvider#acceptsUrl} for more info. + * + * @param connProvider the {@link ConnectionProvider} to use to establish new connections + */ + public static void setCustomConnectionProvider(ConnectionProvider connProvider) { + customConnectionProvider.set(connProvider); + } + + public static ConnectionProvider getCustomConnectionProvider() { + return customConnectionProvider.get(); + } + + /** + * Clears the non-default {@link ConnectionProvider} if it has been set. The default + * ConnectionProvider will be used if the non-default ConnectionProvider has not been set or has + * been cleared. + */ + public static void resetCustomConnectionProvider() { + customConnectionProvider.set(null); + } + + public static void setConnectionInitFunc(final @NonNull ConnectionInitFunc func) { + connectionInitFunc.set(func); + } + + public static ConnectionInitFunc getConnectionInitFunc() { + return connectionInitFunc.get(); + } + + public static void resetConnectionInitFunc() { + connectionInitFunc.set(null); + } + public static void clearCaches() { RdsUtils.clearCache(); RdsHostListProvider.clearAll(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java b/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java index 16440362f..90c031be3 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java @@ -26,6 +26,7 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.AwsWrapperProperty; +import software.amazon.jdbc.Driver; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.util.CacheMap; @@ -40,8 +41,6 @@ public class DialectManager implements DialectProvider { private static final Logger LOGGER = Logger.getLogger(DialectManager.class.getName()); - protected static Dialect customDialect; - public static final AwsWrapperProperty DIALECT = new AwsWrapperProperty( "wrapperDialect", "", "A unique identifier for the supported database dialect."); @@ -88,12 +87,24 @@ public DialectManager(PluginService pluginService) { this.pluginService = pluginService; } + /** + * Sets a custom dialect handler. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void setCustomDialect(final @NonNull Dialect dialect) { - customDialect = dialect; + Driver.setCustomDialect(dialect); } + /** + * Resets a custom dialect handler. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void resetCustomDialect() { - customDialect = null; + Driver.resetCustomDialect(); } public static void resetEndpointCache() { @@ -110,6 +121,7 @@ public Dialect getDialect( this.canUpdate = false; this.dialect = null; + final Dialect customDialect = Driver.getCustomDialect(); if (customDialect != null) { this.dialectCode = DialectCodes.CUSTOM; this.dialect = customDialect; diff --git a/wrapper/src/main/java/software/amazon/jdbc/exceptions/ExceptionManager.java b/wrapper/src/main/java/software/amazon/jdbc/exceptions/ExceptionManager.java index 16f66a8e7..a6badc629 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/exceptions/ExceptionManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/exceptions/ExceptionManager.java @@ -16,18 +16,29 @@ package software.amazon.jdbc.exceptions; +import software.amazon.jdbc.Driver; import software.amazon.jdbc.dialect.Dialect; public class ExceptionManager { - protected static ExceptionHandler customHandler; - + /** + * Sets a custom exception handler. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void setCustomHandler(final ExceptionHandler exceptionHandler) { - customHandler = exceptionHandler; + Driver.setCustomExceptionHandler(exceptionHandler); } + /** + * Resets a custom exception handler. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void resetCustomHandler() { - customHandler = null; + Driver.resetCustomExceptionHandler(); } public boolean isLoginException(final Dialect dialect, final Throwable throwable) { @@ -51,6 +62,7 @@ public boolean isNetworkException(final Dialect dialect, final String sqlState) } private ExceptionHandler getHandler(final Dialect dialect) { + final ExceptionHandler customHandler = Driver.getCustomExceptionHandler(); return customHandler != null ? customHandler : dialect.getExceptionHandler(); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/TargetDriverDialectManager.java b/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/TargetDriverDialectManager.java index 63501bff7..76cbf19eb 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/TargetDriverDialectManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/TargetDriverDialectManager.java @@ -34,8 +34,6 @@ public class TargetDriverDialectManager implements TargetDriverDialectProvider { private static final Logger LOGGER = Logger.getLogger(TargetDriverDialectManager.class.getName()); - protected static TargetDriverDialect customDialect; - public static final AwsWrapperProperty TARGET_DRIVER_DIALECT = new AwsWrapperProperty( "wrapperTargetDriverDialect", "", "A unique identifier for the target driver dialect."); @@ -72,12 +70,24 @@ public class TargetDriverDialectManager implements TargetDriverDialectProvider { PropertyDefinition.registerPluginProperties(TargetDriverDialectManager.class); } + /** + * Sets a custom target driver dialect handler. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void setCustomDialect(final @NonNull TargetDriverDialect targetDriverDialect) { - customDialect = targetDriverDialect; + software.amazon.jdbc.Driver.setCustomTargetDriverDialect(targetDriverDialect); } + /** + * Resets a custom target driver dialect. + * + * @deprecated Use software.amazon.jdbc.Driver instead + */ + @Deprecated public static void resetCustomDialect() { - customDialect = null; + software.amazon.jdbc.Driver.resetCustomTargetDriverDialect(); } @Override @@ -100,6 +110,7 @@ private TargetDriverDialect getDialect( final @NonNull Properties props, Function checkFunc) throws SQLException { + final TargetDriverDialect customDialect = software.amazon.jdbc.Driver.getCustomTargetDriverDialect(); if (customDialect != null) { if (checkFunc.apply(customDialect)) { this.logDialect("custom", customDialect);