Add X-Databricks-Org-Id header to hand-written extensions for SPOG hosts

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
 

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,
         )

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

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"]