feat(kernel): forward socket timeout (#931)

* feat(kernel): forward socket timeout

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

* fix(kernel): validate request timeout

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

* refactor(kernel): rely on timeout binding validation

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

* refactor(kernel): pass request timeout directly

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

---------

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

Author: vuanhphung

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). |

KERNEL_REV

@@ -1 +1 @@
-ad78a5be3dc8bb7fc78ec574492515ab24e23d4c
+dd810d6d0a179886b923c6e22dc785ddca16ebef

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,

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.

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]

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

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()