more changes

Author: hectorcast-db

databricks/sdk/config.py

@@ -406,6 +406,8 @@ def client_type(self) -> ClientType:
             return ClientType.WORKSPACE
 
         if host_type == HostType.UNIFIED:
+            if self.workspace_id:
+                return ClientType.WORKSPACE
             if self.account_id:
                 return ClientType.ACCOUNT
             # Legacy workspace hosts don't have a workspace_id until AFTER the auth is resolved.

databricks/sdk/credentials_provider.py

@@ -431,9 +431,9 @@ def _oidc_credentials_provider(
 
     # Determine the audience for token exchange
     audience = cfg.token_audience
-    if audience is None and cfg.client_type == ClientType.ACCOUNT:
+    if audience is None and cfg.account_id:
         audience = cfg.account_id
-    if audience is None and cfg.client_type != ClientType.ACCOUNT:
+    if audience is None and not cfg.account_id:
         audience = cfg.oidc_endpoints.token_endpoint
 
     # Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.
@@ -590,7 +590,7 @@ def token() -> oauth.Token:
     def refreshed_headers() -> Dict[str, str]:
         credentials.refresh(request)
         headers = {"Authorization": f"Bearer {credentials.token}"}
-        if cfg.client_type == ClientType.ACCOUNT:
+        if cfg.account_id:
             gcp_credentials.refresh(request)
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token
         return headers
@@ -631,7 +631,7 @@ def token() -> oauth.Token:
     def refreshed_headers() -> Dict[str, str]:
         id_creds.refresh(request)
         headers = {"Authorization": f"Bearer {id_creds.token}"}
-        if cfg.client_type == ClientType.ACCOUNT:
+        if cfg.account_id:
             gcp_impersonated_credentials.refresh(request)
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token
         return headers

tests/test_config.py

@@ -361,16 +361,26 @@ def test_is_account_client_backward_compatibility():
     assert config_account.is_account_client
 
 
-def test_is_account_client_raises_on_unified_host():
-    """Test that is_account_client raises ValueError when used with unified hosts."""
+def test_is_account_client_on_unified_host():
+    """Test that is_account_client returns truthiness of account_id for unified hosts."""
     config = Config(
         host="https://unified.databricks.com",
         experimental_is_unified_host=True,
         workspace_id="test-workspace",
         token="test-token",
     )
-    with pytest.raises(ValueError, match="is_account_client cannot be used with unified hosts"):
-        _ = config.is_account_client
+    # Should be falsy since account_id is not set
+    assert not config.is_account_client
+    
+    # With account_id set, should be truthy
+    config_with_account = Config(
+        host="https://unified.databricks.com",
+        experimental_is_unified_host=True,
+        workspace_id="test-workspace",
+        account_id="test-account",
+        token="test-token",
+    )
+    assert config_with_account.is_account_client
 
 
 def test_oidc_endpoints_unified_workspace(mocker, requests_mock):