From 97030ee98b715181d1e41e542d6068d16670321f Mon Sep 17 00:00:00 2001 From: Bogdan Truta Date: Tue, 28 Oct 2025 10:13:07 +0200 Subject: [PATCH 1/2] feat: FIR-50249 support transactions in Core --- .../integration/tests/TransactionTest.java | 17 +- .../FireboltCoreQueryParameterProvider.java | 2 + .../jdbc/connection/FireboltConnection.java | 137 +++++++--- .../FireboltConnectionServiceSecret.java | 95 ------- .../FireboltConnectionUserPassword.java | 51 ++++ .../connection/FireboltCoreConnection.java | 21 +- .../FireboltConnectionServiceSecretTest.java | 50 ++++ .../FireboltCoreConnectionTest.java | 247 +++++++++++++++++- 8 files changed, 485 insertions(+), 135 deletions(-) diff --git a/src/integrationTest/java/integration/tests/TransactionTest.java b/src/integrationTest/java/integration/tests/TransactionTest.java index 8de8e4d22f..85ae41eeae 100644 --- a/src/integrationTest/java/integration/tests/TransactionTest.java +++ b/src/integrationTest/java/integration/tests/TransactionTest.java @@ -7,6 +7,8 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; import java.sql.Connection; import java.sql.PreparedStatement; @@ -24,6 +26,7 @@ @CustomLog @Tag(TestTag.V2) +@Execution(ExecutionMode.SAME_THREAD) class TransactionTest extends IntegrationTest { @BeforeAll @@ -41,6 +44,7 @@ void shouldNotStartTransactionIfAutoCommitIsTrue() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { connection.setAutoCommit(true); //no-op, but ensures auto-commit is true + // In core this breaks Firebolt Core since a transaction is orphaned and Core cannot manage more than 1 transaction statement.execute("BEGIN TRANSACTION;"); statement.execute("INSERT INTO transaction_test VALUES (2, 'test')"); @@ -58,6 +62,7 @@ void shouldNotStartTransactionIfAutoCommitIsTrue() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldCommitTransactionWhenSwitchingToAutoCommit() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -81,6 +86,7 @@ void shouldCommitTransactionWhenSwitchingToAutoCommit() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldHandleSequentialTransactions() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -109,6 +115,7 @@ void shouldHandleSequentialTransactions() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldRollbackTransactionSuccessfully() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -134,6 +141,7 @@ void shouldRollbackTransactionSuccessfully() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldWorkWithPreparedStatements() throws SQLException { try (Connection connection = createConnection(); PreparedStatement ps = connection.prepareStatement("INSERT INTO transaction_test VALUES (?, 'test')")) { @@ -151,6 +159,7 @@ void shouldWorkWithPreparedStatements() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldNotCommitTransactionWhenCommitWasManuallyExecuted() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -174,7 +183,8 @@ void shouldNotCommitTransactionWhenCommitWasManuallyExecuted() throws SQLExcepti } @Test - void shouldNotCommitTransactionWhenConnectionClosesOnAutoCommitFalse() throws SQLException { + @Tag(TestTag.CORE) + void shouldRollbackTransactionWhenConnectionClosesOnAutoCommitFalse() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -188,6 +198,7 @@ void shouldNotCommitTransactionWhenConnectionClosesOnAutoCommitFalse() throws SQ } @Test + @Tag(TestTag.CORE) void shouldThrowExceptionWhenStartingTransactionDuringTransaction() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -202,6 +213,7 @@ void shouldThrowExceptionWhenStartingTransactionDuringTransaction() throws SQLEx } @Test + @Tag(TestTag.CORE) void shouldNotRollbackTransactionWhenRollbackWasManuallyExecuted() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -225,6 +237,7 @@ void shouldNotRollbackTransactionWhenRollbackWasManuallyExecuted() throws SQLExc } @Test + @Tag(TestTag.CORE) void shouldThrowExceptionWhenCommitingWithNoTransaction() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -237,6 +250,7 @@ void shouldThrowExceptionWhenCommitingWithNoTransaction() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldThrowExceptionWhenRollbackWithNoTransaction() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { @@ -249,6 +263,7 @@ void shouldThrowExceptionWhenRollbackWithNoTransaction() throws SQLException { } @Test + @Tag(TestTag.CORE) void shouldNotCommitOnSetStatement() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { diff --git a/src/main/java/com/firebolt/jdbc/client/query/FireboltCoreQueryParameterProvider.java b/src/main/java/com/firebolt/jdbc/client/query/FireboltCoreQueryParameterProvider.java index 0e44874694..2789896791 100644 --- a/src/main/java/com/firebolt/jdbc/client/query/FireboltCoreQueryParameterProvider.java +++ b/src/main/java/com/firebolt/jdbc/client/query/FireboltCoreQueryParameterProvider.java @@ -28,6 +28,8 @@ public Map getQueryParams(FireboltProperties fireboltProperties, addQueryParameterIfNeeded(params, statementInfoWrapper.getPreparedStatementParameters()); addQueryTimeoutIfNeeded(params, queryTimeout); addServerAsyncIfNeeded(params, isServerAsync); + addTransactionIdIfNeeded(params, fireboltProperties.getTransactionId()); + addTransactionSequenceIdIfNeeded(params, fireboltProperties.getTransactionSequenceId()); return params; } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java index 3e46b14f3c..3b421ea69c 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java @@ -92,6 +92,12 @@ public abstract class FireboltConnection extends JdbcBase implements Connection, // Parameter parser is determined by the version we're running on @Getter public final ParserVersion parserVersion; + // Indicates if auto-commit is enabled + private boolean autoCommit = true; + // Indicates if the connection is currently in a transaction + private boolean inTransaction = false; + // Indicates if a transaction command (commit/rollback) is currently being executed + private boolean executingTransactionCommand = false; protected FireboltConnection(@NonNull String url, Properties connectionSettings, @@ -220,16 +226,93 @@ private void addStatement(FireboltStatement statement) throws SQLException { } } - @Override - public boolean getAutoCommit() throws SQLException { - validateConnectionIsNotClose(); - return true; - } - - @Override - @ExcludeFromJacocoGeneratedReport - @NotImplemented - public void setAutoCommit(boolean autoCommit) throws SQLException {} + @Override + public boolean getAutoCommit() throws SQLException { + validateConnectionIsNotClose(); + return autoCommit; + } + + @Override + public void setAutoCommit(boolean autoCommit) throws SQLException { + validateConnectionIsNotClose(); + + if (autoCommit && inTransaction) { + commit(); + } + this.autoCommit = autoCommit; + } + + @Override + public int getTransactionIsolation() throws SQLException { + validateConnectionIsNotClose(); + return Connection.TRANSACTION_REPEATABLE_READ; + } + + @Override + public void setTransactionIsolation(int level) throws SQLException { + validateConnectionIsNotClose(); + if (level != Connection.TRANSACTION_REPEATABLE_READ) { + throw new FireboltSQLFeatureNotSupportedException("Only TRANSACTION_REPEATABLE_READ isolation level is supported"); + } + } + + @Override + public void commit() throws SQLException { + executeTransactionCommand("COMMIT"); + } + + @Override + public void rollback() throws SQLException { + executeTransactionCommand("ROLLBACK"); + } + + /** + * Ensures a transaction is started if auto-commit is disabled and no transaction is active. + * Called automatically before query execution. + * + * @throws SQLException if there's an error starting the transaction + */ + public void ensureTransactionForQueryExecution() throws SQLException { + validateConnectionIsNotClose(); + if (sessionProperties.getTransactionId() == null) { + inTransaction = false; + } + + if (executingTransactionCommand || autoCommit) { + return; + } + + if (!inTransaction) { + executingTransactionCommand = true; + try (Statement statement = createStatement()) { + statement.execute("BEGIN TRANSACTION"); + inTransaction = true; + } catch (SQLException ex) { + throw new FireboltException("Could not start transaction for query execution", ex); + } finally { + executingTransactionCommand = false; + } + } + } + + private void executeTransactionCommand(String sql) throws SQLException { + validateConnectionIsNotClose(); + if (autoCommit) { + throw new FireboltException(String.format("Cannot %s when auto-commit is enabled", sql.toLowerCase())); + } + if (!inTransaction) { + throw new FireboltException("No transaction is currently active"); + } + executingTransactionCommand = true; + try (Statement statement = createStatement()) { + statement.execute(sql); + inTransaction = false; + } catch (SQLException ex) { + throw new FireboltException(String.format("Could not %s the transaction", sql.toLowerCase()), ex); + } finally { + executingTransactionCommand = false; + } + } @Override public boolean isClosed() { @@ -266,19 +349,6 @@ public String getEngine() { return getSessionProperties().getEngine(); } - @Override - public int getTransactionIsolation() throws SQLException { - validateConnectionIsNotClose(); - return Connection.TRANSACTION_NONE; - } - - @Override - public void setTransactionIsolation(int level) throws SQLException { - if (level != Connection.TRANSACTION_NONE) { - throw new FireboltSQLFeatureNotSupportedException(); - } - } - @Override public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException { validateConnectionIsNotClose(); @@ -313,6 +383,13 @@ public void abort(Executor executor) throws SQLException { @Override public void close() { log.debug("Closing connection"); + if (inTransaction) { + try { + rollback(); + } catch (SQLException e) { + log.error("Exception encountered while rolling back transaction on close"); + } + } synchronized (this) { if (isClosed()) { return; @@ -525,20 +602,6 @@ public String getEndpoint() { return httpConnectionUrl; } - public void ensureTransactionForQueryExecution() throws SQLException {} - - @Override - @NotImplemented - public void commit() throws SQLException { - // no-op as transactions are not supported - } - - @Override - @NotImplemented - public void rollback() throws SQLException { - // no-op as transactions are not supported - } - @Override public void setNetworkTimeout(Executor executor, int milliseconds) throws SQLException { validateConnectionIsNotClose(); diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java index 1b0855fe38..0cb3d44850 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecret.java @@ -39,12 +39,6 @@ public class FireboltConnectionServiceSecret extends FireboltConnection { private static final String PROTOCOL_VERSION = "2.4"; private final FireboltGatewayUrlService fireboltGatewayUrlService; private FireboltEngineVersion2Service fireboltEngineVersion2Service; - // Indicates if auto-commit is enabled - private boolean autoCommit = true; - // Indicates if the connection is currently in a transaction - private boolean inTransaction = false; - // Indicates if a transaction command (commit/rollback) is currently being executed - private boolean executingTransactionCommand = false; private final CacheService cacheService; @@ -106,95 +100,6 @@ protected void authenticate() throws SQLException { } } - @Override - public boolean getAutoCommit() throws SQLException { - validateConnectionIsNotClose(); - return autoCommit; - } - - @Override - public void setAutoCommit(boolean autoCommit) throws SQLException { - validateConnectionIsNotClose(); - - if (autoCommit && inTransaction) { - commit(); - } - this.autoCommit = autoCommit; - } - - @Override - public int getTransactionIsolation() throws SQLException { - validateConnectionIsNotClose(); - return Connection.TRANSACTION_REPEATABLE_READ; - } - - @Override - public void setTransactionIsolation(int level) throws SQLException { - validateConnectionIsNotClose(); - if (level != Connection.TRANSACTION_REPEATABLE_READ) { - throw new FireboltSQLFeatureNotSupportedException("Only TRANSACTION_REPEATABLE_READ isolation level is supported"); - } - } - - @Override - public void commit() throws SQLException { - executeTransactionCommand("COMMIT"); - } - - @Override - public void rollback() throws SQLException { - executeTransactionCommand("ROLLBACK"); - } - - /** - * Ensures a transaction is started if auto-commit is disabled and no transaction is active. - * Called automatically before query execution. - * - * @throws SQLException if there's an error starting the transaction - */ - @Override - public void ensureTransactionForQueryExecution() throws SQLException { - validateConnectionIsNotClose(); - if (sessionProperties.getTransactionId() == null) { - inTransaction = false; - } - - if (executingTransactionCommand || autoCommit) { - return; - } - - if (!inTransaction) { - executingTransactionCommand = true; - try (Statement statement = createStatement()) { - statement.execute("BEGIN TRANSACTION"); - inTransaction = true; - } catch (SQLException ex) { - throw new FireboltException("Could not start transaction for query execution", ex); - } finally { - executingTransactionCommand = false; - } - } - } - - private void executeTransactionCommand(String sql) throws SQLException { - validateConnectionIsNotClose(); - if (autoCommit) { - throw new FireboltException(String.format("Cannot %s when auto-commit is enabled", sql.toLowerCase())); - } - if (!inTransaction) { - throw new FireboltException("No transaction is currently active"); - } - executingTransactionCommand = true; - try (Statement statement = createStatement()) { - statement.execute(sql); - inTransaction = false; - } catch (SQLException ex) { - throw new FireboltException(String.format("Could not %s the transaction", sql.toLowerCase()), ex); - } finally { - executingTransactionCommand = false; - } - } - private String getAccessTokenInternal() throws SQLException { if (!isConnectionCachingEnabled()) { return getAccessToken(loginProperties).orElse(""); diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java index 5b0ad8492a..e79f60fbf1 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java @@ -2,18 +2,22 @@ import com.firebolt.jdbc.FireboltBackendType; import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport; +import com.firebolt.jdbc.annotation.NotImplemented; import com.firebolt.jdbc.client.account.FireboltAccountClient; import com.firebolt.jdbc.client.authentication.AuthenticationRequest; import com.firebolt.jdbc.client.authentication.FireboltAuthenticationClient; import com.firebolt.jdbc.client.authentication.UsernamePasswordAuthenticationRequest; import com.firebolt.jdbc.connection.settings.FireboltProperties; import com.firebolt.jdbc.exception.FireboltException; +import com.firebolt.jdbc.exception.FireboltSQLFeatureNotSupportedException; import com.firebolt.jdbc.service.FireboltAuthenticationService; import com.firebolt.jdbc.service.FireboltEngineApiService; import com.firebolt.jdbc.service.FireboltEngineInformationSchemaService; import com.firebolt.jdbc.service.FireboltEngineService; import com.firebolt.jdbc.service.FireboltStatementService; import com.firebolt.jdbc.type.ParserVersion; + +import java.sql.Connection; import java.sql.SQLException; import java.util.Optional; import java.util.Properties; @@ -97,6 +101,53 @@ protected boolean isConnectionCachingEnabled() { return false; } + @Override + public boolean getAutoCommit() throws SQLException { + validateConnectionIsNotClose(); + return true; + } + + @Override + @ExcludeFromJacocoGeneratedReport + @NotImplemented + public void setAutoCommit(boolean autoCommit) throws SQLException {} + + @Override + public int getTransactionIsolation() throws SQLException { + validateConnectionIsNotClose(); + return Connection.TRANSACTION_NONE; + } + + @Override + public void setTransactionIsolation(int level) throws SQLException { + if (level != Connection.TRANSACTION_NONE) { + throw new FireboltSQLFeatureNotSupportedException(); + } + } + + @Override + @NotImplemented + public void commit() throws SQLException { + // no-op as transactions are not supported + } + + @Override + @NotImplemented + public void rollback() throws SQLException { + // no-op as transactions are not supported + } + + /** + * Ensures a transaction is started if auto-commit is disabled and no transaction is active. + * Called automatically before query execution. + * + * + * @throws SQLException if there's an error starting the transaction + */ + @Override + public void ensureTransactionForQueryExecution() throws SQLException { + /* Since v1 does not support transactions this will be no-op */ + } @Override public Optional getConnectionUserAgentHeader() { return Optional.empty(); diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltCoreConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltCoreConnection.java index eff94d4ae3..5245a909f1 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltCoreConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltCoreConnection.java @@ -4,6 +4,8 @@ import com.firebolt.jdbc.client.authentication.AuthenticationRequest; import com.firebolt.jdbc.client.authentication.FireboltAuthenticationClient; import com.firebolt.jdbc.metadata.FireboltDatabaseMetadata; +import com.firebolt.jdbc.service.FireboltAuthenticationService; +import com.firebolt.jdbc.service.FireboltStatementService; import com.firebolt.jdbc.type.ParserVersion; import java.net.MalformedURLException; import java.net.URI; @@ -18,17 +20,34 @@ import lombok.CustomLog; import okhttp3.OkHttpClient; import org.apache.commons.lang3.StringUtils; +import org.jetbrains.annotations.TestOnly; @CustomLog public class FireboltCoreConnection extends FireboltConnection { - private static final String PROTOCOL_VERSION = "2.3"; + private static final String PROTOCOL_VERSION = "2.4"; public FireboltCoreConnection(String url, Properties connectionSettings) throws SQLException { super(url, connectionSettings, PROTOCOL_VERSION, ParserVersion.CURRENT); connect(); } + /** + * Test-only constructor that allows injection of FireboltStatementService for mocking + * @param url the connection URL + * @param connectionSettings connection properties + * @param fireboltAuthenticationService authentication service (can be null for core) + * @param fireboltStatementService statement service to inject + * @throws SQLException if connection fails + */ + @TestOnly + protected FireboltCoreConnection(String url, Properties connectionSettings, + FireboltAuthenticationService fireboltAuthenticationService, + FireboltStatementService fireboltStatementService) throws SQLException { + super(url, connectionSettings, fireboltAuthenticationService, fireboltStatementService, PROTOCOL_VERSION, ParserVersion.CURRENT); + connect(); + } + @Override protected void authenticate() throws SQLException { // there is no authentication for firebolt core diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java index 417b3c6423..28b23743b5 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java @@ -60,6 +60,7 @@ import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.atLeast; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -479,6 +480,55 @@ private void disableCacheConnection() { connectionProperties.put("cache_connection", "none"); } + @Test + void shouldRollbackTransactionWhenClosingConnectionWithActiveTransaction() throws SQLException { + try (FireboltConnection connection = createConnection(url, connectionProperties)) { + when(fireboltStatementService.execute(any(), any(), any())).thenReturn(Optional.empty()); + + // Start a transaction + connection.setAutoCommit(false); + connection.ensureTransactionForQueryExecution(); + + // Verify we're in a transaction + assertFalse(connection.getAutoCommit()); + + // Close the connection - this should trigger a rollback + connection.close(); + + // Verify rollback was called + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + verify(fireboltStatementService, times(2)) + .execute(statementCaptor.capture(), any(), any()); + + List statements = statementCaptor.getAllValues(); + assertEquals("BEGIN TRANSACTION", statements.get(0).getSql()); + assertEquals("ROLLBACK", statements.get(1).getSql()); + // The third call might be from connection cleanup + } + } + + @Test + void shouldNotRollbackWhenClosingConnectionWithoutActiveTransaction() throws SQLException { + try (FireboltConnection connection = createConnection(url, connectionProperties)) { + // Don't start a transaction - keep auto-commit enabled + assertTrue(connection.getAutoCommit()); + + // Close the connection - this should NOT trigger a rollback + connection.close(); + + // Verify no rollback was called (only connection setup calls) + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + verify(fireboltStatementService, atLeast(0)) + .execute(statementCaptor.capture(), any(), any()); + + // Check that no ROLLBACK statement was executed + List statements = statementCaptor.getAllValues(); + boolean hasRollback = statements.stream() + .anyMatch(stmt -> "ROLLBACK".equals(stmt.getSql())); + assertFalse(hasRollback, "No ROLLBACK statement should be executed when closing without active transaction"); + } + } + protected FireboltConnection createConnection(String url, Properties props) throws SQLException { return new FireboltConnectionServiceSecret(Pair.of(url, props), fireboltAuthenticationService, fireboltGatewayUrlService, fireboltStatementService, fireboltEngineVersion2Service, mockConnectionIdGenerator, mockCacheService); diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltCoreConnectionTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltCoreConnectionTest.java index 54f4192683..42012bae20 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltCoreConnectionTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltCoreConnectionTest.java @@ -1,16 +1,23 @@ package com.firebolt.jdbc.connection; import com.firebolt.jdbc.connection.settings.FireboltProperties; +import com.firebolt.jdbc.service.FireboltStatementService; import java.sql.SQLException; import java.sql.Statement; +import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Properties; import java.util.stream.Collectors; + +import com.firebolt.jdbc.exception.FireboltException; +import com.firebolt.jdbc.statement.StatementInfoWrapper; import org.apache.commons.lang3.StringUtils; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -19,6 +26,10 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @@ -30,6 +41,9 @@ class FireboltCoreConnectionTest { @Mock private Statement mockStatement; + @Mock + private FireboltStatementService fireboltStatementService; + @ParameterizedTest(name = "Valid URL: {0}") @ValueSource(strings = { "http://localhost:8080", @@ -155,6 +169,222 @@ void canConnectOverHttps() throws SQLException { } } + @Test + void shouldStartTransactionAndCommitWhenSwitchingToAutoCommit() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + when(fireboltStatementService.execute(any(), any(), any())).thenReturn(Optional.empty()); + + connection.setAutoCommit(false); + assertFalse(connection.getAutoCommit()); + + connection.createStatement().execute("SELECT 1"); + + connection.setAutoCommit(true); + assertTrue(connection.getAutoCommit()); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + + verify(fireboltStatementService, times(4)) + .execute(statementCaptor.capture(), any(), any()); + + List statement = statementCaptor.getAllValues(); + assertEquals("USE DATABASE \"my_db\"", statement.get(0).getSql()); + assertEquals("BEGIN TRANSACTION", statement.get(1).getSql()); + assertEquals("SELECT 1", statement.get(2).getSql()); + assertEquals("COMMIT", statement.get(3).getSql()); + } + } + + @Test + void shouldThrowExceptionWhenCommittingOrRollbackWithAutoCommitEnabled() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParams(connectionParams)) { + assertTrue(connection.getAutoCommit()); + + FireboltException commitException = assertThrows(FireboltException.class, connection::commit); + assertEquals("Cannot commit when auto-commit is enabled", commitException.getMessage()); + + FireboltException rollbackException = assertThrows(FireboltException.class, connection::rollback); + assertEquals("Cannot rollback when auto-commit is enabled", rollbackException.getMessage()); + } + } + + @Test + void shouldThrowExceptionWhenCommittingOrRollbackWithoutActiveTransaction() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParams(connectionParams)) { + connection.setAutoCommit(false); + + FireboltException commitException = assertThrows(FireboltException.class, connection::commit); + assertEquals("No transaction is currently active", commitException.getMessage()); + + FireboltException rollbackException = assertThrows(FireboltException.class, connection::rollback); + assertEquals("No transaction is currently active", rollbackException.getMessage()); + } + } + + @Test + void shouldCommitRollbackAndBeginTransactionSuccessfully() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + when(fireboltStatementService.execute(any(), any(), any())).thenReturn(Optional.empty()); + + connection.setAutoCommit(false); + + Statement statement1 = connection.createStatement(); + statement1.execute("SELECT 1"); + connection.commit(); + + Statement statement2 = connection.createStatement(); + statement2.execute("SELECT 2"); + connection.rollback(); + + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + + verify(fireboltStatementService, times(7)) + .execute(statementCaptor.capture(), any(), any()); + + List statement = statementCaptor.getAllValues(); + assertEquals("USE DATABASE \"my_db\"", statement.get(0).getSql()); + assertEquals("BEGIN TRANSACTION", statement.get(1).getSql()); + assertEquals("SELECT 1", statement.get(2).getSql()); + assertEquals("COMMIT", statement.get(3).getSql()); + assertEquals("BEGIN TRANSACTION", statement.get(4).getSql()); + assertEquals("SELECT 2", statement.get(5).getSql()); + assertEquals("ROLLBACK", statement.get(6).getSql()); + } + } + + @Test + void shouldHandleTransactionErrorsGracefully() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + connection.setAutoCommit(false); + + doThrow(new SQLException("")).when(fireboltStatementService).execute(any(), any(), any()); + + FireboltException exception = assertThrows(FireboltException.class, + connection::ensureTransactionForQueryExecution); + assertEquals("Could not start transaction for query execution", exception.getMessage()); + } + } + + @Test + void shouldHandleCommitErrorsGracefully() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + connection.setAutoCommit(false); + + connection.ensureTransactionForQueryExecution(); + + doThrow(new SQLException("Commit failed")).when(fireboltStatementService).execute(any(), any(), any()); + + FireboltException exception = assertThrows(FireboltException.class, connection::commit); + assertEquals("Could not commit the transaction", exception.getMessage()); + } + } + + @Test + void shouldHandleRollbackErrorsGracefully() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + connection.setAutoCommit(false); + + connection.ensureTransactionForQueryExecution(); + + doThrow(new SQLException("Rollback failed")).when(fireboltStatementService).execute(any(), any(), any()); + + FireboltException exception = assertThrows(FireboltException.class, connection::rollback); + assertEquals("Could not rollback the transaction", exception.getMessage()); + } + } + + @Test + void shouldRollbackTransactionWhenClosingConnectionWithActiveTransaction() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + when(fireboltStatementService.execute(any(), any(), any())).thenReturn(Optional.empty()); + + // Start a transaction + connection.setAutoCommit(false); + connection.ensureTransactionForQueryExecution(); + + // Verify we're in a transaction + assertFalse(connection.getAutoCommit()); + + // Close the connection - this should trigger a rollback + connection.close(); + + // Verify rollback was called + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + verify(fireboltStatementService, times(3)) + .execute(statementCaptor.capture(), any(), any()); + + List statements = statementCaptor.getAllValues(); + //needs to validate db + assertEquals("USE DATABASE \"my_db\"", statements.get(0).getSql()); + assertEquals("BEGIN TRANSACTION", statements.get(1).getSql()); + assertEquals("ROLLBACK", statements.get(2).getSql()); + } + } + + @Test + void shouldNotRollbackWhenClosingConnectionWithoutActiveTransaction() throws SQLException { + Map connectionParams = Map.of( + "url", "https://localhost:3473", + "database", "my_db" + ); + + try (FireboltCoreConnection connection = createConnectionWithParamsAndMockStatementService(connectionParams)) { + // Don't start a transaction - keep auto-commit enabled + assertTrue(connection.getAutoCommit()); + + // Close the connection - this should NOT trigger a rollback + connection.close(); + + // Verify only the USE DATABASE call was made, no ROLLBACK + ArgumentCaptor statementCaptor = ArgumentCaptor.forClass(StatementInfoWrapper.class); + verify(fireboltStatementService, times(1)) + .execute(statementCaptor.capture(), any(), any()); + + List statements = statementCaptor.getAllValues(); + assertEquals("USE DATABASE \"my_db\"", statements.get(0).getSql()); + } + } + private FireboltCoreConnection createConnection(String url) throws SQLException { StringBuilder jdbcUrlBuilder = new StringBuilder(VALID_URL_WITH_DB); @@ -176,12 +406,27 @@ private FireboltCoreConnection createConnectionWithParams(Map pa return aFireboltCoreConnection(jdbcUrlBuilder.toString(), new Properties()); } + private FireboltCoreConnection createConnectionWithParamsAndMockStatementService(Map parameters) throws SQLException { + StringBuilder jdbcUrlBuilder = new StringBuilder(VALID_URL_WITH_DB); + + String params = parameters.entrySet().stream() + .map(entry -> entry.getKey() + "=" + entry.getValue()) + .collect(Collectors.joining("&")); + + jdbcUrlBuilder.append("&").append(params); + return aFireboltCoreConnectionAndMockStatementService(jdbcUrlBuilder.toString(), new Properties()); + } + private FireboltCoreConnection aFireboltCoreConnection(String jdbcUrl, Properties properties) throws SQLException { - return new FireboltCoreConnection(jdbcUrl, properties){ + return new FireboltCoreConnection(jdbcUrl, properties, null, fireboltStatementService){ @Override public Statement createStatement() { return mockStatement; } }; } + + private FireboltCoreConnection aFireboltCoreConnectionAndMockStatementService(String jdbcUrl, Properties properties) throws SQLException { + return new FireboltCoreConnection(jdbcUrl, properties, null, fireboltStatementService); + } } From 6d581a64764399a03ca52552b97f4e0b27fc6271 Mon Sep 17 00:00:00 2001 From: Bogdan Truta Date: Wed, 29 Oct 2025 11:30:23 +0200 Subject: [PATCH 2/2] fix sonar and comms --- .../jdbc/connection/FireboltConnection.java | 14 +++++++------- .../connection/FireboltConnectionUserPassword.java | 4 +++- .../FireboltConnectionServiceSecretTest.java | 1 - 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java index 3b421ea69c..e664884a9a 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java @@ -383,17 +383,17 @@ public void abort(Executor executor) throws SQLException { @Override public void close() { log.debug("Closing connection"); - if (inTransaction) { - try { - rollback(); - } catch (SQLException e) { - log.error("Exception encountered while rolling back transaction on close"); - } - } synchronized (this) { if (isClosed()) { return; } else { + if (inTransaction) { + try { + rollback(); + } catch (SQLException e) { + log.error("Exception encountered while rolling back transaction on close"); + } + } closed = true; } } diff --git a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java index e79f60fbf1..6019d782e1 100644 --- a/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java +++ b/src/main/java/com/firebolt/jdbc/connection/FireboltConnectionUserPassword.java @@ -110,7 +110,9 @@ public boolean getAutoCommit() throws SQLException { @Override @ExcludeFromJacocoGeneratedReport @NotImplemented - public void setAutoCommit(boolean autoCommit) throws SQLException {} + public void setAutoCommit(boolean autoCommit) throws SQLException { + // no-op as transactions are not supported + } @Override public int getTransactionIsolation() throws SQLException { diff --git a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java index 28b23743b5..59c9697d2e 100644 --- a/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java +++ b/src/test/java/com/firebolt/jdbc/connection/FireboltConnectionServiceSecretTest.java @@ -503,7 +503,6 @@ void shouldRollbackTransactionWhenClosingConnectionWithActiveTransaction() throw List statements = statementCaptor.getAllValues(); assertEquals("BEGIN TRANSACTION", statements.get(0).getSql()); assertEquals("ROLLBACK", statements.get(1).getSql()); - // The third call might be from connection cleanup } }