[Feature] Support CLI Profiles for Unified Host

Author: tanmay-db

databricks/sdk/credentials_provider.py

@@ -430,9 +430,9 @@ def _oidc_credentials_provider(
 
     # Determine the audience for token exchange
     audience = cfg.token_audience
-    if audience is None and cfg.is_account_client:
+    if audience is None and cfg.client_type.value == "account":
         audience = cfg.account_id
-    if audience is None and not cfg.is_account_client:
+    if audience is None and cfg.client_type.value != "account":
         audience = cfg.oidc_endpoints.token_endpoint
 
     # Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.
@@ -589,7 +589,7 @@ def token() -> oauth.Token:
     def refreshed_headers() -> Dict[str, str]:
         credentials.refresh(request)
         headers = {"Authorization": f"Bearer {credentials.token}"}
-        if cfg.is_account_client:
+        if cfg.client_type.value == "account":
             gcp_credentials.refresh(request)
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token
         return headers
@@ -630,7 +630,7 @@ def token() -> oauth.Token:
     def refreshed_headers() -> Dict[str, str]:
         id_creds.refresh(request)
         headers = {"Authorization": f"Bearer {id_creds.token}"}
-        if cfg.is_account_client:
+        if cfg.client_type.value == "account":
             gcp_impersonated_credentials.refresh(request)
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token
         return headers
@@ -852,7 +852,14 @@ class DatabricksCliTokenSource(CliTokenSource):
 
     def __init__(self, cfg: "Config"):
         args = ["auth", "token", "--host", cfg.host]
-        if cfg.is_account_client:
+        if cfg.experimental_is_unified_host:
+            # For unified hosts, pass account_id, workspace_id, and experimental flag
+            args += ["--experimental-is-unified-host"]
+            if cfg.account_id:
+                args += ["--account-id", cfg.account_id]
+            if cfg.workspace_id:
+                args += ["--workspace-id", str(cfg.workspace_id)]
+        elif cfg.client_type.value == "account":
             args += ["--account-id", cfg.account_id]
 
         cli_path = cfg.databricks_cli_path