Skip unused Python HTTP client for explicit kernel connections (#1017)

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

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
         )

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,

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
 

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",