Set necessary headers when authenticating via Azure CLI

databricks/sdk/azure.py

@@ -29,3 +29,9 @@ class AzureEnvironment:
                            resource_manager_endpoint="https://management.chinacloudapi.cn/",
                            active_directory_endpoint="https://login.chinacloudapi.cn/"),
 )
+
+
+def add_workspace_id_header(cfg: 'Config', headers: dict[str, str]) -> dict[str, str]:
+    if cfg.azure_workspace_resource_id:
+        headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id
+    return headers

databricks/sdk/core.py

@@ -21,7 +21,7 @@
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
-from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
+from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment, add_workspace_id_header
 from .oauth import (ClientCredentials, OAuthClient, OidcEndpoints, Refreshable,
                     Token, TokenCache, TokenSource)
 from .version import __version__
@@ -210,9 +210,7 @@ def refreshed_headers() -> Dict[str, str]:
             'Authorization': f"Bearer {inner.token().access_token}",
             'X-Databricks-Azure-SP-Management-Token': cloud.token().access_token,
         }
-        if cfg.azure_workspace_resource_id:
-            headers["X-Databricks-Azure-Workspace-Resource-Id"] = cfg.azure_workspace_resource_id
-        return headers
+        return add_workspace_id_header(cfg, headers)
 
     return refreshed_headers
 
@@ -281,7 +279,8 @@ def azure_cli(cfg: 'Config') -> Optional[HeaderFactory]:
 
     def inner() -> Dict[str, str]:
         token = token_source.token()
-        return {'Authorization': f'{token.token_type} {token.access_token}'}
+        headers = {'Authorization': f'{token.token_type} {token.access_token}'}
+        return add_workspace_id_header(cfg, headers)
 
     return inner
 

tests/test_auth_manual_tests.py

@@ -0,0 +1,12 @@
+from databricks.sdk.core import Config
+
+from .conftest import __tests__
+
+def test_azure_cli_workspace_header_present(monkeypatch):
+    monkeypatch.setenv('HOME', __tests__ + '/testdata/azure')
+    monkeypatch.setenv('PATH', __tests__ + '/testdata:/bin')
+    resource_id = '/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123'
+    cfg = Config(auth_type='azure-cli', host='x', azure_workspace_resource_id=resource_id)
+    assert 'X-Databricks-Azure-Workspace-Resource-Id' in cfg.authenticate()
+    assert cfg.authenticate()['X-Databricks-Azure-Workspace-Resource-Id'] == resource_id
+