From e8c071af86638461a0a8fddce760addcbf9280c3 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 25 Aug 2026 20:13:22 +0000 Subject: [PATCH 1/4] feat(kernel): forward socket timeout Signed-off-by: Vu Anh Phung --- CONNECTION_PARAMETERS.md | 2 +- KERNEL_REV | 2 +- src/databricks/sql/backend/kernel/client.py | 9 +++++++ src/databricks/sql/client.py | 6 +++-- src/databricks/sql/session.py | 1 + tests/unit/test_kernel_client.py | 27 +++++++++++++++++++++ tests/unit/test_session.py | 11 +++++---- 7 files changed, 49 insertions(+), 9 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index f63de23a9..b2f429a20 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -99,7 +99,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `_socket_timeout` | `float` (s) | ✅ | ❌ | `900` | Socket send/recv/connect timeout. Not forwarded to the kernel, which manages its own request timeout. | +| `_socket_timeout` | `float` (s) | ✅ | ✅ | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor an immediate timeout. | | `_pool_connections` | `int` | ✅ | ⚠️ | `10` | Number of urllib3 connection pools. Configures the connector's shared Python HTTP client; the kernel's query transport is its own Rust stack. | | `_pool_maxsize` | `int` | ✅ | ⚠️ | `20` | Max connections per pool on the shared Python HTTP client. Same kernel caveat as `_pool_connections`. | | `_proxy_auth_method` | `str` | ✅ | ⚠️ | `None` | `basic` or `negotiate` (Kerberos). Applies to the shared Python HTTP client; not threaded to the kernel query transport. See [`docs/proxy.md`](docs/proxy.md). | diff --git a/KERNEL_REV b/KERNEL_REV index 6cd3da53d..f751c496b 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -ad78a5be3dc8bb7fc78ec574492515ab24e23d4c +dd810d6d0a179886b923c6e22dc785ddca16ebef diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 93b0a98a4..6d30c4457 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -217,6 +217,9 @@ def __init__( # to the kernel ``Session``'s ``retry_*`` kwargs in # ``open_session`` via ``_kernel_retry_kwargs``. self._retry_options = kwargs.get("retry_options") or {} + # The connector's ``_socket_timeout`` is already expressed in + # seconds, matching the kernel's request-timeout binding. + self._request_timeout_secs = kwargs.get("request_timeout_secs") self._catalog = catalog self._schema = schema # ``_use_arrow_native_complex_types`` is the connector-side @@ -330,6 +333,11 @@ def open_session( # Translate the connector's ``_retry_*`` kwargs into the # kernel's ``retry_*`` kwargs. Empty when at defaults. retry_kwargs = _kernel_retry_kwargs(self._retry_options) + request_timeout_kwargs: Dict[str, Any] = {} + if self._request_timeout_secs is not None: + request_timeout_kwargs["request_timeout_secs"] = ( + self._request_timeout_secs + ) # Forward caller / connector HTTP headers. The kernel applies # them on every request; a caller ``User-Agent`` is appended # to the kernel's base UA. Only pass the kwarg when there's @@ -372,6 +380,7 @@ def open_session( **auth_kwargs, **tls_kwargs, **retry_kwargs, + **request_timeout_kwargs, **http_headers_kwargs, ) except Exception as exc: diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 44895954f..914a24dde 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -273,8 +273,10 @@ def read(self) -> Optional[OAuthToken]: # _retry_stop_after_attempts_count # The maximum number of attempts during a request retry sequence (defaults to 24) # _socket_timeout - # The timeout in seconds for socket send, recv and connect operations. Defaults to None for - # no timeout. Should be a positive float or integer. + # On Thrift, the timeout in seconds for socket send, recv and connect + # operations. On the kernel path, a positive value is the total HTTP + # request deadline. Kernel values of None or 0 select its 120-second + # default; 0 is neither unlimited nor an immediate timeout. # _disable_pandas # In case the deserialisation through pandas causes any issues, it can be disabled with # this flag. diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index f35cdf525..19cd1dba3 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -230,6 +230,7 @@ def _create_backend( _use_arrow_native_complex_types=_use_arrow_native_complex_types, auth_options=kernel_auth_options, retry_options=kernel_retry_options, + request_timeout_secs=kwargs.get("_socket_timeout"), ) databricks_client_class: Type[DatabricksClient] diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 3eb5a9006..5964893fa 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -344,6 +344,33 @@ def fake_session(**kw): assert captured.get("complex_types_as_json") is expected_flag +@pytest.mark.parametrize("timeout", [None, 0, 12.5]) +def test_open_session_passes_request_timeout_to_kernel(monkeypatch, timeout): + captured = {} + + def fake_session(**kw): + captured.update(kw) + sess = MagicMock() + sess.session_id = "sess-id" + return sess + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session) + c = kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + request_timeout_secs=timeout, + ) + + c.open_session(session_configuration=None, catalog=None, schema=None) + + if timeout is None: + assert "request_timeout_secs" not in captured + else: + assert captured["request_timeout_secs"] == timeout + + def test_execute_command_forwards_parameters_to_bind_param(): """``execute_command(parameters=[...])`` routes each parameter through ``bind_tspark_params`` onto the kernel statement before diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 6fcefcade..c50650e4d 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -410,10 +410,9 @@ def test_use_kernel_pat_builds_minimal_access_token_provider(self): assert isinstance(sess.auth_provider, AccessTokenAuthProvider) -class TestKernelRetryOptionsThreading: - """The connector's ``_retry_*`` kwargs must be forwarded into the - kernel client's ``retry_options`` on the use_kernel path (the kernel - owns the retry loop). Captures the kwargs session.py passes by +class TestKernelTransportOptionsThreading: + """The connector's retry and socket timeout kwargs must be forwarded + on the use_kernel path. Captures the kwargs session.py passes by patching ``KernelDatabricksClient`` and inspecting its call args. Patching ``KernelDatabricksClient`` requires importing @@ -426,7 +425,7 @@ class TestKernelRetryOptionsThreading: PACKAGE = "databricks.sql" - def test_retry_kwargs_threaded_into_kernel_client(self): + def test_retry_and_socket_timeout_threaded_into_kernel_client(self): import sys import types @@ -466,6 +465,7 @@ def test_retry_kwargs_threaded_into_kernel_client(self): _retry_delay_max=90.0, _retry_stop_after_attempts_count=10, _retry_stop_after_attempts_duration=600.0, + _socket_timeout=12.5, ) try: _, kwargs = mock_kernel_client.call_args @@ -474,6 +474,7 @@ def test_retry_kwargs_threaded_into_kernel_client(self): assert opts["retry_delay_max"] == 90.0 assert opts["retry_stop_after_attempts_count"] == 10 assert opts["retry_stop_after_attempts_duration"] == 600.0 + assert kwargs["request_timeout_secs"] == 12.5 finally: conn.close() From 2ccda76d6010b6b771ea5138f2d3647bb50dfba2 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 25 Aug 2026 20:34:38 +0000 Subject: [PATCH 2/4] fix(kernel): validate request timeout Signed-off-by: Vu Anh Phung --- CONNECTION_PARAMETERS.md | 2 +- src/databricks/sql/backend/kernel/client.py | 19 ++++++++++++++++++- src/databricks/sql/client.py | 3 ++- tests/unit/test_kernel_client.py | 12 ++++++++++++ 4 files changed, 33 insertions(+), 3 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index b2f429a20..55b3480fe 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -99,7 +99,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `_socket_timeout` | `float` (s) | ✅ | ✅ | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor an immediate timeout. | +| `_socket_timeout` | `float` (s) | ✅ | ✅ | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor immediate; negative and non-finite values raise `ValueError`. | | `_pool_connections` | `int` | ✅ | ⚠️ | `10` | Number of urllib3 connection pools. Configures the connector's shared Python HTTP client; the kernel's query transport is its own Rust stack. | | `_pool_maxsize` | `int` | ✅ | ⚠️ | `20` | Max connections per pool on the shared Python HTTP client. Same kernel caveat as `_pool_connections`. | | `_proxy_auth_method` | `str` | ✅ | ⚠️ | `None` | `basic` or `negotiate` (Kerberos). Applies to the shared Python HTTP client; not threaded to the kernel query transport. See [`docs/proxy.md`](docs/proxy.md). | diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 6d30c4457..a7118b114 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -23,6 +23,7 @@ from __future__ import annotations import logging +import math import threading import uuid from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union @@ -219,7 +220,23 @@ def __init__( self._retry_options = kwargs.get("retry_options") or {} # The connector's ``_socket_timeout`` is already expressed in # seconds, matching the kernel's request-timeout binding. - self._request_timeout_secs = kwargs.get("request_timeout_secs") + request_timeout_secs = kwargs.get("request_timeout_secs") + if request_timeout_secs is None: + self._request_timeout_secs = None + else: + try: + self._request_timeout_secs = float(request_timeout_secs) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError( + "_socket_timeout must be a non-negative finite number of seconds" + ) from exc + if ( + not math.isfinite(self._request_timeout_secs) + or self._request_timeout_secs < 0 + ): + raise ValueError( + "_socket_timeout must be a non-negative finite number of seconds" + ) self._catalog = catalog self._schema = schema # ``_use_arrow_native_complex_types`` is the connector-side diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 914a24dde..2edea3d40 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -276,7 +276,8 @@ def read(self) -> Optional[OAuthToken]: # On Thrift, the timeout in seconds for socket send, recv and connect # operations. On the kernel path, a positive value is the total HTTP # request deadline. Kernel values of None or 0 select its 120-second - # default; 0 is neither unlimited nor an immediate timeout. + # default; 0 is neither unlimited nor an immediate timeout. Negative + # and non-finite values are rejected. # _disable_pandas # In case the deserialisation through pandas causes any issues, it can be disabled with # this flag. diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 5964893fa..a99b14b60 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -371,6 +371,18 @@ def fake_session(**kw): assert captured["request_timeout_secs"] == timeout +@pytest.mark.parametrize("timeout", [-1, float("nan"), float("inf")]) +def test_request_timeout_rejects_invalid_values(timeout): + with pytest.raises(ValueError, match="non-negative finite"): + kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + request_timeout_secs=timeout, + ) + + def test_execute_command_forwards_parameters_to_bind_param(): """``execute_command(parameters=[...])`` routes each parameter through ``bind_tspark_params`` onto the kernel statement before From 9151648aaaf0f9761c93c0ff8e58a48b264ea504 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 26 Aug 2026 03:47:39 +0000 Subject: [PATCH 3/4] refactor(kernel): rely on timeout binding validation Signed-off-by: Vu Anh Phung --- CONNECTION_PARAMETERS.md | 2 +- src/databricks/sql/backend/kernel/client.py | 22 ++------------------- src/databricks/sql/client.py | 3 +-- tests/unit/test_kernel_client.py | 12 ----------- 4 files changed, 4 insertions(+), 35 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 55b3480fe..b2f429a20 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -99,7 +99,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `_socket_timeout` | `float` (s) | ✅ | ✅ | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor immediate; negative and non-finite values raise `ValueError`. | +| `_socket_timeout` | `float` (s) | ✅ | ✅ | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor an immediate timeout. | | `_pool_connections` | `int` | ✅ | ⚠️ | `10` | Number of urllib3 connection pools. Configures the connector's shared Python HTTP client; the kernel's query transport is its own Rust stack. | | `_pool_maxsize` | `int` | ✅ | ⚠️ | `20` | Max connections per pool on the shared Python HTTP client. Same kernel caveat as `_pool_connections`. | | `_proxy_auth_method` | `str` | ✅ | ⚠️ | `None` | `basic` or `negotiate` (Kerberos). Applies to the shared Python HTTP client; not threaded to the kernel query transport. See [`docs/proxy.md`](docs/proxy.md). | diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index a7118b114..2cd9c51b9 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -23,7 +23,6 @@ from __future__ import annotations import logging -import math import threading import uuid from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union @@ -218,25 +217,8 @@ def __init__( # to the kernel ``Session``'s ``retry_*`` kwargs in # ``open_session`` via ``_kernel_retry_kwargs``. self._retry_options = kwargs.get("retry_options") or {} - # The connector's ``_socket_timeout`` is already expressed in - # seconds, matching the kernel's request-timeout binding. - request_timeout_secs = kwargs.get("request_timeout_secs") - if request_timeout_secs is None: - self._request_timeout_secs = None - else: - try: - self._request_timeout_secs = float(request_timeout_secs) - except (TypeError, ValueError, OverflowError) as exc: - raise ValueError( - "_socket_timeout must be a non-negative finite number of seconds" - ) from exc - if ( - not math.isfinite(self._request_timeout_secs) - or self._request_timeout_secs < 0 - ): - raise ValueError( - "_socket_timeout must be a non-negative finite number of seconds" - ) + # The kernel binding owns type and range validation. + self._request_timeout_secs = kwargs.get("request_timeout_secs") self._catalog = catalog self._schema = schema # ``_use_arrow_native_complex_types`` is the connector-side diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 2edea3d40..914a24dde 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -276,8 +276,7 @@ def read(self) -> Optional[OAuthToken]: # On Thrift, the timeout in seconds for socket send, recv and connect # operations. On the kernel path, a positive value is the total HTTP # request deadline. Kernel values of None or 0 select its 120-second - # default; 0 is neither unlimited nor an immediate timeout. Negative - # and non-finite values are rejected. + # default; 0 is neither unlimited nor an immediate timeout. # _disable_pandas # In case the deserialisation through pandas causes any issues, it can be disabled with # this flag. diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index a99b14b60..5964893fa 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -371,18 +371,6 @@ def fake_session(**kw): assert captured["request_timeout_secs"] == timeout -@pytest.mark.parametrize("timeout", [-1, float("nan"), float("inf")]) -def test_request_timeout_rejects_invalid_values(timeout): - with pytest.raises(ValueError, match="non-negative finite"): - kernel_client.KernelDatabricksClient( - server_hostname="example.cloud.databricks.com", - http_path="/sql/1.0/warehouses/abc", - auth_provider=AccessTokenAuthProvider("dapi-test"), - ssl_options=None, - request_timeout_secs=timeout, - ) - - def test_execute_command_forwards_parameters_to_bind_param(): """``execute_command(parameters=[...])`` routes each parameter through ``bind_tspark_params`` onto the kernel statement before From dd7c8b4a2eda5d4f40b4e9735cf8c7556ec29d3d Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 26 Aug 2026 04:54:39 +0000 Subject: [PATCH 4/4] refactor(kernel): pass request timeout directly Signed-off-by: Vu Anh Phung --- src/databricks/sql/backend/kernel/client.py | 7 +------ tests/unit/test_kernel_client.py | 5 +---- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 2cd9c51b9..47892c1e8 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -332,11 +332,6 @@ def open_session( # Translate the connector's ``_retry_*`` kwargs into the # kernel's ``retry_*`` kwargs. Empty when at defaults. retry_kwargs = _kernel_retry_kwargs(self._retry_options) - request_timeout_kwargs: Dict[str, Any] = {} - if self._request_timeout_secs is not None: - request_timeout_kwargs["request_timeout_secs"] = ( - self._request_timeout_secs - ) # Forward caller / connector HTTP headers. The kernel applies # them on every request; a caller ``User-Agent`` is appended # to the kernel's base UA. Only pass the kwarg when there's @@ -376,10 +371,10 @@ def open_session( # backend's surface (interval columns arrive as # strings). intervals_as_string=True, + request_timeout_secs=self._request_timeout_secs, **auth_kwargs, **tls_kwargs, **retry_kwargs, - **request_timeout_kwargs, **http_headers_kwargs, ) except Exception as exc: diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 5964893fa..7e8249553 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -365,10 +365,7 @@ def fake_session(**kw): c.open_session(session_configuration=None, catalog=None, schema=None) - if timeout is None: - assert "request_timeout_secs" not in captured - else: - assert captured["request_timeout_secs"] == timeout + assert captured["request_timeout_secs"] == timeout def test_execute_command_forwards_parameters_to_bind_param():