Relocate custom handlers

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<ConnectionProvider> 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<HostSpec> 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<HostSpec> 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 {

wrapper/src/main/java/software/amazon/jdbc/Driver.java

@@ -25,15 +25,19 @@
 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;
 import java.util.logging.Level;
 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> resetSessionStateOnCloseCallable =
+      new AtomicReference<>(null);
+  private static final AtomicReference<TransferSessionStateOnSwitchCallable> transferSessionStateOnSwitchCallable =
+      new AtomicReference<>(null);
+
+  private static final AtomicReference<ExceptionHandler> customExceptionHandler =
+      new AtomicReference<>(null);
+
+  private static final AtomicReference<Dialect> customDialect =
+      new AtomicReference<>(null);
+
+  private static final AtomicReference<TargetDriverDialect> customTargetDriverDialect =
+      new AtomicReference<>(null);
+
+  private static final AtomicReference<ConnectionProvider> customConnectionProvider =
+      new AtomicReference<>(null);
+
+  private static final AtomicReference<ConnectionInitFunc> 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<String, String> 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();

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;

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();
   }
 }