Extract SPOG org-id from cluster http_path for non-Thrift requests (databricks#817)

msrathore-db · web-flow · commit 67233503db95 · 2026-06-01T17:00:17.000Z
* Extract SPOG org-id from cluster http_path for non-Thrift requests

For all-purpose-compute Thrift connections on SPOG (custom-URL) hosts the
http_path is /sql/protocolv1/o/&lt;workspace-id&gt;/&lt;cluster-id&gt; and the
workspace ID is encoded in the path itself. PoPP routes the Thrift
request correctly off the /o/&lt;wsid&gt;/ segment, so the connection succeeds
without an explicit ?o= query parameter.

Other requests on the same connection (telemetry uploads to
/telemetry-ext, feature-flag fetches, SEA REST calls) hit different
paths that don't carry the workspace ID. Previously _extract_spog_headers
only looked at ?o= in the http_path, so the x-databricks-org-id header
was never set for cluster URLs without ?o=. On SPOG hosts PoPP then had
no workspace context for these requests and redirected them to /login,
silently dropping telemetry.

Extend _extract_spog_headers to also extract the workspace ID from the
cluster path segment as a fallback when ?o= is absent. Priority order:
explicit caller header &gt; ?o= query param &gt; /o/&lt;wsid&gt;/ path segment.

Adds five unit tests covering the new cluster-path extraction, leading
slash, query-param-wins priority, explicit-header-wins priority, and a
warehouse-path regression guard.

Signed-off-by: Madhavendra Rathore &lt;madhavendra.rathore@databricks.com&gt;

* Respect mixed-case SPOG org-id headers

Signed-off-by: Madhavendra Rathore &lt;madhavendra.rathore@databricks.com&gt;

* Validate SPOG org id extraction

Signed-off-by: Madhavendra Rathore &lt;madhavendra.rathore@databricks.com&gt;

---------

Signed-off-by: Madhavendra Rathore &lt;madhavendra.rathore@databricks.com&gt;
diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import Dict, Tuple, List, Optional, Any, Type
 
 from databricks.sql.thrift_api.TCLIService import ttypes
@@ -72,10 +73,10 @@ def __init__(
         base_headers = [("User-Agent", self.useragent_header)]
         all_headers = (http_headers or []) + base_headers
 
-        # Extract ?o=<workspaceId> from http_path for SPOG routing.
-        # On SPOG hosts, the httpPath contains ?o=<workspaceId> which routes Thrift
-        # requests via the URL. For SEA, telemetry, and feature flags (which use
-        # separate endpoints), we inject x-databricks-org-id as an HTTP header.
+        # Extract workspace context from http_path for SPOG routing.
+        # On SPOG hosts, the http_path can contain either ?o=<workspaceId> or an
+        # all-purpose-compute /o/<workspaceId>/ path segment. For SEA, telemetry,
+        # and feature flags, we inject x-databricks-org-id as an HTTP header.
         self._spog_headers = self._extract_spog_headers(http_path, all_headers)
         if self._spog_headers:
             all_headers = all_headers + list(self._spog_headers.items())
@@ -170,42 +171,76 @@ def _create_backend(
         }
         return databricks_client_class(**common_args)
 
+    # All-purpose-compute Thrift http_path:
+    # [/]sql/protocolv1/o/<workspace-id>/<cluster-id>[/...][?...]
+    _ORG_ID_RE = re.compile(r"^[0-9]+$")
+    _CLUSTER_PATH_ORG_ID_RE = re.compile(r"^/?sql/protocolv1/o/([0-9]+)/[^/?]+")
+
     @staticmethod
     def _extract_spog_headers(http_path, existing_headers):
-        """Extract ?o=<workspaceId> from http_path and return as a header dict for SPOG routing."""
-        if not http_path or "?" not in http_path:
+        """Extract the workspace ID from http_path for SPOG routing and return it
+        as an ``x-databricks-org-id`` header dict.
+
+        Two sources are inspected, in priority order:
+          1. ``?o=<workspace-id>`` query parameter in http_path (warehouse paths
+             typically encode the workspace this way on SPOG).
+          2. ``/sql/protocolv1/o/<workspace-id>/<cluster-id>`` path segment
+             (all-purpose compute paths embed the workspace in the path itself).
+
+        An explicit ``x-databricks-org-id`` already set by the caller wins over
+        both. Returns an empty dict when no workspace ID can be determined.
+
+        On SPOG (Custom URL) hosts this header is required for non-Thrift
+        endpoints — telemetry, feature flags, SEA — to be routed to the right
+        workspace. Without it, PoPP falls back to default routing and
+        workspace-scoped requests are redirected to ``/login``.
+        """
+        if not http_path:
             return {}
 
-        from urllib.parse import parse_qs
-
-        query_string = http_path.split("?", 1)[1]
-        params = parse_qs(query_string)
-        org_id = params.get("o", [None])[0]
-        if not org_id:
+        # Caller already set the header; never override. Header names are case-insensitive.
+        if any(k.lower() == "x-databricks-org-id" for k, _ in existing_headers):
             logger.debug(
-                "SPOG header extraction: http_path has query string but no ?o= param, "
-                "skipping x-databricks-org-id injection"
+                "SPOG header extraction: x-databricks-org-id already set by caller, "
+                "not extracting from http_path"
             )
             return {}
 
-        # Don't override if explicitly set
-        if any(k == "x-databricks-org-id" for k, _ in existing_headers):
+        org_id = None
+        source = None
+
+        if "?" in http_path:
+            from urllib.parse import parse_qs
+
+            query_string = http_path.split("?", 1)[1]
+            params = parse_qs(query_string)
+            value = params.get("o", [None])[0]
+            if value and Session._ORG_ID_RE.fullmatch(value):
+                org_id = value
+                source = "?o= in http_path"
+
+        if org_id is None:
+            cluster_match = Session._CLUSTER_PATH_ORG_ID_RE.match(http_path)
+            if cluster_match:
+                org_id = cluster_match.group(1)
+                source = "cluster path segment"
+
+        if org_id is None:
             logger.debug(
-                "SPOG header extraction: x-databricks-org-id already set by caller, "
-                "not overriding with ?o=%s from http_path",
-                org_id,
+                "SPOG header extraction: no workspace ID found in http_path, "
+                "skipping x-databricks-org-id injection"
             )
             return {}
 
         logger.debug(
-            "SPOG header extraction: injecting x-databricks-org-id=%s "
-            "(extracted from ?o= in http_path)",
+            "SPOG header extraction: injecting x-databricks-org-id=%s (extracted from %s)",
             org_id,
+            source,
         )
         return {"x-databricks-org-id": org_id}
 
     def get_spog_headers(self):
-        """Returns SPOG routing headers (x-databricks-org-id) if ?o= was in http_path."""
+        """Returns extracted SPOG routing headers (x-databricks-org-id), if any."""
         return dict(self._spog_headers)
 
     def open(self):
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
@@ -162,7 +162,10 @@ def test_socket_timeout_passthrough(self, mock_client_class):
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
-        mock_session_config = {"ANSI_MODE": "FALSE", "QUERY_TAGS": "team:engineering,project:data-pipeline"}
+        mock_session_config = {
+            "ANSI_MODE": "FALSE",
+            "QUERY_TAGS": "team:engineering,project:data-pipeline",
+        }
         databricks.sql.connect(
             session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS
         )
@@ -218,10 +221,15 @@ def test_query_tags_dict_sets_session_config(self, mock_client_class):
         )
 
         call_kwargs = mock_client_class.return_value.open_session.call_args[1]
-        assert call_kwargs["session_configuration"]["QUERY_TAGS"] == "team:data-eng,project:etl"
+        assert (
+            call_kwargs["session_configuration"]["QUERY_TAGS"]
+            == "team:data-eng,project:etl"
+        )
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
-    def test_query_tags_dict_takes_precedence_over_session_config(self, mock_client_class):
+    def test_query_tags_dict_takes_precedence_over_session_config(
+        self, mock_client_class
+    ):
         databricks.sql.connect(
             query_tags={"team": "new-team"},
             session_configuration={"QUERY_TAGS": "team:old-team,other:value"},
@@ -242,9 +250,7 @@ def test_extracts_org_id_from_query_param(self):
         assert result == {"x-databricks-org-id": "6051921418418893"}
 
     def test_no_query_param_returns_empty(self):
-        result = Session._extract_spog_headers(
-            "/sql/1.0/warehouses/abc123", []
-        )
+        result = Session._extract_spog_headers("/sql/1.0/warehouses/abc123", [])
         assert result == {}
 
     def test_no_o_param_returns_empty(self):
@@ -268,8 +274,79 @@ def test_explicit_header_takes_precedence(self):
         )
         assert result == {}
 
+    def test_explicit_header_takes_precedence_case_insensitively(self):
+        existing = [("X-Databricks-Org-Id", "explicit-value")]
+        result = Session._extract_spog_headers(
+            "/sql/1.0/warehouses/abc123?o=6051921418418893", existing
+        )
+        assert result == {}
+
     def test_multiple_query_params(self):
         result = Session._extract_spog_headers(
             "/sql/1.0/warehouses/abc123?o=12345&extra=val", []
         )
         assert result == {"x-databricks-org-id": "12345"}
+
+    def test_non_numeric_query_param_returns_empty(self):
+        result = Session._extract_spog_headers(
+            "/sql/1.0/warehouses/abc123?o=abc123", []
+        )
+        assert result == {}
+
+    def test_control_char_query_param_returns_empty(self):
+        result = Session._extract_spog_headers(
+            "/sql/1.0/warehouses/abc123?o=123%0D%0AX-Injected:%20yes", []
+        )
+        assert result == {}
+
+    def test_empty_query_param_returns_empty(self):
+        result = Session._extract_spog_headers(
+            "/sql/1.0/warehouses/abc123?o=", []
+        )
+        assert result == {}
+
+    def test_extracts_org_id_from_cluster_path_segment(self):
+        # All-purpose-compute path embeds workspace ID in /o/<wsid>/<cluster>.
+        # Without ?o=, the driver must still set x-databricks-org-id so that
+        # telemetry and other non-Thrift requests route to the right workspace
+        # on SPOG hosts.
+        result = Session._extract_spog_headers(
+            "sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt", []
+        )
+        assert result == {"x-databricks-org-id": "6051921418418893"}
+
+    def test_extracts_org_id_from_cluster_path_with_leading_slash(self):
+        result = Session._extract_spog_headers(
+            "/sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt", []
+        )
+        assert result == {"x-databricks-org-id": "6051921418418893"}
+
+    def test_query_param_wins_over_cluster_path_segment(self):
+        # When both forms are present, ?o= takes precedence.
+        result = Session._extract_spog_headers(
+            "sql/protocolv1/o/111/0528-220959-uzmcn1qt?o=222", []
+        )
+        assert result == {"x-databricks-org-id": "222"}
+
+    def test_explicit_header_wins_over_cluster_path_segment(self):
+        existing = [("x-databricks-org-id", "from-caller")]
+        result = Session._extract_spog_headers(
+            "sql/protocolv1/o/111/0528-220959-uzmcn1qt", existing
+        )
+        assert result == {}
+
+    def test_nested_cluster_path_prefix_returns_empty(self):
+        result = Session._extract_spog_headers(
+            "evil/sql/protocolv1/o/999/0528-220959-uzmcn1qt", []
+        )
+        assert result == {}
+
+    def test_incomplete_cluster_path_returns_empty(self):
+        result = Session._extract_spog_headers("sql/protocolv1/o/999/", [])
+        assert result == {}
+
+    def test_warehouse_path_without_query_param_returns_empty(self):
+        # Regression guard: the new cluster-path regex must not accidentally
+        # match warehouse paths (which never embed the workspace ID).
+        result = Session._extract_spog_headers("/sql/1.0/warehouses/abc123", [])
+        assert result == {}
