Validate SPOG org id extraction

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

src/databricks/sql/session.py

@@ -173,7 +173,8 @@ def _create_backend(
 
     # All-purpose-compute Thrift http_path:
     # [/]sql/protocolv1/o/<workspace-id>/<cluster-id>[/...][?...]
-    _CLUSTER_PATH_ORG_ID_RE = re.compile(r"(?:^|/)sql/protocolv1/o/(\d+)/[^/?]+")
+    _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):
@@ -214,12 +215,12 @@ def _extract_spog_headers(http_path, existing_headers):
             query_string = http_path.split("?", 1)[1]
             params = parse_qs(query_string)
             value = params.get("o", [None])[0]
-            if value:
+            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.search(http_path)
+            cluster_match = Session._CLUSTER_PATH_ORG_ID_RE.match(http_path)
             if cluster_match:
                 org_id = cluster_match.group(1)
                 source = "cluster path segment"

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):
@@ -281,6 +287,24 @@ def test_multiple_query_params(self):
         )
         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
@@ -311,10 +335,18 @@ def test_explicit_header_wins_over_cluster_path_segment(self):
         )
         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", []
-        )
+        result = Session._extract_spog_headers("/sql/1.0/warehouses/abc123", [])
         assert result == {}