Skip to content

Commit 3792886

Browse files
committed
feat(kernel): forward _pool_maxsize
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
1 parent a22ac63 commit 3792886

5 files changed

Lines changed: 39 additions & 4 deletions

File tree

CONNECTION_PARAMETERS.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ to change without notice.
102102
| ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- |
103103
| `_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. |
104104
| `_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. |
105-
| `_pool_maxsize` | `int` || ⚠️ | `20` | Max connections per pool on the shared Python HTTP client. Same kernel caveat as `_pool_connections`. |
105+
| `_pool_maxsize` | `int` || | `20` (Thrift); `100` (kernel when unset) | Max idle connections retained per host. An explicit value configures both the shared Python HTTP client and the kernel's Rust HTTP pool. |
106106
| `_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). |
107107
| `_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. |
108108
| `_retry_stop_after_attempts_duration`| `float` (s) ||| `900` | Max total wall-clock seconds spent retrying. Forwarded to the kernel. |
@@ -202,9 +202,9 @@ None — the kernel's parameter surface is currently a subset of Thrift's.
202202

203203
### Behavioral divergences to watch
204204

205-
- **Connection pooling / proxy** (`_pool_connections`, `_pool_maxsize`,
206-
`_proxy_auth_method`) configure the connector's shared Python HTTP client
207-
(auth/telemetry); the kernel's query traffic uses its own Rust transport.
205+
- **Connection pooling / proxy**: `_pool_connections` and `_proxy_auth_method`
206+
configure only the shared Python HTTP client. `_pool_maxsize` also configures
207+
the kernel's Rust HTTP pool when explicitly set.
208208
- **`use_inline_params`** renders parameters inline on Thrift; the kernel uses
209209
native parameter binding.
210210

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,7 @@ def __init__(
260260
self._retry_options = kwargs.get("retry_options") or {}
261261
# The kernel binding owns type and range validation.
262262
self._request_timeout_secs = kwargs.get("request_timeout_secs")
263+
self._max_connections = kwargs.get("max_connections")
263264
# Kernel telemetry phase 7 adds binding/runtime identity and
264265
# telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
265266
self._telemetry_options = kwargs.get("telemetry_options") or {}
@@ -399,6 +400,9 @@ def open_session(
399400
]
400401
if forwarded:
401402
http_headers_kwargs["http_headers"] = forwarded
403+
pool_kwargs: Dict[str, Any] = {}
404+
if self._max_connections is not None:
405+
pool_kwargs["max_connections"] = self._max_connections
402406
self._kernel_session = _kernel.Session(
403407
host=self._server_hostname,
404408
http_path=self._http_path,
@@ -417,6 +421,7 @@ def open_session(
417421
# strings).
418422
intervals_as_string=True,
419423
request_timeout_secs=self._request_timeout_secs,
424+
**pool_kwargs,
420425
**auth_kwargs,
421426
**tls_kwargs,
422427
**retry_kwargs,

src/databricks/sql/session.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,7 @@ def _create_backend(
313313
auth_options=kernel_auth_options,
314314
retry_options=kernel_retry_options,
315315
request_timeout_secs=kwargs.get("_socket_timeout"),
316+
max_connections=kwargs.get("_pool_maxsize"),
316317
telemetry_options=kernel_telemetry_options,
317318
)
318319

tests/unit/test_kernel_client.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,33 @@ def fake_session(**kw):
368368
assert captured["request_timeout_secs"] == timeout
369369

370370

371+
@pytest.mark.parametrize("max_connections", [None, 41])
372+
def test_open_session_passes_max_connections_to_kernel(monkeypatch, max_connections):
373+
captured = {}
374+
375+
def fake_session(**kw):
376+
captured.update(kw)
377+
sess = MagicMock()
378+
sess.session_id = "sess-id"
379+
return sess
380+
381+
monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
382+
c = kernel_client.KernelDatabricksClient(
383+
server_hostname="example.cloud.databricks.com",
384+
http_path="/sql/1.0/warehouses/abc",
385+
auth_provider=AccessTokenAuthProvider("dapi-test"),
386+
ssl_options=None,
387+
max_connections=max_connections,
388+
)
389+
390+
c.open_session(session_configuration=None, catalog=None, schema=None)
391+
392+
if max_connections is None:
393+
assert "max_connections" not in captured
394+
else:
395+
assert captured["max_connections"] == max_connections
396+
397+
371398
def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
372399
"""Kernel telemetry phase 7 added binding/runtime identity and
373400
telemetry config kwargs to ``databricks_sql_kernel.Session``."""

tests/unit/test_session.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
466466
_retry_stop_after_attempts_count=10,
467467
_retry_stop_after_attempts_duration=600.0,
468468
_socket_timeout=12.5,
469+
_pool_maxsize=41,
469470
)
470471
try:
471472
_, kwargs = mock_kernel_client.call_args
@@ -475,6 +476,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self):
475476
assert opts["retry_stop_after_attempts_count"] == 10
476477
assert opts["retry_stop_after_attempts_duration"] == 600.0
477478
assert kwargs["request_timeout_secs"] == 12.5
479+
assert kwargs["max_connections"] == 41
478480
finally:
479481
conn.close()
480482

0 commit comments

Comments
 (0)