transfer session state during failover (aws#814)

Bit-Quill · Jan 3, 2024 · 774aaee · 774aaee
1 parent 7e1d15e
commit 774aaee
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 20 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java
@@ -29,8 +29,6 @@
 import java.util.logging.Handler;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
-import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.profile.ConfigurationProfile;
 import software.amazon.jdbc.profile.DriverConfigurationProfiles;
@@ -112,7 +110,7 @@ public Connection connect(final String url, final Properties info) throws SQLExc
       PropertyDefinition.DATABASE.set(props, databaseName);
     }
 
-    LOGGER.finest(() -> PropertyUtils.logProperties(props, "DEBUG: Connecting with properties: \n"));
+    LOGGER.finest(() -> PropertyUtils.logProperties(props, "Connecting with properties: \n"));
 
     final String profileName = PropertyDefinition.PROFILE_NAME.getString(props);
     ConfigurationProfile configurationProfile = null;

diff --git a/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java
@@ -113,8 +113,7 @@ public Connection connect(
       final @NonNull Properties props)
       throws SQLException {
 
-    LOGGER.finest(() -> "DEBUG: Connecting with properties:\n"
-        + PropertyUtils.logProperties(props, ""));
+    LOGGER.finest(() -> PropertyUtils.logProperties(props, "Connecting with properties: \n"));
 
     final Properties copy = PropertyUtils.copyProperties(props);
     dialect.prepareConnectProperties(copy, protocol, hostSpec);

diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java
@@ -482,13 +482,17 @@ private boolean canUpdateTopology(final String methodName) {
   /**
    * Connects this dynamic failover connection proxy to the host pointed out by the given host
    * index.
+   * <p></p>
+   * The method assumes that current connection is not setup. If it's not true, a session state
+   * transfer from the current connection to a new one may be necessary. This should be handled by callee.
    *
    * @param host The host.
    * @throws SQLException if an error occurs
    */
   private void connectTo(final HostSpec host) throws SQLException {
     try {
-      switchCurrentConnectionTo(host, createConnectionForHost(host));
+      this.pluginService.setCurrentConnection(createConnectionForHost(host), host);
+
       LOGGER.fine(
           () -> Messages.get(
               "Failover.establishedConnection",
@@ -536,26 +540,20 @@ private boolean shouldAttemptReaderConnection() {
   }
 
   /**
-   * Replaces the previous underlying connection by the connection given. State from previous
-   * connection, if any, is synchronized with the new one.
+   * Synchronizes state from current connection, if any, is the new one.
    *
-   * @param host       The host that matches the given connection.
-   * @param connection The connection instance to switch to.
+   * @param newHostSpec   The host that matches the given connection.
+   * @param newConnection The connection instance to switch to.
    * @throws SQLException if an error occurs
    */
-  private void switchCurrentConnectionTo(final HostSpec host, final Connection connection) throws SQLException {
+  private void transferSessionStateToNewConnection(final HostSpec newHostSpec, final Connection newConnection)
+      throws SQLException {
+
     Connection currentConnection = this.pluginService.getCurrentConnection();
     HostSpec currentHostSpec = this.pluginService.getCurrentHostSpec();
 
-    if (currentConnection != connection) {
-      transferSessionState(currentConnection, currentHostSpec, connection, host);
-      invalidateCurrentConnection();
-    }
-
-    this.pluginService.setCurrentConnection(connection, host);
-
-    if (this.pluginManagerService != null) {
-      this.pluginManagerService.setInTransaction(false);
+    if (currentConnection != newConnection) {
+      transferSessionState(currentConnection, currentHostSpec, newConnection, newHostSpec);
     }
   }
 
@@ -744,7 +742,9 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException
 
       if (keepSessionStateOnFailover) {
         restoreSessionState(result.getConnection());
+        transferSessionStateToNewConnection(result.getHost(), result.getConnection());
       }
+      this.invalidateCurrentConnection();
       this.pluginService.setCurrentConnection(result.getConnection(), result.getHost());
 
       this.pluginService.getCurrentHostSpec().removeAlias(oldAliases.toArray(new String[]{}));
@@ -801,7 +801,9 @@ protected void failoverWriter() throws SQLException {
       final HostSpec writerHostSpec = getWriter(failoverResult.getTopology());
       if (keepSessionStateOnFailover) {
         restoreSessionState(failoverResult.getNewConnection());
+        transferSessionStateToNewConnection(writerHostSpec, failoverResult.getNewConnection());
       }
+      this.invalidateCurrentConnection();
       this.pluginService.setCurrentConnection(failoverResult.getNewConnection(), writerHostSpec);
 
       LOGGER.fine(