feat(kernel): forward _pool_maxsize (#940)

* feat(kernel): forward _pool_maxsize

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

* fix(kernel): gate max connections kwarg

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

* chore(kernel): bump pinned revision

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

* refactor(kernel): pass max connections directly

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

* fix(kernel): gate max connections by wheel support

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

* docs(kernel): clarify max connections version

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

* fix(kernel): normalize zero pool maxsize

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

@@ -102,7 +102,7 @@ to change without notice.
 | ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- |
 | `_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`.                                          |
+| `_pool_maxsize`                      | `int`       |   ✅   |   ⚠️   | `20` (Thrift); `100` (kernel when unset) | Max idle connections retained per host. A positive value always configures the shared Python HTTP client and also configures the kernel's Rust HTTP pool with kernel ≥ 1.1.0. Unset or `0` keeps each client's default. |
 | `_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). |
 | `_retry_stop_after_attempts_count`   | `int`       |   ✅   |   ✅   | `30`          | Max attempts in a retry sequence. Bounded to `[1, 60]` on Thrift; forwarded to the kernel's retry policy.                                       |
 | `_retry_stop_after_attempts_duration`| `float` (s) |   ✅   |   ✅   | `900`         | Max total wall-clock seconds spent retrying. Forwarded to the kernel.                                                                           |
@@ -202,9 +202,10 @@ None — the kernel's parameter surface is currently a subset of Thrift's.
 
 ### Behavioral divergences to watch
 
-- **Connection pooling / proxy** (`_pool_connections`, `_pool_maxsize`,
-  `_proxy_auth_method`) configure the connector's shared Python HTTP client
-  (auth/telemetry); the kernel's query traffic uses its own Rust transport.
+- **Connection pooling / proxy**: `_pool_connections` and `_proxy_auth_method`
+  configure only the shared Python HTTP client. `_pool_maxsize` also configures
+  the kernel's Rust HTTP pool when positive and using kernel ≥ 1.1.0. A value of
+  `0` is treated as unset: the shared client keeps 20 and the kernel keeps 100.
 - **`use_inline_params`** renders parameters inline on Thrift; the kernel uses
   native parameter binding.
 

KERNEL_REV

@@ -1 +1 @@
-d64009eb59404c1b082cb020296337f96dc0d4d7
+167651ecc67143ef258ad70fd2682ced00beaa99

src/databricks/sql/backend/kernel/client.py

@@ -156,7 +156,7 @@ def _kernel_session_accepts_kwarg(name: str) -> bool:
     ``**kwargs`` catch-all), so forwarding a kwarg it doesn't declare raises
     ``TypeError`` at construction, so we gate kwargs on what the installed
     wheel supports. Falls **closed** (returns ``False``) when the signature
-    can't be introspected because omitting an accepted telemetry kwarg is safer
+    can't be introspected because omitting an accepted optional kwarg is safer
     than forwarding an unsupported one.
     """
     try:
@@ -260,6 +260,7 @@ def __init__(
         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._max_connections = kwargs.get("max_connections")
         # Kernel telemetry phase 7 adds binding/runtime identity and
         # telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
         self._telemetry_options = kwargs.get("telemetry_options") or {}
@@ -377,6 +378,9 @@ def open_session(
             # kernel's ``retry_*`` kwargs. Empty when at defaults.
             retry_kwargs = _kernel_retry_kwargs(self._retry_options)
             telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
+            max_connections_kwargs: Dict[str, Any] = {}
+            if _kernel_session_accepts_kwarg("max_connections"):
+                max_connections_kwargs["max_connections"] = self._max_connections
             # 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
@@ -421,6 +425,7 @@ def open_session(
                 **tls_kwargs,
                 **retry_kwargs,
                 **telemetry_kwargs,
+                **max_connections_kwargs,
                 **http_headers_kwargs,
             )
         except Exception as exc:

src/databricks/sql/session.py

@@ -313,6 +313,7 @@ def _create_backend(
                 auth_options=kernel_auth_options,
                 retry_options=kernel_retry_options,
                 request_timeout_secs=kwargs.get("_socket_timeout"),
+                max_connections=kwargs.get("_pool_maxsize") or None,
                 telemetry_options=kernel_telemetry_options,
             )
 

tests/unit/test_kernel_client.py

@@ -368,6 +368,30 @@ def fake_session(**kw):
     assert captured["request_timeout_secs"] == timeout
 
 
+@pytest.mark.parametrize("max_connections", [None, 41])
+def test_open_session_passes_max_connections_to_kernel(monkeypatch, max_connections):
+    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,
+        max_connections=max_connections,
+    )
+
+    c.open_session(session_configuration=None, catalog=None, schema=None)
+
+    assert captured["max_connections"] == max_connections
+
+
 def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
     """Kernel telemetry phase 7 added binding/runtime identity and
     telemetry config kwargs to ``databricks_sql_kernel.Session``."""
@@ -428,25 +452,20 @@ def fake_session(**kw):
     assert captured["telemetry_circuit_breaker_enabled"] is False
 
 
-def test_open_session_omits_phase_7_kwargs_kernel_does_not_accept(monkeypatch):
-    """Phase-7 identity/telemetry kwargs must NOT be forwarded to a kernel
-    ``Session`` whose (fixed, no-``**kwargs``) constructor doesn't declare
-    them.
+def test_open_session_omits_optional_kwargs_kernel_does_not_accept(monkeypatch):
+    """Optional kwargs must NOT be forwarded to a kernel ``Session`` whose
+    fixed constructor doesn't declare them.
 
     The real ``databricks_sql_kernel.Session`` is a PyO3 class with a fixed
-    signature; the pinned ``^0.2.0`` wheel predates phase 7 and accepts none
-    of these kwargs, so forwarding them unconditionally raises ``TypeError``
-    and breaks every ``use_kernel=True`` connection. The other tests here use
-    a ``**kwargs`` MagicMock that silently swallows the kwargs and hides the
-    break; this one uses a fixed-signature fake mirroring the real 0.2.0
-    surface to prove the client gates on what the installed Session supports.
+    signature. The other tests here use a ``**kwargs`` MagicMock that silently
+    swallows unsupported kwargs; this fixed-signature fake proves the client
+    filters optional kwargs.
     """
     captured = {}
 
-    # Fixed signature mirroring the pinned 0.2.0 kernel Session: it accepts
-    # the base connection/tls/retry kwargs but NONE of the phase-7 identity
-    # or telemetry kwargs, and has no **kwargs catch-all.
-    def fake_session_v0_2_0(
+    # Accept baseline connection/tls/retry kwargs but no max-connections,
+    # phase-7 identity, or telemetry kwargs, and no **kwargs catch-all.
+    def fake_session_without_optional_kwargs(
         host,
         http_path,
         *,
@@ -480,7 +499,9 @@ def fake_session_v0_2_0(
         sess.session_id = "sess-id"
         return sess
 
-    monkeypatch.setattr(kernel_client._kernel, "Session", fake_session_v0_2_0)
+    monkeypatch.setattr(
+        kernel_client._kernel, "Session", fake_session_without_optional_kwargs
+    )
     monkeypatch.setattr(
         kernel_client.TelemetryHelper,
         "get_driver_system_configuration",
@@ -504,17 +525,18 @@ def fake_session_v0_2_0(
     kwargs = kernel_client._kernel_telemetry_kwargs(
         {"enable_telemetry": True, "telemetry_batch_size": 17}
     )
-    assert kwargs == {}, f"expected no phase-7 kwargs on 0.2.0 Session, got {kwargs}"
+    assert kwargs == {}, f"expected no unsupported phase-7 kwargs, got {kwargs}"
 
     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,
+        max_connections=41,
         telemetry_options={"enable_telemetry": True, "telemetry_batch_size": 17},
     )
-    # Would raise TypeError: unexpected keyword argument if the client
-    # forwarded phase-7 kwargs the fixed-signature Session doesn't declare.
+    # Would raise TypeError if the client forwarded max-connections or
+    # phase-7 kwargs the fixed-signature Session doesn't declare.
     c.open_session(session_configuration=None, catalog=None, schema=None)
     assert captured["host"] == "example.cloud.databricks.com"
 

tests/unit/test_session.py

@@ -466,6 +466,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
                 _retry_stop_after_attempts_count=10,
                 _retry_stop_after_attempts_duration=600.0,
                 _socket_timeout=12.5,
+                _pool_maxsize=41,
             )
             try:
                 _, kwargs = mock_kernel_client.call_args
@@ -475,6 +476,21 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
                 assert opts["retry_stop_after_attempts_count"] == 10
                 assert opts["retry_stop_after_attempts_duration"] == 600.0
                 assert kwargs["request_timeout_secs"] == 12.5
+                assert kwargs["max_connections"] == 41
+            finally:
+                conn.close()
+
+            conn = databricks.sql.connect(
+                server_hostname="foo",
+                http_path="/sql/1.0/warehouses/abc",
+                use_kernel=True,
+                access_token="dapi-xyz",
+                enable_telemetry=False,
+                _pool_maxsize=0,
+            )
+            try:
+                _, kwargs = mock_kernel_client.call_args
+                assert kwargs["max_connections"] is None
             finally:
                 conn.close()