diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 1ad5681f1..01fa7c443 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,6 +7,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. * Add `X-Databricks-Org-Id` header to `WorkspaceExt.upload()` and `WorkspaceExt.download()` for SPOG host compatibility. * `WorkspaceClient.get_workspace_id()` now returns `Config.workspace_id` directly when set, instead of calling `/api/2.0/preview/scim/v2/Me`. This removes an API round-trip on every call where the workspace ID is already known (profile, `?o=` query param, env var, or host metadata) and fixes a failure on SPOG hosts where the unauthenticated probe request was rejected with `Unable to load OAuth Config`. 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/workspace.py b/databricks/sdk/mixins/workspace.py index db9e44ea8..a26a325be 100644 --- a/databricks/sdk/mixins/workspace.py +++ b/databricks/sdk/mixins/workspace.py @@ -93,6 +93,7 @@ def upload( return self._api.do( "POST", "/api/2.0/workspace/import", + headers=headers, files={"content": content}, data=data, headers=headers,