fix(bigquery-jdbc): abort session when connection is closed

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

@@ -1073,6 +1073,10 @@ private void closeImpl() throws SQLException {
         }
       }
 
+      if (this.sessionInfoConnectionProperty != null) {
+        abortSession();
+      }
+
       boolean interrupted = Thread.currentThread().isInterrupted();
 
       try {
@@ -1467,6 +1471,38 @@ private void commitTransaction() {
     }
   }
 
+  private void abortSession() {
+    try {
+      LOG.fine(
+          "Aborting session on connection close: " + this.sessionInfoConnectionProperty.getValue());
+      QueryJobConfiguration abortSessionJobConfig =
+          QueryJobConfiguration.newBuilder("CALL BQ.ABORT_SESSION();")
+              .setConnectionProperties(this.queryProperties)
+              .build();
+      Job abortJob = this.bigQuery.create(JobInfo.of(abortSessionJobConfig));
+      abortJob.waitFor();
+    } catch (InterruptedException ex) {
+      Thread.currentThread().interrupt();
+      throw new BigQueryJdbcRuntimeException("Interrupted during session abort", ex);
+    } catch (BigQueryException ex) {
+      LOG.warning(
+          "Failed to abort session during session abort (session may have already ended): "
+              + ex.getMessage());
+    } finally {
+      this.sessionInfoConnectionProperty = null;
+      if (this.queryProperties != null) {
+        List<ConnectionProperty> updated = new ArrayList<>();
+        for (ConnectionProperty cp : this.queryProperties) {
+          if (!"session_id".equalsIgnoreCase(cp.getKey())) {
+            updated.add(cp);
+          }
+        }
+        this.queryProperties = Collections.unmodifiableList(updated);
+      }
+      this.transactionStarted = false;
+    }
+  }
+
   @Override
   public CallableStatement prepareCall(String sql) throws SQLException {
     checkClosed();

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

@@ -16,7 +16,6 @@
 
 package com.google.cloud.bigquery.jdbc;
 
-import static org.junit.jupiter.api.Assertions.*;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -44,7 +43,10 @@
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.JobInfo;
 import com.google.cloud.bigquery.Project;
+import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
 import com.google.cloud.bigquery.exception.BigQueryJdbcException;
 import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
@@ -70,6 +72,7 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
+import org.mockito.ArgumentCaptor;
 import org.mockito.MockedStatic;
 
 public class BigQueryConnectionTest extends BigQueryJdbcLoggingBaseTest {
@@ -819,4 +822,39 @@ public void testUserSuppliedSessionId() throws Exception {
           "user_supplied_session_999", connection.getSessionInfoConnectionProperty().getValue());
     }
   }
+
+  @Test
+  public void testCloseWithActiveSessionAbortsSession() throws Exception {
+    try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
+      BigQuery mockBigQuery = mock(BigQuery.class);
+      Job mockJob = mock(Job.class);
+      when(mockBigQuery.create(any(JobInfo.class))).thenReturn(mockJob);
+      when(mockJob.waitFor()).thenReturn(mockJob);
+      connection.bigQuery = mockBigQuery;
+
+      connection.updateSessionInfo("test_session_id_to_abort");
+      connection.close();
+
+      ArgumentCaptor<JobInfo> jobCaptor = ArgumentCaptor.forClass(JobInfo.class);
+      verify(mockBigQuery).create(jobCaptor.capture());
+      QueryJobConfiguration config =
+          (QueryJobConfiguration) jobCaptor.getValue().getConfiguration();
+      assertEquals("CALL BQ.ABORT_SESSION();", config.getQuery());
+      assertNull(connection.getSessionInfoConnectionProperty());
+      assertTrue(connection.isClosed());
+    }
+  }
+
+  @Test
+  public void testCloseWithoutSessionDoesNotAbortSession() throws Exception {
+    try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
+      BigQuery mockBigQuery = mock(BigQuery.class);
+      connection.bigQuery = mockBigQuery;
+
+      connection.close();
+
+      verify(mockBigQuery, never()).create(any(JobInfo.class));
+      assertTrue(connection.isClosed());
+    }
+  }
 }

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

@@ -2880,4 +2880,31 @@ public void testPerConnectionLoggingE2E() throws SQLException, IOException {
       }
     }
   }
+
+  @Test
+  public void testSessionAbortedOnConnectionClose() throws SQLException {
+    String sessionId;
+    try (Connection connection = DriverManager.getConnection(session_enabled_connection_uri)) {
+      try (Statement statement = connection.createStatement()) {
+        statement.execute("CREATE TEMP TABLE session_temp_table (id INT64);");
+      }
+      BigQueryConnection bqConn = connection.unwrap(BigQueryConnection.class);
+      assertNotNull(bqConn.getSessionInfoConnectionProperty());
+      sessionId = bqConn.getSessionInfoConnectionProperty().getValue();
+      assertNotNull(sessionId);
+    }
+
+    // After connection is closed, the session is aborted on the BigQuery server.
+    // Attaching to the same session_id in a new connection should fail when running a query.
+    String urlWithAbortedSession =
+        connection_uri + "EnableSession=1;QueryProperties=session_id=" + sessionId + ";";
+    try (Connection newConnection = DriverManager.getConnection(urlWithAbortedSession)) {
+      try (Statement statement = newConnection.createStatement()) {
+        SQLException ex =
+            assertThrows(
+                SQLException.class, () -> statement.execute("SELECT * FROM session_temp_table;"));
+        assertTrue(ex.getMessage().toLowerCase().contains("not found".toLowerCase()));
+      }
+    }
+  }
 }