From f42828d46e2777b6aa4f56d46197781b2aa7e236 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 4 Sep 2026 14:56:08 -0400 Subject: [PATCH 1/5] fix(bigquery-jdbc): abort session when connection is closed --- .../bigquery/jdbc/BigQueryConnection.java | 36 +++++++++++++++++ .../bigquery/jdbc/BigQueryConnectionTest.java | 40 ++++++++++++++++++- .../bigquery/jdbc/it/ITBigQueryJDBCTest.java | 30 ++++++++++++++ 3 files changed, 105 insertions(+), 1 deletion(-) 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 2f5863054903..7adc6d1afd46 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 @@ -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 close", ex); + } catch (BigQueryException ex) { + LOG.warning( + "Failed to abort session during connection close (session may have already ended): " + + ex.getMessage()); + } finally { + this.sessionInfoConnectionProperty = null; + if (this.queryProperties != null) { + List 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(); diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java index 7b01f9ac760e..f206dc8a7651 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java +++ b/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 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()); + } + } } 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 81d172a70018..8e1ea9686d5b 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 @@ -2880,4 +2880,34 @@ 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("session ended") + || ex.getMessage().toLowerCase().contains("not found"), + "Expected session ended error but got: " + ex.getMessage()); + } + } + } } From 488ea7020707304c8b4aae7ae2f7bf40b458ac3b Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 4 Sep 2026 15:08:04 -0400 Subject: [PATCH 2/5] nit --- .../com/google/cloud/bigquery/jdbc/BigQueryConnection.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 7adc6d1afd46..c39665fcd8f3 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 @@ -1483,10 +1483,10 @@ private void abortSession() { abortJob.waitFor(); } catch (InterruptedException ex) { Thread.currentThread().interrupt(); - throw new BigQueryJdbcRuntimeException("Interrupted during close", ex); + throw new BigQueryJdbcRuntimeException("Interrupted during session abort", ex); } catch (BigQueryException ex) { LOG.warning( - "Failed to abort session during connection close (session may have already ended): " + "Failed to abort session during session abort (session may have already ended): " + ex.getMessage()); } finally { this.sessionInfoConnectionProperty = null; From 4840c9ed9c0bf11af469cc3ee0d18e50fccebb50 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 4 Sep 2026 15:53:26 -0400 Subject: [PATCH 3/5] fix test --- .../google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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 8e1ea9686d5b..282a776121b3 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 @@ -2904,9 +2904,11 @@ public void testSessionAbortedOnConnectionClose() throws SQLException { assertThrows( SQLException.class, () -> statement.execute("SELECT * FROM session_temp_table;")); assertTrue( - ex.getMessage().toLowerCase().contains("session ended") - || ex.getMessage().toLowerCase().contains("not found"), - "Expected session ended error but got: " + ex.getMessage()); + ex.getMessage().contains(sessionId), + "Expected exception message to not contain session ID: " + + sessionId + + ", but got: " + + ex.getMessage()); } } } From f12d6f4eee1e1312ebab60ab981b88b3fa0363e6 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 4 Sep 2026 20:54:59 -0400 Subject: [PATCH 4/5] nit --- .../google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) 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 282a776121b3..625836aa3302 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 @@ -2903,12 +2903,7 @@ public void testSessionAbortedOnConnectionClose() throws SQLException { SQLException ex = assertThrows( SQLException.class, () -> statement.execute("SELECT * FROM session_temp_table;")); - assertTrue( - ex.getMessage().contains(sessionId), - "Expected exception message to not contain session ID: " - + sessionId - + ", but got: " - + ex.getMessage()); + assertTrue(ex.getMessage().toLowerCase().contains("not found".toLowerCase())); } } } From ad5f561f9c689a0810586642dc34143c41875ee9 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Tue, 8 Sep 2026 14:27:07 -0400 Subject: [PATCH 5/5] ony abort driver created session --- java-bigquery-jdbc/docs/USER_GUIDE.md | 94 +++++++++---------- .../bigquery/jdbc/BigQueryConnection.java | 10 +- .../bigquery/jdbc/BigQueryStatement.java | 1 + .../bigquery/jdbc/BigQueryConnectionTest.java | 23 +++++ .../bigquery/jdbc/it/ITBigQueryJDBCTest.java | 1 - 5 files changed, 76 insertions(+), 53 deletions(-) diff --git a/java-bigquery-jdbc/docs/USER_GUIDE.md b/java-bigquery-jdbc/docs/USER_GUIDE.md index 0f38e0fe252e..cb0291275e54 100644 --- a/java-bigquery-jdbc/docs/USER_GUIDE.md +++ b/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=` 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=` 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[]`). 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 c39665fcd8f3..d13568a370b7 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 @@ -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,7 +1080,7 @@ private void closeImpl() throws SQLException { } } - if (this.sessionInfoConnectionProperty != null) { + if (this.sessionInfoConnectionProperty != null && this.isSessionCreatedByDriver) { abortSession(); } @@ -1499,6 +1506,7 @@ private void abortSession() { } this.queryProperties = Collections.unmodifiableList(updated); } + this.isSessionCreatedByDriver = false; this.transactionStarted = false; } } 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 3fd724c820d9..f3107579add0 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 @@ -1512,6 +1512,7 @@ QueryJobConfiguration.Builder getJobConfig(String query) { } } else if (isSessionEnabled) { queryConfigBuilder.setCreateSession(true); + this.connection.isSessionCreatedByDriver = true; } if (!props.isEmpty()) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java index f206dc8a7651..341341dcc5b4 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java @@ -816,6 +816,7 @@ 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( @@ -833,6 +834,8 @@ public void testCloseWithActiveSessionAbortsSession() throws Exception { connection.bigQuery = mockBigQuery; connection.updateSessionInfo("test_session_id_to_abort"); + connection.isSessionCreatedByDriver = true; + assertTrue(connection.isSessionCreatedByDriver()); connection.close(); ArgumentCaptor jobCaptor = ArgumentCaptor.forClass(JobInfo.class); @@ -841,6 +844,26 @@ public void testCloseWithActiveSessionAbortsSession() throws Exception { (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()); } } 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 625836aa3302..1e4e49ab6a5b 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 @@ -2903,7 +2903,6 @@ public void testSessionAbortedOnConnectionClose() throws SQLException { SQLException ex = assertThrows( SQLException.class, () -> statement.execute("SELECT * FROM session_temp_table;")); - assertTrue(ex.getMessage().toLowerCase().contains("not found".toLowerCase())); } } }