Merge remote-tracking branch 'origin/main' into migrate-java-cloud-bom

Author: jinseopkim0

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -992,6 +992,21 @@ 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 {

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -526,9 +526,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);
     }
   }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java

@@ -492,6 +492,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 {

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java

@@ -2211,6 +2211,96 @@ 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 =