From b03afc8642ed8017b5019ed9c32462503ceb78dc Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 16 Apr 2026 14:18:16 +0200 Subject: [PATCH] Add X-Databricks-Org-Id header to hand-written extension methods for SPOG host compatibility Hand-written mixin methods that make direct _api.do() calls were missing the X-Databricks-Org-Id header that SPOG hosts require for routing. The generated service methods already include this header per-call, but Workspace.upload(), Workspace.download(), Shares.list(), and ServingEndpoints.http_request() were missing it. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + databricks/sdk/mixins/open_ai_client.py | 6 +++++- databricks/sdk/mixins/sharing.py | 3 +++ databricks/sdk/mixins/workspace.py | 11 ++++++++++- 4 files changed, 19 insertions(+), 2 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 1b9a3ae70..80cf40b1a 100755 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -9,6 +9,7 @@ ### Security ### Bug Fixes +* Add `X-Databricks-Org-Id` header to hand-written extension methods (`Workspace.upload()`, `Workspace.download()`, `Shares.list()`, `ServingEndpoints.http_request()`) for SPOG host compatibility. ### Documentation diff --git a/databricks/sdk/mixins/open_ai_client.py b/databricks/sdk/mixins/open_ai_client.py index 127e3522b..86a269bd2 100644 --- a/databricks/sdk/mixins/open_ai_client.py +++ b/databricks/sdk/mixins/open_ai_client.py @@ -170,6 +170,10 @@ def http_request( # This is a temporary fix to get the headers we need for the MCP session id # TODO: Remove this once we have a better way to get back the response headers headers_to_capture = ["mcp-session-id"] + request_headers = {"Accept": "text/plain", "Content-Type": "application/json"} + cfg = self._api._cfg + if cfg.workspace_id: + request_headers["X-Databricks-Org-Id"] = cfg.workspace_id res = self._api.do( "POST", "/api/2.0/external-function", @@ -181,7 +185,7 @@ def http_request( "json": js.dumps(json) if json is not None else None, "params": js.dumps(params) if params is not None else None, }, - headers={"Accept": "text/plain", "Content-Type": "application/json"}, + headers=request_headers, raw=True, response_headers=headers_to_capture, ) diff --git a/databricks/sdk/mixins/sharing.py b/databricks/sdk/mixins/sharing.py index 65e03d665..be9ed9552 100644 --- a/databricks/sdk/mixins/sharing.py +++ b/databricks/sdk/mixins/sharing.py @@ -31,6 +31,9 @@ def list(self, *, max_results: Optional[int] = None, page_token: Optional[str] = headers = { "Accept": "application/json", } + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Org-Id"] = cfg.workspace_id if "max_results" not in query: query["max_results"] = 0 diff --git a/databricks/sdk/mixins/workspace.py b/databricks/sdk/mixins/workspace.py index f62ad5bff..78cbdf859 100644 --- a/databricks/sdk/mixins/workspace.py +++ b/databricks/sdk/mixins/workspace.py @@ -85,10 +85,15 @@ def upload( data["language"] = language.value if overwrite: data["overwrite"] = "true" + headers = {} + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Org-Id"] = cfg.workspace_id try: return self._api.do( "POST", "/api/2.0/workspace/import", + headers=headers, files={"content": content}, data=data, ) @@ -113,5 +118,9 @@ def download(self, path: str, *, format: Optional[ExportFormat] = None) -> Binar query = {"path": path, "direct_download": "true"} if format: query["format"] = format.value - response = self._api.do("GET", "/api/2.0/workspace/export", query=query, raw=True) + headers = {} + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Org-Id"] = cfg.workspace_id + response = self._api.do("GET", "/api/2.0/workspace/export", query=query, headers=headers, raw=True) return response["contents"]