Skip to content

Commit 3a66597

Browse files
committed
Respect mixed-case SPOG org-id headers
Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
1 parent d1e0c47 commit 3a66597

2 files changed

Lines changed: 14 additions & 7 deletions

File tree

src/databricks/sql/session.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -73,10 +73,10 @@ def __init__(
7373
base_headers = [("User-Agent", self.useragent_header)]
7474
all_headers = (http_headers or []) + base_headers
7575

76-
# Extract ?o=<workspaceId> from http_path for SPOG routing.
77-
# On SPOG hosts, the httpPath contains ?o=<workspaceId> which routes Thrift
78-
# requests via the URL. For SEA, telemetry, and feature flags (which use
79-
# separate endpoints), we inject x-databricks-org-id as an HTTP header.
76+
# Extract workspace context from http_path for SPOG routing.
77+
# On SPOG hosts, the http_path can contain either ?o=<workspaceId> or an
78+
# all-purpose-compute /o/<workspaceId>/ path segment. For SEA, telemetry,
79+
# and feature flags, we inject x-databricks-org-id as an HTTP header.
8080
self._spog_headers = self._extract_spog_headers(http_path, all_headers)
8181
if self._spog_headers:
8282
all_headers = all_headers + list(self._spog_headers.items())
@@ -197,8 +197,8 @@ def _extract_spog_headers(http_path, existing_headers):
197197
if not http_path:
198198
return {}
199199

200-
# Caller already set the headernever override.
201-
if any(k == "x-databricks-org-id" for k, _ in existing_headers):
200+
# Caller already set the header; never override. Header names are case-insensitive.
201+
if any(k.lower() == "x-databricks-org-id" for k, _ in existing_headers):
202202
logger.debug(
203203
"SPOG header extraction: x-databricks-org-id already set by caller, "
204204
"not extracting from http_path"
@@ -239,7 +239,7 @@ def _extract_spog_headers(http_path, existing_headers):
239239
return {"x-databricks-org-id": org_id}
240240

241241
def get_spog_headers(self):
242-
"""Returns SPOG routing headers (x-databricks-org-id) if ?o= was in http_path."""
242+
"""Returns extracted SPOG routing headers (x-databricks-org-id), if any."""
243243
return dict(self._spog_headers)
244244

245245
def open(self):

tests/unit/test_session.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,13 @@ def test_explicit_header_takes_precedence(self):
268268
)
269269
assert result == {}
270270

271+
def test_explicit_header_takes_precedence_case_insensitively(self):
272+
existing = [("X-Databricks-Org-Id", "explicit-value")]
273+
result = Session._extract_spog_headers(
274+
"/sql/1.0/warehouses/abc123?o=6051921418418893", existing
275+
)
276+
assert result == {}
277+
271278
def test_multiple_query_params(self):
272279
result = Session._extract_spog_headers(
273280
"/sql/1.0/warehouses/abc123?o=12345&extra=val", []

0 commit comments

Comments
 (0)