diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..d9b43313d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) # 4.9.0 (2026-11-04) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 37736ed1a..f3185a73d 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -149,7 +149,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | ------------------------------------- | ------ | :----: | :----: | ------------- | ------------------------------------------------------------------------------------------------------------- | -| `use_cloud_fetch` | `bool` | ✅ | ❌ | `True` | Download large result sets in parallel from cloud storage. The kernel manages result transport internally. | +| `use_cloud_fetch` | `bool` | ✅ | ✅ | `True` | Set `False` at connection creation to fetch complete results without cloud-storage downloads. On the kernel path, changing this attribute after connecting does not change result delivery. | | `max_download_threads` | `int` | ✅ | ❌ | `10` | Worker threads for cloud-fetch downloads. Not forwarded to the kernel. | | `enable_query_result_lz4_compression` | `bool` | ✅ | ❌ | `True` | LZ4-compress result payloads. Not forwarded; the kernel handles compression internally. | | `_disable_pandas` | `bool` | ✅ | ✅ | `False` | Skip the pandas-based Arrow→row deserialization and materialize rows directly with PyArrow. This is a **Python-side** result-conversion toggle, not a wire option: the kernel returns results as Arrow (`RecordBatch`es) and the connector runs the *same* `_convert_arrow_table` for both backends, so the flag is honored on the kernel path too. Affects only row fetches (`fetchone`/`fetchmany`/`fetchall`); the `fetch*_arrow` methods return the Arrow table unchanged regardless of this flag. | @@ -202,7 +202,7 @@ installed kernel binding. `force_enable_telemetry` remains Python-driver-only. 2. `experimental_oauth_persistence` (custom OAuth token store). 3. TLS-client-cert *authentication* (`_use_cert_as_auth`) — note the TLS *transport* options (`_tls_*`) themselves **are** honored on both backends. -4. Result-transport tuning: `use_cloud_fetch`, `max_download_threads`, +4. Result-transport tuning: `max_download_threads`, `enable_query_result_lz4_compression`. 5. Arrow-native rendering for `_use_arrow_native_decimals` / `_use_arrow_native_timestamps` (complex types **are** forwarded). diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..567b98450 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -293,6 +293,7 @@ def __init__( # The kernel binding owns type and range validation. self._request_timeout_secs = kwargs.get("request_timeout_secs") self._max_connections = kwargs.get("max_connections") + self._use_cloud_fetch = kwargs.get("use_cloud_fetch", True) # Opt-in native GEOMETRY / GEOGRAPHY support. True requests # the canonical Arrow struct and surfaces as ``{"srid": int, "wkb": # bytes}`` through pyarrow; False requests WKT / EWKT strings. This is @@ -396,6 +397,13 @@ def open_session( session_conf: Optional[Dict[str, str]] = None if session_configuration: session_conf = {k: str(v) for k, v in session_configuration.items()} + if not self._use_cloud_fetch: + session_conf = { + k: v + for k, v in (session_conf or {}).items() + if k.lower() != "can_cloud_download" + } + session_conf["can_cloud_download"] = "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 365b06f8f..9856e6771 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -332,6 +332,7 @@ def _create_backend( retry_options=kernel_retry_options, request_timeout_secs=kwargs.get("_socket_timeout"), max_connections=kwargs.get("_pool_maxsize") or None, + use_cloud_fetch=kwargs.get("use_cloud_fetch", True), enable_geospatial_support=kwargs.get( "enable_geospatial_support", False ), diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index ae0301fb6..d84f6347a 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -21,7 +21,12 @@ from __future__ import annotations +import hashlib +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import select +import socket import sys +import threading from uuid import uuid4 import pytest @@ -136,6 +141,87 @@ def conn(kernel_conn_params): c.close() +@pytest.fixture +def workspace_only_proxy(kernel_conn_params, monkeypatch): + connections = {"workspace": 0, "blocked": 0} + workspace = kernel_conn_params["server_hostname"] + ":443" + + class Handler(BaseHTTPRequestHandler): + def do_CONNECT(self): + self.close_connection = True + if self.path != workspace: + connections["blocked"] += 1 + self.send_error(403) + return + with socket.create_connection( + (kernel_conn_params["server_hostname"], 443), timeout=30 + ) as upstream: + connections["workspace"] += 1 + self.send_response(200) + self.end_headers() + while True: + ready, _, _ = select.select([self.connection, upstream], [], [], 30) + if not ready: + return + for source in ready: + data = source.recv(65536) + if not data: + return + target = ( + upstream if source is self.connection else self.connection + ) + target.sendall(data) + + def log_message(self, *_args): + pass + + proxy = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=proxy.serve_forever, daemon=True) + thread.start() + for name in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy"): + monkeypatch.setenv(name, f"http://127.0.0.1:{proxy.server_port}") + for name in ("NO_PROXY", "no_proxy"): + monkeypatch.setenv(name, "") + try: + yield connections + finally: + proxy.shutdown() + proxy.server_close() + thread.join() + + +def test_cloudfetch_disabled_returns_complete_large_result( + kernel_conn_params, workspace_only_proxy +): + from databricks.sql.backend.kernel.client import KernelDatabricksClient + + with sql.connect(**kernel_conn_params, use_cloud_fetch=False) as connection: + assert isinstance(connection.session.backend, KernelDatabricksClient) + with connection.cursor() as cursor: + cursor.execute( + "SELECT id, sha2(CAST(id AS STRING), 512) AS payload " + "FROM range(250000) ORDER BY id" + ) + count = 0 + payload_bytes = 0 + digest = hashlib.sha256() + while rows := cursor.fetchmany(8192): + for identity, payload in rows: + assert identity == count + assert payload == hashlib.sha512(str(identity).encode()).hexdigest() + payload_bytes += len(payload) + digest.update(f"{identity}\t{payload}\n".encode()) + count += 1 + assert count == 250000 + assert payload_bytes == 32000000 + assert digest.hexdigest() == ( + "a25ada4c6c476b53e97eb6e8a945e9ce4c14cc432c7ad672dc6fb607d660c86a" + ) + + assert workspace_only_proxy["workspace"] > 0 + assert workspace_only_proxy["blocked"] == 0 + + def test_connect_with_use_kernel_opens_a_session(conn): assert conn.open, "connection should report open after connect()" diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..150eeed51 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -344,6 +344,58 @@ def fake_session(**kw): assert captured.get("complex_types_as_json") is expected_flag +@pytest.mark.parametrize( + "kwargs, configuration, expected", + [ + ({}, {"TIMEZONE": "UTC"}, {"TIMEZONE": "UTC"}), + ( + {"use_cloud_fetch": True}, + {"can_cloud_download": "false"}, + {"can_cloud_download": "false"}, + ), + ({"use_cloud_fetch": False}, None, {"can_cloud_download": "false"}), + ( + {"use_cloud_fetch": False}, + {"TIMEZONE": "UTC"}, + {"TIMEZONE": "UTC", "can_cloud_download": "false"}, + ), + ( + {"use_cloud_fetch": False}, + {"can_cloud_download": "true"}, + {"can_cloud_download": "false"}, + ), + ( + {"use_cloud_fetch": False}, + {"Can_Cloud_Download": "true"}, + {"can_cloud_download": "false"}, + ), + ], +) +def test_open_session_honors_cloudfetch_option( + monkeypatch, kwargs, configuration, expected +): + captured = {} + + def fake_session(**options): + captured.update(options) + session = MagicMock() + session.session_id = "sess-id" + return session + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session) + client = kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + **kwargs, + ) + original = dict(configuration) if configuration is not None else None + client.open_session(configuration, catalog=None, schema=None) + assert captured["session_conf"] == expected + assert configuration == original + + @pytest.mark.parametrize("timeout", [None, 0, 12.5]) def test_open_session_passes_request_timeout_to_kernel(monkeypatch, timeout): captured = {} diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index a9f774208..25d617af6 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -484,6 +484,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self): _retry_stop_after_attempts_duration=600.0, _socket_timeout=12.5, _pool_maxsize=41, + use_cloud_fetch=False, enable_geospatial_support=False, ) try: @@ -495,6 +496,7 @@ def test_retry_and_socket_timeout_threaded_into_kernel_client(self): assert opts["retry_stop_after_attempts_duration"] == 600.0 assert kwargs["request_timeout_secs"] == 12.5 assert kwargs["max_connections"] == 41 + assert kwargs["use_cloud_fetch"] is False assert kwargs["enable_geospatial_support"] is False finally: conn.close()