From 1adcb29c44cc6ac3e5738d8fcf8e274a66dfba11 Mon Sep 17 00:00:00 2001
From: eric-wang-1990 <e.wang@databricks.com>
Date: Mon, 17 Aug 2026 17:45:26 -0700
Subject: [PATCH 01/10] feat(kernel): honor _connection_uri and _port on the
 use_kernel path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 src/databricks/sql/session.py | 58 ++++++++++++++++++++++++++++--
 tests/unit/test_session.py    | 68 ++++++++++++++++++++++++++++++++++-
 2 files changed, 123 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index a83d62db1..27977c2ed 100644
--- a/src/databricks/sql/session.py
+++ b/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,
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index ba008b103..21fa7adf9 100644
--- a/tests/unit/test_session.py
+++ b/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"

From 7fc2d5fcbc8a7876a28777a75d19fb79c57c7379 Mon Sep 17 00:00:00 2001
From: eric-wang-1990 <e.wang@databricks.com>
Date: Mon, 17 Aug 2026 18:17:05 -0700
Subject: [PATCH 02/10] refactor(kernel): simplify _port folding to the
 bare-host case

server_hostname reaches the backend as a bare host on this path, so drop
the defensive scheme peel/re-add: just append the port when the host has
none, and let the kernel's normalise_host add the scheme. Removes the
now-moot scheme-preservation test.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
---
 src/databricks/sql/session.py | 18 +++++++-----------
 tests/unit/test_session.py    |  7 -------
 2 files changed, 7 insertions(+), 18 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index 27977c2ed..f959869f0 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -56,17 +56,13 @@ def _kernel_host_and_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
+        # server_hostname is a bare host on this path (e.g.
+        # ``dbc-123.cloud.databricks.com``); the kernel adds the scheme.
+        # Append the port unless the host already carries one.
+        host = server_hostname.rstrip("/")
+        if ":" not in host:
+            host = "{}:{}".format(host, port)
+        return host, http_path
 
     return server_hostname, http_path
 
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 21fa7adf9..48b196d6f 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -643,13 +643,6 @@ def test_port_folded_into_bare_host(self):
         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"

From a5e0a08897f1403bc82a0d07faea9c3355471974 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 07:35:17 +0000
Subject: [PATCH 03/10] ai: apply changes for #915 (5 review threads)

Addresses:
  - #3800053148 at src/databricks/sql/session.py:48
  - #3800053178 at src/databricks/sql/session.py:52
  - #3800053209 at src/databricks/sql/session.py:69
  - #3800056285 at src/databricks/sql/session.py:245
  - #3800209433 at src/databricks/sql/session.py:245

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py |  40 +++++++++-
 tests/unit/test_session.py    | 134 ++++++++++++++++++++++++++++++++++
 2 files changed, 170 insertions(+), 4 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index f959869f0..c42e56e5e 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -44,10 +44,19 @@ def _kernel_host_and_path(
     """
     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.
+        # Ensure a scheme so urlsplit populates netloc rather than path. A
+        # scheme-less URI is ambiguous here, so default to https to match the
+        # connector's default transport.
         uri = connection_uri if "://" in connection_uri else "https://" + connection_uri
         parts = urlsplit(uri)
+        if not parts.netloc:
+            # A missing authority (e.g. a value like ``//foo`` or ``https:///p``)
+            # would otherwise yield a scheme-only host such as ``https://`` and
+            # silently connect to the wrong endpoint. Fail loudly instead.
+            raise ValueError(
+                "Invalid _connection_uri {!r}: could not determine host "
+                "authority (expected scheme://host[:port]/path)".format(connection_uri)
+            )
         host = "{}://{}".format(parts.scheme, parts.netloc)
         path = parts.path or http_path
         if parts.query:
@@ -58,9 +67,15 @@ def _kernel_host_and_path(
     if port is not None:
         # server_hostname is a bare host on this path (e.g.
         # ``dbc-123.cloud.databricks.com``); the kernel adds the scheme.
-        # Append the port unless the host already carries one.
+        # Append the port unless the host already carries one. Detect an
+        # existing port via ``urlsplit`` rather than a naive ``":" in host``
+        # check, so IPv6 literals (whose authority legitimately contains ``:``
+        # even without a port, e.g. ``[::1]``) are handled correctly.
+        # ``urlsplit`` only populates ``netloc``/``port`` when a scheme is
+        # present, so add a temporary one when the host is scheme-less.
         host = server_hostname.rstrip("/")
-        if ":" not in host:
+        probe = host if "://" in host else "https://" + host
+        if urlsplit(probe).port is None:
             host = "{}:{}".format(host, port)
         return host, http_path
 
@@ -245,6 +260,23 @@ def _create_backend(
             kernel_host, kernel_http_path = _kernel_host_and_path(
                 server_hostname, http_path, kwargs
             )
+            # The SPOG ``x-databricks-org-id`` header in ``all_headers`` was
+            # derived in ``__init__`` from the *original* ``http_path``. A
+            # ``_connection_uri`` override can rewrite the kernel path to a
+            # different workspace (a different ``?o=`` or cluster path), so
+            # re-derive the routing header from the resolved path and swap it
+            # in — otherwise the kernel would receive an org-id that points at
+            # the pre-override workspace (a silent mis-routing). A caller-set
+            # header still wins: ``_spog_headers`` is empty in that case, so the
+            # explicit header is left untouched below.
+            if kernel_http_path != http_path and self._spog_headers:
+                base_headers = [
+                    h for h in all_headers if h not in self._spog_headers.items()
+                ]
+                self._spog_headers = self._extract_spog_headers(
+                    kernel_http_path, base_headers
+                )
+                all_headers = base_headers + list(self._spog_headers.items())
             return KernelDatabricksClient(
                 server_hostname=kernel_host,
                 http_path=kernel_http_path,
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 48b196d6f..3322d37bb 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -477,6 +477,52 @@ def test_retry_kwargs_threaded_into_kernel_client(self):
             finally:
                 conn.close()
 
+    def test_remapped_host_and_path_threaded_into_kernel_client(self):
+        """The remapped ``server_hostname``/``http_path`` from
+        ``_kernel_host_and_path`` must reach ``KernelDatabricksClient``.
+        Guards against a regression that dropped or reordered the call so
+        the raw (pre-override) values leaked through — the silent-ignore
+        bug this PR fixes, which the pure-function tests alone can't catch.
+        """
+        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(
+            "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
+        ):
+            instance = mock_kernel_client.return_value
+            instance.open_session.return_value = SessionId(
+                BackendType.SEA, "sess-id", None
+            )
+
+            conn = databricks.sql.connect(
+                server_hostname="foo.cloud.databricks.com",
+                http_path="/sql/1.0/warehouses/abc",
+                use_kernel=True,
+                access_token="dapi-xyz",
+                enable_telemetry=False,
+                _connection_uri="https://direct.example.com:8443/sql/1.0/warehouses/xyz",
+            )
+            try:
+                _, kwargs = mock_kernel_client.call_args
+                # The _connection_uri override must be split and remapped
+                # onto the kernel client, not passed through raw.
+                assert kwargs["server_hostname"] == "https://direct.example.com:8443"
+                assert kwargs["http_path"] == "/sql/1.0/warehouses/xyz"
+            finally:
+                conn.close()
+
 
 class TestKernelUserAgentForwarding:
     """user_agent_entry must reach the kernel on the use_kernel path —
@@ -527,6 +573,79 @@ def test_user_agent_entry_reaches_kernel_client_http_headers(self):
                 conn.close()
 
 
+class TestKernelSpogHeaderReDerivedFromResolvedPath:
+    """On the kernel path a ``_connection_uri`` override can rewrite the
+    http_path to a different workspace. The SPOG ``x-databricks-org-id``
+    header is computed in ``__init__`` from the *original* http_path, so it
+    must be re-derived from the resolved kernel path — otherwise the kernel
+    would receive an org-id pointing at the pre-override workspace."""
+
+    PACKAGE = "databricks.sql"
+
+    def _connect_and_get_kernel_headers(self, connect_kwargs):
+        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(
+            "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
+        ):
+            instance = mock_kernel_client.return_value
+            instance.open_session.return_value = SessionId(
+                BackendType.SEA, "sess-id", None
+            )
+
+            conn = databricks.sql.connect(
+                server_hostname="foo.cloud.databricks.com",
+                use_kernel=True,
+                access_token="dapi-xyz",
+                enable_telemetry=False,
+                **connect_kwargs,
+            )
+            try:
+                _, kwargs = mock_kernel_client.call_args
+                return dict(kwargs["http_headers"])
+            finally:
+                conn.close()
+
+    def test_org_id_re_derived_from_connection_uri_override(self):
+        # Original path routes to workspace 111; the _connection_uri override
+        # points at workspace 222 — the kernel must see org-id 222.
+        headers = self._connect_and_get_kernel_headers(
+            {
+                "http_path": "/sql/1.0/warehouses/abc?o=111",
+                "_connection_uri": "https://direct.example.com/sql/1.0/warehouses/xyz?o=222",
+            }
+        )
+        assert headers.get("x-databricks-org-id") == "222"
+
+    def test_org_id_dropped_when_override_has_no_workspace(self):
+        # The override points at a path with no workspace routing info, so the
+        # stale org-id (from the original path) must not be carried over.
+        headers = self._connect_and_get_kernel_headers(
+            {
+                "http_path": "/sql/1.0/warehouses/abc?o=111",
+                "_connection_uri": "https://direct.example.com/sql/1.0/warehouses/xyz",
+            }
+        )
+        assert "x-databricks-org-id" not in headers
+
+    def test_org_id_preserved_when_no_override(self):
+        headers = self._connect_and_get_kernel_headers(
+            {"http_path": "/sql/1.0/warehouses/abc?o=111"}
+        )
+        assert headers.get("x-databricks-org-id") == "111"
+
+
 @pytest.mark.realkernel
 class TestUseKernelRoutesThroughRealWheel:
     """No-network proof that ``sql.connect(use_kernel=True)`` actually
@@ -638,6 +757,11 @@ def test_connection_uri_wins_over_port(self):
         assert host == "https://direct.example.com:9999"
         assert path == "/p"
 
+    def test_connection_uri_without_authority_raises(self):
+        for bad in ("//no-scheme-authority", "https:///only-path"):
+            with pytest.raises(ValueError, match="could not determine host authority"):
+                _kernel_host_and_path(self.HOST, self.PATH, {"_connection_uri": bad})
+
     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)
@@ -646,3 +770,13 @@ def test_port_folded_into_bare_host(self):
     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"
+
+    def test_port_folded_into_ipv6_literal_without_port(self):
+        # An IPv6 authority contains ``:`` even without a port, so a naive
+        # ``":" in host`` check would wrongly skip appending _port.
+        host, _ = _kernel_host_and_path("[::1]", self.PATH, {"_port": 8443})
+        assert host == "[::1]:8443"
+
+    def test_port_not_double_appended_for_ipv6_literal_with_port(self):
+        host, _ = _kernel_host_and_path("[::1]:7000", self.PATH, {"_port": 8443})
+        assert host == "[::1]:7000"

From 043df782d6dd5a12ced737bf42a96c0a5113d4cd Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 07:42:07 +0000
Subject: [PATCH 04/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3801996184 at src/databricks/sql/session.py:272

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py | 14 ++++++++++----
 tests/unit/test_session.py    | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index c42e56e5e..a0fcfbf87 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -266,10 +266,16 @@ def _create_backend(
             # different workspace (a different ``?o=`` or cluster path), so
             # re-derive the routing header from the resolved path and swap it
             # in — otherwise the kernel would receive an org-id that points at
-            # the pre-override workspace (a silent mis-routing). A caller-set
-            # header still wins: ``_spog_headers`` is empty in that case, so the
-            # explicit header is left untouched below.
-            if kernel_http_path != http_path and self._spog_headers:
+            # the pre-override workspace (a silent mis-routing). This must fire
+            # whenever the path changed, in *both* directions: when the original
+            # path carried routing and the override drops or changes it, and
+            # when the original had none but the override introduces one — the
+            # latter is skipped if we also gate on ``self._spog_headers``. Any
+            # stale extracted header is stripped first, then re-derived; when
+            # ``_spog_headers`` is empty the filter strips nothing. A caller-set
+            # header still wins because ``_extract_spog_headers`` re-checks the
+            # existing headers and returns ``{}`` in that case.
+            if kernel_http_path != http_path:
                 base_headers = [
                     h for h in all_headers if h not in self._spog_headers.items()
                 ]
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 3322d37bb..1b9108937 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -639,6 +639,20 @@ def test_org_id_dropped_when_override_has_no_workspace(self):
         )
         assert "x-databricks-org-id" not in headers
 
+    def test_org_id_added_when_override_introduces_workspace(self):
+        # The original path carries no workspace routing (so no org-id header is
+        # set in __init__), but the _connection_uri override introduces one —
+        # the kernel must see org-id 222 rather than falling back to default
+        # routing. This exercises the none -> o=222 direction, which is skipped
+        # if re-derivation is gated on the original path having had a header.
+        headers = self._connect_and_get_kernel_headers(
+            {
+                "http_path": "/sql/1.0/warehouses/abc",
+                "_connection_uri": "https://direct.example.com/sql/1.0/warehouses/xyz?o=222",
+            }
+        )
+        assert headers.get("x-databricks-org-id") == "222"
+
     def test_org_id_preserved_when_no_override(self):
         headers = self._connect_and_get_kernel_headers(
             {"http_path": "/sql/1.0/warehouses/abc?o=111"}

From 80a6dd931b6adff6ee09f0356d40eac946ab0e59 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 07:51:24 +0000
Subject: [PATCH 05/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802045805 at src/databricks/sql/session.py:62

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py |  9 +++++++++
 tests/unit/test_session.py    | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index a0fcfbf87..7e68a8943 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -36,6 +36,12 @@ def _kernel_host_and_path(
       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.
+      When the URI omits a path (e.g. ``https://host:8443`` or
+      ``https://host:8443?o=222``) the connection's original ``http_path`` is
+      retained, and any query on the URI is applied to that retained path — so
+      a path-less, query-bearing URI overrides only the host and query while
+      keeping the original warehouse path. The fragment (``#...``) is dropped
+      since it never goes on the wire.
     - ``_port`` is otherwise folded into the host authority, unless the
       hostname already carries a port.
 
@@ -58,6 +64,9 @@ def _kernel_host_and_path(
                 "authority (expected scheme://host[:port]/path)".format(connection_uri)
             )
         host = "{}://{}".format(parts.scheme, parts.netloc)
+        # A path-less URI keeps the connection's original http_path; any query
+        # on the URI is then applied to that retained path (host + query
+        # override, path preserved). See the docstring for the rationale.
         path = parts.path or http_path
         if parts.query:
             path = "{}?{}".format(path, parts.query)
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 1b9108937..ba67c420d 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -762,6 +762,24 @@ def test_connection_uri_preserves_query(self):
         assert host == "https://h.example.com"
         assert path == "/sql/1.0/warehouses/xyz?o=123"
 
+    def test_connection_uri_without_path_retains_original_path(self):
+        # A path-less URI overrides only the host; the connection's original
+        # http_path is retained.
+        host, path = _kernel_host_and_path(
+            self.HOST, self.PATH, {"_connection_uri": "https://h.example.com:8443"}
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == self.PATH
+
+    def test_connection_uri_without_path_applies_query_to_retained_path(self):
+        # A path-less, query-bearing URI overrides the host and applies the
+        # query to the retained original path (documented fallback semantics).
+        host, path = _kernel_host_and_path(
+            self.HOST, self.PATH, {"_connection_uri": "https://h.example.com:8443?o=222"}
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == "{}?o=222".format(self.PATH)
+
     def test_connection_uri_wins_over_port(self):
         host, path = _kernel_host_and_path(
             self.HOST,

From fde990f9662395a9cbab4467b364d1779ee4b9a4 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 07:56:54 +0000
Subject: [PATCH 06/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802102078 at src/databricks/sql/session.py:72

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py | 16 +++++++++++-----
 tests/unit/test_session.py    | 12 ++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index 7e68a8943..07fb0bcc6 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -38,10 +38,11 @@ def _kernel_host_and_path(
       ``_connection_uri`` wins over ``_port``, matching the Thrift backend.
       When the URI omits a path (e.g. ``https://host:8443`` or
       ``https://host:8443?o=222``) the connection's original ``http_path`` is
-      retained, and any query on the URI is applied to that retained path — so
-      a path-less, query-bearing URI overrides only the host and query while
-      keeping the original warehouse path. The fragment (``#...``) is dropped
-      since it never goes on the wire.
+      retained, and any query on the URI is applied to that retained path,
+      replacing any query the retained path already carried — so a path-less,
+      query-bearing URI overrides only the host and query while keeping the
+      original warehouse path. The fragment (``#...``) is dropped since it
+      never goes on the wire.
     - ``_port`` is otherwise folded into the host authority, unless the
       hostname already carries a port.
 
@@ -69,7 +70,12 @@ def _kernel_host_and_path(
         # override, path preserved). See the docstring for the rationale.
         path = parts.path or http_path
         if parts.query:
-            path = "{}?{}".format(path, parts.query)
+            # Drop any query already on the retained path before applying the
+            # URI's query, so the URI's query fully overrides it. Appending
+            # unconditionally would otherwise yield a malformed double-``?``
+            # path (e.g. ``/warehouses/abc?o=111?o=222``) that mis-parses
+            # downstream and silently drops the org-id header.
+            path = "{}?{}".format(path.split("?", 1)[0], parts.query)
         return host, path
 
     port = kwargs.get("_port")
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index ba67c420d..bab45b12d 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -780,6 +780,18 @@ def test_connection_uri_without_path_applies_query_to_retained_path(self):
         assert host == "https://h.example.com:8443"
         assert path == "{}?o=222".format(self.PATH)
 
+    def test_connection_uri_without_path_query_replaces_original_path_query(self):
+        # When the retained original http_path already carries a query, the
+        # URI's query replaces it rather than being appended after a second
+        # ``?`` (which would produce a malformed double-query path).
+        host, path = _kernel_host_and_path(
+            self.HOST,
+            "/sql/1.0/warehouses/abc?o=111",
+            {"_connection_uri": "https://h.example.com:8443?o=222"},
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == "/sql/1.0/warehouses/abc?o=222"
+
     def test_connection_uri_wins_over_port(self):
         host, path = _kernel_host_and_path(
             self.HOST,

From 9d06e483eaa6486e6988bd70b8eeb66677531146 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 08:03:15 +0000
Subject: [PATCH 07/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802164062 at src/databricks/sql/session.py:71

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py |  2 +-
 tests/unit/test_session.py    | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index 07fb0bcc6..f31e133a9 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -75,7 +75,7 @@ def _kernel_host_and_path(
             # unconditionally would otherwise yield a malformed double-``?``
             # path (e.g. ``/warehouses/abc?o=111?o=222``) that mis-parses
             # downstream and silently drops the org-id header.
-            path = "{}?{}".format(path.split("?", 1)[0], parts.query)
+            path = "{}?{}".format((path or "").split("?", 1)[0], parts.query)
         return host, path
 
     port = kwargs.get("_port")
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index bab45b12d..8d31164f4 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -792,6 +792,16 @@ def test_connection_uri_without_path_query_replaces_original_path_query(self):
         assert host == "https://h.example.com:8443"
         assert path == "/sql/1.0/warehouses/abc?o=222"
 
+    def test_connection_uri_without_path_and_none_http_path_applies_query(self):
+        # http_path is legitimately nullable on the connect path. A path-less,
+        # query-bearing URI combined with http_path=None must not raise; the
+        # retained path coerces to empty before the query is applied.
+        host, path = _kernel_host_and_path(
+            self.HOST, None, {"_connection_uri": "https://h.example.com:8443?o=222"}
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == "?o=222"
+
     def test_connection_uri_wins_over_port(self):
         host, path = _kernel_host_and_path(
             self.HOST,

From a699cb66b14dcc01e0a80ee6d5b19385059598f2 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 08:08:54 +0000
Subject: [PATCH 08/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802212055 at src/databricks/sql/session.py:24

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index f31e133a9..c48e59d98 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -22,8 +22,8 @@
 
 
 def _kernel_host_and_path(
-    server_hostname: str, http_path: str, kwargs: dict
-) -> Tuple[str, str]:
+    server_hostname: str, http_path: Optional[str], kwargs: dict
+) -> Tuple[str, Optional[str]]:
     """Resolve the ``(host, http_path)`` the kernel ``Session`` should use,
     honoring the Thrift-style ``_connection_uri`` / ``_port`` overrides.
 

From e0e90c1ef93f1facf9dc657d210d3d096b80546d Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 08:18:03 +0000
Subject: [PATCH 09/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802255225 at src/databricks/sql/session.py:93

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py | 13 ++++++++++++-
 tests/unit/test_session.py    |  9 +++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index c48e59d98..0b4aa8ed0 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -90,7 +90,18 @@ def _kernel_host_and_path(
         # present, so add a temporary one when the host is scheme-less.
         host = server_hostname.rstrip("/")
         probe = host if "://" in host else "https://" + host
-        if urlsplit(probe).port is None:
+        try:
+            existing_port = urlsplit(probe).port
+        except ValueError as exc:
+            # ``SplitResult.port`` raises when the authority carries a
+            # non-numeric or out-of-range port. Re-raise as a clear,
+            # connector-side error consistent with the ``_connection_uri``
+            # validation above, instead of leaking an opaque URL-parsing error.
+            raise ValueError(
+                "Invalid server_hostname {!r}: could not parse its port "
+                "({})".format(server_hostname, exc)
+            )
+        if existing_port is None:
             host = "{}:{}".format(host, port)
         return host, http_path
 
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 8d31164f4..ce2d145dd 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -834,3 +834,12 @@ def test_port_folded_into_ipv6_literal_without_port(self):
     def test_port_not_double_appended_for_ipv6_literal_with_port(self):
         host, _ = _kernel_host_and_path("[::1]:7000", self.PATH, {"_port": 8443})
         assert host == "[::1]:7000"
+
+    def test_malformed_port_in_hostname_raises_clear_error(self):
+        # A non-numeric/out-of-range port on server_hostname makes
+        # ``SplitResult.port`` raise; surface a clear connector-side error
+        # instead of leaking an opaque URL-parsing ValueError.
+        with pytest.raises(ValueError, match="could not parse its port"):
+            _kernel_host_and_path(
+                self.HOST + ":notaport", self.PATH, {"_port": 8443}
+            )

From dda710f64b66aa3232d1137f5f464d0da9e1a14b Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Tue, 18 Aug 2026 08:26:54 +0000
Subject: [PATCH 10/10] ai: apply changes for #915 (1 review thread)

Addresses:
  - #3802314928 at src/databricks/sql/session.py:71

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/session.py | 11 ++++++++---
 tests/unit/test_session.py    | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index 0b4aa8ed0..73af70080 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -37,7 +37,8 @@ def _kernel_host_and_path(
       (returned as ``host``) and its path+query (returned as ``http_path``).
       ``_connection_uri`` wins over ``_port``, matching the Thrift backend.
       When the URI omits a path (e.g. ``https://host:8443`` or
-      ``https://host:8443?o=222``) the connection's original ``http_path`` is
+      ``https://host:8443?o=222``, and likewise a lone trailing slash such as
+      ``https://host:8443/``) the connection's original ``http_path`` is
       retained, and any query on the URI is applied to that retained path,
       replacing any query the retained path already carried — so a path-less,
       query-bearing URI overrides only the host and query while keeping the
@@ -67,8 +68,12 @@ def _kernel_host_and_path(
         host = "{}://{}".format(parts.scheme, parts.netloc)
         # A path-less URI keeps the connection's original http_path; any query
         # on the URI is then applied to that retained path (host + query
-        # override, path preserved). See the docstring for the rationale.
-        path = parts.path or http_path
+        # override, path preserved). A lone ``"/"`` path (common from a
+        # copy-pasted base URL like ``https://host:8443/``) is treated the same
+        # as an absent path so it does not silently point the kernel at the host
+        # root instead of the warehouse. See the docstring for the rationale.
+        uri_path = parts.path if parts.path not in ("", "/") else None
+        path = uri_path or http_path
         if parts.query:
             # Drop any query already on the retained path before applying the
             # URI's query, so the URI's query fully overrides it. Appending
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index ce2d145dd..a4840555d 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -771,6 +771,26 @@ def test_connection_uri_without_path_retains_original_path(self):
         assert host == "https://h.example.com:8443"
         assert path == self.PATH
 
+    def test_connection_uri_trailing_slash_retains_original_path(self):
+        # A lone trailing slash (common when a base URL is copy-pasted, e.g.
+        # ``https://host:8443/``) must be treated the same as an absent path so
+        # the original warehouse http_path is retained rather than silently
+        # replaced with the host root ``"/"``.
+        host, path = _kernel_host_and_path(
+            self.HOST, self.PATH, {"_connection_uri": "https://h.example.com:8443/"}
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == self.PATH
+
+    def test_connection_uri_trailing_slash_applies_query_to_retained_path(self):
+        # A trailing-slash, query-bearing URI overrides the host and applies the
+        # query to the retained original path, just like the path-less case.
+        host, path = _kernel_host_and_path(
+            self.HOST, self.PATH, {"_connection_uri": "https://h.example.com:8443/?o=222"}
+        )
+        assert host == "https://h.example.com:8443"
+        assert path == "{}?o=222".format(self.PATH)
+
     def test_connection_uri_without_path_applies_query_to_retained_path(self):
         # A path-less, query-bearing URI overrides the host and applies the
         # query to the retained original path (documented fallback semantics).