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

java-bigquery-jdbc/docs/USER_GUIDE.md

@@ -14,10 +14,11 @@ This guide provides comprehensive instructions for configuring, developing with,
 4. [Connection Properties Reference](#4-connection-properties-reference)
 5. [Data Type Mapping Reference](#5-data-type-mapping-reference)
 6. [JDBC Driver Architecture & Core Features](#6-jdbc-driver-architecture--core-features)
-   - [Transaction Management & Multi-Statement Sessions](#transaction-management--multi-statement-sessions)
+   - [Multi-Statement Sessions & Transaction Management](#multi-statement-sessions--transaction-management)
    - [High-Throughput Storage Read & Write APIs](#high-throughput-storage-read--write-apis)
 7. [Feature Examples & Code Snippets](#7-feature-examples--code-snippets)
    - [Transactions (Manual Commit & Rollback)](#transactions-manual-commit--rollback)
+   - [Connecting to a Pre-Existing Session](#connecting-to-a-pre-existing-session)
    - [Prepared Statements & Parameter Binding](#prepared-statements--parameter-binding)
    - [Callable Statements & Stored Procedures](#callable-statements--stored-procedures)
    - [Batch Ingestion with Storage Write API](#batch-ingestion-with-storage-write-api)
@@ -205,7 +206,8 @@ String url = "jdbc:bigquery://https://bigquery.googleapis.com:443"
 
 | Property Name | Default Value | Description |
 | :--- | :---: | :--- |
-| `EnableSession` | `false` | Enables multi-statement session creation and transaction support (`BEGIN`, `COMMIT`, `ROLLBACK`). |
+| `EnableSession` | `false` | Enables BigQuery multi-statement session creation and transaction support (`BEGIN`, `COMMIT`, `ROLLBACK`). |
+| `QueryProperties` | `null` | Comma- or semicolon-separated key-value pairs passed as connection-level job properties (e.g., `QueryProperties=session_id=<session_id>` to connect to a pre-existing session). |
 
 ### High-Throughput Storage & Write API Properties
 
@@ -291,58 +293,26 @@ When running queries through the JDBC driver for BigQuery, data types map as spe
 
 ## 6. JDBC Driver Architecture & Core Features
 
-### Transaction Management & Multi-Statement Sessions
+### Multi-Statement Sessions & Transaction Management
 
-BigQuery supports **Multi-Statement Transactions** across tables using standard SQL primitives (`BEGIN TRANSACTION`, `COMMIT TRANSACTION`, `ROLLBACK TRANSACTION`). The driver bridges standard JDBC methods (`setAutoCommit`, `commit`, `rollback`) directly to BigQuery's underlying session engine.
+BigQuery supports **Multi-Statement Sessions**, which preserve state across multiple SQL statements executed on the same connection.
 
-#### Session Lifecycle Flow:
+1. **Enabling Sessions (`EnableSession=true`)**:
+   - Add `;EnableSession=true` (or `EnableSession=1`) to the JDBC connection URL or DataSource properties.
+   - Under default auto-commit mode (`autoCommit=true`), statements execute and commit individually while sharing session state.
 
-```
-[DriverManager.getConnection()]
-             │
-   (EnableSession=true)
-             │
-  ┌──────────▼──────────┐
-  │ setAutoCommit(false)│ ──────► Begins transaction block in session
-  └──────────┬──────────┘
-             │
-  ┌──────────▼──────────┐
-  │ Execute DML & SQL   │ ──────► Runs queries within active session
-  │ Statements          │
-  └──────────┬──────────┘
-             │
-     ┌───────┴───────┐
-     │               │
-     ▼               ▼
-┌─────────┐     ┌──────────┐
-│commit() │     │rollback()│
-└────┬────┘     └────┬─────┘
-     │               │
-     ▼               ▼
-Executes:        Executes:
-COMMIT           ROLLBACK
-TRANSACTION;     TRANSACTION;
-     │               │
-     └───────┬───────┘
-             │
-             ▼
-(Auto-re-executes BEGIN TRANSACTION; if setAutoCommit remains false)
-```
+2. **Multi-Statement Transactions (`setAutoCommit(false)`)**:
+   - Multi-statement transactions require `;EnableSession=true` (calling `setAutoCommit(false)`, `commit()`, or `rollback()` with sessions disabled throws an exception).
+   - Calling `conn.setAutoCommit(false)` begins a multi-statement transaction in BigQuery. Statements executed within the transaction block remain uncommitted until `conn.commit()` is explicitly called (or discarded via `conn.rollback()`).
+   - If `autoCommit` remains `false`, the driver automatically starts the next transaction block for subsequent statements.
+   - Isolation level: `Connection.TRANSACTION_SERIALIZABLE` (BigQuery snapshot isolation).
 
-1. **Pre-requisite Check**: Calling `setAutoCommit(false)`, `commit()`, or `rollback()` requires `;EnableSession=true` in the connection URL. If disabled or invoked without an active transaction, an exception is thrown by the driver.
-2. **Session & Transaction Start**:
-   - `setAutoCommit(false)` initiates a multi-statement transaction session in BigQuery.
-3. **Statement Propagation**:
-   - All `Statement` or `PreparedStatement` instances created on the connection execute within the scope of the active session.
-4. **Commit & Rollback**:
-   - `commit()` executes `COMMIT TRANSACTION;` to commit changes.
-   - `rollback()` executes `ROLLBACK TRANSACTION;` to discard changes.
-   - If `autoCommit` remains `false`, the driver automatically starts the next transaction block.
-5. **Connection Close Safety**:
-   - If an uncommitted transaction is pending when `conn.close()` is invoked, the driver automatically rolls back the transaction to prevent uncommitted changes from persisting.
-6. **Isolation Level & Holdability**:
-   - Isolation level: `Connection.TRANSACTION_SERIALIZABLE` (BigQuery multi-statement snapshot isolation).
-   - Holdability: `ResultSet.CLOSE_CURSORS_AT_COMMIT`.
+3. **Using an Existing Session (`QueryProperties=session_id=...`)**:
+   - You can attach to a pre-existing BigQuery session by specifying `;QueryProperties=session_id=<existing_session_id>` in the connection URL.
+
+4. **Connection Closure & Lifecycle**:
+   - When `conn.close()` is called, sessions created by the driver are automatically terminated to release BigQuery server resources.
+   - If a pre-existing session ID was supplied by the user (`QueryProperties=session_id=...`), the session is preserved when the connection is closed.
 
 ---
 
@@ -360,7 +330,7 @@ For enterprise data ingestion and analytics extraction, the driver integrates wi
 ## 7. Feature Examples & Code Snippets
 
 ### Transactions (Manual Commit & Rollback)
-Transactions require `;EnableSession=true` in the connection URL to enable multi-statement sessions in BigQuery.
+Transactions require `;EnableSession=true` in the connection URL to enable multi-statement transactions in BigQuery:
 
 ```java
 String url = "jdbc:bigquery://https://bigquery.googleapis.com:443;ProjectId=my-project;EnableSession=true;OAuthType=3";
@@ -389,6 +359,28 @@ try (Connection conn = DriverManager.getConnection(url)) {
 
 ---
 
+### Connecting to a Pre-Existing Session
+To attach to an existing BigQuery session created outside the driver:
+
+```java
+String existingSessionId = "your_existing_session_id_here";
+String url = "jdbc:bigquery://https://bigquery.googleapis.com:443;ProjectId=my-project;EnableSession=true;"
+           + "QueryProperties=session_id=" + existingSessionId + ";OAuthType=3";
+
+try (Connection conn = DriverManager.getConnection(url);
+     Statement stmt = conn.createStatement()) {
+
+    // Query tables or temporary objects in the pre-existing session
+    try (ResultSet rs = stmt.executeQuery("SELECT * FROM ExistingTempTable")) {
+        while (rs.next()) {
+            // Process rows...
+        }
+    }
+}
+```
+
+---
+
 ### Prepared Statements & Parameter Binding
 Use `PreparedStatement` to safely bind parameters including primitive types, decimals (`BigDecimal`), temporal values (`Date`, `Timestamp`), and byte arrays (`byte[]`).
 

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

@@ -179,6 +179,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
   // when autocommit is false transaction starts and session is initialized.
   boolean transactionStarted;
   volatile ConnectionProperty sessionInfoConnectionProperty;
+  // isSessionCreatedByDriver is false by default.
+  boolean isSessionCreatedByDriver = false;
   boolean isClosed;
   DatasetId defaultDataset;
   String location;
@@ -683,6 +685,7 @@ private void beginTransaction() {
         transactionBeginJobConfig.setConnectionProperties(this.queryProperties);
       } else {
         transactionBeginJobConfig.setCreateSession(true);
+        this.isSessionCreatedByDriver = true;
       }
       Job job = this.bigQuery.create(JobInfo.of(transactionBeginJobConfig.build()));
       job = job.waitFor();
@@ -744,6 +747,10 @@ public ConnectionProperty getSessionInfoConnectionProperty() {
     return this.sessionInfoConnectionProperty;
   }
 
+  boolean isSessionCreatedByDriver() {
+    return this.isSessionCreatedByDriver;
+  }
+
   boolean isEnableHighThroughputAPI() {
     return this.enableHighThroughputAPI;
   }
@@ -1073,6 +1080,10 @@ private void closeImpl() throws SQLException {
         }
       }
 
+      if (this.sessionInfoConnectionProperty != null && this.isSessionCreatedByDriver) {
+        abortSession();
+      }
+
       boolean interrupted = Thread.currentThread().isInterrupted();
 
       try {
@@ -1467,6 +1478,39 @@ 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.isSessionCreatedByDriver = false;
+      this.transactionStarted = false;
+    }
+  }
+
   @Override
   public CallableStatement prepareCall(String sql) throws SQLException {
     checkClosed();

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

@@ -1512,6 +1512,7 @@ QueryJobConfiguration.Builder getJobConfig(String query) {
       }
     } else if (isSessionEnabled) {
       queryConfigBuilder.setCreateSession(true);
+      this.connection.isSessionCreatedByDriver = true;
     }
 
     if (!props.isEmpty()) {

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 {
@@ -813,10 +816,68 @@ public void testUserSuppliedSessionId() throws Exception {
         BASE_URL + ";EnableSession=1;QueryProperties=session_id=user_supplied_session_999";
     try (BigQueryConnection connection = new BigQueryConnection(urlWithSessionId)) {
       assertTrue(connection.isSessionEnabled());
+      assertFalse(connection.isSessionCreatedByDriver());
       assertNotNull(connection.getSessionInfoConnectionProperty());
       assertEquals("session_id", connection.getSessionInfoConnectionProperty().getKey());
       assertEquals(
           "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.isSessionCreatedByDriver = true;
+      assertTrue(connection.isSessionCreatedByDriver());
+      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());
+      assertFalse(connection.isSessionCreatedByDriver());
+      assertTrue(connection.isClosed());
+    }
+  }
+
+  @Test
+  public void testCloseWithUserSuppliedSessionDoesNotAbortSession() throws Exception {
+    String urlWithSessionId =
+        BASE_URL + ";EnableSession=1;QueryProperties=session_id=user_supplied_session_999";
+    try (BigQueryConnection connection = new BigQueryConnection(urlWithSessionId)) {
+      BigQuery mockBigQuery = mock(BigQuery.class);
+      connection.bigQuery = mockBigQuery;
+
+      assertFalse(connection.isSessionCreatedByDriver());
+      assertEquals(
+          "user_supplied_session_999", connection.getSessionInfoConnectionProperty().getValue());
+
+      connection.close();
+
+      verify(mockBigQuery, never()).create(any(JobInfo.class));
+      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,30 @@ 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;"));
+      }
+    }
+  }
 }