-

Author: tanmay-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java

@@ -5,15 +5,9 @@
 /** Represents the type of Databricks client being used for API operations. */
 @InternalApi
 public enum ClientType {
-  /** Traditional workspace client. */
+  /** Workspace client (traditional or unified host with workspaceId). */
   WORKSPACE,
 
-  /** Traditional account client. */
-  ACCOUNT,
-
-  /** Workspace-scoped client on unified host. */
-  WORKSPACE_ON_UNIFIED,
-
-  /** Account-scoped client on unified host. */
-  ACCOUNT_ON_UNIFIED
+  /** Account client (traditional or unified host without workspaceId). */
+  ACCOUNT
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java

@@ -31,8 +31,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
     }
     List<String> cmd =
         new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
-    if (config.getClientType() == ClientType.ACCOUNT
-        || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) {
+    if (config.getClientType() == ClientType.ACCOUNT) {
       cmd.add("--account-id");
       cmd.add(config.getAccountId());
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -28,19 +28,15 @@ public class DatabricksConfig {
   private String accountId;
 
   /**
-   * Workspace ID for unified host operations.
-   *
-   * <p><b>Note:</b> This API is experimental and may change or be removed in future releases
-   * without notice.
+   * Workspace ID for unified host operations. Note: This API is experimental and may change or be
+   * removed in future releases without notice.
    */
   @ConfigAttribute(env = "DATABRICKS_WORKSPACE_ID")
   private String workspaceId;
 
   /**
-   * Flag to explicitly mark a host as a unified host.
-   *
-   * <p><b>Note:</b> This API is experimental and may change or be removed in future releases
-   * without notice.
+   * Flag to explicitly mark a host as a unified host. Note: This API is experimental and may change
+   * or be removed in future releases without notice.
    */
   @ConfigAttribute(env = "DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST")
   private Boolean experimentalIsUnifiedHost;
@@ -61,10 +57,8 @@ public class DatabricksConfig {
   private String redirectUrl;
 
   /**
-   * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints.
-   *
-   * <p><b>Note:</b> This API is experimental and may change or be removed in future releases
-   * without notice.
+   * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints. Note: This
+   * API is experimental and may change or be removed in future releases without notice.
    */
   @ConfigAttribute(env = "DATABRICKS_DISCOVERY_URL")
   private String discoveryUrl;
@@ -256,7 +250,7 @@ public synchronized Map<String, String> authenticate() throws DatabricksExceptio
 
         // For unified hosts with workspace operations, wrap the header factory
         // to inject the X-Databricks-Org-Id header
-        if (getClientType() == ClientType.WORKSPACE_ON_UNIFIED) {
+        if (getHostType() == HostType.UNIFIED && workspaceId != null && !workspaceId.isEmpty()) {
           headerFactory = new UnifiedHostHeaderFactory(rawHeaderFactory, workspaceId);
         } else {
           headerFactory = rawHeaderFactory;
@@ -754,9 +748,10 @@ public ClientType getClientType() {
     HostType hostType = getHostType();
     switch (hostType) {
       case UNIFIED:
+        // For unified hosts, client type depends on whether workspaceId is set
         return (workspaceId != null && !workspaceId.isEmpty())
-            ? ClientType.WORKSPACE_ON_UNIFIED
-            : ClientType.ACCOUNT_ON_UNIFIED;
+            ? ClientType.WORKSPACE
+            : ClientType.ACCOUNT;
       case ACCOUNTS:
         return ClientType.ACCOUNT;
       case WORKSPACE:

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java

@@ -151,10 +151,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
                   config.getHttpClient())
               .audience(config.getTokenAudience())
               .accountId(
-                  (config.getClientType() == ClientType.ACCOUNT
-                          || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED)
-                      ? config.getAccountId()
-                      : null)
+                  config.getClientType() == ClientType.ACCOUNT ? config.getAccountId() : null)
               .scopes(config.getScopes())
               .build();
 

databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java

@@ -66,8 +66,7 @@ public HeaderFactory configure(DatabricksConfig config) {
       Map<String, String> headers = new HashMap<>();
       headers.put("Authorization", String.format("Bearer %s", idToken.getTokenValue()));
 
-      if (config.getClientType() == ClientType.ACCOUNT
-          || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) {
+      if (config.getClientType() == ClientType.ACCOUNT) {
         AccessToken token;
         try {
           token = finalServiceAccountCredentials.createScoped(GCP_SCOPES).refreshAccessToken();

databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java

@@ -69,8 +69,7 @@ public HeaderFactory configure(DatabricksConfig config) {
         throw new DatabricksException(message, e);
       }
 
-      if (config.getClientType() == ClientType.ACCOUNT
-          || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) {
+      if (config.getClientType() == ClientType.ACCOUNT) {
         try {
           headers.put(
               SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue());

databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java

@@ -55,8 +55,11 @@ public void testGetWorkspaceClientForUnifiedHost() {
     // Should have workspace ID set
     assertEquals("123456", workspaceClient.config().getWorkspaceId());
 
-    // Should be workspace-on-unified client type
-    assertEquals(ClientType.WORKSPACE_ON_UNIFIED, workspaceClient.config().getClientType());
+    // Should be workspace client type (on unified host)
+    assertEquals(ClientType.WORKSPACE, workspaceClient.config().getClientType());
+
+    // Host type should still be unified
+    assertEquals(HostType.UNIFIED, workspaceClient.config().getHostType());
   }
 
   @Test

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -401,8 +401,9 @@ public void testGetClientTypeAccount() {
 
   @Test
   public void testGetClientTypeWorkspaceOnUnified() {
+    // For unified hosts with workspaceId, client type is WORKSPACE
     assertEquals(
-        ClientType.WORKSPACE_ON_UNIFIED,
+        ClientType.WORKSPACE,
         new DatabricksConfig()
             .setHost("https://unified.databricks.com")
             .setExperimentalIsUnifiedHost(true)
@@ -412,8 +413,9 @@ public void testGetClientTypeWorkspaceOnUnified() {
 
   @Test
   public void testGetClientTypeAccountOnUnified() {
+    // For unified hosts without workspaceId, client type is ACCOUNT
     assertEquals(
-        ClientType.ACCOUNT_ON_UNIFIED,
+        ClientType.ACCOUNT,
         new DatabricksConfig()
             .setHost("https://unified.databricks.com")
             .setExperimentalIsUnifiedHost(true)

databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java

@@ -89,13 +89,13 @@ private static Stream<Arguments> provideClientTypeTestCases() {
             "https://unified.databricks.com",
             null,
             true,
-            ClientType.ACCOUNT_ON_UNIFIED),
+            ClientType.ACCOUNT),
         Arguments.of(
             "Unified with workspace ID",
             "https://unified.databricks.com",
             "123456",
             true,
-            ClientType.WORKSPACE_ON_UNIFIED));
+            ClientType.WORKSPACE));
   }
 
   @ParameterizedTest(name = "{0}")
@@ -182,7 +182,7 @@ public void testUnifiedHostFromEnvironmentVariables() {
     assertEquals(HostType.UNIFIED, config.getHostType());
     assertEquals("987654321", config.getWorkspaceId());
     assertEquals("account-abc", config.getAccountId());
-    assertEquals(ClientType.WORKSPACE_ON_UNIFIED, config.getClientType());
+    assertEquals(ClientType.WORKSPACE, config.getClientType());
   }
 
   // --- UnifiedHostHeaderFactory Tests ---