[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,17 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_ACCOUNT_ID")
   private String accountId;
 
+  /** Workspace ID for unified host operations. */
+  @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;
 
@@ -236,7 +247,14 @@ public synchronized Map<String, String> authenticate() throws DatabricksExceptio
         headerFactory = credentialsProvider.configure(this);
         setAuthType(credentialsProvider.authType());
       }
-      return headerFactory.headers();
+      Map<String, String> headers = new HashMap<>(headerFactory.headers());
+
+      // For unified hosts with workspace operations, add the X-Databricks-Org-Id header
+      if (getHostType() == HostType.UNIFIED && workspaceId != null && !workspaceId.isEmpty()) {
+        headers.put("X-Databricks-Org-Id", workspaceId);
+      }
+
+      return headers;
     } catch (DatabricksException e) {
       String msg = String.format("%s auth: %s", credentialsProvider.authType(), e.getMessage());
       DatabricksException wrapperException = new DatabricksException(msg, e);
@@ -298,6 +316,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 +715,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 +778,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

@@ -98,6 +98,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
     // TODO: refactor the code so that the IdTokenSources are created within the
     // configure call of their corresponding CredentialsProvider. This will allow
     // us to simplify the code by validating IdTokenSources when they are created.
+    // This would also need to be updated to support unified hosts.
     OpenIDConnectEndpoints endpoints = null;
     try {
       endpoints = config.getOidcEndpoints();
@@ -150,7 +151,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/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());
+  }
+}

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

@@ -358,4 +358,67 @@ public void testConfigFileScopes(String testName, String profile, List<String> e
     List<String> scopes = config.getScopes();
     assertIterableEquals(expectedScopes, scopes);
   }
+
+  // --- Unified Host Tests (added for SPOG support) ---
+
+  @Test
+  public void testGetHostTypeWorkspace() {
+    assertEquals(
+        HostType.WORKSPACE,
+        new DatabricksConfig().setHost("https://adb-123.azuredatabricks.net").getHostType());
+  }
+
+  @Test
+  public void testGetHostTypeAccounts() {
+    assertEquals(
+        HostType.ACCOUNTS,
+        new DatabricksConfig().setHost("https://accounts.cloud.databricks.com").getHostType());
+  }
+
+  @Test
+  public void testGetHostTypeUnified() {
+    assertEquals(
+        HostType.UNIFIED,
+        new DatabricksConfig()
+            .setHost("https://unified.databricks.com")
+            .setExperimentalIsUnifiedHost(true)
+            .getHostType());
+  }
+
+  @Test
+  public void testGetClientTypeWorkspace() {
+    assertEquals(
+        ClientType.WORKSPACE,
+        new DatabricksConfig().setHost("https://adb-123.azuredatabricks.net").getClientType());
+  }
+
+  @Test
+  public void testGetClientTypeAccount() {
+    assertEquals(
+        ClientType.ACCOUNT,
+        new DatabricksConfig().setHost("https://accounts.cloud.databricks.com").getClientType());
+  }
+
+  @Test
+  public void testGetClientTypeWorkspaceOnUnified() {
+    // For unified hosts with workspaceId, client type is WORKSPACE
+    assertEquals(
+        ClientType.WORKSPACE,
+        new DatabricksConfig()
+            .setHost("https://unified.databricks.com")
+            .setExperimentalIsUnifiedHost(true)
+            .setWorkspaceId("123456")
+            .getClientType());
+  }
+
+  @Test
+  public void testGetClientTypeAccountOnUnified() {
+    // For unified hosts without workspaceId, client type is ACCOUNT
+    assertEquals(
+        ClientType.ACCOUNT,
+        new DatabricksConfig()
+            .setHost("https://unified.databricks.com")
+            .setExperimentalIsUnifiedHost(true)
+            .getClientType());
+  }
 }