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>

Author: peco-engineer-bot[bot]

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()
                 ]

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