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
21 changes: 15 additions & 6 deletions src/databricks/sql/backend/kernel/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,9 @@ def __init__(
self._use_arrow_native_complex_types = kwargs.get(
"_use_arrow_native_complex_types", True
)
# This is a connection option: the kernel fixes the SEA result
# disposition policy for the lifetime of its session.
self._use_cloud_fetch = bool(kwargs.get("use_cloud_fetch", True))
Comment thread
vuanhphung marked this conversation as resolved.
# NB: don't call ``kernel_auth_kwargs`` here. That call
# materialises the bearer token in-process; keeping a
# cleartext copy on a long-lived connector object that may
Expand Down Expand Up @@ -293,12 +296,18 @@ def open_session(
) -> SessionId:
if self._kernel_session is not None:
raise InterfaceError("KernelDatabricksClient already has an open session.")
# ``session_configuration`` flows through to the kernel's
# ``session_conf`` map verbatim; the SEA endpoint enforces
# its own allow-list and rejects unknown keys.
session_conf: Optional[Dict[str, str]] = None
if session_configuration:
session_conf = {k: str(v) for k, v in session_configuration.items()}
# Convert server session confs to strings, then add the kernel's
# client-side CloudFetch knob to the same boundary map.
session_conf = (
{k: str(v) for k, v in session_configuration.items()}
if session_configuration
else {}
)
# The kernel consumes this before filtering the server confs and
# selects INLINE when CloudFetch is disabled.
session_conf["cloudfetch_enabled"] = (
"true" if self._use_cloud_fetch else "false"
)
# The kwarg builds run INSIDE the try so the ``finally`` scrub
# below always fires — including when ``kernel_auth_kwargs``
# itself raises mid-build (e.g. an OAuth token-exchange failure
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 @@ -204,6 +204,7 @@ def _create_backend(
http_client=self.http_client,
catalog=kwargs.get("catalog"),
schema=kwargs.get("schema"),
use_cloud_fetch=kwargs.get("use_cloud_fetch", True),
_use_arrow_native_complex_types=_use_arrow_native_complex_types,
auth_options=kernel_auth_options,
retry_options=kernel_retry_options,
Expand Down
27 changes: 25 additions & 2 deletions tests/e2e/test_kernel_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

from __future__ import annotations

import logging
import sys
from uuid import uuid4

Expand Down Expand Up @@ -163,6 +164,30 @@ def test_drain_large_range_to_arrow(conn):
assert len(rows) == 10000


@pytest.mark.realkernel
def test_use_cloud_fetch_false_uses_inline_results(kernel_conn_params, caplog):
"""The real wheel consumes the client knob and selects inline results."""
params = dict(kernel_conn_params)
params["use_cloud_fetch"] = False

with caplog.at_level(logging.INFO, logger="databricks.sql.kernel"):
with sql.connect(**params) as c:
with c.cursor() as cur:
# Large enough to exercise multi-chunk inline delivery.
cur.execute("SELECT * FROM range(5000000)")
assert cur.fetchmany(1)[0][0] == 0

messages = [
record.getMessage()
for record in caplog.records
if record.name.startswith("databricks.sql.kernel")
]
assert any("Using inline" in message for message in messages), messages
assert not any(
"Using CloudFetch reader" in message for message in messages
), messages


def test_fetchmany_pacing(conn):
"""fetchmany honours the requested size and stops cleanly at
end-of-stream — covers the buffer-slicing logic in
Expand Down Expand Up @@ -194,8 +219,6 @@ def test_fetchall_arrow(conn):
# `databricks.sql.kernel.pyo3`. If the kernel's tracing target or the
# pyo3-log wiring ever drifts, these fail.

import logging


def test_kernel_logs_reach_python_logging(kernel_conn_params, caplog):
"""A query at DEBUG produces records on the `databricks.sql.kernel`
Expand Down
42 changes: 42 additions & 0 deletions tests/unit/test_kernel_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,48 @@ def fake_session(**kw):
assert captured.get("complex_types_as_json") is expected_flag


@pytest.mark.parametrize(
"client_kwargs, expected",
[
({}, "true"),
({"use_cloud_fetch": True}, "true"),
({"use_cloud_fetch": False}, "false"),
({"use_cloud_fetch": None}, "false"),
({"use_cloud_fetch": "false"}, "true"),
],
)
def test_open_session_passes_cloud_fetch_setting_to_kernel(
monkeypatch, client_kwargs, expected
):
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,
**client_kwargs,
)
c.open_session(
session_configuration={"ANSI_MODE": "false"},
catalog=None,
schema=None,
)

assert captured["session_conf"] == {
"ANSI_MODE": "false",
"cloudfetch_enabled": expected,
}


def test_execute_command_forwards_parameters_to_bind_param():
"""``execute_command(parameters=[...])`` routes each parameter
through ``bind_tspark_params`` onto the kernel statement before
Expand Down
38 changes: 38 additions & 0 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,44 @@ def test_retry_kwargs_threaded_into_kernel_client(self):
conn.close()


class TestKernelCloudFetchThreading:
def test_use_cloud_fetch_threaded_into_kernel_client(self):
import sys
import types

pytest.importorskip(
"pyarrow",
reason="kernel client module imports pyarrow at load",
)

fake = types.ModuleType("databricks_sql_kernel")
fake.KernelError = type("KernelError", (Exception,), {})
fake.Session = MagicMock()

with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
"databricks.sql.backend.kernel.client.KernelDatabricksClient"
) as mock_kernel_client, patch(
"databricks.sql.session.get_python_sql_connector_auth_provider"
):
instance = mock_kernel_client.return_value
instance.open_session.return_value = SessionId(
BackendType.SEA, "sess-id", None
)

conn = databricks.sql.connect(
server_hostname="foo",
http_path="/sql/1.0/warehouses/abc",
use_kernel=True,
use_cloud_fetch=False,
access_token="dapi-xyz",
enable_telemetry=False,
)
try:
assert mock_kernel_client.call_args.kwargs["use_cloud_fetch"] is False
finally:
conn.close()


class TestKernelUserAgentForwarding:
"""user_agent_entry must reach the kernel on the use_kernel path —
session.py folds it into the composed User-Agent and includes it in
Expand Down
Loading