code review

aws · Nov 7, 2023 · b374fae · b374fae
1 parent b383aa0
commit b374fae
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/...rc/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/...rc/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java
@@ -445,8 +445,8 @@ protected void transferSessionStateOnReadWriteSplit(
       final HostSpec destHostSpec)
       throws SQLException {
 
-    final Connection from = this.pluginService.getCurrentConnection();
-    if (from == null || dest == null) {
+    final Connection src = this.pluginService.getCurrentConnection();
+    if (src == null || dest == null) {
       return;
     }
 
@@ -456,7 +456,7 @@ protected void transferSessionStateOnReadWriteSplit(
     if (callableCopy != null) {
       final boolean isHandled = callableCopy.transferSessionState(
           sessionState,
-          from,
+          src,
           this.pluginService.getCurrentHostSpec(),
           dest,
           destHostSpec);
@@ -469,7 +469,7 @@ protected void transferSessionStateOnReadWriteSplit(
     sessionState = this.pluginService.getCurrentConnectionState();
     sessionState.remove(SessionDirtyFlag.READONLY); // We don't want to change READONLY flag of the connection
     final SessionStateHelper helper = new SessionStateHelper();
-    helper.transferSessionState(sessionState, from, dest);
+    helper.transferSessionState(sessionState, src, dest);
   }
 
   private synchronized void switchToReaderConnection(final List<HostSpec> hosts)

diff --git a/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java b/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java
@@ -180,6 +180,7 @@ public void abort(final Executor executor) throws SQLException {
         () -> {
           this.pluginService.getCurrentConnection().abort(executor);
           this.pluginManagerService.setInTransaction(false);
+          this.pluginService.resetCurrentConnectionStates();
         },
         executor);
   }
@@ -205,6 +206,7 @@ public void close() throws SQLException {
           this.pluginService.getCurrentConnection().close();
           this.openConnectionStacktrace = null;
           this.pluginManagerService.setInTransaction(false);
+          this.pluginService.resetCurrentConnectionStates();
         });
     this.releaseResources();
   }
@@ -218,9 +220,9 @@ public void commit() throws SQLException {
         "Connection.commit",
         () -> {
           this.pluginService.getCurrentConnection().commit();
-          final boolean autoCommit = this.pluginService.getAutoCommit();
+          final boolean isInTransaction = this.pluginService.isInTransaction();
           this.pluginManagerService.setInTransaction(false);
-          if (!autoCommit
+          if (isInTransaction
               && this.pluginService.getCurrentConnectionState().contains(SessionDirtyFlag.AUTO_COMMIT)) {
             this.pluginService.resetCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT);
           }
@@ -682,9 +684,9 @@ public void rollback() throws SQLException {
         "Connection.rollback",
         () -> {
           this.pluginService.getCurrentConnection().rollback();
-          final boolean autoCommit = this.pluginService.getAutoCommit();
+          final boolean isInTransaction = this.pluginService.isInTransaction();
           this.pluginManagerService.setInTransaction(false);
-          if (!autoCommit
+          if (isInTransaction
               && this.pluginService.getCurrentConnectionState().contains(SessionDirtyFlag.AUTO_COMMIT)) {
             this.pluginService.resetCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT);
           }
@@ -717,11 +719,12 @@ public void setAutoCommit(final boolean autoCommit) throws SQLException {
         this.pluginService.getCurrentConnection(),
         "Connection.setAutoCommit",
         () -> {
+          final boolean currentAutoCommit = this.pluginService.getAutoCommit();
           this.pluginService.getCurrentConnection().setAutoCommit(autoCommit);
-          if (this.pluginService.getAutoCommit() != autoCommit) {
+          this.pluginService.setAutoCommit(autoCommit);
+          if (currentAutoCommit != autoCommit) {
             this.pluginService.setCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT);
           }
-          this.pluginService.setAutoCommit(autoCommit);
         },
         autoCommit);
   }