Skip to content

Commit 388a3e7

Browse files
committed
fix(kernel): gate max connections by wheel support
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
1 parent dcb6b0b commit 388a3e7

2 files changed

Lines changed: 14 additions & 12 deletions

File tree

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ def _kernel_session_accepts_kwarg(name: str) -> bool:
156156
``**kwargs`` catch-all), so forwarding a kwarg it doesn't declare raises
157157
``TypeError`` at construction, so we gate kwargs on what the installed
158158
wheel supports. Falls **closed** (returns ``False``) when the signature
159-
can't be introspected because omitting an accepted telemetry kwarg is safer
159+
can't be introspected because omitting an accepted optional kwarg is safer
160160
than forwarding an unsupported one.
161161
"""
162162
try:
@@ -378,6 +378,9 @@ def open_session(
378378
# kernel's ``retry_*`` kwargs. Empty when at defaults.
379379
retry_kwargs = _kernel_retry_kwargs(self._retry_options)
380380
telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
381+
max_connections_kwargs: Dict[str, Any] = {}
382+
if _kernel_session_accepts_kwarg("max_connections"):
383+
max_connections_kwargs["max_connections"] = self._max_connections
381384
# Forward caller / connector HTTP headers. The kernel applies
382385
# them on every request; a caller ``User-Agent`` is appended
383386
# to the kernel's base UA. Only pass the kwarg when there's
@@ -418,11 +421,11 @@ def open_session(
418421
# strings).
419422
intervals_as_string=True,
420423
request_timeout_secs=self._request_timeout_secs,
421-
max_connections=self._max_connections,
422424
**auth_kwargs,
423425
**tls_kwargs,
424426
**retry_kwargs,
425427
**telemetry_kwargs,
428+
**max_connections_kwargs,
426429
**http_headers_kwargs,
427430
)
428431
except Exception as exc:

tests/unit/test_kernel_client.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -452,20 +452,19 @@ def fake_session(**kw):
452452
assert captured["telemetry_circuit_breaker_enabled"] is False
453453

454454

455-
def test_open_session_omits_phase_7_kwargs_kernel_does_not_accept(monkeypatch):
456-
"""Phase-7 identity/telemetry kwargs must NOT be forwarded to a kernel
457-
``Session`` whose (fixed, no-``**kwargs``) constructor doesn't declare
458-
them.
455+
def test_open_session_omits_optional_kwargs_kernel_does_not_accept(monkeypatch):
456+
"""Optional kwargs must NOT be forwarded to a kernel ``Session`` whose
457+
fixed constructor doesn't declare them.
459458
460459
The real ``databricks_sql_kernel.Session`` is a PyO3 class with a fixed
461460
signature. The other tests here use a ``**kwargs`` MagicMock that silently
462461
swallows unsupported kwargs; this fixed-signature fake proves the client
463-
filters optional telemetry kwargs.
462+
filters optional kwargs.
464463
"""
465464
captured = {}
466465

467-
# Accept baseline connection/tls/retry/pool kwargs but no phase-7 identity
468-
# or telemetry kwargs, and no **kwargs catch-all.
466+
# Accept baseline connection/tls/retry kwargs but no max-connections,
467+
# phase-7 identity, or telemetry kwargs, and no **kwargs catch-all.
469468
def fake_session_v0_2_0(
470469
host,
471470
http_path,
@@ -494,7 +493,6 @@ def fake_session_v0_2_0(
494493
complex_types_as_json=False,
495494
intervals_as_string=False,
496495
request_timeout_secs=None,
497-
max_connections=None,
498496
):
499497
captured["host"] = host
500498
sess = MagicMock()
@@ -532,10 +530,11 @@ def fake_session_v0_2_0(
532530
http_path="/sql/1.0/warehouses/abc",
533531
auth_provider=AccessTokenAuthProvider("dapi-test"),
534532
ssl_options=None,
533+
max_connections=41,
535534
telemetry_options={"enable_telemetry": True, "telemetry_batch_size": 17},
536535
)
537-
# Would raise TypeError: unexpected keyword argument if the client
538-
# forwarded phase-7 kwargs the fixed-signature Session doesn't declare.
536+
# Would raise TypeError if the client forwarded max-connections or
537+
# phase-7 kwargs the fixed-signature Session doesn't declare.
539538
c.open_session(session_configuration=None, catalog=None, schema=None)
540539
assert captured["host"] == "example.cloud.databricks.com"
541540

0 commit comments

Comments
 (0)