[Feature] Support for Unified Host

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### New Features and Improvements
 
+* Add support for unified hosts with experimental flag.
+
 ### Bug Fixes
 
 ### Security Vulnerabilities

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

@@ -0,0 +1,13 @@
+package com.databricks.sdk.core;
+
+import com.databricks.sdk.support.InternalApi;
+
+/** Represents the type of Databricks client being used for API operations. */
+@InternalApi
+public enum ClientType {
+  /** Workspace client (traditional or unified host with workspaceId). */
+  WORKSPACE,
+
+  /** Account client (traditional or unified host without workspaceId). */
+  ACCOUNT
+}

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

@@ -31,7 +31,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
     }
     List<String> cmd =
         new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
-    if (config.isAccountClient()) {
+    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

@@ -27,6 +27,20 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_ACCOUNT_ID")
   private String accountId;
 
+  /**
+   * 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. 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;
+
   @ConfigAttribute(env = "DATABRICKS_TOKEN", auth = "pat", sensitive = true)
   private String token;
 
@@ -43,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;
@@ -233,8 +245,16 @@ public synchronized Map<String, String> authenticate() throws DatabricksExceptio
       if (headerFactory == null) {
         // Calling authenticate without resolve
         ConfigLoader.fixHostIfNeeded(this);
-        headerFactory = credentialsProvider.configure(this);
+        HeaderFactory rawHeaderFactory = credentialsProvider.configure(this);
         setAuthType(credentialsProvider.authType());
+
+        // For unified hosts with workspace operations, wrap the header factory
+        // to inject the X-Databricks-Org-Id header
+        if (getHostType() == HostType.UNIFIED && workspaceId != null && !workspaceId.isEmpty()) {
+          headerFactory = new UnifiedHostHeaderFactory(rawHeaderFactory, workspaceId);
+        } else {
+          headerFactory = rawHeaderFactory;
+        }
       }
       return headerFactory.headers();
     } catch (DatabricksException e) {
@@ -298,6 +318,24 @@ public DatabricksConfig setAccountId(String accountId) {
     return this;
   }
 
+  public String getWorkspaceId() {
+    return workspaceId;
+  }
+
+  public DatabricksConfig setWorkspaceId(String workspaceId) {
+    this.workspaceId = workspaceId;
+    return this;
+  }
+
+  public Boolean getExperimentalIsUnifiedHost() {
+    return experimentalIsUnifiedHost;
+  }
+
+  public DatabricksConfig setExperimentalIsUnifiedHost(Boolean experimentalIsUnifiedHost) {
+    this.experimentalIsUnifiedHost = experimentalIsUnifiedHost;
+    return this;
+  }
+
   public String getDatabricksCliPath() {
     return this.databricksCliPath;
   }
@@ -679,12 +717,49 @@ public boolean isAws() {
   }
 
   public boolean isAccountClient() {
+    if (getHostType() == HostType.UNIFIED) {
+      throw new DatabricksException(
+          "Cannot determine account client status for unified hosts. "
+              + "Use getHostType() or getClientType() instead. "
+              + "For unified hosts, client type depends on whether workspaceId is set.");
+    }
     if (host == null) {
       return false;
     }
     return host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod.");
   }
 
+  /** Returns the host type based on configuration settings and host URL. */
+  public HostType getHostType() {
+    if (experimentalIsUnifiedHost != null && experimentalIsUnifiedHost) {
+      return HostType.UNIFIED;
+    }
+    if (host == null) {
+      return HostType.WORKSPACE;
+    }
+    if (host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod.")) {
+      return HostType.ACCOUNTS;
+    }
+    return HostType.WORKSPACE;
+  }
+
+  /** Returns the client type based on host type and workspace ID configuration. */
+  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
+            : ClientType.ACCOUNT;
+      case ACCOUNTS:
+        return ClientType.ACCOUNT;
+      case WORKSPACE:
+      default:
+        return ClientType.WORKSPACE;
+    }
+  }
+
   public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
     if (discoveryUrl == null) {
       return fetchDefaultOidcEndpoints();
@@ -705,10 +780,25 @@ private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
     return null;
   }
 
+  private OpenIDConnectEndpoints getUnifiedOidcEndpoints(String accountId) throws IOException {
+    if (accountId == null || accountId.isEmpty()) {
+      throw new DatabricksException(
+          "account_id is required for unified host OIDC endpoint discovery");
+    }
+    String prefix = getHost() + "/oidc/accounts/" + accountId;
+    return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize");
+  }
+
   private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
     if (getHost() == null) {
       return null;
     }
+
+    // For unified hosts, use account-based OIDC endpoints
+    if (getHostType() == HostType.UNIFIED) {
+      return getUnifiedOidcEndpoints(getAccountId());
+    }
+
     if (isAzure() && getAzureClientId() != null) {
       Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize");
       request.setRedirectionBehavior(false);

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

@@ -150,7 +150,8 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
                   namedIdTokenSource.idTokenSource,
                   config.getHttpClient())
               .audience(config.getTokenAudience())
-              .accountId(config.isAccountClient() ? config.getAccountId() : null)
+              .accountId(
+                  config.getClientType() == ClientType.ACCOUNT ? config.getAccountId() : null)
               .scopes(config.getScopes())
               .build();
 

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

@@ -66,7 +66,7 @@ public HeaderFactory configure(DatabricksConfig config) {
       Map<String, String> headers = new HashMap<>();
       headers.put("Authorization", String.format("Bearer %s", idToken.getTokenValue()));
 
-      if (config.isAccountClient()) {
+      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,7 +69,7 @@ public HeaderFactory configure(DatabricksConfig config) {
         throw new DatabricksException(message, e);
       }
 
-      if (config.isAccountClient()) {
+      if (config.getClientType() == ClientType.ACCOUNT) {
         try {
           headers.put(
               SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue());

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

@@ -0,0 +1,16 @@
+package com.databricks.sdk.core;
+
+import com.databricks.sdk.support.InternalApi;
+
+/** Represents the type of Databricks host being used. */
+@InternalApi
+public enum HostType {
+  /** Traditional workspace host. */
+  WORKSPACE,
+
+  /** Traditional accounts host. */
+  ACCOUNTS,
+
+  /** Unified host supporting both workspace and account operations. */
+  UNIFIED
+}

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

@@ -0,0 +1,36 @@
+package com.databricks.sdk.core;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * HeaderFactory wrapper that adds X-Databricks-Org-Id header for unified host workspace operations.
+ */
+class UnifiedHostHeaderFactory implements HeaderFactory {
+  private final HeaderFactory delegate;
+  private final String workspaceId;
+
+  /**
+   * Creates a new unified host header factory.
+   *
+   * @param delegate The underlying header factory (e.g., OAuth, PAT)
+   * @param workspaceId The workspace ID to inject in the X-Databricks-Org-Id header
+   */
+  public UnifiedHostHeaderFactory(HeaderFactory delegate, String workspaceId) {
+    if (delegate == null) {
+      throw new IllegalArgumentException("delegate cannot be null");
+    }
+    if (workspaceId == null || workspaceId.isEmpty()) {
+      throw new IllegalArgumentException("workspaceId cannot be null or empty");
+    }
+    this.delegate = delegate;
+    this.workspaceId = workspaceId;
+  }
+
+  @Override
+  public Map<String, String> headers() {
+    Map<String, String> headers = new HashMap<>(delegate.headers());
+    headers.put("X-Databricks-Org-Id", workspaceId);
+    return headers;
+  }
+}

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

@@ -0,0 +1,75 @@
+package com.databricks.sdk;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import com.databricks.sdk.core.ClientType;
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.HostType;
+import com.databricks.sdk.service.provisioning.Workspace;
+import org.junit.jupiter.api.Test;
+
+public class AccountClientTest {
+
+  @Test
+  public void testGetWorkspaceClientForTraditionalAccount() {
+    DatabricksConfig accountConfig =
+        new DatabricksConfig()
+            .setHost("https://accounts.cloud.databricks.com")
+            .setAccountId("test-account")
+            .setToken("test-token");
+
+    AccountClient accountClient = new AccountClient(accountConfig);
+
+    Workspace workspace = new Workspace();
+    workspace.setWorkspaceId(123L);
+    workspace.setDeploymentName("test-workspace");
+
+    WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace);
+
+    // Should have a different host
+    assertNotEquals(accountConfig.getHost(), workspaceClient.config().getHost());
+    assertTrue(workspaceClient.config().getHost().contains("test-workspace"));
+  }
+
+  @Test
+  public void testGetWorkspaceClientForUnifiedHost() {
+    String unifiedHost = "https://unified.databricks.com";
+    DatabricksConfig accountConfig =
+        new DatabricksConfig()
+            .setHost(unifiedHost)
+            .setExperimentalIsUnifiedHost(true)
+            .setAccountId("test-account")
+            .setToken("test-token");
+
+    AccountClient accountClient = new AccountClient(accountConfig);
+
+    Workspace workspace = new Workspace();
+    workspace.setWorkspaceId(123456L);
+    workspace.setDeploymentName("test-workspace");
+
+    WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace);
+
+    // Should have the same host
+    assertEquals(unifiedHost, workspaceClient.config().getHost());
+
+    // Should have workspace ID set
+    assertEquals("123456", workspaceClient.config().getWorkspaceId());
+
+    // 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
+  public void testGetWorkspaceClientForUnifiedHostType() {
+    // Verify unified host type is correctly detected
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("https://unified.databricks.com")
+            .setExperimentalIsUnifiedHost(true);
+
+    assertEquals(HostType.UNIFIED, config.getHostType());
+  }
+}