-

Author: tanmay-db

databricks/sdk/config.py

@@ -383,13 +383,25 @@ def client_type(self) -> ClientType:
 
         This is separate from host_type. For example, a unified host can support both
         workspace and account client types.
+
+        Returns ClientType.ACCOUNT or ClientType.WORKSPACE based on the configuration.
+
+        For unified hosts, account_id must be set. If workspace_id is also set,
+        returns WORKSPACE, otherwise returns ACCOUNT.
         """
-        # If workspace_id is set, this is a workspace client
-        if self.workspace_id:
+        host_type = self.host_type
+
+        if host_type == HostType.ACCOUNTS:
+            return ClientType.ACCOUNT
+
+        if host_type == HostType.WORKSPACE:
             return ClientType.WORKSPACE
 
-        # If account_id is set and no workspace_id, this is an account client
-        if self.account_id:
+        if host_type == HostType.UNIFIED:
+            if not self.account_id:
+                raise ValueError("Unified host requires account_id to be set")
+            if self.workspace_id:
+                return ClientType.WORKSPACE
             return ClientType.ACCOUNT
 
         # Default to workspace for backward compatibility
@@ -463,9 +475,7 @@ def oidc_endpoints(self) -> Optional[OidcEndpoints]:
         # Handle unified hosts
         if self.host_type == HostType.UNIFIED:
             if not self.account_id:
-                raise ValueError(
-                    "Unified host requires account_id to be set for OAuth endpoints"
-                )
+                raise ValueError("Unified host requires account_id to be set for OAuth endpoints")
             return get_unified_endpoints(self.host, self.account_id)
 
         # Handle traditional account hosts

tests/test_config.py

@@ -293,10 +293,11 @@ def test_host_type_unified():
 
 
 def test_client_type_workspace():
-    """Test that client type is workspace when workspace_id is set."""
+    """Test that client type is workspace when workspace_id is set on unified host."""
     config = Config(
         host="https://unified.databricks.com",
         workspace_id="test-workspace",
+        account_id="test-account",
         experimental_is_unified_host=True,
         token="test-token",
     )
@@ -320,6 +321,27 @@ def test_client_type_workspace_default():
     assert config.client_type == ClientType.WORKSPACE
 
 
+def test_client_type_accounts_host():
+    """Test that client type is account for accounts host."""
+    config = Config(
+        host="https://accounts.cloud.databricks.com",
+        account_id="test-account",
+        token="test-token",
+    )
+    assert config.client_type == ClientType.ACCOUNT
+
+
+def test_client_type_unified_without_account_id():
+    """Test that client type raises error for unified host without account_id."""
+    config = Config(
+        host="https://unified.databricks.com",
+        experimental_is_unified_host=True,
+        token="test-token",
+    )
+    with pytest.raises(ValueError, match="Unified host requires account_id"):
+        _ = config.client_type
+
+
 def test_is_account_client_backward_compatibility():
     """Test that is_account_client property still works for backward compatibility."""
     config_workspace = Config(host="https://test.databricks.com", token="test-token")