feat(kernel): honor _connection_uri and _port on the use_kernel path

The kernel branch of Session._create_backend forwarded only server_hostname
and http_path, so _connection_uri and _port were silently ignored on
use_kernel=True (connection reached server_hostname/http_path with no error).

Add _kernel_host_and_path(): decompose _connection_uri into the kernel host
(scheme+authority) + http_path, and fold _port into the host authority.
No kernel change needed — the kernel Session host accepts a fully-qualified
https://host:port and its normalise_host preserves scheme and port.

PECOBLR-4151.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

Author: eric-wang-1990

src/databricks/sql/session.py

@@ -1,6 +1,7 @@
 import logging
 import re
 from typing import Dict, Tuple, List, Optional, Any, Type
+from urllib.parse import urlsplit
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
@@ -20,6 +21,56 @@
 logger = logging.getLogger(__name__)
 
 
+def _kernel_host_and_path(
+    server_hostname: str, http_path: str, kwargs: dict
+) -> Tuple[str, str]:
+    """Resolve the ``(host, http_path)`` the kernel ``Session`` should use,
+    honoring the Thrift-style ``_connection_uri`` / ``_port`` overrides.
+
+    The kernel derives its endpoint from ``host`` + ``http_path`` and accepts a
+    fully-qualified ``host`` (its ``normalise_host`` preserves the scheme and
+    any port), so both overrides can be expressed connector-side without a
+    kernel change:
+
+    - ``_connection_uri`` (a full ``scheme://host[:port]/path`` URI, mirroring
+      the Thrift backend's direct-URI override) is split into its authority
+      (returned as ``host``) and its path+query (returned as ``http_path``).
+      ``_connection_uri`` wins over ``_port``, matching the Thrift backend.
+    - ``_port`` is otherwise folded into the host authority, unless the
+      hostname already carries a port.
+
+    Neither override is set on the common path, so the connection's
+    ``server_hostname`` / ``http_path`` pass through unchanged.
+    """
+    connection_uri = kwargs.get("_connection_uri")
+    if connection_uri:
+        # Ensure a scheme so urlsplit populates netloc rather than path; the
+        # Thrift backend defaults a scheme-less URI to https, so do the same.
+        uri = connection_uri if "://" in connection_uri else "https://" + connection_uri
+        parts = urlsplit(uri)
+        host = "{}://{}".format(parts.scheme, parts.netloc)
+        path = parts.path or http_path
+        if parts.query:
+            path = "{}?{}".format(path, parts.query)
+        return host, path
+
+    port = kwargs.get("_port")
+    if port is not None:
+        # Split off any scheme so we can inspect the authority; the kernel
+        # re-adds https:// when it is absent. Only append the port when the
+        # authority does not already carry one.
+        scheme_match = re.match(r"^(https?://)(.*)$", server_hostname)
+        scheme = scheme_match.group(1) if scheme_match else ""
+        authority = (scheme_match.group(2) if scheme_match else server_hostname).rstrip(
+            "/"
+        )
+        if ":" not in authority:
+            authority = "{}:{}".format(authority, port)
+        return "{}{}".format(scheme, authority), http_path
+
+    return server_hostname, http_path
+
+
 class Session:
     def __init__(
         self,
@@ -195,9 +246,12 @@ def _create_backend(
                     "_retry_stop_after_attempts_duration"
                 ),
             }
+            kernel_host, kernel_http_path = _kernel_host_and_path(
+                server_hostname, http_path, kwargs
+            )
             return KernelDatabricksClient(
-                server_hostname=server_hostname,
-                http_path=http_path,
+                server_hostname=kernel_host,
+                http_path=kernel_http_path,
                 http_headers=all_headers,
                 auth_provider=auth_provider,
                 ssl_options=self.ssl_options,

tests/unit/test_session.py

@@ -9,7 +9,7 @@
 )
 from databricks.sql.backend.types import SessionId, BackendType
 from databricks.sql.common.agent import KNOWN_AGENTS
-from databricks.sql.session import Session
+from databricks.sql.session import Session, _kernel_host_and_path
 
 import databricks.sql
 
@@ -587,3 +587,69 @@ def test_connect_use_kernel_instantiates_real_kernel_backend(self):
                 )
             finally:
                 conn.close()
+
+
+class TestKernelHostAndPathOverrides:
+    """``_kernel_host_and_path`` maps the Thrift-style ``_connection_uri`` /
+    ``_port`` overrides onto the kernel ``Session``'s ``host`` + ``http_path``.
+
+    Pure-function tests (no kernel wheel / pyarrow needed) covering the
+    connector-side handling that lets these overrides work on use_kernel=True
+    without any kernel change.
+    """
+
+    HOST = "foo.cloud.databricks.com"
+    PATH = "/sql/1.0/warehouses/abc"
+
+    def test_no_override_passes_through(self):
+        assert _kernel_host_and_path(self.HOST, self.PATH, {}) == (self.HOST, self.PATH)
+
+    def test_connection_uri_split_into_authority_and_path(self):
+        host, path = _kernel_host_and_path(
+            self.HOST,
+            self.PATH,
+            {"_connection_uri": "https://direct.example.com:8443/sql/1.0/warehouses/xyz"},
+        )
+        assert host == "https://direct.example.com:8443"
+        assert path == "/sql/1.0/warehouses/xyz"
+
+    def test_connection_uri_without_scheme_defaults_to_https(self):
+        host, path = _kernel_host_and_path(
+            self.HOST, self.PATH, {"_connection_uri": "direct.example.com/sql/1.0/warehouses/xyz"}
+        )
+        assert host == "https://direct.example.com"
+        assert path == "/sql/1.0/warehouses/xyz"
+
+    def test_connection_uri_preserves_query(self):
+        host, path = _kernel_host_and_path(
+            self.HOST,
+            self.PATH,
+            {"_connection_uri": "https://h.example.com/sql/1.0/warehouses/xyz?o=123"},
+        )
+        assert host == "https://h.example.com"
+        assert path == "/sql/1.0/warehouses/xyz?o=123"
+
+    def test_connection_uri_wins_over_port(self):
+        host, path = _kernel_host_and_path(
+            self.HOST,
+            self.PATH,
+            {"_connection_uri": "https://direct.example.com:9999/p", "_port": 8443},
+        )
+        assert host == "https://direct.example.com:9999"
+        assert path == "/p"
+
+    def test_port_folded_into_bare_host(self):
+        host, path = _kernel_host_and_path(self.HOST, self.PATH, {"_port": 8443})
+        assert host == "{}:8443".format(self.HOST)
+        assert path == self.PATH
+
+    def test_port_preserves_existing_scheme(self):
+        host, path = _kernel_host_and_path(
+            "https://" + self.HOST, self.PATH, {"_port": 8443}
+        )
+        assert host == "https://{}:8443".format(self.HOST)
+        assert path == self.PATH
+
+    def test_port_not_double_appended_when_host_has_port(self):
+        host, _ = _kernel_host_and_path(self.HOST + ":7000", self.PATH, {"_port": 8443})
+        assert host == self.HOST + ":7000"