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>

Author: peco-engineer-bot[bot]

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,

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"