fix(kernel): honor use_cloud_fetch

Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>

Author: vuanhphung

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

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,

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,9 +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`
     logger — proving the tracing -> log -> pyo3-log -> logging chain."""

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

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