Update hand-written files for X-Databricks-Workspace-Id header

Generator can't touch these (no DO NOT EDIT marker), or need
post-generation reformatting because the new literal is longer:

- config.py: SQL HTTP path resolver header read + widen workspace_id
  field comment (now accepts non-classic workspace identifier formats)
- mixins/workspace.py: upload/download header injection (2 sites)
- mixins/sharing.py: shares-list header injection
- mixins/files.py: storage-proxy workspace probe (2 sites)
- tests/test_config.py: renamed three header tests
  (test_workspace_org_id_header_* -> test_workspace_id_header_*,
   test_no_org_id_header_* -> test_no_workspace_id_header_*)
  and swapped asserted literals/docstrings
- tests/test_files.py: updated mock response header fixtures
- tests/integration/test_unified_profile.py: comment refresh
- NEXT_CHANGELOG.md: internal-changes entry
- __init__.py: black-wrap of the get_workspace_id() probe line that
  exceeded 120 chars after the longer literal substitution

Signed-off-by: Divyansh Vijayvergia <divyansh.vijayvergia@databricks.com>

Author: Divyansh-db

NEXT_CHANGELOG.md

@@ -14,4 +14,6 @@
 
 ### Internal Changes
 
+* Switch workspace addressing header on workspace-scoped API calls from `X-Databricks-Org-Id` to `X-Databricks-Workspace-Id`. The value continues to come from `Config.workspace_id` (`DATABRICKS_WORKSPACE_ID`), and now accepts either a classic numeric workspace ID or another workspace identifier format (server disambiguates).
+
 ### API Changes

databricks/sdk/__init__.py

@@ -89,6 +89,9 @@ def with_user_agent_extra(key: str, value: str):
 class Config:
     host: str = ConfigAttribute(env="DATABRICKS_HOST")
     account_id: str = ConfigAttribute(env="DATABRICKS_ACCOUNT_ID")
+    # Workspace identifier sent on workspace-scoped API calls so unified hosts
+    # can route to the right workspace. Accepts a classic numeric workspace ID
+    # or another workspace identifier format that the server understands.
     workspace_id: str = ConfigAttribute(env="DATABRICKS_WORKSPACE_ID")
 
     # Cloud provider. When set, is_aws/is_azure/is_gcp use this value directly
@@ -596,7 +599,7 @@ def sql_http_path(self) -> Optional[str]:
         if self.cluster_id:
             response = requests.get(f"{self.host}/api/2.0/preview/scim/v2/Me", headers=headers)
             # get workspace ID from the response header
-            workspace_id = response.headers.get("x-databricks-org-id")
+            workspace_id = response.headers.get("x-databricks-workspace-id")
             return f"sql/protocolv1/o/{workspace_id}/{self.cluster_id}"
         if self.warehouse_id:
             return f"/sql/1.0/warehouses/{self.warehouse_id}"

databricks/sdk/config.py

@@ -996,8 +996,10 @@ def _create_request_builder(self):
     def _resolve_workspace_id(self) -> Optional[str]:
         """Resolves the workspace ID by querying the SCIM Me endpoint."""
         try:
-            response = self._api.do("GET", "/api/2.0/preview/scim/v2/Me", response_headers=["X-Databricks-Org-Id"])
-            workspace_id = response.get("X-Databricks-Org-Id")
+            response = self._api.do(
+                "GET", "/api/2.0/preview/scim/v2/Me", response_headers=["X-Databricks-Workspace-Id"]
+            )
+            workspace_id = response.get("X-Databricks-Workspace-Id")
             if workspace_id:
                 return str(workspace_id)
         except Exception:

databricks/sdk/mixins/files.py

@@ -33,7 +33,7 @@ def list(self, *, max_results: Optional[int] = None, page_token: Optional[str] =
         }
         cfg = self._api._cfg
         if cfg.workspace_id:
-            headers["X-Databricks-Org-Id"] = cfg.workspace_id
+            headers["X-Databricks-Workspace-Id"] = cfg.workspace_id
 
         if "max_results" not in query:
             query["max_results"] = 0

databricks/sdk/mixins/sharing.py

@@ -88,7 +88,7 @@ def upload(
         headers = {}
         cfg = self._api._cfg
         if cfg.workspace_id:
-            headers["X-Databricks-Org-Id"] = cfg.workspace_id
+            headers["X-Databricks-Workspace-Id"] = cfg.workspace_id
         try:
             return self._api.do(
                 "POST",
@@ -121,6 +121,6 @@ def download(self, path: str, *, format: Optional[ExportFormat] = None) -> Binar
         headers = {}
         cfg = self._api._cfg
         if cfg.workspace_id:
-            headers["X-Databricks-Org-Id"] = cfg.workspace_id
+            headers["X-Databricks-Workspace-Id"] = cfg.workspace_id
         response = self._api.do("GET", "/api/2.0/workspace/export", query=query, headers=headers, raw=True)
         return response["contents"]

databricks/sdk/mixins/workspace.py

@@ -90,7 +90,7 @@ def test_spog_workspace_google_credentials(isolated_env):
     # google-credentials uses a GCP ID token with target_audience=cfg.host.
     # On the unified host this produces the same token for both account and workspace
     # requests (identical OIDC exchange, identical audience). Account-level APIs accept
-    # this token, but workspace-level APIs return 401. The X-Databricks-Org-Id header
+    # this token, but workspace-level APIs return 401. The X-Databricks-Workspace-Id header
     # is set correctly. This appears to be a server-side limitation on unified hosts.
     pytest.skip("google-credentials ID token is rejected for workspace operations on unified hosts")
     env = isolated_env("ucacct")

tests/integration/test_unified_profile.py

@@ -453,8 +453,8 @@ def test_oidc_endpoints_falls_back_to_databricks_when_no_azure_client_id(mocker,
     assert "https://adb-123.4.azuredatabricks.net/oidc/v1/token" == endpoints.token_endpoint
 
 
-def test_workspace_org_id_header_on_unified_host(requests_mock):
-    """Test that X-Databricks-Org-Id header is added for workspace clients on unified hosts."""
+def test_workspace_id_header_on_unified_host(requests_mock):
+    """Test that X-Databricks-Workspace-Id header is added for workspace clients on unified hosts."""
 
     requests_mock.get("https://unified.databricks.com/api/2.0/preview/scim/v2/Me", json={"result": "success"})
 
@@ -468,12 +468,12 @@ def test_workspace_org_id_header_on_unified_host(requests_mock):
     workspace_client = WorkspaceClient(config=config)
     workspace_client.current_user.me()
 
-    # Verify the request was made with the X-Databricks-Org-Id header
-    assert requests_mock.last_request.headers.get("X-Databricks-Org-Id") == "test-workspace-123"
+    # Verify the request was made with the X-Databricks-Workspace-Id header
+    assert requests_mock.last_request.headers.get("X-Databricks-Workspace-Id") == "test-workspace-123"
 
 
-def test_not_workspace_org_id_header_on_unified_host_on_account_endpoint(requests_mock):
-    """Test that X-Databricks-Org-Id header is added for workspace clients on unified hosts."""
+def test_not_workspace_id_header_on_unified_host_on_account_endpoint(requests_mock):
+    """Test that X-Databricks-Workspace-Id header is added for workspace clients on unified hosts."""
 
     requests_mock.get(
         "https://unified.databricks.com/api/2.0/accounts/test-account/scim/v2/Groups/test-group-123",
@@ -490,12 +490,12 @@ def test_not_workspace_org_id_header_on_unified_host_on_account_endpoint(request
     account_client = AccountClient(config=config)
     account_client.groups.get("test-group-123")
 
-    # Verify the request was made without the X-Databricks-Org-Id header
-    assert "X-Databricks-Org-Id" not in requests_mock.last_request.headers
+    # Verify the request was made without the X-Databricks-Workspace-Id header
+    assert "X-Databricks-Workspace-Id" not in requests_mock.last_request.headers
 
 
-def test_no_org_id_header_on_regular_workspace(requests_mock):
-    """Test that X-Databricks-Org-Id header is NOT added for regular workspace hosts."""
+def test_no_workspace_id_header_on_regular_workspace(requests_mock):
+    """Test that X-Databricks-Workspace-Id header is NOT added for regular workspace hosts."""
     from databricks.sdk.core import ApiClient
 
     requests_mock.get("https://test.databricks.com/api/2.0/test", json={"result": "success"})
@@ -505,8 +505,8 @@ def test_no_org_id_header_on_regular_workspace(requests_mock):
     api_client = ApiClient(config)
     api_client.do("GET", "/api/2.0/test")
 
-    # Verify the X-Databricks-Org-Id header was NOT added
-    assert "X-Databricks-Org-Id" not in requests_mock.last_request.headers
+    # Verify the X-Databricks-Workspace-Id header was NOT added
+    assert "X-Databricks-Workspace-Id" not in requests_mock.last_request.headers
 
 
 def test_disable_oauth_refresh_token_from_env(monkeypatch, mocker):

tests/test_config.py

@@ -926,7 +926,7 @@ def match_request_to_response(
             resp = requests.Response()
             resp.status_code = 200
             resp._content = b"{}"
-            resp.headers["X-Databricks-Org-Id"] = "12345"
+            resp.headers["X-Databricks-Workspace-Id"] = "12345"
             resp.request = request
             resp.url = request.url
             return resp
@@ -1478,7 +1478,7 @@ def custom_matcher(request: requests.Request) -> Optional[requests.Response]:
                             resp = requests.Response()
                             resp.status_code = 200
                             resp._content = b"{}"
-                            resp.headers["X-Databricks-Org-Id"] = "12345"
+                            resp.headers["X-Databricks-Workspace-Id"] = "12345"
                             resp.request = request
                             resp.url = request.url
                             return resp