[WIP] Draft unified host changes

Author: jh-db

CHANGELOG.md

@@ -1,5 +1,22 @@
 # Version changelog
 
+## Unreleased
+
+### New Features and Improvements
+* Add support for unified hosts, i.e. hosts that support both workspace-level and account-level operations
+* Add `HostType` and `ConfigType` enums to `DatabricksConfig` for better host type management
+* Add `workspaceId` field to `DatabricksConfig` for workspace clients on unified hosts
+* Add `experimentalIsUnifiedHost` field to `DatabricksConfig` to mark unified hosts
+* Add `getHostType()` and `getConfigType()` methods to `DatabricksConfig`
+* Add X-Databricks-Org-Id header support for unified host workspace requests
+
+### Deprecations
+* Deprecate `isAccountClient()` method in `DatabricksConfig`. Use `getHostType()` or `getConfigType()` instead.
+
+### Internal Changes
+* Update OAuth endpoint discovery to support unified hosts
+* Update credential providers to use `getHostType()` instead of `isAccountClient()`
+
 ## Release v0.68.0 (2025-10-30)
 
 ### Documentation

databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java

@@ -40,6 +40,8 @@ public static class Builder {
     private Integer debugTruncateBytes;
     private HttpClient httpClient;
     private String accountId;
+    private String workspaceId;
+    private DatabricksConfig.HostType hostType;
     private RetryStrategyPicker retryStrategyPicker;
     private boolean isDebugHeaders;
 
@@ -50,6 +52,8 @@ public Builder withDatabricksConfig(DatabricksConfig config) {
       this.httpClient = config.getHttpClient();
       this.debugTruncateBytes = config.getDebugTruncateBytes();
       this.accountId = config.getAccountId();
+      this.workspaceId = config.getWorkspaceId();
+      this.hostType = config.getHostType();
       this.isDebugHeaders = config.isDebugHeaders();
 
       if (config.getDisableRetries()) {
@@ -112,6 +116,8 @@ public ApiClient build() {
   private final Function<Void, String> getHostFunc;
   private final Function<Void, String> getAuthTypeFunc;
   private final String accountId;
+  private final String workspaceId;
+  private final DatabricksConfig.HostType hostType;
   private final boolean isDebugHeaders;
   private static final String RETRY_AFTER_HEADER = "retry-after";
 
@@ -141,6 +147,8 @@ private ApiClient(Builder builder) {
     this.getAuthTypeFunc = builder.getAuthTypeFunc != null ? builder.getAuthTypeFunc : v -> "";
     this.httpClient = builder.httpClient;
     this.accountId = builder.accountId;
+    this.workspaceId = builder.workspaceId;
+    this.hostType = builder.hostType;
     this.retryStrategyPicker =
         builder.retryStrategyPicker != null
             ? builder.retryStrategyPicker
@@ -240,6 +248,16 @@ private Response executeInner(Request in, String path, RequestOptions options) {
       }
       in.withHeader("User-Agent", userAgent);
 
+      // Unified hosts use X-Databricks-Org-Id header to determine which workspace to route the
+      // request to. The header must not be set for account-level API requests, otherwise the
+      // request will fail. This relies on the assumption that workspaceId is only set for
+      // workspace client configs.
+      if (hostType == DatabricksConfig.HostType.UNIFIED_HOST
+          && workspaceId != null
+          && !workspaceId.isEmpty()) {
+        in.withHeader("X-Databricks-Org-Id", workspaceId);
+      }
+
       options.applyOptions(in);
 
       // Make the request, catching any exceptions, as we may want to retry.

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

@@ -29,7 +29,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
     }
     List<String> cmd =
         new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
-    if (config.isAccountClient()) {
+    if (config.getHostType() != DatabricksConfig.HostType.WORKSPACE_HOST) {
       cmd.add("--account-id");
       cmd.add(config.getAccountId());
     }

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

@@ -19,6 +19,28 @@
 import org.apache.http.HttpMessage;
 
 public class DatabricksConfig {
+  /** HostType represents the type of API the configured host supports. */
+  public enum HostType {
+    /** WorkspaceHost supports only workspace-level APIs. */
+    WORKSPACE_HOST,
+    /** AccountHost supports only account-level APIs. */
+    ACCOUNT_HOST,
+    /** UnifiedHost supports both workspace-level and account-level APIs. */
+    UNIFIED_HOST
+  }
+
+  /** ConfigType represents the type of API this config is valid for. */
+  public enum ConfigType {
+    /** WorkspaceConfig is valid for workspace-level API requests. */
+    WORKSPACE_CONFIG,
+    /** AccountConfig is valid for account-level API requests. */
+    ACCOUNT_CONFIG,
+    /**
+     * InvalidConfig is returned when the config is not valid for either workspace-level or
+     * account-level APIs.
+     */
+    INVALID_CONFIG
+  }
   private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider();
 
   @ConfigAttribute(env = "DATABRICKS_HOST")
@@ -27,6 +49,14 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_ACCOUNT_ID")
   private String accountId;
 
+  /** Databricks Workspace ID for Workspace clients when working with unified hosts. */
+  @ConfigAttribute(env = "DATABRICKS_WORKSPACE_ID")
+  private String workspaceId;
+
+  /** Marker for unified hosts. Will be redundant once we can recognize unified hosts by their hostname. */
+  @ConfigAttribute(env = "DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST")
+  private Boolean experimentalIsUnifiedHost;
+
   @ConfigAttribute(env = "DATABRICKS_TOKEN", auth = "pat", sensitive = true)
   private String token;
 
@@ -290,6 +320,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 != null && experimentalIsUnifiedHost;
+  }
+
+  public DatabricksConfig setExperimentalIsUnifiedHost(boolean experimentalIsUnifiedHost) {
+    this.experimentalIsUnifiedHost = experimentalIsUnifiedHost;
+    return this;
+  }
+
   public String getDatabricksCliPath() {
     return this.databricksCliPath;
   }
@@ -670,13 +718,67 @@ public boolean isAws() {
     return this.getDatabricksEnvironment().getCloud() == Cloud.AWS;
   }
 
+  /**
+   * Returns true if client is configured for Accounts API. Panics if the config has the unified
+   * host flag set.
+   *
+   * @deprecated Use {@link #getHostType()} if possible, or {@link #getConfigType()} if necessary.
+   */
+  @Deprecated
   public boolean isAccountClient() {
+    if (getExperimentalIsUnifiedHost()) {
+      throw new IllegalStateException(
+          "isAccountClient cannot be used with unified hosts; use getHostType() instead");
+    }
     if (host == null) {
       return false;
     }
     return host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod.");
   }
 
+  /** Returns the type of host that the client is configured for. */
+  public HostType getHostType() {
+    if (getExperimentalIsUnifiedHost()) {
+      return HostType.UNIFIED_HOST;
+    }
+
+    if (host == null) {
+      return HostType.WORKSPACE_HOST;
+    }
+
+    if (host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod.")) {
+      return HostType.ACCOUNT_HOST;
+    }
+
+    return HostType.WORKSPACE_HOST;
+  }
+
+  /**
+   * Returns the type of config that the client is configured for. Returns InvalidConfig if the
+   * config is invalid. Use of this method should be avoided where possible, because we plan to
+   * remove WorkspaceClient and AccountClient in favor of a single unified client in the future.
+   */
+  public ConfigType getConfigType() {
+    HostType hostType = getHostType();
+    switch (hostType) {
+      case ACCOUNT_HOST:
+        return ConfigType.ACCOUNT_CONFIG;
+      case WORKSPACE_HOST:
+        return ConfigType.WORKSPACE_CONFIG;
+      case UNIFIED_HOST:
+        if (accountId == null || accountId.isEmpty()) {
+          // All unified host configs must have an account ID
+          return ConfigType.INVALID_CONFIG;
+        }
+        if (workspaceId != null && !workspaceId.isEmpty()) {
+          return ConfigType.WORKSPACE_CONFIG;
+        }
+        return ConfigType.ACCOUNT_CONFIG;
+      default:
+        return ConfigType.INVALID_CONFIG;
+    }
+  }
+
   public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
     if (discoveryUrl == null) {
       return fetchDefaultOidcEndpoints();
@@ -712,23 +814,49 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
       return new OpenIDConnectEndpoints(
           realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl);
     }
-    if (isAccountClient() && getAccountId() != null) {
-      String prefix = getHost() + "/oidc/accounts/" + getAccountId();
-      return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize");
-    }
 
-    ApiClient apiClient =
-        new ApiClient.Builder()
-            .withHttpClient(getHttpClient())
-            .withGetHostFunc(v -> getHost())
-            .build();
-    try {
-      return apiClient.execute(
-          new Request("GET", "/oidc/.well-known/oauth-authorization-server"),
-          OpenIDConnectEndpoints.class);
-    } catch (IOException e) {
-      throw new DatabricksException("IO error: " + e.getMessage(), e);
+    HostType hostType = getHostType();
+    switch (hostType) {
+      case ACCOUNT_HOST:
+        if (getAccountId() != null) {
+          String prefix = getHost() + "/oidc/accounts/" + getAccountId();
+          return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize");
+        }
+        break;
+      case UNIFIED_HOST:
+        if (getAccountId() != null) {
+          ApiClient apiClient =
+              new ApiClient.Builder()
+                  .withHttpClient(getHttpClient())
+                  .withGetHostFunc(v -> getHost())
+                  .build();
+          try {
+            String discoveryPath =
+                "/oidc/accounts/" + getAccountId() + "/.well-known/oauth-authorization-server";
+            return apiClient.execute(
+                new Request("GET", discoveryPath), OpenIDConnectEndpoints.class);
+          } catch (IOException e) {
+            throw new DatabricksException("IO error: " + e.getMessage(), e);
+          }
+        }
+        break;
+      case WORKSPACE_HOST:
+        ApiClient apiClient =
+            new ApiClient.Builder()
+                .withHttpClient(getHttpClient())
+                .withGetHostFunc(v -> getHost())
+                .build();
+        try {
+          return apiClient.execute(
+              new Request("GET", "/oidc/.well-known/oauth-authorization-server"),
+              OpenIDConnectEndpoints.class);
+        } catch (IOException e) {
+          throw new DatabricksException("IO error: " + e.getMessage(), e);
+        }
+      default:
+        break;
     }
+    return null;
   }
 
   @Override
@@ -795,9 +923,10 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
             Arrays.asList(
                 // The config for WorkspaceClient has a different host and Azure Workspace resource
                 // ID, and also omits
-                // the account ID.
+                // the account ID and workspace ID.
                 "host",
                 "accountId",
+                "workspaceId",
                 "azureWorkspaceResourceId",
                 // For cloud-native OAuth, we need to reauthenticate as the audience has changed, so
                 // don't cache the

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

@@ -148,7 +148,10 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
                   namedIdTokenSource.idTokenSource,
                   config.getHttpClient())
               .audience(config.getTokenAudience())
-              .accountId(config.isAccountClient() ? config.getAccountId() : null)
+              .accountId(
+                  config.getHostType() != DatabricksConfig.HostType.WORKSPACE_HOST
+                      ? config.getAccountId()
+                      : null)
               .scopes(config.getScopes())
               .build();
 

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

@@ -64,7 +64,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.getHostType() != DatabricksConfig.HostType.WORKSPACE_HOST) {
         AccessToken token;
         try {
           token = finalServiceAccountCredentials.createScoped(GCP_SCOPES).refreshAccessToken();

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

@@ -67,7 +67,7 @@ public HeaderFactory configure(DatabricksConfig config) {
         throw new DatabricksException(message, e);
       }
 
-      if (config.isAccountClient()) {
+      if (config.getHostType() != DatabricksConfig.HostType.WORKSPACE_HOST) {
         try {
           headers.put(
               SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue());