From f42828d46e2777b6aa4f56d46197781b2aa7e236 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 4 Sep 2026 14:56:08 -0400 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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 05/10] 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())); } } } From 5182a533f5194bbf86f1bcd13d310148a22f65cc Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Thu, 17 Sep 2026 15:48:18 -0400 Subject: [PATCH 06/10] publish session state as a n atomic snapshot --- .../bigquery/jdbc/BigQueryConnection.java | 244 ++++++++++++------ .../bigquery/jdbc/BigQueryStatement.java | 16 +- .../jdbc/BigQueryCallableStatementTest.java | 4 + .../bigquery/jdbc/BigQueryConnectionTest.java | 67 +++-- .../jdbc/BigQueryJdbcContextProxyTest.java | 5 + .../BigQueryPreparedStatementSettersTest.java | 3 + .../bigquery/jdbc/BigQueryStatementTest.java | 4 + 7 files changed, 240 insertions(+), 103 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 c3ea68baf644..735450c9a93f 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 @@ -78,6 +78,8 @@ import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.logging.Level; /** * An implementation of {@link java.sql.Connection} for establishing a connection with BigQuery and @@ -92,6 +94,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { private final String connectionId; private static final String DEFAULT_JDBC_TOKEN_VALUE = "Google-BigQuery-JDBC-Driver"; private static final String DEFAULT_VERSION = "0.0.0"; + // Canonical spelling of the BigQuery session_id connection property key. + static final String SESSION_ID_KEY = "session_id"; private static final Set SAFE_TO_LOG_PROPERTIES = ImmutableSortedSet.orderedBy(String.CASE_INSENSITIVE_ORDER) .add( @@ -180,9 +184,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { // transactionStarted is false by default. // 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; @@ -203,7 +204,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { long destinationDatasetExpirationTime; String kmsKeyName; String universeDomain; - private volatile List queryProperties; Map authProperties; Map overrideProperties; Map proxyProperties; @@ -247,6 +247,14 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { private final ExecutorService metadataExecutor; private final ExecutorService queryExecutor; + /** + * All session-scoped state, published as a single immutable snapshot. + * + *

Never null; starts out as an empty, session-less state. + */ + private final AtomicReference sessionState = + new AtomicReference<>(new SessionState(null, Collections.emptyList(), false)); + BigQueryConnection(String url) throws IOException { this(url, DataSource.fromUrl(url)); } @@ -260,7 +268,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.otelContext = Context.current().with(baggage); try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(this.connectionId)) { this.connectionUrl = url; - if (LOG.isLoggable(java.util.logging.Level.CONFIG)) { + if (LOG.isLoggable(Level.CONFIG)) { Properties connectionProps = ds.createProperties(); Properties maskedProps = new Properties(); for (String name : connectionProps.stringPropertyNames()) { @@ -366,9 +374,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { this.unsupportedHTAPIFallback = ds.getUnsupportedHTAPIFallback(); this.maxResults = ds.getMaxResults(); Map queryPropertiesMap = ds.getQueryProperties(); - this.sessionInfoConnectionProperty = - getSessionPropertyFromQueryProperties(queryPropertiesMap); - this.queryProperties = convertMapToConnectionPropertiesList(queryPropertiesMap); + this.sessionState.set( + new SessionState( + getSessionPropertyFromQueryProperties(queryPropertiesMap), + convertMapToConnectionPropertiesList(queryPropertiesMap), + false)); this.enableWriteAPI = ds.getEnableWriteAPI(); this.writeAPIActivationRowCount = ds.getSwaActivationRowCount(); this.writeAPIAppendRowCount = ds.getSwaAppendRowCount(); @@ -602,10 +612,6 @@ String getKmsKeyName() { return this.kmsKeyName; } - List getQueryProperties() { - return this.queryProperties; - } - public String getLocation() { checkClosed(); return this.location; @@ -656,6 +662,14 @@ Map getLabels() { return this.labels; } + List getQueryProperties() { + return this.sessionState.get().queryProperties; + } + + SessionState getSessionStateSnapshot() { + return this.sessionState.get(); + } + /** * Begins a transaction.
* The transaction ends when a {@link BigQueryConnection#commit()} or {@link @@ -665,19 +679,21 @@ Map getLabels() { */ private void beginTransaction() { LOG.finer("++enter++"); + SessionState snapshot = this.sessionState.get(); QueryJobConfiguration.Builder transactionBeginJobConfig = QueryJobConfiguration.newBuilder("BEGIN TRANSACTION;"); try { - if (this.sessionInfoConnectionProperty != null) { - transactionBeginJobConfig.setConnectionProperties(this.queryProperties); + + if (snapshot.sessionInfo != null) { + transactionBeginJobConfig.setConnectionProperties(snapshot.queryProperties); } else { transactionBeginJobConfig.setCreateSession(true); - this.isSessionCreatedByDriver = true; + markSessionCreatedByDriver(); } TableResult transactionResult = this.bigQuery.query(transactionBeginJobConfig.build()); - if (this.sessionInfoConnectionProperty == null + if (this.sessionState.get().sessionInfo == null && transactionResult != null - && transactionResult.getSessionInfo()!= null) { + && transactionResult.getSessionInfo() != null) { updateSessionInfo(transactionResult.getSessionInfo().getSessionId()); } this.transactionStarted = true; @@ -686,33 +702,22 @@ private void beginTransaction() { } } - synchronized void updateSessionInfo(String sessionId) { - LOG.fine("++enter++ "); - if (sessionId != null && !sessionId.isEmpty()) { - if (this.sessionInfoConnectionProperty == null - || !sessionId.equals(this.sessionInfoConnectionProperty.getValue())) { - ConnectionProperty sessionProperty = - ConnectionProperty.newBuilder().setKey("session_id").setValue(sessionId).build(); - this.sessionInfoConnectionProperty = sessionProperty; - List updated = - this.queryProperties != null - ? new ArrayList<>(this.queryProperties) - : new ArrayList<>(); - boolean found = false; - for (int i = 0; i < updated.size(); i++) { - if ("session_id".equalsIgnoreCase(updated.get(i).getKey())) { - updated.set(i, sessionProperty); - found = true; - break; - } - } - if (!found) { - updated.add(sessionProperty); - } - LOG.info("Updated session info: " + sessionId); - this.queryProperties = Collections.unmodifiableList(updated); - } + void updateSessionInfo(String sessionId) { + LOG.finer("++enter++"); + if (sessionId == null || sessionId.isEmpty()) { + return; } + this.sessionState.updateAndGet( + current -> current.withSessionId(sessionId, current.createdByDriver)); + } + + // Marks the session as driver-owned, so it is aborted when the connection closes. + void markSessionCreatedByDriver() { + this.sessionState.updateAndGet( + current -> + current.sessionInfo == null + ? current + : new SessionState(current.sessionInfo, current.queryProperties, true)); } public boolean isTransactionStarted() { @@ -732,11 +737,11 @@ boolean isUnsupportedHTAPIFallback() { } public ConnectionProperty getSessionInfoConnectionProperty() { - return this.sessionInfoConnectionProperty; + return this.sessionState.get().sessionInfo; } boolean isSessionCreatedByDriver() { - return this.isSessionCreatedByDriver; + return this.sessionState.get().createdByDriver; } boolean isEnableHighThroughputAPI() { @@ -958,7 +963,7 @@ private void rollbackImpl() throws SQLException { try { QueryJobConfiguration transactionRollbackJobConfig = QueryJobConfiguration.newBuilder("ROLLBACK TRANSACTION;") - .setConnectionProperties(this.queryProperties) + .setConnectionProperties(this.sessionState.get().queryProperties) .build(); Job rollbackJob = this.bigQuery.create(JobInfo.of(transactionRollbackJobConfig)); rollbackJob.waitFor(); @@ -1068,8 +1073,9 @@ private void closeImpl() throws SQLException { } } - if (this.sessionInfoConnectionProperty != null && this.isSessionCreatedByDriver) { - abortSession(); + SessionState snapshot = this.sessionState.get(); + if (snapshot.sessionInfo != null && snapshot.createdByDriver) { + abortSession(snapshot); } boolean interrupted = Thread.currentThread().isInterrupted(); @@ -1181,13 +1187,38 @@ private void checkIfEnabledSession(String methodName) { private ConnectionProperty getSessionPropertyFromQueryProperties( Map queryPropertiesMap) { LOG.finer("++enter++"); - if (queryPropertiesMap != null && queryPropertiesMap.containsKey("session_id")) { - return ConnectionProperty.newBuilder() - .setKey("session_id") - .setValue(queryPropertiesMap.get("session_id")) - .build(); + + if (queryPropertiesMap == null) { + return null; } - return null; + Map.Entry match = null; + for (Map.Entry entry : queryPropertiesMap.entrySet()) { + if (!isSessionIdKey(entry.getKey())) { + continue; + } + if (match != null) { + // HashMap iteration order is undefined, so picking one would be non-deterministic. + throw new BigQueryJdbcRuntimeException( + String.format( + "QueryProperties contains multiple '%s' entries differing only by case ('%s' and" + + " '%s'). Specify exactly one.", + SESSION_ID_KEY, match.getKey(), entry.getKey())); + } + match = entry; + } + if (match == null) { + return null; + } + if (!SESSION_ID_KEY.equals(match.getKey())) { + LOG.warning( + "Normalizing QueryProperties key '%s' to '%s'; BigQuery connection property keys are" + + " case-sensitive.", + match.getKey(), SESSION_ID_KEY); + } + return ConnectionProperty.newBuilder() + .setKey(SESSION_ID_KEY) + .setValue(match.getValue()) + .build(); } private List convertMapToConnectionPropertiesList( @@ -1196,11 +1227,9 @@ private List convertMapToConnectionPropertiesList( List connectionProperties = new ArrayList(); if (queryPropertiesMap != null) { for (Map.Entry entry : queryPropertiesMap.entrySet()) { + String key = isSessionIdKey(entry.getKey()) ? SESSION_ID_KEY : entry.getKey(); connectionProperties.add( - ConnectionProperty.newBuilder() - .setKey(entry.getKey()) - .setValue(entry.getValue()) - .build()); + ConnectionProperty.newBuilder().setKey(key).setValue(entry.getValue()).build()); } } return Collections.unmodifiableList(connectionProperties); @@ -1456,7 +1485,7 @@ private void commitTransaction() { try { QueryJobConfiguration transactionCommitJobConfig = QueryJobConfiguration.newBuilder("COMMIT TRANSACTION;") - .setConnectionProperties(this.queryProperties) + .setConnectionProperties(this.sessionState.get().queryProperties) .build(); Job commitJob = this.bigQuery.create(JobInfo.of(transactionCommitJobConfig)); commitJob.waitFor(); @@ -1466,13 +1495,12 @@ private void commitTransaction() { } } - private void abortSession() { + private void abortSession(SessionState snapshot) { try { - LOG.fine( - "Aborting session on connection close: " + this.sessionInfoConnectionProperty.getValue()); + LOG.fine("Aborting session on connection close: %s", snapshot.sessionInfo.getValue()); QueryJobConfiguration abortSessionJobConfig = QueryJobConfiguration.newBuilder("CALL BQ.ABORT_SESSION();") - .setConnectionProperties(this.queryProperties) + .setConnectionProperties(snapshot.queryProperties) .build(); this.bigQuery.query(abortSessionJobConfig); } catch (InterruptedException ex) { @@ -1483,17 +1511,7 @@ private void abortSession() { "Failed to abort session during session abort (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.isSessionCreatedByDriver = false; + this.sessionState.updateAndGet(SessionState::withoutSession); this.transactionStarted = false; } } @@ -1594,4 +1612,82 @@ public T unwrap(Class iface) throws SQLException { public boolean isWrapperFor(Class iface) throws SQLException { return iface != null && iface.isInstance(this); } + + /** Returns whether {@code key} is the {@code session_id} property, ignoring case. */ + private static boolean isSessionIdKey(String key) { + return SESSION_ID_KEY.equalsIgnoreCase(key); + } + + /** + * Immutable snapshot of session-scoped connection state. + * + *

Grouping these values into a single object lets them be published with one atomic write, so + * a reader can never observe the session property updated without the matching {@code session_id} + * entry in the query property list. + */ + static final class SessionState { + + /** The active {@code session_id} property, or {@code null} when no session is active. */ + final ConnectionProperty sessionInfo; + + /** Unmodifiable properties, including {@code session_id} when a session is active. */ + final List queryProperties; + + /** + * Whether the driver created the session and is therefore responsible for aborting it on close. + * False by default, and for user-supplied sessions. + */ + final boolean createdByDriver; + + SessionState( + ConnectionProperty sessionInfo, + List queryProperties, + boolean createdByDriver) { + this.sessionInfo = sessionInfo; + this.queryProperties = queryProperties; + this.createdByDriver = createdByDriver; + } + + /** A state with no session and no query properties. Also the default for mocked connections. */ + static SessionState empty() { + return new SessionState(null, Collections.emptyList(), false); + } + + /** + * Returns a copy of this state with {@code session_id} set to {@code sessionId}, collapsing any + * pre-existing entries that differ only by case. + */ + SessionState withSessionId(String sessionId, boolean createdByDriver) { + ConnectionProperty session = + ConnectionProperty.newBuilder().setKey(SESSION_ID_KEY).setValue(sessionId).build(); + List updated = new ArrayList<>(this.queryProperties.size() + 1); + boolean replaced = false; + for (ConnectionProperty existing : this.queryProperties) { + if (isSessionIdKey(existing.getKey())) { + if (!replaced) { + updated.add(session); + replaced = true; + } + // Any further case-variant duplicates are dropped. + } else { + updated.add(existing); + } + } + if (!replaced) { + updated.add(session); + } + return new SessionState(session, Collections.unmodifiableList(updated), createdByDriver); + } + + /** Returns a copy of this state with every {@code session_id} entry removed. */ + SessionState withoutSession() { + List updated = new ArrayList<>(this.queryProperties.size()); + for (ConnectionProperty existing : this.queryProperties) { + if (!isSessionIdKey(existing.getKey())) { + updated.add(existing); + } + } + return new SessionState(null, Collections.unmodifiableList(updated), 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 fc6e6e751ce1..e56d84831fe9 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 @@ -196,7 +196,8 @@ private BigQuerySettings generateBigQuerySettings() { querySettings.setUseQueryCache(this.connection.isUseQueryCache()); querySettings.setQueryDialect(this.connection.getQueryDialect()); querySettings.setKmsKeyName(this.connection.getKmsKeyName()); - querySettings.setQueryProperties(this.connection.getQueryProperties()); + BigQueryConnection.SessionState snapshot = this.connection.getSessionStateSnapshot(); + querySettings.setQueryProperties(snapshot.queryProperties); querySettings.setAllowLargeResults(this.connection.isAllowLargeResults()); if (this.connection.getJobTimeoutInSeconds() > 0) { querySettings.setJobTimeoutMs(this.connection.getJobTimeoutInSeconds() * 1000L); @@ -212,8 +213,7 @@ private BigQuerySettings generateBigQuerySettings() { // only create session if enable session and session info is null if (this.connection.isSessionEnabled()) { querySettings.setEnableSession(this.connection.isSessionEnabled()); - querySettings.setSessionInfoConnectionProperty( - this.connection.getSessionInfoConnectionProperty()); + querySettings.setSessionInfoConnectionProperty(snapshot.sessionInfo); } querySettings.setUseWriteAPI(this.connection.isEnableWriteAPI()); querySettings.setWriteAPIActivationRowCount(this.connection.getWriteAPIActivationRowCount()); @@ -1508,9 +1508,10 @@ QueryJobConfiguration.Builder getJobConfig(String query) { queryConfigBuilder.setUseQueryCache(this.querySettings.getUseQueryCache()); queryConfigBuilder.setMaxResults(this.querySettings.getMaxResultPerPage()); + BigQueryConnection.SessionState snapshot = this.connection.getSessionStateSnapshot(); ConnectionProperty sessionProperty = this.connection != null - ? this.connection.getSessionInfoConnectionProperty() + ? snapshot.sessionInfo : this.querySettings.getSessionInfoConnectionProperty(); boolean isSessionEnabled = this.connection != null @@ -1518,7 +1519,7 @@ QueryJobConfiguration.Builder getJobConfig(String query) { : this.querySettings.isEnableSession(); List queryProperties = this.connection != null - ? this.connection.getQueryProperties() + ? snapshot.queryProperties : this.querySettings.getQueryProperties(); List props = @@ -1526,13 +1527,14 @@ QueryJobConfiguration.Builder getJobConfig(String query) { if (sessionProperty != null) { boolean hasSessionId = - props.stream().anyMatch(cp -> "session_id".equalsIgnoreCase(cp.getKey())); + props.stream() + .anyMatch(cp -> BigQueryConnection.SESSION_ID_KEY.equalsIgnoreCase(cp.getKey())); if (!hasSessionId) { props.add(sessionProperty); } } else if (isSessionEnabled) { queryConfigBuilder.setCreateSession(true); - this.connection.isSessionCreatedByDriver = true; + this.connection.markSessionCreatedByDriver(); } if (!props.isEmpty()) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryCallableStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryCallableStatementTest.java index 72f8ee067067..bfe907b934f4 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryCallableStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryCallableStatementTest.java @@ -20,6 +20,7 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import com.google.cloud.bigquery.StandardSQLTypeName; @@ -44,6 +45,9 @@ public class BigQueryCallableStatementTest { @BeforeEach public void setUp() throws IOException, SQLException { bigQueryConnection = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()) + .when(bigQueryConnection) + .getSessionStateSnapshot(); } @Test 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 45998a2379af..0c9fccfedea9 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 @@ -43,12 +43,11 @@ 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.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; import com.google.cloud.logging.Logging; @@ -828,27 +827,23 @@ public void testUserSuppliedSessionId() throws Exception { 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; + connection.markSessionCreatedByDriver(); assertTrue(connection.isSessionCreatedByDriver()); 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()); + ArgumentCaptor jobCaptor = + ArgumentCaptor.forClass(QueryJobConfiguration.class); + verify(mockBigQuery).query(jobCaptor.capture()); + assertEquals("CALL BQ.ABORT_SESSION();", jobCaptor.getValue().getQuery()); assertNull(connection.getSessionInfoConnectionProperty()); assertFalse(connection.isSessionCreatedByDriver()); assertTrue(connection.isClosed()); - } } + @Test public void testEnableTimestampPicosDefault() throws Exception { try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { @@ -870,29 +865,57 @@ public void testCloseWithUserSuppliedSessionDoesNotAbortSession() throws Excepti connection.close(); - verify(mockBigQuery, never()).create(any(JobInfo.class)); + verify(mockBigQuery, never()).query(any(QueryJobConfiguration.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; + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { + BigQuery mockBigQuery = mock(BigQuery.class); + connection.bigQuery = mockBigQuery; - connection.close(); + connection.close(); - verify(mockBigQuery, never()).create(any(JobInfo.class)); - assertTrue(connection.isClosed()); - } - } + verify(mockBigQuery, never()).query(any(QueryJobConfiguration.class)); + assertTrue(connection.isClosed()); + } + } - @Test + @Test public void testEnableTimestampPicosConfigured() throws Exception { String url = BASE_URL + "EnableTimestampPicos=1;"; try (BigQueryConnection connection = new BigQueryConnection(url)) { assertTrue(connection.isEnableTimestampPicos()); } } + + @Test + public void testSessionIdIsMatchedRegardlessOfCase() throws Exception { + String url = BASE_URL + ";QueryProperties=Session_Id=abc123"; + try (BigQueryConnection connection = new BigQueryConnection(url)) { + assertNotNull(connection.getSessionInfoConnectionProperty()); + assertEquals("abc123", connection.getSessionInfoConnectionProperty().getValue()); + } + } + + @Test + public void testSessionIdKeyIsNormalizedInQueryProperties() throws Exception { + String url = BASE_URL + ";QueryProperties=SESSION_ID=abc123"; + try (BigQueryConnection connection = new BigQueryConnection(url)) { + // The key reaches BigQuery lowercased + assertTrue( + connection.getQueryProperties().stream() + .anyMatch(cp -> "session_id".equals(cp.getKey()))); + } + } + + @Test + public void testDuplicateSessionIdKeysAreRejected() { + String url = BASE_URL + ";QueryProperties=session_id=a,Session_Id=b"; + BigQueryJdbcRuntimeException ex = + assertThrows(BigQueryJdbcRuntimeException.class, () -> new BigQueryConnection(url)); + assertTrue(ex.getMessage().contains("multiple 'session_id' entries")); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java index 6cb46cb8de21..eb0e06390b09 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java @@ -22,6 +22,7 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -60,6 +61,8 @@ public void testExtractConnectionIdFromConnection() throws SQLException { @Test public void testExtractConnectionIdFromStatement() throws SQLException { BigQueryConnection mockConn = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()).when(mockConn).getSessionStateSnapshot(); + when(mockConn.getBigQuery()).thenReturn(mock(com.google.cloud.bigquery.BigQuery.class)); BigQueryStatement stmt = new BigQueryStatement(mockConn); @@ -95,6 +98,8 @@ public void testExtractConnectionIdFromDatabaseMetaData() throws SQLException { @Test public void testExtractConnectionIdFromResultSetMetaData() throws SQLException { BigQueryConnection mockConn = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()).when(mockConn).getSessionStateSnapshot(); + BigQueryStatement stmt = new BigQueryStatement(mockConn); stmt.connectionId = "conn-uuid-999"; diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java index 466d21b870d4..6b02d5754d09 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java @@ -22,6 +22,7 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -60,6 +61,8 @@ public class BigQueryPreparedStatementSettersTest { public void setUp() throws Exception { connection = mock(BigQueryConnection.class); when(connection.getQueryDialect()).thenReturn("SQL"); + doReturn(BigQueryConnection.SessionState.empty()).when(connection).getSessionStateSnapshot(); + preparedStatement = new BigQueryPreparedStatement(connection, "SELECT ?, ?, ?, ?, ?"); } 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 f8dcb3ac6e67..78d493be9344 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 @@ -223,6 +223,10 @@ public void setUp() throws IOException, SQLException { doReturn(1000L).when(bigQueryConnection).getMaxResults(); testExecutorService = Executors.newSingleThreadExecutor(); doReturn(testExecutorService).when(bigQueryConnection).getExecutorService(); + doReturn(BigQueryConnection.SessionState.empty()) + .when(bigQueryConnection) + .getSessionStateSnapshot(); + bigQueryStatement = new BigQueryStatement(bigQueryConnection); VectorSchemaRoot vectorSchemaRoot = getTestVectorSchemaRoot(); arrowSchema = From a72f736c35fd7e34dd206eda8ae3b59c2917c882 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Thu, 17 Sep 2026 20:03:15 -0400 Subject: [PATCH 07/10] make the session write onec per connection --- .../bigquery/jdbc/BigQueryConnection.java | 41 +++++++++++++--- .../bigquery/jdbc/BigQueryStatement.java | 19 +++---- .../bigquery/jdbc/BigQueryConnectionTest.java | 49 ++++++++++++++----- .../bigquery/jdbc/BigQueryStatementTest.java | 2 +- 4 files changed, 79 insertions(+), 32 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 735450c9a93f..b3e0b5e9a97f 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 @@ -694,7 +694,7 @@ private void beginTransaction() { if (this.sessionState.get().sessionInfo == null && transactionResult != null && transactionResult.getSessionInfo() != null) { - updateSessionInfo(transactionResult.getSessionInfo().getSessionId()); + initSessionInfo(transactionResult.getSessionInfo().getSessionId()); } this.transactionStarted = true; } catch (InterruptedException ex) { @@ -702,20 +702,49 @@ private void beginTransaction() { } } - void updateSessionInfo(String sessionId) { + /** + * Establishes the session for this connection. + * + *

A connection's session is write-once: it is fixed either from a user-supplied {@code + * session_id} at construction or by the first job that creates one, and does not change until the + * connection closes. Repeat calls with the same id are no-ops. A call with a different + * id is ignored and logged, because silently swapping the session would strand the original and + * invalidate any open transaction. + */ + void initSessionInfo(String sessionId) { LOG.finer("++enter++"); if (sessionId == null || sessionId.isEmpty()) { return; } - this.sessionState.updateAndGet( - current -> current.withSessionId(sessionId, current.createdByDriver)); + SessionState previous = + this.sessionState.getAndUpdate( + current -> + current.sessionInfo == null + ? current.withSessionId(sessionId, current.createdByDriver) + : current); + if (previous.sessionInfo == null) { + LOG.info("Established session: %s", sessionId); + } else if (!sessionId.equals(previous.sessionInfo.getValue())) { + LOG.warning( + "Ignoring attempt to change session from '%s' to '%s'; a connection's session is" + + " immutable for its lifetime.", + previous.sessionInfo.getValue(), sessionId); + } } - // Marks the session as driver-owned, so it is aborted when the connection closes. + /** + * Marks this connection's session as driver-owned, so that it is aborted when the connection + * closes. + * + *

This is called when a job is configured with {@code createSession=true}, which is + * necessarily before the session id is known. {@link #initSessionInfo} carries the flag + * forward onto the session once BigQuery returns its id. If the job fails and no session is ever + * established, the flag is harmless: {@link #close} aborts only when an id is also present. + */ void markSessionCreatedByDriver() { this.sessionState.updateAndGet( current -> - current.sessionInfo == null + current.createdByDriver ? current : new SessionState(current.sessionInfo, current.queryProperties, true)); } 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 e56d84831fe9..203c7e30fc02 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 @@ -614,7 +614,7 @@ private void saveSessionIdIfPresent(TableResult tableResult) { if (tableResult.getSessionInfo() != null) { String sessionId = tableResult.getSessionInfo().getSessionId(); if (sessionId != null && !sessionId.isEmpty()) { - this.connection.updateSessionInfo(sessionId); + this.connection.initSessionInfo(sessionId); } } } @@ -1508,19 +1508,12 @@ QueryJobConfiguration.Builder getJobConfig(String query) { queryConfigBuilder.setUseQueryCache(this.querySettings.getUseQueryCache()); queryConfigBuilder.setMaxResults(this.querySettings.getMaxResultPerPage()); + // Only reachable from execute paths, which call checkClosed() first, so this.connection is + // non-null here; close() is the only thing that nulls it. BigQueryConnection.SessionState snapshot = this.connection.getSessionStateSnapshot(); - ConnectionProperty sessionProperty = - this.connection != null - ? snapshot.sessionInfo - : this.querySettings.getSessionInfoConnectionProperty(); - boolean isSessionEnabled = - this.connection != null - ? this.connection.isSessionEnabled() - : this.querySettings.isEnableSession(); - List queryProperties = - this.connection != null - ? snapshot.queryProperties - : this.querySettings.getQueryProperties(); + ConnectionProperty sessionProperty = snapshot.sessionInfo; + boolean isSessionEnabled = this.connection.isSessionEnabled(); + List queryProperties = snapshot.queryProperties; List props = queryProperties != null ? new ArrayList<>(queryProperties) : new ArrayList<>(); 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 0c9fccfedea9..6dd4b4b64302 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 @@ -28,6 +28,7 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.never; @@ -43,9 +44,13 @@ 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.JobStatistics.SessionInfo; import com.google.cloud.bigquery.Project; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; +import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; @@ -780,11 +785,11 @@ public void testGetDiscoveredProjects_OtherExceptionThrown() throws Exception { } @Test - public void testUpdateSessionInfo() throws Exception { + public void testSessionIdIsWriteOnce() throws Exception { try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { assertNull(connection.getSessionInfoConnectionProperty()); - connection.updateSessionInfo("test_session_id_1"); + connection.initSessionInfo("test_session_id_1"); assertNotNull(connection.getSessionInfoConnectionProperty()); assertEquals("session_id", connection.getSessionInfoConnectionProperty().getKey()); assertEquals("test_session_id_1", connection.getSessionInfoConnectionProperty().getValue()); @@ -798,9 +803,10 @@ public void testUpdateSessionInfo() throws Exception { && "test_session_id_1".equals(cp.getValue())); assertTrue(found, "queryProperties should contain session_id property"); - // Update to a new session ID and ensure it updates without creating duplicates - connection.updateSessionInfo("test_session_id_2"); - assertEquals("test_session_id_2", connection.getSessionInfoConnectionProperty().getValue()); + // A connection's session is write-once: a second, different id is ignored rather than + // silently swapping the session and stranding the original. + connection.initSessionInfo("test_session_id_2"); + assertEquals("test_session_id_1", connection.getSessionInfoConnectionProperty().getValue()); long count = connection.getQueryProperties().stream() .filter(cp -> "session_id".equalsIgnoreCase(cp.getKey())) @@ -824,20 +830,39 @@ public void testUserSuppliedSessionId() throws Exception { } @Test - public void testCloseWithActiveSessionAbortsSession() throws Exception { - try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { + public void testCloseAbortsSessionCreatedByDriver() throws Exception { + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL + ";EnableSession=1")) { BigQuery mockBigQuery = mock(BigQuery.class); connection.bigQuery = mockBigQuery; - connection.updateSessionInfo("test_session_id_to_abort"); - connection.markSessionCreatedByDriver(); + // BEGIN TRANSACTION asks BigQuery to create the session, and the result carries the new id. + SessionInfo sessionInfo = mock(SessionInfo.class); + when(sessionInfo.getSessionId()).thenReturn("driver_created_session"); + TableResult beginResult = mock(TableResult.class); + when(beginResult.getSessionInfo()).thenReturn(sessionInfo); + when(mockBigQuery.query(any(QueryJobConfiguration.class))).thenReturn(beginResult); + + // close() rolls the open transaction back before it aborts the session. + Job rollbackJob = mock(Job.class); + when(mockBigQuery.create(any(JobInfo.class))).thenReturn(rollbackJob); + when(rollbackJob.waitFor()).thenReturn(rollbackJob); + + // Drives beginTransaction(), which claims ownership before the id is known. + connection.setAutoCommit(false); + assertTrue(connection.isSessionCreatedByDriver()); + assertEquals( + "driver_created_session", connection.getSessionInfoConnectionProperty().getValue()); + connection.close(); + // close() also rolls back the open transaction, so match on the abort specifically. ArgumentCaptor jobCaptor = ArgumentCaptor.forClass(QueryJobConfiguration.class); - verify(mockBigQuery).query(jobCaptor.capture()); - assertEquals("CALL BQ.ABORT_SESSION();", jobCaptor.getValue().getQuery()); + verify(mockBigQuery, atLeastOnce()).query(jobCaptor.capture()); + assertTrue( + jobCaptor.getAllValues().stream() + .anyMatch(config -> "CALL BQ.ABORT_SESSION();".equals(config.getQuery()))); assertNull(connection.getSessionInfoConnectionProperty()); assertFalse(connection.isSessionCreatedByDriver()); assertTrue(connection.isClosed()); @@ -865,7 +890,7 @@ public void testCloseWithUserSuppliedSessionDoesNotAbortSession() throws Excepti connection.close(); - verify(mockBigQuery, never()).query(any(QueryJobConfiguration.class)); + verify(mockBigQuery, never()).create(any(JobInfo.class)); assertTrue(connection.isClosed()); } } 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 78d493be9344..4411eab3b7b6 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 @@ -1168,7 +1168,7 @@ public void testSessionIdSavedFromTableResult() throws Exception { QueryJobConfiguration.newBuilder("CREATE TEMP TABLE t1 (id INT64)").build(); bigQueryStatement.executeJob(jobConfig); - verify(bigQueryConnection).updateSessionInfo("session_xyz_123"); + verify(bigQueryConnection).initSessionInfo("session_xyz_123"); } @Test From b74de6de33968e320bdb04b0eec4d908ca96195a Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Thu, 17 Sep 2026 20:40:20 -0400 Subject: [PATCH 08/10] fix tests --- .../cloud/bigquery/jdbc/BigQueryConnectionTest.java | 2 +- .../jdbc/BigQueryPreparedStatementSettersTest.java | 12 ++++++++++++ .../cloud/bigquery/jdbc/BigQueryStatementTest.java | 3 +++ 3 files changed, 16 insertions(+), 1 deletion(-) 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 8a8f68493e79..62d01229efea 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 @@ -44,10 +44,10 @@ import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQueryException; +import com.google.cloud.bigquery.DataFormatOptions; import com.google.cloud.bigquery.Job; import com.google.cloud.bigquery.JobInfo; import com.google.cloud.bigquery.JobStatistics.SessionInfo; -import com.google.cloud.bigquery.DataFormatOptions; import com.google.cloud.bigquery.Project; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java index 8c85a117adb7..b9a14147a428 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatementSettersTest.java @@ -332,6 +332,9 @@ public void testCreateJsonRowWithSetObjectNull() throws Exception { @Test public void testSetObjectWithTimestampStringAndTypesTimestamp_picosEnabled() throws Exception { BigQueryConnection picosConnection = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()) + .when(picosConnection) + .getSessionStateSnapshot(); doReturn(true).when(picosConnection).isEnableTimestampPicos(); doReturn(BigQueryJdbcUrlUtility.DEFAULT_QUERY_DIALECT_VALUE) .when(picosConnection) @@ -356,6 +359,9 @@ public void testSetObjectWithTimestampStringAndTypesTimestamp_picosEnabled() thr @Test public void testSetTimestamp_picosEnabledPreservesNanoseconds() throws Exception { BigQueryConnection picosConnection = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()) + .when(picosConnection) + .getSessionStateSnapshot(); doReturn(true).when(picosConnection).isEnableTimestampPicos(); doReturn(BigQueryJdbcUrlUtility.DEFAULT_QUERY_DIALECT_VALUE) .when(picosConnection) @@ -378,6 +384,9 @@ public void testSetTimestamp_picosEnabledPreservesNanoseconds() throws Exception @Test public void testSetTimestamp_picosDisabledTruncatesToMicroseconds() throws Exception { BigQueryConnection nonPicosConnection = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()) + .when(nonPicosConnection) + .getSessionStateSnapshot(); doReturn(false).when(nonPicosConnection).isEnableTimestampPicos(); doReturn(BigQueryJdbcUrlUtility.DEFAULT_QUERY_DIALECT_VALUE) .when(nonPicosConnection) @@ -399,6 +408,9 @@ public void testSetTimestamp_picosDisabledTruncatesToMicroseconds() throws Excep @Test public void testBatchConfiguration_withEnableTimestampPicos() throws Exception { BigQueryConnection picosConnection = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()) + .when(picosConnection) + .getSessionStateSnapshot(); doReturn(true).when(picosConnection).isEnableTimestampPicos(); doReturn(BigQueryJdbcUrlUtility.DEFAULT_QUERY_DIALECT_VALUE) .when(picosConnection) 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 4c406ed584ba..6a20a1c4dfde 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 @@ -510,6 +510,7 @@ public void testGetJobConfigWithExtraLabels() { @Test public void testExecute_legacySqlWithEnableTimestampPicos_throwsException() { BigQueryConnection mockConn = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()).when(mockConn).getSessionStateSnapshot(); doReturn("BIG_QUERY").when(mockConn).getQueryDialect(); doReturn(true).when(mockConn).isEnableTimestampPicos(); @@ -524,6 +525,7 @@ public void testExecute_legacySqlWithEnableTimestampPicos_throwsException() { @Test public void testGetJobConfig_standardSql_setsUseLegacySqlFalse() { BigQueryConnection mockConn = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()).when(mockConn).getSessionStateSnapshot(); doReturn("SQL").when(mockConn).getQueryDialect(); BigQueryStatement statement = new BigQueryStatement(mockConn); @@ -536,6 +538,7 @@ public void testGetJobConfig_standardSql_setsUseLegacySqlFalse() { @Test public void testGetJobConfig_legacySql_setsUseLegacySqlTrue() { BigQueryConnection mockConn = mock(BigQueryConnection.class); + doReturn(BigQueryConnection.SessionState.empty()).when(mockConn).getSessionStateSnapshot(); doReturn("BIG_QUERY").when(mockConn).getQueryDialect(); BigQueryStatement statement = new BigQueryStatement(mockConn); From 69310bf01a449fe58803b2c620d02b6653930cdc Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 18 Sep 2026 09:26:52 -0400 Subject: [PATCH 09/10] fix exception --- .../java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java | 2 +- 1 file changed, 1 insertion(+), 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 73de3132d0af..005054a35e05 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 @@ -1228,7 +1228,7 @@ private ConnectionProperty getSessionPropertyFromQueryProperties( } if (match != null) { // HashMap iteration order is undefined, so picking one would be non-deterministic. - throw new BigQueryJdbcRuntimeException( + throw new BigQueryJdbcException( String.format( "QueryProperties contains multiple '%s' entries differing only by case ('%s' and" + " '%s'). Specify exactly one.", From 17694562cb51b432bf9a0febef2b951e96457b74 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Fri, 18 Sep 2026 10:19:13 -0400 Subject: [PATCH 10/10] fix exceptions --- .../bigquery/jdbc/BigQueryConnection.java | 6 ++-- .../bigquery/jdbc/BigQueryConnectionTest.java | 31 +++++++++---------- 2 files changed, 18 insertions(+), 19 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 005054a35e05..b4d2f0a4f956 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 @@ -256,11 +256,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { private final AtomicReference sessionState = new AtomicReference<>(new SessionState(null, Collections.emptyList(), false)); - BigQueryConnection(String url) throws IOException { + BigQueryConnection(String url) throws SQLException { this(url, DataSource.fromUrl(url)); } - BigQueryConnection(String url, DataSource ds) throws IOException { + BigQueryConnection(String url, DataSource ds) throws SQLException { this.connectionId = UUID.randomUUID().toString(); Baggage baggage = Baggage.builder() @@ -1215,7 +1215,7 @@ private void checkIfEnabledSession(String methodName) { } private ConnectionProperty getSessionPropertyFromQueryProperties( - Map queryPropertiesMap) { + Map queryPropertiesMap) throws SQLException { LOG.finer("++enter++"); if (queryPropertiesMap == null) { 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 62d01229efea..5f3846873b53 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 @@ -53,7 +53,6 @@ import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; -import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; import com.google.cloud.logging.Logging; @@ -247,7 +246,7 @@ public void testWriteAPIConnectionProperties() throws SQLException { assertFalse(connectionDefault.enableWriteAPI); assertEquals(3, connectionDefault.writeAPIActivationRowCount); assertEquals(1000, connectionDefault.writeAPIAppendRowCount); - } catch (IOException | SQLException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } @@ -262,13 +261,13 @@ public void testWriteAPIConnectionProperties() throws SQLException { assertTrue(connection.enableWriteAPI); assertEquals(6, connection.writeAPIActivationRowCount); assertEquals(500, connection.writeAPIAppendRowCount); - } catch (IOException | SQLException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } } @Test - public void testTimestampPicosControlsDataFormatOptions() throws IOException, SQLException { + public void testTimestampPicosControlsDataFormatOptions() throws SQLException { try (BigQueryConnection connection = new BigQueryConnection(BASE_URL + "EnableTimestampPicos=1;")) { assertEquals( @@ -299,13 +298,13 @@ public void testGetWriteClient() throws SQLException { BigQueryWriteClient writeClient = connectionDefault.getBigQueryWriteClient(); assertNotNull(writeClient); assertFalse(writeClient.isShutdown()); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } } @Test - public void testAdditionalProjects() throws IOException, BigQueryJdbcException { + public void testAdditionalProjects() throws BigQueryJdbcException { String url1 = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + "OAuthType=2;ProjectId=MyBigQueryProject;" @@ -316,7 +315,7 @@ public void testAdditionalProjects() throws IOException, BigQueryJdbcException { String additionalProjects1 = conn1.getAdditionalProjects(); assertNotNull(additionalProjects1); assertEquals("projA,projB", additionalProjects1); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } String url2 = @@ -329,13 +328,13 @@ public void testAdditionalProjects() throws IOException, BigQueryJdbcException { String additionalProjects2 = conn2.getAdditionalProjects(); assertNotNull(additionalProjects2); assertEquals("projX", additionalProjects2); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } } @Test - public void testFilterTablesOnDefaultDatasetProperty() throws SQLException, IOException { + public void testFilterTablesOnDefaultDatasetProperty() throws SQLException { // Test default value String urlDefault = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" @@ -346,7 +345,7 @@ public void testFilterTablesOnDefaultDatasetProperty() throws SQLException, IOEx assertFalse( connectionDefault.isFilterTablesOnDefaultDataset(), "Default value for FilterTablesOnDefaultDataset should be false"); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } @@ -361,13 +360,13 @@ public void testFilterTablesOnDefaultDatasetProperty() throws SQLException, IOEx assertTrue( connectionTrue.isFilterTablesOnDefaultDataset(), "FilterTablesOnDefaultDataset should be true when set to 1"); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } } @Test - public void testRequestGoogleDriveScopeProperty() throws IOException, SQLException { + public void testRequestGoogleDriveScopeProperty() throws SQLException { // Test enabled String urlEnabled = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" @@ -380,7 +379,7 @@ public void testRequestGoogleDriveScopeProperty() throws IOException, SQLExcepti 1, connectionEnabled.isRequestGoogleDriveScope(), "RequestGoogleDriveScope should be enabled when set to 1"); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } @@ -396,7 +395,7 @@ public void testRequestGoogleDriveScopeProperty() throws IOException, SQLExcepti 0, connectionDisabled.isRequestGoogleDriveScope(), "RequestGoogleDriveScope should be disabled when set to 0"); - } catch (SQLException | IOException e) { + } catch (SQLException e) { throw new BigQueryJdbcException(e); } } @@ -958,8 +957,8 @@ public void testSessionIdKeyIsNormalizedInQueryProperties() throws Exception { @Test public void testDuplicateSessionIdKeysAreRejected() { String url = BASE_URL + ";QueryProperties=session_id=a,Session_Id=b"; - BigQueryJdbcRuntimeException ex = - assertThrows(BigQueryJdbcRuntimeException.class, () -> new BigQueryConnection(url)); + BigQueryJdbcException ex = + assertThrows(BigQueryJdbcException.class, () -> new BigQueryConnection(url)); assertTrue(ex.getMessage().contains("multiple 'session_id' entries")); } }