Allow optional failover in the connect pipeline

CHANGELOG.md

@@ -3,6 +3,9 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/#semantic-versioning-200).
 
+### :magic_wand: Added
+- Logic and a connection property to enable driver failover when network exceptions occur in the connect pipeline (PR #1099)[https://github.com/aws/aws-advanced-jdbc-wrapper/pull/1099]
+
 ## [2.3.9] - 2024-08-09
 
 ### :bug: Fixed

docs/using-the-jdbc-driver/using-plugins/UsingTheFailoverPlugin.md

@@ -30,6 +30,7 @@ In addition to the parameters that you can configure for the underlying driver,
 | `failoverReaderConnectTimeoutMs`       | Integer |                                        No                                        | Maximum allowed time in milliseconds to attempt to connect to a reader instance during a reader failover process.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    | `30000`                                                                                                                                                                                             |
 | `failoverTimeoutMs`                    | Integer |                                        No                                        | Maximum allowed time in milliseconds to attempt reconnecting to a new writer or reader instance after a cluster failover is initiated.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               | `300000`                                                                                                                                                                                            |
 | `failoverWriterReconnectIntervalMs`    | Integer |                                        No                                        | Interval of time in milliseconds to wait between attempts to reconnect to a failed writer during a writer failover process.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          | `2000`                                                                                                                                                                                              |
+| `enableConnectFailover`                | Boolean |                                        No                                        | Enables/disables cluster-aware failover if the initial connection to the database fails due to a network exception. Note that this may result in a connection to a different instance in the cluster than was specified by the URL.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  | `false`                                                                                                                                                                                             |
 | ~~`keepSessionStateOnFailover`~~       | Boolean |                                        No                                        | This parameter is no longer available. If specified, it will be ignored by the driver. See [Session State](../SessionState.md) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     | `false`                                                                                                                                                                                             |
 | ~~`enableFailoverStrictReader`~~       | Boolean |                                        No                                        | This parameter is no longer available and, if specified, it will be ignored by the driver. See `failoverMode` (`reader-or-writer` or `strict-reader`) for more details.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |                                                                                                                                                                                                     |
 

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

@@ -0,0 +1,35 @@
+/*
+ * 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 java.util.Properties;
+
+@FunctionalInterface
+public interface AcceptsUrlFunc {
+
+  /**
+   * This function can be passed to a {@link ConnectionProvider} constructor to specify when the
+   * {@link ConnectionProvider} should be used to open a connection to the given {@link HostSpec} with the
+   * given {@link Properties}.
+   *
+   * @param hostSpec      the host details for the requested connection
+   * @param props         the properties for the requested connection
+   * @return a boolean indicating whether a {@link ConnectionProvider} should be used to open a connection to the given
+   *     {@link HostSpec} with the given {@link Properties}.
+   */
+  boolean acceptsUrl(HostSpec hostSpec, Properties props);
+}

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

@@ -69,6 +69,7 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
   private static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);
   private final HikariPoolConfigurator poolConfigurator;
   private final HikariPoolMapping poolMapping;
+  private final AcceptsUrlFunc acceptsUrlFunc;
   private final LeastConnectionsHostSelector leastConnectionsHostSelector;
 
   /**
@@ -112,6 +113,7 @@ public HikariPooledConnectionProvider(
       HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
+    this.acceptsUrlFunc = null;
     this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
   }
 
@@ -147,14 +149,62 @@ public HikariPooledConnectionProvider(
       long poolCleanupNanos) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
+    this.acceptsUrlFunc = null;
     poolExpirationCheckNanos = poolExpirationNanos;
     databasePools.setCleanupIntervalNanos(poolCleanupNanos);
     this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
   }
 
+  /**
+   * {@link HikariPooledConnectionProvider} constructor. This class can be passed to
+   * {@link ConnectionProviderManager#setConnectionProvider} to enable internal connection pools for
+   * each database instance in a cluster. By maintaining internal connection pools, the driver can
+   * improve performance by reusing old {@link Connection} objects.
+   *
+   * @param hikariPoolConfigurator a function that returns a {@link HikariConfig} with specific
+   *                               Hikari configurations. By default, the
+   *                               {@link HikariPooledConnectionProvider} will configure the
+   *                               jdbcUrl, exceptionOverrideClassName, username, and password. Any
+   *                               additional configuration should be defined by passing in this
+   *                               parameter. If no additional configuration is desired, pass in a
+   *                               {@link HikariPoolConfigurator} that returns an empty
+   *                               HikariConfig.
+   * @param mapping                a function that returns a String key used for the internal
+   *                               connection pool keys. An internal connection pool will be
+   *                               generated for each unique key returned by this function.
+   * @param acceptsUrlFunc         a function that defines when an internal connection pool should be created for a
+   *                               requested connection. An internal connection pool will be created when the connect
+   *                               pipeline is being executed and this function returns <code>true</code>.
+   * @param poolExpirationNanos    the amount of time that a pool should sit in the cache before
+   *                               being marked as expired for cleanup, in nanoseconds. Expired
+   *                               pools can still be used and will not be closed unless there
+   *                               are no active connections.
+   * @param poolCleanupNanos       the interval defining how often expired connection pools
+   *                               should be cleaned up, in nanoseconds. Note that expired pools
+   *                               will not be closed unless there are no active connections.
+   */
+  public HikariPooledConnectionProvider(
+      HikariPoolConfigurator hikariPoolConfigurator,
+      HikariPoolMapping mapping,
+      AcceptsUrlFunc acceptsUrlFunc,
+      long poolExpirationNanos,
+      long poolCleanupNanos) {
+    this.poolConfigurator = hikariPoolConfigurator;
+    this.poolMapping = mapping;
+    this.acceptsUrlFunc = acceptsUrlFunc;
+    poolExpirationCheckNanos = poolExpirationNanos;
+    databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
+  }
+
+
   @Override
   public boolean acceptsUrl(
       @NonNull String protocol, @NonNull HostSpec hostSpec, @NonNull Properties props) {
+    if (this.acceptsUrlFunc != null) {
+      return this.acceptsUrlFunc.acceptsUrl(hostSpec, props);
+    }
+
     final RdsUrlType urlType = rdsUtils.identifyRdsType(hostSpec.getHost());
     return RdsUrlType.RDS_INSTANCE.equals(urlType);
   }

wrapper/src/main/java/software/amazon/jdbc/exceptions/MySQLExceptionHandler.java

@@ -21,14 +21,29 @@
 
 public class MySQLExceptionHandler implements ExceptionHandler {
   public static final String SQLSTATE_ACCESS_ERROR = "28000";
+  public static final String SQLSTATE_SYNTAX_ERROR_OR_ACCESS_VIOLATION = "42000";
+  public static final String SET_NETWORK_TIMEOUT_ON_CLOSED_CONNECTION =
+      "setNetworkTimeout cannot be called on a closed connection";
 
   @Override
   public boolean isNetworkException(final Throwable throwable) {
     Throwable exception = throwable;
 
     while (exception != null) {
       if (exception instanceof SQLException) {
-        return isNetworkException(((SQLException) exception).getSQLState());
+        SQLException sqlException = (SQLException) exception;
+
+        // Hikari throws a network exception with SQL state 42000 if all the following points are true:
+        // - HikariDataSource#getConnection is called and the cached connection that was grabbed is broken due to server
+        // failover.
+        // - the MariaDB driver is being used (the underlying driver determines the SQL state of the Hikari exception).
+        //
+        // The check for the Hikari MariaDB exception is added here because the exception handler is determined by the
+        // database dialect. Consequently, this exception handler is used when using the MariaDB driver against a MySQL
+        // database engine.
+        if (isNetworkException(sqlException.getSQLState()) || isHikariMariaDbNetworkException(sqlException)) {
+          return true;
+        }
       } else if (exception instanceof CJException) {
         return isNetworkException(((CJException) exception).getSQLState());
       }
@@ -76,4 +91,9 @@ public boolean isLoginException(final String sqlState) {
 
     return SQLSTATE_ACCESS_ERROR.equals(sqlState);
   }
+
+  private boolean isHikariMariaDbNetworkException(final SQLException sqlException) {
+    return sqlException.getSQLState().equals(SQLSTATE_SYNTAX_ERROR_OR_ACCESS_VIOLATION)
+        && sqlException.getMessage().contains(SET_NETWORK_TIMEOUT_ON_CLOSED_CONNECTION);
+  }
 }

wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java

@@ -84,6 +84,7 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin {
   private final PluginService pluginService;
   protected final Properties properties;
   protected boolean enableFailoverSetting;
+  protected boolean enableConnectFailover;
   protected int failoverTimeoutMsSetting;
   protected int failoverClusterTopologyRefreshRateMsSetting;
   protected int failoverWriterReconnectIntervalMsSetting;
@@ -137,6 +138,13 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin {
           "enableClusterAwareFailover", "true",
           "Enable/disable cluster-aware failover logic");
 
+  public static final AwsWrapperProperty ENABLE_CONNECT_FAILOVER =
+      new AwsWrapperProperty(
+          "enableConnectFailover", "false",
+          "Enable/disable cluster-aware failover if the initial connection to the database fails due to a "
+              + "network exception. Note that this may result in a connection to a different instance in the cluster "
+              + "than was specified by the URL.");
+
   public static final AwsWrapperProperty FAILOVER_MODE =
       new AwsWrapperProperty(
           "failoverMode", null,
@@ -353,6 +361,7 @@ public boolean isFailoverEnabled() {
 
   private void initSettings() {
     this.enableFailoverSetting = ENABLE_CLUSTER_AWARE_FAILOVER.getBoolean(this.properties);
+    this.enableConnectFailover = ENABLE_CONNECT_FAILOVER.getBoolean(this.properties);
     this.failoverTimeoutMsSetting = FAILOVER_TIMEOUT_MS.getInteger(this.properties);
     this.failoverClusterTopologyRefreshRateMsSetting =
         FAILOVER_CLUSTER_TOPOLOGY_REFRESH_RATE_MS.getInteger(this.properties);
@@ -767,15 +776,34 @@ public Connection connect(
       final boolean isInitialConnection,
       final JdbcCallable<Connection, SQLException> connectFunc)
       throws SQLException {
-    return connectInternal(driverProtocol, hostSpec, props, isInitialConnection, connectFunc);
+    return connectInternal(driverProtocol, hostSpec, props, isInitialConnection, connectFunc, false);
   }
 
   private Connection connectInternal(String driverProtocol, HostSpec hostSpec, Properties props,
-      boolean isInitialConnection, JdbcCallable<Connection, SQLException> connectFunc)
+      boolean isInitialConnection, JdbcCallable<Connection, SQLException> connectFunc, boolean isForceConnect)
       throws SQLException {
-    final Connection conn =
-        this.staleDnsHelper.getVerifiedConnection(isInitialConnection, this.hostListProviderService,
-            driverProtocol, hostSpec, props, connectFunc);
+
+    Connection conn = null;
+    try {
+      conn =
+          this.staleDnsHelper.getVerifiedConnection(isInitialConnection, this.hostListProviderService,
+              driverProtocol, hostSpec, props, connectFunc);
+    } catch (final SQLException e) {
+      if (!this.enableConnectFailover || isForceConnect || !shouldExceptionTriggerConnectionSwitch(e)) {
+        throw e;
+      }
+
+      try {
+        failover(this.pluginService.getCurrentHostSpec());
+      } catch (FailoverSuccessSQLException failoverSuccessException) {
+        conn = this.pluginService.getCurrentConnection();
+      }
+    }
+
+    if (conn == null) {
+      // This should be unreachable, the above logic will either get a connection successfully or throw an exception.
+      throw new SQLException(Messages.get("Failover.unableToConnect"));
+    }
 
     if (isInitialConnection) {
       this.pluginService.refreshHostList(conn);
@@ -792,6 +820,6 @@ public Connection forceConnect(
       final boolean isInitialConnection,
       final JdbcCallable<Connection, SQLException> forceConnectFunc)
       throws SQLException {
-    return connectInternal(driverProtocol, hostSpec, props, isInitialConnection, forceConnectFunc);
+    return connectInternal(driverProtocol, hostSpec, props, isInitialConnection, forceConnectFunc, true);
   }
 }