From 8b39570a118e96c7bbeb0129ca0509d67b05a957 Mon Sep 17 00:00:00 2001 From: Sergey Chernov Date: Thu, 18 Jun 2026 18:00:07 -0700 Subject: [PATCH 1/3] Fixed clearing session when do cancel on Statement in jdbc --- CHANGELOG.md | 4 +++- .../com/clickhouse/client/api/Session.java | 13 ++++++++-- .../client/api/insert/InsertSettings.java | 20 ++++++++++++++++ .../client/api/internal/CommonSettings.java | 24 ++++++++++++++----- .../client/api/query/QuerySettings.java | 15 ++++++++++++ .../com/clickhouse/client/ClientTests.java | 19 +++++++++++++++ docs/features.md | 2 +- .../com/clickhouse/jdbc/StatementTest.java | 14 +++++++++-- 8 files changed, 99 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dd347a3a8..e219cca49 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,8 @@ ### New Features +- **[client-v2]** Added `Session` API to encapsulate and manage ClickHouse session settings (`session_id`, `session_check`, `session_timeout`, `session_timezone`) as a reusable object. The `Session` instance can be applied to any request settings using `applyTo()`, and session state can be cleared via `clearSession()`. Additionally, added `resetOption(String)` and `suppressOption(String)` to `InsertSettings`, `QuerySettings`, and `CommonSettings` to allow removing or suppressing specific settings. Suppressed settings (set to `null`) will not be sent to the server, which is useful for overriding global settings. + - **[client-v2]** Added runtime credential update APIs on `Client`: `updateUserAndPassword(String, String)`, `updateAccessToken(String)`, and `updateBearerToken(String)`. Subsequent requests on the same `Client` instance use the new credentials without rebuilding the client. The authentication method is fixed at construction time; calling a runtime updater that does not match the configured method throws `ClientMisconfigurationException`. See `docs/authentication.md` for details and migration guidance. - **[jdbc-v2]** Added `cluster_name` configuration property to specify a target cluster for statements like `KILL QUERY` that require an `ON CLUSTER` clause to execute across all nodes. (https://github.com/ClickHouse/clickhouse-java/issues/2837) @@ -39,7 +41,7 @@ ### Bug Fixes -- **[jdbc-v2]** Fixed `Statement.cancel()` throwing `SESSION_IS_LOCKED` when the statement was running inside a ClickHouse session (e.g. via `clickhouse_setting_session_id`). The `KILL QUERY` request issued by `cancel()` now runs outside the session, so it no longer contends with the running query for the session lock. (https://github.com/ClickHouse/clickhouse-java/issues/2690) +- **[jdbc-v2]** Fixed `Statement.cancel()` throwing `SESSION_IS_LOCKED` when the statement was running inside a ClickHouse session. The driver now accepts `session_id`, `session_check`, and `session_timeout` as first-class connection properties and correctly suppresses them when issuing a `KILL QUERY` during cancellation. This ensures the cancellation request runs outside the session and no longer contends with the running query for the session lock. (https://github.com/ClickHouse/clickhouse-java/issues/2690, https://github.com/ClickHouse/clickhouse-java/issues/2881) - **[client-v2]** Fixed inconsistent use of `executionTimeout` parameter in `Client` component. The timeout was previously set in milliseconds but mistakenly retrieved and used in seconds in some places. Now it correctly uses milliseconds consistently. (https://github.com/ClickHouse/clickhouse-java/issues/2358) diff --git a/client-v2/src/main/java/com/clickhouse/client/api/Session.java b/client-v2/src/main/java/com/clickhouse/client/api/Session.java index d06e3d314..50a0e9016 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/Session.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/Session.java @@ -3,6 +3,7 @@ import com.clickhouse.client.api.http.ClickHouseHttpProto; import com.clickhouse.client.api.internal.ValidationUtils; +import java.util.HashMap; import java.util.Map; /** @@ -83,7 +84,7 @@ public synchronized void updateSessionId(String sessionId) { setSessionId(sessionId); } - public synchronized void applyTo(Map requestSettings) { + public synchronized void applyTo(Map requestSettings) { putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_ID, sessionId); putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_CHECK, sessionCheck == null ? null : (sessionCheck ? "1" : "0")); @@ -92,7 +93,15 @@ public synchronized void applyTo(Map requestSettings) { putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE, sessionTimezone); } - private static void putIfSet(Map settings, String key, String value) { + public static void clearSession(HashMap settings) { + settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID), null); + settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT), null); + settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK), null); + // Do not clean `session_timezone` setting because it is not related to session management and used to + // set timezone for consequent queries in some multi-user applications. + } + + private static void putIfSet(Map settings, String key, String value) { if (value != null) { settings.put(ClientConfigProperties.serverSetting(key), value); } diff --git a/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java index 7a2001eda..445a31e7e 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java @@ -60,6 +60,26 @@ public InsertSettings setOption(String option, Object value) { return this; } + /** + * Removes options from the settings. + * @param option - configuration option name + */ + public InsertSettings resetOption(String option) { + settings.resetOption(option); + return this; + } + + /** + * Makes option value to null that makes agent to remove it from final collection. + * This is useful to override even global settings when they need to be removed. + * @param option - option key + * @return current settings instance + */ + public InsertSettings suppressOption(String option) { + settings.suppressOption(option); + return this; + } + /** * Get all settings as an unmodifiable map. * diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java index e55d1429e..06f55f459 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java @@ -20,7 +20,7 @@ public class CommonSettings { private String operationId; private String logComment; - protected Map settings; + protected HashMap settings; // using hashmap to store null values public CommonSettings() { settings = new HashMap<>(); @@ -66,11 +66,27 @@ public CommonSettings setOption(String option, Object value) { return this; } + /** + * Removes option from the setting + * @param option - option key + * @return current settings instance + */ public CommonSettings resetOption(String option) { settings.remove(option); return this; } + /** + * Makes option value to null that makes agent to remove it from final collection. + * This is useful to override even global settings when they need to be removed. + * @param option - option key + * @return current settings instance + */ + public CommonSettings suppressOption(String option) { + settings.put(option, null); + return this; + } + /** * Get all settings as an unmodifiable map. * @@ -140,11 +156,7 @@ public CommonSettings use(Session session) { } public void clearSession() { - resetOption(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID)); - resetOption(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK)); - resetOption(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT)); - // Do not clean `session_timezone` setting because it is not related to session management and used to - // set timezone for consequent queries in some multi-user applications. + Session.clearSession(settings); } /** diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java index 9df39f407..0d3f43dfc 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java @@ -50,11 +50,26 @@ public QuerySettings setOption(String option, Object value) { return this; } + /** + * Removes options from the settings. + * @param option - configuration option name + */ public QuerySettings resetOption(String option) { settings.resetOption(option); return this; } + /** + * Makes option value to null that makes agent to remove it from final collection. + * This is useful to override even global settings when they need to be removed. + * @param option - option key + * @return current settings instance + */ + public QuerySettings suppressOption(String option) { + settings.suppressOption(option); + return this; + } + /** * Gets a configuration option. * diff --git a/client-v2/src/test/java/com/clickhouse/client/ClientTests.java b/client-v2/src/test/java/com/clickhouse/client/ClientTests.java index eaa675349..3c7529abd 100644 --- a/client-v2/src/test/java/com/clickhouse/client/ClientTests.java +++ b/client-v2/src/test/java/com/clickhouse/client/ClientTests.java @@ -50,6 +50,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Consumer; import java.util.function.Supplier; @@ -729,6 +730,24 @@ public void testInvalidAuthConfiguration() throws Exception { Assert.assertTrue(e.getMessage().contains("Trust store and certificates cannot be used together"), e.getMessage())); } + @Test(groups = {"integration"}) + public void testOverrideSettings() throws Exception { + final String clientTimezone = "America/Los_Angeles"; + try (Client client = newClient().setSessionTimezone(clientTimezone).build()) { + + final BiConsumer checkTzStmt = (settings, timezone) -> { + GenericRecord firstRecord = client.queryAll("SELECT timezone()", settings).stream().findFirst().get(); + Assert.assertEquals(firstRecord.getString(1), timezone); + }; + + checkTzStmt.accept(new QuerySettings(), clientTimezone); + + final String altTimezone = "America/New_York"; + checkTzStmt.accept(new QuerySettings().setSessionTimezone(altTimezone), altTimezone); + checkTzStmt.accept(new QuerySettings().setOption(ClientConfigProperties.serverSetting("session_timezone"), null), "UTC"); + } + } + public boolean isVersionMatch(String versionExpression, Client client) { List serverVersion = client.queryAll("SELECT version()"); return ClickHouseVersion.of(serverVersion.get(0).getString(1)).check(versionExpression); diff --git a/docs/features.md b/docs/features.md index be63e9f99..eb232ebdf 100644 --- a/docs/features.md +++ b/docs/features.md @@ -11,7 +11,7 @@ This document lists stable, user-visible behavior in `client-v2` and `jdbc-v2` t - Proxy support: Can send requests through configured HTTP proxies, including proxy credentials. - Connection and socket tuning: Exposes pool sizing, keep-alive, reuse strategy, connect/request/socket timeouts, and low-level socket options. - Query execution: Executes SQL asynchronously and returns streaming query responses with response metadata and metrics. -- Query settings: Supports per-query database selection, output format, execution limits, roles, log comments, headers, reusable `Session` objects, session settings, server settings, and network timeout overrides. +- Query settings: Supports per-query database selection, output format, execution limits, roles, log comments, headers, reusable `Session` objects, session settings, server settings, and network timeout overrides. Settings explicitly set to `null` are suppressed and will not be sent to the server. - Parameterized SQL: Accepts named query parameters and can send them through supported HTTP request encodings. - Result materialization helpers: Provides streaming `Records`, generic row access, and convenience APIs that materialize all rows into generic records or typed POJOs. - Binary format readers: Reads ClickHouse binary result formats including `Native`, `RowBinary`, `RowBinaryWithNames`, and `RowBinaryWithNamesAndTypes`. diff --git a/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java b/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java index eed401a14..daab1afd2 100644 --- a/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java +++ b/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java @@ -1,6 +1,7 @@ package com.clickhouse.jdbc; import com.clickhouse.client.api.ClientConfigProperties; +import com.clickhouse.client.api.Session; import com.clickhouse.client.api.internal.ServerSettings; import com.clickhouse.client.api.query.GenericRecord; import com.clickhouse.data.ClickHouseVersion; @@ -758,7 +759,11 @@ public void testCancelQueryWithSession() throws Exception { // "Session is locked by a concurrent client" (SESSION_IS_LOCKED). The KILL QUERY request issued by // cancel() must not carry the session id of the query being cancelled. String sessionId = "test-session-" + UUID.randomUUID(); - try (Connection conn = getJdbcConnection()) { + Properties properties = new Properties(); + Session session = new Session(); + session.setSessionId(sessionId); + session.applyTo(properties); + try (Connection conn = getJdbcConnection(properties)) { try (StatementImpl stmt = (StatementImpl) conn.createStatement()) { stmt.getLocalSettings().setSessionId(sessionId); stmt.setQueryTimeout(30); // safety net so a failed cancel cannot hang the test @@ -799,7 +804,12 @@ public void testCancelInsertWithSession() throws Exception { // Regression test for #2690 covering a long-running INSERT executed inside a session. String tableName = getDatabase() + ".cancel_insert_with_session"; String sessionId = "test-session-" + UUID.randomUUID(); - try (Connection conn = getJdbcConnection(Map.of(ASYNC_INSERT_SETTING_KEY, ServerSettings.OFF))) { + Properties properties = new Properties(); + properties.put(ASYNC_INSERT_SETTING_KEY, ServerSettings.OFF); + Session session = new Session(); + session.setSessionId(sessionId); + session.applyTo(properties); + try (Connection conn = getJdbcConnection(properties)) { try (Statement setup = conn.createStatement()) { setup.execute("DROP TABLE IF EXISTS " + tableName); setup.execute("CREATE TABLE " + tableName + " (num UInt64) ENGINE = MergeTree ORDER BY ()"); From a504d68dbb3935104c77d5b7d4070b86f723921b Mon Sep 17 00:00:00 2001 From: Sergey Chernov Date: Thu, 18 Jun 2026 18:12:07 -0700 Subject: [PATCH 2/3] fixed --- CHANGELOG.md | 2 +- .../com/clickhouse/client/api/Session.java | 3 +-- .../client/api/insert/InsertSettings.java | 11 ----------- .../client/api/internal/CommonSettings.java | 18 +----------------- .../client/api/query/QuerySettings.java | 15 --------------- docs/features.md | 2 +- 6 files changed, 4 insertions(+), 47 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e219cca49..320805a55 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,7 +29,7 @@ ### New Features -- **[client-v2]** Added `Session` API to encapsulate and manage ClickHouse session settings (`session_id`, `session_check`, `session_timeout`, `session_timezone`) as a reusable object. The `Session` instance can be applied to any request settings using `applyTo()`, and session state can be cleared via `clearSession()`. Additionally, added `resetOption(String)` and `suppressOption(String)` to `InsertSettings`, `QuerySettings`, and `CommonSettings` to allow removing or suppressing specific settings. Suppressed settings (set to `null`) will not be sent to the server, which is useful for overriding global settings. +- **[client-v2]** Added `Session` API to encapsulate and manage ClickHouse session settings (`session_id`, `session_check`, `session_timeout`, `session_timezone`) as a reusable object. The `Session` instance can be applied to any request settings using `applyTo()`, and session state can be cleared via `clearSession()`. Additionally, added `resetOption(String)` to `InsertSettings`, `QuerySettings`, and `CommonSettings` to allow removing specific settings. Settings explicitly set to `null` will not be sent to the server, which is useful for overriding global settings. - **[client-v2]** Added runtime credential update APIs on `Client`: `updateUserAndPassword(String, String)`, `updateAccessToken(String)`, and `updateBearerToken(String)`. Subsequent requests on the same `Client` instance use the new credentials without rebuilding the client. The authentication method is fixed at construction time; calling a runtime updater that does not match the configured method throws `ClientMisconfigurationException`. See `docs/authentication.md` for details and migration guidance. diff --git a/client-v2/src/main/java/com/clickhouse/client/api/Session.java b/client-v2/src/main/java/com/clickhouse/client/api/Session.java index 50a0e9016..ab00f6ad0 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/Session.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/Session.java @@ -3,7 +3,6 @@ import com.clickhouse.client.api.http.ClickHouseHttpProto; import com.clickhouse.client.api.internal.ValidationUtils; -import java.util.HashMap; import java.util.Map; /** @@ -93,7 +92,7 @@ public synchronized void applyTo(Map requestSettings) { putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE, sessionTimezone); } - public static void clearSession(HashMap settings) { + public static void clearSession(Map settings) { settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID), null); settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT), null); settings.put(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK), null); diff --git a/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java index 445a31e7e..078019b17 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java @@ -69,17 +69,6 @@ public InsertSettings resetOption(String option) { return this; } - /** - * Makes option value to null that makes agent to remove it from final collection. - * This is useful to override even global settings when they need to be removed. - * @param option - option key - * @return current settings instance - */ - public InsertSettings suppressOption(String option) { - settings.suppressOption(option); - return this; - } - /** * Get all settings as an unmodifiable map. * diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java index 06f55f459..e84873fb5 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java @@ -20,7 +20,7 @@ public class CommonSettings { private String operationId; private String logComment; - protected HashMap settings; // using hashmap to store null values + protected Map settings; public CommonSettings() { settings = new HashMap<>(); @@ -66,27 +66,11 @@ public CommonSettings setOption(String option, Object value) { return this; } - /** - * Removes option from the setting - * @param option - option key - * @return current settings instance - */ public CommonSettings resetOption(String option) { settings.remove(option); return this; } - /** - * Makes option value to null that makes agent to remove it from final collection. - * This is useful to override even global settings when they need to be removed. - * @param option - option key - * @return current settings instance - */ - public CommonSettings suppressOption(String option) { - settings.put(option, null); - return this; - } - /** * Get all settings as an unmodifiable map. * diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java index 0d3f43dfc..9df39f407 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java @@ -50,26 +50,11 @@ public QuerySettings setOption(String option, Object value) { return this; } - /** - * Removes options from the settings. - * @param option - configuration option name - */ public QuerySettings resetOption(String option) { settings.resetOption(option); return this; } - /** - * Makes option value to null that makes agent to remove it from final collection. - * This is useful to override even global settings when they need to be removed. - * @param option - option key - * @return current settings instance - */ - public QuerySettings suppressOption(String option) { - settings.suppressOption(option); - return this; - } - /** * Gets a configuration option. * diff --git a/docs/features.md b/docs/features.md index eb232ebdf..c3ab24541 100644 --- a/docs/features.md +++ b/docs/features.md @@ -11,7 +11,7 @@ This document lists stable, user-visible behavior in `client-v2` and `jdbc-v2` t - Proxy support: Can send requests through configured HTTP proxies, including proxy credentials. - Connection and socket tuning: Exposes pool sizing, keep-alive, reuse strategy, connect/request/socket timeouts, and low-level socket options. - Query execution: Executes SQL asynchronously and returns streaming query responses with response metadata and metrics. -- Query settings: Supports per-query database selection, output format, execution limits, roles, log comments, headers, reusable `Session` objects, session settings, server settings, and network timeout overrides. Settings explicitly set to `null` are suppressed and will not be sent to the server. +- Query settings: Supports per-query database selection, output format, execution limits, roles, log comments, headers, reusable `Session` objects, session settings, server settings, and network timeout overrides. Settings explicitly set to `null` will not be sent to the server. - Parameterized SQL: Accepts named query parameters and can send them through supported HTTP request encodings. - Result materialization helpers: Provides streaming `Records`, generic row access, and convenience APIs that materialize all rows into generic records or typed POJOs. - Binary format readers: Reads ClickHouse binary result formats including `Native`, `RowBinary`, `RowBinaryWithNames`, and `RowBinaryWithNamesAndTypes`. From 22f15ad83a2c7eaa1fe9bf65ed9a962a3de2f005 Mon Sep 17 00:00:00 2001 From: Sergey Chernov Date: Fri, 19 Jun 2026 11:19:19 -0700 Subject: [PATCH 3/3] added verification for custom_http_header --- .../com/clickhouse/client/SettingsTests.java | 2 + .../com/clickhouse/jdbc/StatementTest.java | 64 +++++++++++-------- 2 files changed, 38 insertions(+), 28 deletions(-) diff --git a/client-v2/src/test/java/com/clickhouse/client/SettingsTests.java b/client-v2/src/test/java/com/clickhouse/client/SettingsTests.java index 81261dc33..00c69e0b2 100644 --- a/client-v2/src/test/java/com/clickhouse/client/SettingsTests.java +++ b/client-v2/src/test/java/com/clickhouse/client/SettingsTests.java @@ -177,6 +177,8 @@ public void testInsertSettingsSpecific() throws Exception { final InsertSettings settings = new InsertSettings(); settings.setDatabase("test_db1"); Assert.assertEquals(settings.getDatabase(), "test_db1"); + settings.resetOption(ClientConfigProperties.DATABASE.getKey()); + Assert.assertNull(settings.getDatabase()); } { diff --git a/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java b/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java index daab1afd2..ec32e9ec6 100644 --- a/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java +++ b/jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java @@ -755,44 +755,52 @@ public void testCancelQueryWithSession() throws Exception { throw new SkipException("Cloud + HTTP doesn't work well. Enough to test locally"); } - // Regression test for #2690: cancelling a query that runs inside a session must not fail with - // "Session is locked by a concurrent client" (SESSION_IS_LOCKED). The KILL QUERY request issued by - // cancel() must not carry the session id of the query being cancelled. String sessionId = "test-session-" + UUID.randomUUID(); Properties properties = new Properties(); Session session = new Session(); session.setSessionId(sessionId); session.applyTo(properties); try (Connection conn = getJdbcConnection(properties)) { - try (StatementImpl stmt = (StatementImpl) conn.createStatement()) { - stmt.getLocalSettings().setSessionId(sessionId); - stmt.setQueryTimeout(30); // safety net so a failed cancel cannot hang the test + testCancelQueryWithSessionValidation(conn, sessionId); + } - final AtomicReference threadError = new AtomicReference<>(); - final CountDownLatch started = new CountDownLatch(1); - Thread worker = new Thread(() -> { - started.countDown(); - // Long-running query that only completes when killed. - try (ResultSet rs = stmt.executeQuery("SELECT count() FROM system.numbers_mt")) { - rs.next(); - } catch (Throwable t) { - System.out.println("Error: " + t.getMessage()); - threadError.set(t); - } - }); - worker.start(); - started.await(); + // Test case when session id is in custom_http_header + properties = new Properties(); + properties.put(DriverProperties.CUSTOM_HTTP_PARAMS.getKey(), "session_id=" + sessionId); + try (Connection conn = getJdbcConnection(properties)) { + testCancelQueryWithSessionValidation(conn, sessionId); + } + } - String queryId = waitForQueryId(stmt, 15); - assertNotNull(queryId, "Query id was not assigned in time"); - assertTrue(waitForQueryToStart(queryId, 15), "Query did not start on the server in time"); + private void testCancelQueryWithSessionValidation(Connection conn, String sessionId) throws Exception { + try (StatementImpl stmt = (StatementImpl) conn.createStatement()) { + stmt.getLocalSettings().setSessionId(sessionId); + stmt.setQueryTimeout(30); // safety net so a failed cancel cannot hang the test + + final AtomicReference threadError = new AtomicReference<>(); + final CountDownLatch started = new CountDownLatch(1); + Thread worker = new Thread(() -> { + started.countDown(); + // Long-running query that only completes when killed. + try (ResultSet rs = stmt.executeQuery("SELECT count() FROM system.numbers_mt")) { + rs.next(); + } catch (Throwable t) { + System.out.println("Error: " + t.getMessage()); + threadError.set(t); + } + }); + worker.start(); + started.await(); - // Cancel from the main thread - must not throw SESSION_IS_LOCKED. - stmt.cancel(); + String queryId = waitForQueryId(stmt, 15); + assertNotNull(queryId, "Query id was not assigned in time"); + assertTrue(waitForQueryToStart(queryId, 15), "Query did not start on the server in time"); - worker.join(TimeUnit.SECONDS.toMillis(20)); - assertFalse(worker.isAlive(), "Query was not cancelled and is still running"); - } + // Cancel from the main thread - must not throw SESSION_IS_LOCKED. + stmt.cancel(); + + worker.join(TimeUnit.SECONDS.toMillis(20)); + assertFalse(worker.isAlive(), "Query was not cancelled and is still running"); } }