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..47892c1e8 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -217,6 +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 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 @@ -369,6 +371,7 @@ 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, 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..7e8249553 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -344,6 +344,30 @@ 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) + + 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()