From 004727c7e9247365ec83d1bbaa4fc20728160836 Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Mon, 6 Nov 2023 17:40:30 -0800 Subject: [PATCH 1/3] - track session states - allow user to provide custom function to handle session state transfers --- .../software/amazon/jdbc/PluginService.java | 13 +++ .../amazon/jdbc/PluginServiceImpl.java | 26 +++++ .../failover/FailoverConnectionPlugin.java | 72 ++++++++++++-- .../ReadWriteSplittingPlugin.java | 41 +++++++- .../states/RestoreSessionStateCallable.java | 42 ++++++++ .../amazon/jdbc/states/SessionDirtyFlag.java | 33 +++++++ .../jdbc/states/SessionStateHelper.java | 99 +++++++++++++++++++ .../states/SessionStateTransferCallable.java | 45 +++++++++ .../jdbc/wrapper/ConnectionWrapper.java | 62 +++++++++--- .../jdbc/plugin/efm/ConcurrencyTests.java | 31 ++++++ .../FailoverConnectionPluginTest.java | 8 +- .../ReadWriteSplittingPluginTest.java | 2 + 12 files changed, 441 insertions(+), 33 deletions(-) create mode 100644 wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/states/SessionDirtyFlag.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java create mode 100644 wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java index a75c6ca2a..7f6a22d00 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java @@ -27,6 +27,7 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.exceptions.ExceptionHandler; import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.telemetry.TelemetryFactory; /** @@ -47,6 +48,18 @@ EnumSet setCurrentConnection( @Nullable ConnectionPlugin skipNotificationForThisPlugin) throws SQLException; + EnumSet getCurrentConnectionState(); + + void setCurrentConnectionState(SessionDirtyFlag flag); + + void resetCurrentConnectionState(SessionDirtyFlag flag); + + void resetCurrentConnectionStates(); + + boolean getAutoCommit(); + + void setAutoCommit(final boolean autoCommit); + List getHosts(); HostSpec getInitialConnectionHostSpec(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java index 299360aee..2cd701359 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java @@ -43,6 +43,7 @@ import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.HostAvailabilityStrategyFactory; import software.amazon.jdbc.hostlistprovider.StaticHostListProvider; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.CacheMap; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -68,6 +69,8 @@ public class PluginServiceImpl implements PluginService, CanReleaseResources, private final ExceptionManager exceptionManager; protected final DialectProvider dialectProvider; protected Dialect dialect; + protected EnumSet currentConnectionSessionState = EnumSet.noneOf(SessionDirtyFlag.class); + protected boolean isAutoCommit = false; public PluginServiceImpl( @NonNull final ConnectionPluginManager pluginManager, @@ -568,4 +571,27 @@ public String getTargetName() { return this.pluginManager.getDefaultConnProvider().getTargetName(); } + public EnumSet getCurrentConnectionState() { + return this.currentConnectionSessionState.clone(); + } + + public void setCurrentConnectionState(SessionDirtyFlag flag) { + this.currentConnectionSessionState.add(flag); + } + + public void resetCurrentConnectionState(SessionDirtyFlag flag) { + this.currentConnectionSessionState.remove(flag); + } + + public void resetCurrentConnectionStates() { + this.currentConnectionSessionState.clear(); + } + + public boolean getAutoCommit() { + return this.isAutoCommit; + } + + public void setAutoCommit(final boolean autoCommit) { + this.isAutoCommit = autoCommit; + } } 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 index f7577bfe7..da69917fc 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -42,6 +42,10 @@ import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsHelper; +import software.amazon.jdbc.states.RestoreSessionStateCallable; +import software.amazon.jdbc.states.SessionDirtyFlag; +import software.amazon.jdbc.states.SessionStateHelper; +import software.amazon.jdbc.states.SessionStateTransferCallable; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.RdsUtils; @@ -85,6 +89,10 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { static final String METHOD_ABORT = "Connection.abort"; static final String METHOD_CLOSE = "Connection.close"; static final String METHOD_IS_CLOSED = "Connection.isClosed"; + + protected static SessionStateTransferCallable sessionStateTransferCallable; + protected static RestoreSessionStateCallable restoreSessionStateCallable; + private final PluginService pluginService; protected final Properties properties; protected boolean enableFailoverSetting; @@ -199,6 +207,22 @@ public FailoverConnectionPlugin(final PluginService pluginService, final Propert this.failoverReaderFailedCounter = telemetryFactory.createCounter("readerFailover.completed.failed.count"); } + public static void setSessionStateTransferFunc(SessionStateTransferCallable callable) { + sessionStateTransferCallable = callable; + } + + public static void resetSessionStateTransferFunc() { + sessionStateTransferCallable = null; + } + + public static void setRestoreSessionStateFunc(RestoreSessionStateCallable callable) { + restoreSessionStateCallable = callable; + } + + public static void resetRestoreSessionStateFunc() { + restoreSessionStateCallable = null; + } + @Override public Set getSubscribedMethods() { return subscribedMethods; @@ -521,9 +545,10 @@ private boolean shouldAttemptReaderConnection() { */ private void switchCurrentConnectionTo(final HostSpec host, final Connection connection) throws SQLException { Connection currentConnection = this.pluginService.getCurrentConnection(); + HostSpec currentHostSpec = this.pluginService.getCurrentHostSpec(); if (currentConnection != connection) { - transferSessionState(currentConnection, connection); + transferSessionState(currentConnection, currentHostSpec, connection, host); invalidateCurrentConnection(); } @@ -538,22 +563,38 @@ private void switchCurrentConnectionTo(final HostSpec host, final Connection con * Transfers basic session state from one connection to another. * * @param from The connection to transfer state from + * @param fromHostSpec The connection {@link HostSpec} to transfer state from * @param to The connection to transfer state to + * @param toHostSpec The connection {@link HostSpec} to transfer state to * @throws SQLException if a database access error occurs, this method is called on a closed connection, this * method is called during a distributed transaction, or this method is called during a * transaction */ protected void transferSessionState( final Connection from, - final Connection to) throws SQLException { + final HostSpec fromHostSpec, + final Connection to, + final HostSpec toHostSpec) throws SQLException { if (from == null || to == null) { return; } - to.setReadOnly(from.isReadOnly()); - to.setAutoCommit(from.getAutoCommit()); - to.setTransactionIsolation(from.getTransactionIsolation()); + EnumSet sessionState = this.pluginService.getCurrentConnectionState(); + + SessionStateTransferCallable callableCopy = sessionStateTransferCallable; + if (callableCopy != null) { + final boolean isHandled = callableCopy.transferSessionState(sessionState, from, fromHostSpec, to, toHostSpec); + if (isHandled) { + // Custom function has handled session transfer + return; + } + } + + // Otherwise, lets run default logic. + sessionState = this.pluginService.getCurrentConnectionState(); + final SessionStateHelper helper = new SessionStateHelper(); + helper.transferSessionState(sessionState, from, to); } /** @@ -569,12 +610,23 @@ protected void restoreSessionState(final Connection to) throws SQLException { return; } - if (savedReadOnlyStatus != null) { - to.setReadOnly(savedReadOnlyStatus); - } - if (savedAutoCommitStatus != null) { - to.setAutoCommit(savedAutoCommitStatus); + final RestoreSessionStateCallable callableCopy = restoreSessionStateCallable; + if (callableCopy != null) { + final boolean isHandled = callableCopy.restoreSessionState( + this.pluginService.getCurrentConnectionState(), + to, + this.savedReadOnlyStatus, + this.savedAutoCommitStatus + ); + if (isHandled) { + // Custom function has handled everything. + return; + } } + + // Otherwise, lets run default logic. + final SessionStateHelper helper = new SessionStateHelper(); + helper.restoreSessionState(to, this.savedReadOnlyStatus, this.savedAutoCommitStatus); } private void dealWithOriginalException( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index d79dfd776..063a2a73e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java @@ -40,6 +40,9 @@ import software.amazon.jdbc.cleanup.CanReleaseResources; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.failover.FailoverSQLException; +import software.amazon.jdbc.states.SessionDirtyFlag; +import software.amazon.jdbc.states.SessionStateHelper; +import software.amazon.jdbc.states.SessionStateTransferCallable; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.WrapperUtils; @@ -61,6 +64,9 @@ public class ReadWriteSplittingPlugin extends AbstractConnectionPlugin static final String METHOD_SET_READ_ONLY = "Connection.setReadOnly"; static final String METHOD_CLEAR_WARNINGS = "Connection.clearWarnings"; + protected static SessionStateTransferCallable sessionStateTransferCallable; + + private final PluginService pluginService; private final Properties properties; private final String readerSelectorStrategy; @@ -105,6 +111,14 @@ public class ReadWriteSplittingPlugin extends AbstractConnectionPlugin this.readerConnection = readerConnection; } + public static void setSessionStateTransferFunc(SessionStateTransferCallable callable) { + sessionStateTransferCallable = callable; + } + + public static void resetSessionStateTransferFunc() { + sessionStateTransferCallable = null; + } + @Override public Set getSubscribedMethods() { return subscribedMethods; @@ -408,7 +422,7 @@ private void switchCurrentConnectionTo( return; } - transferSessionStateOnReadWriteSplit(newConnection); + transferSessionStateOnReadWriteSplit(newConnection, newConnectionHost); this.pluginService.setCurrentConnection(newConnection, newConnectionHost); LOGGER.finest(() -> Messages.get( "ReadWriteSplittingPlugin.settingCurrentConnection", @@ -422,18 +436,37 @@ private void switchCurrentConnectionTo( * will be updated when the setReadOnly call continues down the plugin chain * * @param to The connection to transfer state to + * @param toHostSpec The connection {@link HostSpec} to transfer state to * @throws SQLException if a database access error occurs, this method is called on a closed * connection, or this method is called during a distributed transaction */ protected void transferSessionStateOnReadWriteSplit( - final Connection to) throws SQLException { + final Connection to, final HostSpec toHostSpec) throws SQLException { final Connection from = this.pluginService.getCurrentConnection(); if (from == null || to == null) { return; } - to.setAutoCommit(from.getAutoCommit()); - to.setTransactionIsolation(from.getTransactionIsolation()); + EnumSet sessionState = this.pluginService.getCurrentConnectionState(); + + SessionStateTransferCallable callableCopy = sessionStateTransferCallable; + if (callableCopy != null) { + final boolean result = callableCopy.transferSessionState( + sessionState, + from, + this.pluginService.getCurrentHostSpec(), + to, + toHostSpec); + if (result) { + // Custom function has handled session transfer + return; + } + } + + 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, to); } private synchronized void switchToReaderConnection(final List hosts) diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java b/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java new file mode 100644 index 000000000..f8641b036 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java @@ -0,0 +1,42 @@ +/* + * 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.states; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.EnumSet; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; + +public interface RestoreSessionStateCallable { + /** + * Restores partial session state from saved values to a connection. + * + * @param sessionState Session state flags for from-connection + * @param to The connection to transfer state to + * @param readOnly ReadOnly flag to set to + * @param autoCommit AutoCommit flag to set to + * @return true, if session state is restored successful and no default logic should be executed after. + * False, if default logic should be executed. + */ + boolean restoreSessionState( + final @NonNull EnumSet sessionState, + final @NonNull Connection to, + final @Nullable Boolean readOnly, + final @Nullable Boolean autoCommit) + throws SQLException; +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/SessionDirtyFlag.java b/wrapper/src/main/java/software/amazon/jdbc/states/SessionDirtyFlag.java new file mode 100644 index 000000000..985da0204 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/states/SessionDirtyFlag.java @@ -0,0 +1,33 @@ +/* + * 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.states; + + +import java.util.EnumSet; + +public enum SessionDirtyFlag { + READONLY, + AUTO_COMMIT, + TRANSACTION_ISOLATION, + CATALOG, + NETWORK_TIMEOUT, + SCHEMA, + TYPE_MAP, + HOLDABILITY; + + public static final EnumSet ALL = EnumSet.allOf(SessionDirtyFlag.class); +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java new file mode 100644 index 000000000..9b260ec34 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java @@ -0,0 +1,99 @@ +/* + * 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.states; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.EnumSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +public class SessionStateHelper { + + /** + * Transfers basic session state from one connection to another. + * + * @param sessionState Session state of from-connection + * @param from The connection to transfer state from + * @param to The connection to transfer state to + * @throws SQLException if a database access error occurs, this method is called on a closed connection, this + * method is called during a distributed transaction, or this method is called during a + * transaction + */ + public void transferSessionState( + final EnumSet sessionState, + final Connection from, + final Connection to) throws SQLException { + + if (from == null || to == null) { + return; + } + + if (sessionState.contains(SessionDirtyFlag.READONLY)) { + to.setReadOnly(from.isReadOnly()); + } + if (sessionState.contains(SessionDirtyFlag.AUTO_COMMIT)) { + to.setAutoCommit(from.getAutoCommit()); + } + if (sessionState.contains(SessionDirtyFlag.TRANSACTION_ISOLATION)) { + to.setTransactionIsolation(from.getTransactionIsolation()); + } + if (sessionState.contains(SessionDirtyFlag.CATALOG)) { + to.setCatalog(from.getCatalog()); + } + if (sessionState.contains(SessionDirtyFlag.SCHEMA)) { + to.setSchema(from.getSchema()); + } + if (sessionState.contains(SessionDirtyFlag.TYPE_MAP)) { + to.setTypeMap(from.getTypeMap()); + } + if (sessionState.contains(SessionDirtyFlag.HOLDABILITY)) { + to.setHoldability(from.getHoldability()); + } + if (sessionState.contains(SessionDirtyFlag.NETWORK_TIMEOUT)) { + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + to.setNetworkTimeout(executorService, from.getNetworkTimeout()); + executorService.shutdown(); + } + } + + /** + * Restores partial session state from saved values to a connection. + * + * @param to The connection to transfer state to + * @param readOnly ReadOnly flag to set to + * @param autoCommit AutoCommit flag to set to + * @throws SQLException if a database access error occurs, this method is called on a closed connection, this + * method is called during a distributed transaction, or this method is called during a + * transaction + */ + public void restoreSessionState(final Connection to, final Boolean readOnly, final Boolean autoCommit) + throws SQLException { + + if (to == null) { + return; + } + + if (readOnly != null) { + to.setReadOnly(readOnly); + } + if (autoCommit != null) { + to.setAutoCommit(autoCommit); + } + } + +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java new file mode 100644 index 000000000..90d822447 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java @@ -0,0 +1,45 @@ +/* + * 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.states; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.EnumSet; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import software.amazon.jdbc.HostSpec; + +public interface SessionStateTransferCallable { + + /** + * Transfers session state from one connection to another. + * + * @param sessionState Session state flags for from-connection + * @param from The connection to transfer state from + * @param fromHostSpec The connection {@link HostSpec} to transfer state from + * @param to The connection to transfer state to + * @param toHostSpec The connection {@link HostSpec} to transfer state to + * @return true, if session state transfer is successful and no default logic should be executed after. + * False, if default logic should be executed. + */ + boolean transferSessionState( + final @NonNull EnumSet sessionState, + final @NonNull Connection from, + final @Nullable HostSpec fromHostSpec, + final @NonNull Connection to, + final @Nullable HostSpec toHostSpec) throws SQLException; +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java b/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java index 16cb8a8f8..e92fd17fa 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java +++ b/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java @@ -47,16 +47,12 @@ import software.amazon.jdbc.PluginServiceImpl; import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.cleanup.CanReleaseResources; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.dialect.DialectManager; -import software.amazon.jdbc.dialect.DialectProvider; import software.amazon.jdbc.dialect.HostListProviderSupplier; -import software.amazon.jdbc.hostlistprovider.ConnectionStringHostListProvider; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.WrapperUtils; -import software.amazon.jdbc.util.telemetry.DefaultTelemetryFactory; import software.amazon.jdbc.util.telemetry.TelemetryFactory; public class ConnectionWrapper implements Connection, CanReleaseResources { @@ -222,7 +218,12 @@ public void commit() throws SQLException { "Connection.commit", () -> { this.pluginService.getCurrentConnection().commit(); + final boolean autoCommit = this.pluginService.getAutoCommit(); this.pluginManagerService.setInTransaction(false); + if (!autoCommit + && this.pluginService.getCurrentConnectionState().contains(SessionDirtyFlag.AUTO_COMMIT)) { + this.pluginService.resetCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT); + } }); } @@ -320,10 +321,9 @@ public Statement createStatement( this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.createStatement", - () -> - this.pluginService - .getCurrentConnection() - .createStatement(resultSetType, resultSetConcurrency, resultSetHoldability), + () -> this.pluginService + .getCurrentConnection() + .createStatement(resultSetType, resultSetConcurrency, resultSetHoldability), resultSetType, resultSetConcurrency, resultSetHoldability); @@ -352,6 +352,7 @@ public void setReadOnly(final boolean readOnly) throws SQLException { () -> { this.pluginService.getCurrentConnection().setReadOnly(readOnly); this.pluginManagerService.setReadOnly(readOnly); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.READONLY); }, readOnly); } @@ -681,7 +682,12 @@ public void rollback() throws SQLException { "Connection.rollback", () -> { this.pluginService.getCurrentConnection().rollback(); + final boolean autoCommit = this.pluginService.getAutoCommit(); this.pluginManagerService.setInTransaction(false); + if (!autoCommit + && this.pluginService.getCurrentConnectionState().contains(SessionDirtyFlag.AUTO_COMMIT)) { + this.pluginService.resetCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT); + } }); } @@ -710,7 +716,13 @@ public void setAutoCommit(final boolean autoCommit) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setAutoCommit", - () -> this.pluginService.getCurrentConnection().setAutoCommit(autoCommit), + () -> { + this.pluginService.getCurrentConnection().setAutoCommit(autoCommit); + if (this.pluginService.getAutoCommit() != autoCommit) { + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.AUTO_COMMIT); + } + this.pluginService.setAutoCommit(autoCommit); + }, autoCommit); } @@ -732,7 +744,10 @@ public void setCatalog(final String catalog) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setCatalog", - () -> this.pluginService.getCurrentConnection().setCatalog(catalog), + () -> { + this.pluginService.getCurrentConnection().setCatalog(catalog); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.CATALOG); + }, catalog); } @@ -766,7 +781,10 @@ public void setHoldability(final int holdability) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setHoldability", - () -> this.pluginService.getCurrentConnection().setHoldability(holdability), + () -> { + this.pluginService.getCurrentConnection().setHoldability(holdability); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.HOLDABILITY); + }, holdability); } @@ -777,7 +795,10 @@ public void setNetworkTimeout(final Executor executor, final int milliseconds) t this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setNetworkTimeout", - () -> this.pluginService.getCurrentConnection().setNetworkTimeout(executor, milliseconds), + () -> { + this.pluginService.getCurrentConnection().setNetworkTimeout(executor, milliseconds); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.NETWORK_TIMEOUT); + }, executor, milliseconds); } @@ -812,7 +833,10 @@ public void setSchema(final String schema) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setSchema", - () -> this.pluginService.getCurrentConnection().setSchema(schema), + () -> { + this.pluginService.getCurrentConnection().setSchema(schema); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.SCHEMA); + }, schema); } @@ -823,7 +847,10 @@ public void setTransactionIsolation(final int level) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setTransactionIsolation", - () -> this.pluginService.getCurrentConnection().setTransactionIsolation(level), + () -> { + this.pluginService.getCurrentConnection().setTransactionIsolation(level); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.TRANSACTION_ISOLATION); + }, level); } @@ -834,7 +861,10 @@ public void setTypeMap(final Map> map) throws SQLException { this.pluginManager, this.pluginService.getCurrentConnection(), "Connection.setTypeMap", - () -> this.pluginService.getCurrentConnection().setTypeMap(map), + () -> { + this.pluginService.getCurrentConnection().setTypeMap(map); + this.pluginService.setCurrentConnectionState(SessionDirtyFlag.TYPE_MAP); + }, map); } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java index aae93ae7d..62007ac73 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/efm/ConcurrencyTests.java @@ -69,6 +69,7 @@ import software.amazon.jdbc.dialect.UnknownDialect; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @Disabled @@ -288,6 +289,36 @@ public EnumSet setCurrentConnection(@NonNull Connection conne return null; } + @Override + public EnumSet getCurrentConnectionState() { + return EnumSet.noneOf(SessionDirtyFlag.class); + } + + @Override + public void setCurrentConnectionState(SessionDirtyFlag flag) { + + } + + @Override + public void resetCurrentConnectionState(SessionDirtyFlag flag) { + + } + + @Override + public void resetCurrentConnectionStates() { + + } + + @Override + public boolean getAutoCommit() { + return false; + } + + @Override + public void setAutoCommit(boolean autoCommit) { + + } + @Override public List getHosts() { return null; diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index d02e07306..551f21173 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -62,6 +62,7 @@ import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; import software.amazon.jdbc.hostlistprovider.DynamicHostListProvider; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.telemetry.GaugeCallable; @@ -111,6 +112,7 @@ void init() throws SQLException { when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); + when(mockPluginService.getCurrentConnectionState()).thenReturn(EnumSet.allOf(SessionDirtyFlag.class)); when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); @@ -197,10 +199,10 @@ void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLEx void test_syncSessionState_withNullConnections() throws SQLException { initializePlugin(); - plugin.transferSessionState(null, mockConnection); + plugin.transferSessionState(null, null, mockConnection, null); verify(mockConnection, never()).getAutoCommit(); - plugin.transferSessionState(mockConnection, null); + plugin.transferSessionState(mockConnection, null, null, null); verify(mockConnection, never()).getAutoCommit(); } @@ -214,7 +216,7 @@ void test_syncSessionState() throws SQLException { initializePlugin(); - plugin.transferSessionState(mockConnection, mockConnection); + plugin.transferSessionState(mockConnection, null, mockConnection, null); verify(target).setReadOnly(eq(false)); verify(target).getAutoCommit(); verify(target).getTransactionIsolation(); diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java index 40aa20c38..18ba25f21 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java @@ -58,6 +58,7 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.plugin.failover.FailoverSuccessSQLException; +import software.amazon.jdbc.states.SessionDirtyFlag; import software.amazon.jdbc.util.SqlState; public class ReadWriteSplittingPluginTest { @@ -140,6 +141,7 @@ void mockDefaultBehavior() throws SQLException { when(this.mockPluginService.connect(eq(readerHostSpec3), any(Properties.class))) .thenReturn(mockReaderConn3); when(this.mockPluginService.acceptsStrategy(any(), eq("random"))).thenReturn(true); + when(mockPluginService.getCurrentConnectionState()).thenReturn(EnumSet.allOf(SessionDirtyFlag.class)); when(this.mockConnectFunc.call()).thenReturn(mockWriterConn); when(mockWriterConn.createStatement()).thenReturn(mockStatement); when(mockReaderConn1.createStatement()).thenReturn(mockStatement); From c88645d1db568b0c2d3c76470436eca2a7218320 Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Tue, 7 Nov 2023 09:40:36 -0800 Subject: [PATCH 2/3] code review --- .../failover/FailoverConnectionPlugin.java | 34 ++++++++-------- .../ReadWriteSplittingPlugin.java | 21 +++++----- .../states/RestoreSessionStateCallable.java | 4 +- .../jdbc/states/SessionStateHelper.java | 40 +++++++++---------- .../states/SessionStateTransferCallable.java | 16 ++++---- 5 files changed, 59 insertions(+), 56 deletions(-) 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 index da69917fc..890589021 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -560,23 +560,23 @@ private void switchCurrentConnectionTo(final HostSpec host, final Connection con } /** - * Transfers basic session state from one connection to another. + * Transfers session state from one connection to another. * - * @param from The connection to transfer state from - * @param fromHostSpec The connection {@link HostSpec} to transfer state from - * @param to The connection to transfer state to - * @param toHostSpec The connection {@link HostSpec} to transfer state to + * @param src The connection to transfer state from + * @param srcHostSpec The connection {@link HostSpec} to transfer state from + * @param dest The connection to transfer state to + * @param destHostSpec The connection {@link HostSpec} to transfer state to * @throws SQLException if a database access error occurs, this method is called on a closed connection, this * method is called during a distributed transaction, or this method is called during a * transaction */ protected void transferSessionState( - final Connection from, - final HostSpec fromHostSpec, - final Connection to, - final HostSpec toHostSpec) throws SQLException { + final Connection src, + final HostSpec srcHostSpec, + final Connection dest, + final HostSpec destHostSpec) throws SQLException { - if (from == null || to == null) { + if (src == null || dest == null) { return; } @@ -584,7 +584,7 @@ protected void transferSessionState( SessionStateTransferCallable callableCopy = sessionStateTransferCallable; if (callableCopy != null) { - final boolean isHandled = callableCopy.transferSessionState(sessionState, from, fromHostSpec, to, toHostSpec); + final boolean isHandled = callableCopy.transferSessionState(sessionState, src, srcHostSpec, dest, destHostSpec); if (isHandled) { // Custom function has handled session transfer return; @@ -594,19 +594,19 @@ protected void transferSessionState( // Otherwise, lets run default logic. sessionState = this.pluginService.getCurrentConnectionState(); final SessionStateHelper helper = new SessionStateHelper(); - helper.transferSessionState(sessionState, from, to); + helper.transferSessionState(sessionState, src, dest); } /** * Restores partial session state from saved values to a connection. * - * @param to The connection to transfer state to + * @param dest The connection to transfer state to * @throws SQLException if a database access error occurs, this method is called on a closed connection, this * method is called during a distributed transaction, or this method is called during a * transaction */ - protected void restoreSessionState(final Connection to) throws SQLException { - if (to == null) { + protected void restoreSessionState(final Connection dest) throws SQLException { + if (dest == null) { return; } @@ -614,7 +614,7 @@ protected void restoreSessionState(final Connection to) throws SQLException { if (callableCopy != null) { final boolean isHandled = callableCopy.restoreSessionState( this.pluginService.getCurrentConnectionState(), - to, + dest, this.savedReadOnlyStatus, this.savedAutoCommitStatus ); @@ -626,7 +626,7 @@ protected void restoreSessionState(final Connection to) throws SQLException { // Otherwise, lets run default logic. final SessionStateHelper helper = new SessionStateHelper(); - helper.restoreSessionState(to, this.savedReadOnlyStatus, this.savedAutoCommitStatus); + helper.restoreSessionState(dest, this.savedReadOnlyStatus, this.savedAutoCommitStatus); } private void dealWithOriginalException( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index 063a2a73e..ee394c2d1 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java @@ -435,15 +435,18 @@ private void switchCurrentConnectionTo( * status. This method is only called when setReadOnly is being called; the read-only status * will be updated when the setReadOnly call continues down the plugin chain * - * @param to The connection to transfer state to - * @param toHostSpec The connection {@link HostSpec} to transfer state to + * @param dest The destination connection to transfer state to + * @param destHostSpec The destination connection {@link HostSpec} * @throws SQLException if a database access error occurs, this method is called on a closed * connection, or this method is called during a distributed transaction */ protected void transferSessionStateOnReadWriteSplit( - final Connection to, final HostSpec toHostSpec) throws SQLException { + final Connection dest, + final HostSpec destHostSpec) + throws SQLException { + final Connection from = this.pluginService.getCurrentConnection(); - if (from == null || to == null) { + if (from == null || dest == null) { return; } @@ -451,13 +454,13 @@ protected void transferSessionStateOnReadWriteSplit( SessionStateTransferCallable callableCopy = sessionStateTransferCallable; if (callableCopy != null) { - final boolean result = callableCopy.transferSessionState( + final boolean isHandled = callableCopy.transferSessionState( sessionState, from, this.pluginService.getCurrentHostSpec(), - to, - toHostSpec); - if (result) { + dest, + destHostSpec); + if (isHandled) { // Custom function has handled session transfer return; } @@ -466,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, to); + helper.transferSessionState(sessionState, from, dest); } private synchronized void switchToReaderConnection(final List hosts) diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java b/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java index f8641b036..b1e3a236e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java +++ b/wrapper/src/main/java/software/amazon/jdbc/states/RestoreSessionStateCallable.java @@ -27,7 +27,7 @@ public interface RestoreSessionStateCallable { * Restores partial session state from saved values to a connection. * * @param sessionState Session state flags for from-connection - * @param to The connection to transfer state to + * @param dest The destination connection to transfer state to * @param readOnly ReadOnly flag to set to * @param autoCommit AutoCommit flag to set to * @return true, if session state is restored successful and no default logic should be executed after. @@ -35,7 +35,7 @@ public interface RestoreSessionStateCallable { */ boolean restoreSessionState( final @NonNull EnumSet sessionState, - final @NonNull Connection to, + final @NonNull Connection dest, final @Nullable Boolean readOnly, final @Nullable Boolean autoCommit) throws SQLException; diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java index 9b260ec34..d4b203593 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java +++ b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateHelper.java @@ -25,48 +25,48 @@ public class SessionStateHelper { /** - * Transfers basic session state from one connection to another. + * Transfers session state from source connection to destination connection. * - * @param sessionState Session state of from-connection - * @param from The connection to transfer state from - * @param to The connection to transfer state to + * @param sessionState Session state of source connection + * @param src The source connection to transfer state from + * @param dest The destination connection to transfer state to * @throws SQLException if a database access error occurs, this method is called on a closed connection, this * method is called during a distributed transaction, or this method is called during a * transaction */ public void transferSessionState( final EnumSet sessionState, - final Connection from, - final Connection to) throws SQLException { + final Connection src, + final Connection dest) throws SQLException { - if (from == null || to == null) { + if (src == null || dest == null) { return; } if (sessionState.contains(SessionDirtyFlag.READONLY)) { - to.setReadOnly(from.isReadOnly()); + dest.setReadOnly(src.isReadOnly()); } if (sessionState.contains(SessionDirtyFlag.AUTO_COMMIT)) { - to.setAutoCommit(from.getAutoCommit()); + dest.setAutoCommit(src.getAutoCommit()); } if (sessionState.contains(SessionDirtyFlag.TRANSACTION_ISOLATION)) { - to.setTransactionIsolation(from.getTransactionIsolation()); + dest.setTransactionIsolation(src.getTransactionIsolation()); } if (sessionState.contains(SessionDirtyFlag.CATALOG)) { - to.setCatalog(from.getCatalog()); + dest.setCatalog(src.getCatalog()); } if (sessionState.contains(SessionDirtyFlag.SCHEMA)) { - to.setSchema(from.getSchema()); + dest.setSchema(src.getSchema()); } if (sessionState.contains(SessionDirtyFlag.TYPE_MAP)) { - to.setTypeMap(from.getTypeMap()); + dest.setTypeMap(src.getTypeMap()); } if (sessionState.contains(SessionDirtyFlag.HOLDABILITY)) { - to.setHoldability(from.getHoldability()); + dest.setHoldability(src.getHoldability()); } if (sessionState.contains(SessionDirtyFlag.NETWORK_TIMEOUT)) { final ExecutorService executorService = Executors.newSingleThreadExecutor(); - to.setNetworkTimeout(executorService, from.getNetworkTimeout()); + dest.setNetworkTimeout(executorService, src.getNetworkTimeout()); executorService.shutdown(); } } @@ -74,25 +74,25 @@ public void transferSessionState( /** * Restores partial session state from saved values to a connection. * - * @param to The connection to transfer state to + * @param dest The destination connection to transfer state to * @param readOnly ReadOnly flag to set to * @param autoCommit AutoCommit flag to set to * @throws SQLException if a database access error occurs, this method is called on a closed connection, this * method is called during a distributed transaction, or this method is called during a * transaction */ - public void restoreSessionState(final Connection to, final Boolean readOnly, final Boolean autoCommit) + public void restoreSessionState(final Connection dest, final Boolean readOnly, final Boolean autoCommit) throws SQLException { - if (to == null) { + if (dest == null) { return; } if (readOnly != null) { - to.setReadOnly(readOnly); + dest.setReadOnly(readOnly); } if (autoCommit != null) { - to.setAutoCommit(autoCommit); + dest.setAutoCommit(autoCommit); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java index 90d822447..6872bb66c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java +++ b/wrapper/src/main/java/software/amazon/jdbc/states/SessionStateTransferCallable.java @@ -29,17 +29,17 @@ public interface SessionStateTransferCallable { * Transfers session state from one connection to another. * * @param sessionState Session state flags for from-connection - * @param from The connection to transfer state from - * @param fromHostSpec The connection {@link HostSpec} to transfer state from - * @param to The connection to transfer state to - * @param toHostSpec The connection {@link HostSpec} to transfer state to + * @param src The source connection to transfer state from + * @param srcHostSpec The source connection {@link HostSpec} + * @param dest The destination connection to transfer state to + * @param destHostSpec The destination connection {@link HostSpec} * @return true, if session state transfer is successful and no default logic should be executed after. * False, if default logic should be executed. */ boolean transferSessionState( final @NonNull EnumSet sessionState, - final @NonNull Connection from, - final @Nullable HostSpec fromHostSpec, - final @NonNull Connection to, - final @Nullable HostSpec toHostSpec) throws SQLException; + final @NonNull Connection src, + final @Nullable HostSpec srcHostSpec, + final @NonNull Connection dest, + final @Nullable HostSpec destHostSpec) throws SQLException; } From 898bcc74ec87208cc92e306b9053c406d935930a Mon Sep 17 00:00:00 2001 From: sergiyvamz Date: Tue, 7 Nov 2023 10:03:59 -0800 Subject: [PATCH 3/3] code review --- .../ReadWriteSplittingPlugin.java | 8 ++++---- .../amazon/jdbc/wrapper/ConnectionWrapper.java | 15 +++++++++------ 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index ee394c2d1..4ea8d6345 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/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 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 index e92fd17fa..a48cdbc55 100644 --- 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); }