From 42f0c48c8a0e7a43e53e8d13dd45ad40140b0ebb Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Tue, 10 Nov 2026 05:44:48 +0000 Subject: [PATCH] Skip unused Python HTTP client for explicit kernel connections Signed-off-by: Aakash Saravanan --- src/databricks/sql/client.py | 19 +++- src/databricks/sql/session.py | 2 +- tests/unit/test_kernel_async_readers.py | 18 ++- tests/unit/test_session.py | 140 ++++++++++++++++++++++++ 4 files changed, 174 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 6e03ed227..613b62386 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -413,7 +413,12 @@ def read(self) -> Optional[OAuthToken]: ) client_context = build_client_context(server_hostname, __version__, **kwargs) - self.http_client = UnifiedHttpClient(client_context) + # The kernel owns its HTTP transport. + self.http_client: Optional[UnifiedHttpClient] = ( + None + if kwargs.get("use_kernel", False) + else UnifiedHttpClient(client_context) + ) try: self.session = self._open_session_with_reyden_fallback( @@ -483,10 +488,13 @@ def read(self) -> Optional[OAuthToken]: ) # Determine proxy usage - use_proxy = self.http_client.using_proxy() + use_proxy = ( + self.http_client.using_proxy() if self.http_client is not None else False + ) proxy_host_info = None if ( - use_proxy + self.http_client is not None + and use_proxy and self.http_client.proxy_uri and isinstance(self.http_client.proxy_uri, str) ): @@ -1386,6 +1394,8 @@ def _handle_staging_put( session_id_hex=self.connection.get_session_id_hex(), ) + # Kernel backends reject staging operations before reaching these handlers. + assert self.connection.http_client is not None with open(local_file, "rb") as fh: r = self.connection.http_client.request( HttpMethod.PUT, presigned_url, body=fh.read(), headers=headers @@ -1444,6 +1454,7 @@ def _handle_staging_put_stream( session_id_hex=self.connection.get_session_id_hex(), ) + assert self.connection.http_client is not None r = self.connection.http_client.request( HttpMethod.PUT, presigned_url, body=stream.read(), headers=headers ) @@ -1466,6 +1477,7 @@ def _handle_staging_get( session_id_hex=self.connection.get_session_id_hex(), ) + assert self.connection.http_client is not None r = self.connection.http_client.request( HttpMethod.GET, presigned_url, headers=headers ) @@ -1490,6 +1502,7 @@ def _handle_staging_remove( ): """Make an HTTP DELETE request to the presigned_url""" + assert self.connection.http_client is not None r = self.connection.http_client.request( HttpMethod.DELETE, presigned_url, headers=headers ) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 365b06f8f..e2b27a980 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -52,7 +52,7 @@ def __init__( self, server_hostname: str, http_path: str, - http_client: UnifiedHttpClient, + http_client: Optional[UnifiedHttpClient], http_headers: Optional[List[Tuple[str, str]]] = None, session_configuration: Optional[Dict[str, Any]] = None, catalog: Optional[str] = None, diff --git a/tests/unit/test_kernel_async_readers.py b/tests/unit/test_kernel_async_readers.py index bb770368d..93268f883 100644 --- a/tests/unit/test_kernel_async_readers.py +++ b/tests/unit/test_kernel_async_readers.py @@ -6,6 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import SimpleNamespace +from unittest.mock import patch from urllib.parse import urlsplit import pytest @@ -132,7 +133,10 @@ def handle_request(self): def log_message(self, *_args): pass - with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server, patch( + "databricks.sql.client.UnifiedHttpClient", + side_effect=AssertionError("Explicit kernel must not create a Python client"), + ): worker = threading.Thread(target=server.serve_forever, daemon=True) worker.start() conn = None @@ -144,6 +148,7 @@ def log_message(self, *_args): use_kernel=True, enable_telemetry=False, ) + assert conn.http_client is None yield conn, state finally: state.release.set() @@ -180,6 +185,17 @@ def _assert_read_routes(state): assert execute_bodies[0]["wait_timeout"] == "0s" +def test_sync_cursor_fetch_without_python_http_client(async_connection): + conn, state = async_connection + with conn.cursor() as cursor: + cursor.execute("SELECT id AS n FROM range(1, 5)") + assert cursor.fetchall() == [(1,), (2,), (3,), (4,)] + assert any( + method == "GET" and path.endswith("/result/chunks/1") + for method, path, _ in state.requests + ) + + def test_async_cursor_refetch_preserves_session(async_connection): from databricks.sql.backend.types import CommandState diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index a9f774208..2ca944dc5 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -33,6 +33,60 @@ class TestSession: "enable_telemetry": False, } + @pytest.mark.parametrize( + "backend_kwargs, backend_name", + [ + ({}, "ThriftDatabricksClient"), + ({"use_kernel": False}, "ThriftDatabricksClient"), + ({"use_sea": True}, "SeaDatabricksClient"), + ], + ) + @pytest.mark.parametrize("proxy_uri", [None, "http://proxy.example.com:8081"]) + def test_python_backends_keep_http_client( + self, backend_kwargs, backend_name, proxy_uri + ): + with patch(f"{self.PACKAGE_NAME}.session.{backend_name}") as backend, patch( + f"{self.PACKAGE_NAME}.client.UnifiedHttpClient" + ) as http_client, patch( + f"{self.PACKAGE_NAME}.client.TelemetryClientFactory.get_telemetry_client" + ) as telemetry: + http_client.return_value.using_proxy.return_value = proxy_uri is not None + http_client.return_value.proxy_uri = proxy_uri + conn = databricks.sql.connect( + **self.DUMMY_CONNECTION_ARGS, **backend_kwargs + ) + try: + http_client.assert_called_once() + assert conn.http_client is http_client.return_value + assert conn.session.http_client is conn.http_client + assert backend.call_args.kwargs["http_client"] is conn.http_client + params = telemetry.return_value.export_initial_telemetry_log.call_args.kwargs[ + "driver_connection_params" + ] + assert params.use_proxy is (proxy_uri is not None) + assert params.use_system_proxy is (proxy_uri is not None) + if proxy_uri: + assert params.proxy_host_info.host_url == "proxy.example.com" + assert params.proxy_host_info.port == 8081 + else: + assert params.proxy_host_info is None + finally: + conn.close() + backend.return_value.close_session.assert_called_once() + http_client.return_value.close.assert_called_once() + + def test_conflicting_kernel_and_sea_flags_preserve_error(self): + with patch( + f"{self.PACKAGE_NAME}.client.UnifiedHttpClient", + side_effect=AssertionError( + "Explicit kernel must not create a Python client" + ), + ) as http_client, pytest.raises(ValueError, match="mutually exclusive"): + databricks.sql.connect( + **self.DUMMY_CONNECTION_ARGS, use_kernel=True, use_sea=True + ) + http_client.assert_not_called() + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) def test_close_uses_the_correct_session_id(self, mock_client_class): instance = mock_client_class.return_value @@ -877,6 +931,92 @@ def _connect(self, **overrides): args.update(overrides) return databricks.sql.connect(**args) + @pytest.mark.parametrize("route", ["explicit", "recovery", "cached"]) + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_kernel_http_client_lifecycle(self, mock_thrift, route): + from databricks.sql.backend import reyden_warehouse_cache + from databricks.sql.result_set import ResultSet + + explicit = route == "explicit" + mock_thrift.return_value.open_session.side_effect = self._reject() + if route == "cached": + reyden_warehouse_cache.mark_reyden(self.HOST, "wh-reyden") + + with self._fake_kernel() as kernel, patch( + f"{self.PACKAGE}.client.UnifiedHttpClient" + ) as http_client, patch( + f"{self.PACKAGE}.client.TelemetryClientFactory.get_telemetry_client" + ) as telemetry: + if explicit: + http_client.side_effect = AssertionError( + "Explicit kernel must not create a Python client" + ) + http_client.return_value.using_proxy.return_value = True + http_client.return_value.proxy_uri = "http://proxy.example.com:8081" + result = Mock(spec=ResultSet, is_staging_operation=False) + result.fetchall.return_value = [(1,)] + kernel.return_value.execute_command.return_value = result + conn = self._connect( + **({"use_kernel": True} if explicit else {}), + enable_telemetry=True, + force_enable_telemetry=True, + ) + try: + assert conn.session.use_kernel is True + assert conn.telemetry_enabled is False + assert conn.http_client is ( + None if explicit else http_client.return_value + ) + assert conn.session.http_client is conn.http_client + assert kernel.call_args.kwargs["http_client"] is conn.http_client + params = telemetry.return_value.export_initial_telemetry_log.call_args.kwargs[ + "driver_connection_params" + ] + assert params.use_proxy is (not explicit) + assert params.use_system_proxy is (not explicit) + if explicit: + assert params.proxy_host_info is None + with conn.cursor() as cursor: + cursor.execute("SELECT 1") + assert cursor.fetchall() == [(1,)] + kernel.return_value.execute_command.assert_called_once() + finally: + conn.close() + kernel.return_value.close_session.assert_called_once() + if explicit: + http_client.assert_not_called() + else: + http_client.assert_called_once() + http_client.return_value.close.assert_called_once() + assert mock_thrift.return_value.open_session.call_count == ( + 1 if route == "recovery" else 0 + ) + + @pytest.mark.parametrize("failure_stage", ["backend", "open"]) + def test_explicit_kernel_failure_preserves_error_and_telemetry(self, failure_stage): + with self._fake_kernel() as kernel, patch( + f"{self.PACKAGE}.client.UnifiedHttpClient", + side_effect=AssertionError( + "Explicit kernel must not create a Python client" + ), + ) as http_client, patch( + f"{self.PACKAGE}.client.TelemetryClientFactory.connection_failure_log" + ) as failure_log: + if failure_stage == "backend": + error = ImportError("kernel package unavailable") + kernel.side_effect = error + else: + error = OperationalError("kernel could not open session") + kernel.return_value.open_session.side_effect = error + with pytest.raises(type(error)) as excinfo: + self._connect(use_kernel=True, enable_telemetry=True) + assert excinfo.value is error + http_client.assert_not_called() + failure_log.assert_called_once() + assert failure_log.call_args.kwargs["enable_telemetry"] is ( + failure_stage == "backend" + ) + @pytest.mark.parametrize("route", ["explicit", "recovery", "cached"]) @pytest.mark.parametrize( "geospatial_kwargs, expected",