From 1dd601e5f3913ad9a576b888453603db022f6a1d Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 18 Aug 2026 06:03:33 +0000 Subject: [PATCH] fix(kernel): honor use_cloud_fetch Signed-off-by: Vu Anh Phung --- src/databricks/sql/backend/kernel/client.py | 21 ++++++++--- src/databricks/sql/session.py | 1 + tests/e2e/test_kernel_backend.py | 27 ++++++++++++- tests/unit/test_kernel_client.py | 42 +++++++++++++++++++++ tests/unit/test_session.py | 38 +++++++++++++++++++ 5 files changed, 121 insertions(+), 8 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index b1a1d5b3e..4e3d658ca 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -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)) # 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 @@ -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 diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a83d62db1..c79ed6083 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -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, diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index 8b532a56a..49b2679b0 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -21,6 +21,7 @@ from __future__ import annotations +import logging import sys from uuid import uuid4 @@ -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 @@ -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` diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 79be53e64..c896cf8a6 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -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 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index ba008b103..81c2f9182 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -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