diff --git a/java-bigquery-jdbc/docs/USER_GUIDE.md b/java-bigquery-jdbc/docs/USER_GUIDE.md index 64ed0852a079..1549d1b5fb4a 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). | ### Data Types & Extended Precision Properties @@ -313,58 +315,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. --- @@ -382,7 +352,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"; @@ -411,6 +381,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 423c4f6dd65c..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 @@ -37,6 +37,7 @@ 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.exception.BigQueryJdbcSqlFeatureNotSupportedException; @@ -78,6 +79,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 +95,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,7 +185,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; boolean isClosed; DatasetId defaultDataset; String location; @@ -201,7 +205,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { long destinationDatasetExpirationTime; String kmsKeyName; String universeDomain; - private volatile List queryProperties; Map authProperties; Map overrideProperties; Map proxyProperties; @@ -245,11 +248,19 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { private final ExecutorService metadataExecutor; private final ExecutorService queryExecutor; - BigQueryConnection(String url) throws IOException { + /** + * 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 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() @@ -258,7 +269,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()) { @@ -364,9 +375,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(); @@ -600,10 +613,6 @@ String getKmsKeyName() { return this.kmsKeyName; } - List getQueryProperties() { - return this.queryProperties; - } - public String getLocation() { checkClosed(); return this.location; @@ -654,6 +663,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 @@ -663,22 +680,22 @@ 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); + markSessionCreatedByDriver(); } - Job job = this.bigQuery.create(JobInfo.of(transactionBeginJobConfig.build())); - job = job.waitFor(); - Job transactionBeginJob = this.bigQuery.getJob(job.getJobId()); - if (this.sessionInfoConnectionProperty == null - && transactionBeginJob != null - && transactionBeginJob.getStatistics() != null - && transactionBeginJob.getStatistics().getSessionInfo() != null) { - updateSessionInfo(transactionBeginJob.getStatistics().getSessionInfo().getSessionId()); + TableResult transactionResult = this.bigQuery.query(transactionBeginJobConfig.build()); + if (this.sessionState.get().sessionInfo == null + && transactionResult != null + && transactionResult.getSessionInfo() != null) { + initSessionInfo(transactionResult.getSessionInfo().getSessionId()); } this.transactionStarted = true; } catch (InterruptedException ex) { @@ -686,35 +703,53 @@ 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); - } + /** + * 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; + } + 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 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.createdByDriver + ? current + : new SessionState(current.sessionInfo, current.queryProperties, true)); + } + public boolean isTransactionStarted() { return this.transactionStarted; } @@ -732,7 +767,11 @@ boolean isUnsupportedHTAPIFallback() { } public ConnectionProperty getSessionInfoConnectionProperty() { - return this.sessionInfoConnectionProperty; + return this.sessionState.get().sessionInfo; + } + + boolean isSessionCreatedByDriver() { + return this.sessionState.get().createdByDriver; } boolean isEnableHighThroughputAPI() { @@ -954,7 +993,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(); @@ -1064,6 +1103,11 @@ private void closeImpl() throws SQLException { } } + SessionState snapshot = this.sessionState.get(); + if (snapshot.sessionInfo != null && snapshot.createdByDriver) { + abortSession(snapshot); + } + boolean interrupted = Thread.currentThread().isInterrupted(); try { @@ -1171,15 +1215,40 @@ private void checkIfEnabledSession(String methodName) { } private ConnectionProperty getSessionPropertyFromQueryProperties( - Map queryPropertiesMap) { + Map queryPropertiesMap) throws SQLException { 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; + } + 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 BigQueryJdbcException( + 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 null; + return ConnectionProperty.newBuilder() + .setKey(SESSION_ID_KEY) + .setValue(match.getValue()) + .build(); } private List convertMapToConnectionPropertiesList( @@ -1188,11 +1257,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); @@ -1458,7 +1525,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(); @@ -1468,6 +1535,27 @@ private void commitTransaction() { } } + private void abortSession(SessionState snapshot) { + try { + LOG.fine("Aborting session on connection close: %s", snapshot.sessionInfo.getValue()); + QueryJobConfiguration abortSessionJobConfig = + QueryJobConfiguration.newBuilder("CALL BQ.ABORT_SESSION();") + .setConnectionProperties(snapshot.queryProperties) + .build(); + this.bigQuery.query(abortSessionJobConfig); + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + throw new BigQueryJdbcRuntimeException("Interrupted during session abort", ex); + } catch (BigQueryException ex) { + LOG.warning( + "Failed to abort session during session abort (session may have already ended): " + + ex.getMessage()); + } finally { + this.sessionState.updateAndGet(SessionState::withoutSession); + this.transactionStarted = false; + } + } + @Override public CallableStatement prepareCall(String sql) throws SQLException { checkClosed(); @@ -1564,4 +1652,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 2ee127a26af1..98880e9f6a14 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()); @@ -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,30 +1508,26 @@ QueryJobConfiguration.Builder getJobConfig(String query) { queryConfigBuilder.setUseQueryCache(this.querySettings.getUseQueryCache()); queryConfigBuilder.setMaxResults(this.querySettings.getMaxResultPerPage()); - ConnectionProperty sessionProperty = - this.connection != null - ? this.connection.getSessionInfoConnectionProperty() - : this.querySettings.getSessionInfoConnectionProperty(); - boolean isSessionEnabled = - this.connection != null - ? this.connection.isSessionEnabled() - : this.querySettings.isEnableSession(); - List queryProperties = - this.connection != null - ? this.connection.getQueryProperties() - : this.querySettings.getQueryProperties(); + // 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 = snapshot.sessionInfo; + boolean isSessionEnabled = this.connection.isSessionEnabled(); + List queryProperties = snapshot.queryProperties; List props = queryProperties != null ? new ArrayList<>(queryProperties) : new ArrayList<>(); 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.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 2eb66c9a192c..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 @@ -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; @@ -29,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; @@ -45,8 +45,13 @@ 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.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.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; @@ -71,6 +76,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 { @@ -240,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); } @@ -255,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( @@ -292,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;" @@ -309,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 = @@ -322,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;" @@ -339,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); } @@ -354,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;" @@ -373,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); } @@ -389,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); } } @@ -797,11 +803,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()); @@ -815,9 +821,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())) @@ -832,6 +839,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( @@ -839,6 +847,46 @@ public void testUserSuppliedSessionId() throws Exception { } } + @Test + public void testCloseAbortsSessionCreatedByDriver() throws Exception { + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL + ";EnableSession=1")) { + BigQuery mockBigQuery = mock(BigQuery.class); + connection.bigQuery = mockBigQuery; + + // 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, 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()); + } + } + @Test public void testEnableTimestampPicosDefault() throws Exception { try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { @@ -846,6 +894,38 @@ public void testEnableTimestampPicosDefault() throws Exception { } } + @Test + public void testCloseWithUserSuppliedSessionDoesNotAbortSession() throws Exception { + String urlWithSessionId = + BASE_URL + ";EnableSession=1;QueryProperties=session_id=user_supplied_session_999"; + try (BigQueryConnection connection = new BigQueryConnection(urlWithSessionId)) { + BigQuery mockBigQuery = mock(BigQuery.class); + connection.bigQuery = mockBigQuery; + + assertFalse(connection.isSessionCreatedByDriver()); + assertEquals( + "user_supplied_session_999", connection.getSessionInfoConnectionProperty().getValue()); + + connection.close(); + + verify(mockBigQuery, never()).create(any(JobInfo.class)); + assertTrue(connection.isClosed()); + } + } + + @Test + public void testCloseWithoutSessionDoesNotAbortSession() throws Exception { + try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) { + BigQuery mockBigQuery = mock(BigQuery.class); + connection.bigQuery = mockBigQuery; + + connection.close(); + + verify(mockBigQuery, never()).query(any(QueryJobConfiguration.class)); + assertTrue(connection.isClosed()); + } + } + @Test public void testEnableTimestampPicosConfigured() throws Exception { String url = BASE_URL + "EnableTimestampPicos=1;"; @@ -853,4 +933,32 @@ public void testEnableTimestampPicosConfigured() throws Exception { 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"; + BigQueryJdbcException ex = + assertThrows(BigQueryJdbcException.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 7e36b974b3b4..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 @@ -61,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 ?, ?, ?, ?, ?"); } @@ -330,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) @@ -354,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) @@ -376,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) @@ -397,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 317e0b0da921..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 @@ -224,6 +224,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 = @@ -506,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(); @@ -520,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); @@ -532,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); @@ -1203,7 +1210,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 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 999f2a7ad5f8..ea0799e2aea3 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 @@ -2935,4 +2935,30 @@ public void testPerConnectionLoggingE2E() throws SQLException, IOException { } } } + + @Test + public void testSessionAbortedOnConnectionClose() throws SQLException { + String sessionId; + try (Connection connection = DriverManager.getConnection(session_enabled_connection_uri)) { + try (Statement statement = connection.createStatement()) { + statement.execute("CREATE TEMP TABLE session_temp_table (id INT64);"); + } + BigQueryConnection bqConn = connection.unwrap(BigQueryConnection.class); + assertNotNull(bqConn.getSessionInfoConnectionProperty()); + sessionId = bqConn.getSessionInfoConnectionProperty().getValue(); + assertNotNull(sessionId); + } + + // After connection is closed, the session is aborted on the BigQuery server. + // Attaching to the same session_id in a new connection should fail when running a query. + String urlWithAbortedSession = + connection_uri + "EnableSession=1;QueryProperties=session_id=" + sessionId + ";"; + try (Connection newConnection = DriverManager.getConnection(urlWithAbortedSession)) { + try (Statement statement = newConnection.createStatement()) { + SQLException ex = + assertThrows( + SQLException.class, () -> statement.execute("SELECT * FROM session_temp_table;")); + } + } + } }