diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 586a5c329405..0f18e692c643 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -970,6 +970,20 @@ private void closeImpl() throws SQLException { } this.openStatements.clear(); + if (isTransactionStarted()) { + try { + // It looks like there's no need to start a new transaction after a rollback, + // but the commit behavior is preserved since close() may still fail before isClosed is updated. + rollbackImpl(); + } catch (SQLException e) { + if (exceptionToThrow == null) { + exceptionToThrow = e; + } else { + exceptionToThrow.addSuppressed(e); + } + } + } + boolean interrupted = Thread.currentThread().isInterrupted(); try { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 09cc6cda6940..0c0598ba2381 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -525,9 +525,6 @@ private void closeStatementResources() throws SQLException { this.currentUpdateCount = -1; this.currentJobIdIndex = -1; if (this.connection != null) { - if (this.connection.isTransactionStarted()) { - this.connection.rollback(); - } this.connection.removeStatement(this); } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index 674eb0df64e0..67dfbc71be48 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -486,6 +486,20 @@ public void testCancelWithJoblessQuery() throws SQLException, InterruptedExcepti verify(bigquery, Mockito.never()).cancel(any(JobId.class)); } + @Test + public void testCancelDoesNotRollbackTransaction() throws SQLException { + doReturn(true).when(bigQueryConnection).isTransactionStarted(); + BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement); + statementSpy.jobIds.add(jobId); + + statementSpy.cancel(); + + // Cancel should call bigquery.cancel() but not rollback the transaction + verify(bigquery).cancel(eq(jobId)); + verify(bigQueryConnection, Mockito.never()).rollback(); + verify(bigQueryConnection).removeStatement(statementSpy); + } + @ParameterizedTest @ValueSource(booleans = {true, false}) public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java index 9fe1c4f0f2fc..49f68a0feeed 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java @@ -2324,6 +2324,94 @@ public void testConnectionWithMultipleTransactionCommits() throws SQLException { connection.close(); } + @Test + public void testPreparedStatementCloseDoesNotRollbackTransaction() throws SQLException { + String TRANSACTION_TABLE = "JDBC_PS_CLOSE_TABLE" + randomNumber; + String createTransactionTable = + String.format( + "CREATE OR REPLACE TABLE %s.%s (`id` INTEGER, `name` STRING, `age` INTEGER);", + DATASET, TRANSACTION_TABLE); + String insertQuery = + String.format("INSERT INTO %s.%s (id, name, age) VALUES (?, ?, ?);", DATASET, TRANSACTION_TABLE); + String selectQuery = + String.format("SELECT id, name, age FROM %s.%s ORDER BY id;", DATASET, TRANSACTION_TABLE); + + bigQueryStatement.execute(createTransactionTable); + + try (Connection connection = DriverManager.getConnection(session_enabled_connection_uri)) { + connection.setAutoCommit(false); + try (PreparedStatement ps1 = connection.prepareStatement(insertQuery); + PreparedStatement ps2 = connection.prepareStatement(insertQuery)) { + ps1.setInt(1, 1); + ps1.setString(2, "DwightShrute"); + ps1.setInt(3, 10); + assertEquals(1, ps1.executeUpdate()); + + ps2.setInt(1, 2); + ps2.setString(2, "MichaelScott"); + ps2.setInt(3, 20); + assertEquals(1, ps2.executeUpdate()); + + ps1.close(); + connection.commit(); + + try (ResultSet resultSet = bigQueryStatement.executeQuery(selectQuery)) { + int rowCount = 0; + while (resultSet.next()) { + rowCount++; + assertEquals(rowCount, resultSet.getInt(1)); + } + assertEquals(2, rowCount); + } + } finally { + bigQueryStatement.execute( + String.format("DROP TABLE IF EXISTS %s.%s", DATASET, TRANSACTION_TABLE)); + } + } + } + + @Test + public void testClosingUnusedPreparedStatementDoesNotRollbackPreviousExecute() + throws SQLException { + String TRANSACTION_TABLE = "JDBC_PS_UNUSED_CLOSE_TABLE" + randomNumber; + String createTransactionTable = + String.format( + "CREATE OR REPLACE TABLE %s.%s (`id` INTEGER, `name` STRING, `age` INTEGER);", + DATASET, TRANSACTION_TABLE); + String insertQuery = + String.format("INSERT INTO %s.%s (id, name, age) VALUES (?, ?, ?);", DATASET, TRANSACTION_TABLE); + String selectQuery = + String.format("SELECT id, name, age FROM %s.%s ORDER BY id;", DATASET, TRANSACTION_TABLE); + + bigQueryStatement.execute(createTransactionTable); + + try (Connection connection = DriverManager.getConnection(session_enabled_connection_uri)) { + connection.setAutoCommit(false); + try (PreparedStatement ps1 = connection.prepareStatement(insertQuery); + PreparedStatement ps2 = connection.prepareStatement(insertQuery)) { + + ps2.setInt(1, 1); + ps2.setString(2, "MichaelScott"); + ps2.setInt(3, 20); + assertEquals(1, ps2.executeUpdate()); + + ps1.close(); + connection.commit(); + + try (ResultSet resultSet = bigQueryStatement.executeQuery(selectQuery)) { + assertTrue(resultSet.next()); + assertEquals(1, resultSet.getInt(1)); + assertEquals("MichaelScott", resultSet.getString(2)); + assertEquals(20, resultSet.getInt(3)); + assertFalse(resultSet.next()); + } + } + } finally { + bigQueryStatement.execute( + String.format("DROP TABLE IF EXISTS %s.%s", DATASET, TRANSACTION_TABLE)); + } + } + // Private Helper functions private String getSessionId() throws InterruptedException { QueryJobConfiguration stubJobConfig =