Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions CONNECTION_PARAMETERS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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. |
Expand Down Expand Up @@ -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.

Expand Down
2 changes: 1 addition & 1 deletion KERNEL_REV
Original file line number Diff line number Diff line change
@@ -1 +1 @@
d64009eb59404c1b082cb020296337f96dc0d4d7
167651ecc67143ef258ad70fd2682ced00beaa99
7 changes: 6 additions & 1 deletion src/databricks/sql/backend/kernel/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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 {}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -421,6 +425,7 @@ def open_session(
**tls_kwargs,
**retry_kwargs,
**telemetry_kwargs,
**max_connections_kwargs,
**http_headers_kwargs,
)
except Exception as exc:
Expand Down
1 change: 1 addition & 0 deletions src/databricks/sql/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — kwargs.get("_pool_maxsize") or None only normalizes None and 0 to None; a negative value (e.g. _pool_maxsize=-1) is truthy and is forwarded verbatim to the kernel's max_connections. The PR contract is "positive value is forwarded; unset or 0 keeps the default," so a negative should arguably also collapse to the default rather than reaching the kernel's Rust HTTP-pool config (where it may error at Session construction or be interpreted unexpectedly). Consider gating on > 0 instead:

max_connections=(_v if (_v := kwargs.get("_pool_maxsize")) and _v > 0 else None)

or validating upstream. Low severity since negative values are user error on an internal underscore-prefixed param and would fail loud rather than corrupt data.

telemetry_options=kernel_telemetry_options,
)

Expand Down
58 changes: 40 additions & 18 deletions tests/unit/test_kernel_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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``."""
Expand Down Expand Up @@ -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,
*,
Expand Down Expand Up @@ -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",
Expand All @@ -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"

Expand Down
16 changes: 16 additions & 0 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()

Expand Down
Loading