From e465191a12774989e92849ab3fc78f5a2a524253 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 12 Jan 2026 13:49:16 +0100 Subject: [PATCH 01/13] [Feature] Support for Unified Host --- .../com/databricks/sdk/AccountClient.java | 19 ++ .../com/databricks/sdk/core/ClientType.java | 30 ++ .../DatabricksCliCredentialsProvider.java | 3 +- .../databricks/sdk/core/DatabricksConfig.java | 153 +++++++++- .../sdk/core/DefaultCredentialsProvider.java | 6 +- .../GoogleCredentialsCredentialsProvider.java | 3 +- .../sdk/core/GoogleIdCredentialsProvider.java | 3 +- .../com/databricks/sdk/core/HostType.java | 26 ++ .../sdk/core/UnifiedHostHeaderFactory.java | 39 +++ .../com/databricks/sdk/AccountClientTest.java | 72 +++++ .../sdk/core/DatabricksConfigTest.java | 61 ++++ .../databricks/sdk/core/UnifiedHostTest.java | 265 ++++++++++++++++++ 12 files changed, 675 insertions(+), 5 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java index 5461ba07e..13d160422 100755 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java @@ -5,6 +5,7 @@ import com.databricks.sdk.core.ApiClient; import com.databricks.sdk.core.ConfigLoader; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.HostType; import com.databricks.sdk.core.utils.AzureUtils; import com.databricks.sdk.service.billing.BillableUsageAPI; import com.databricks.sdk.service.billing.BillableUsageService; @@ -1110,7 +1111,25 @@ public DatabricksConfig config() { return config; } + /** + * Creates a WorkspaceClient configured for the specified workspace. + * + *

For unified hosts, this sets the workspace ID on the config instead of changing the host. + * For traditional account hosts, this resolves the workspace deployment URL and creates a config + * with the workspace host. + * + * @param workspace The workspace to create a client for + * @return A configured WorkspaceClient for the specified workspace + */ public WorkspaceClient getWorkspaceClient(Workspace workspace) { + // For unified hosts, reuse the same host and set workspace ID + if (this.config.getHostType() == HostType.UNIFIED) { + DatabricksConfig workspaceConfig = this.config.clone(); + workspaceConfig.setWorkspaceId(String.valueOf(workspace.getWorkspaceId())); + return new WorkspaceClient(workspaceConfig); + } + + // For traditional account hosts, get workspace deployment URL String host = this.config.getDatabricksEnvironment().getDeploymentUrl(workspace.getDeploymentName()); DatabricksConfig config = this.config.newWithWorkspaceHost(host); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java new file mode 100644 index 000000000..e9c603e69 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java @@ -0,0 +1,30 @@ +package com.databricks.sdk.core; + +import com.databricks.sdk.support.InternalApi; + +/** + * Represents the type of Databricks client being used for API operations. + * + *

This is determined by the combination of host type and workspace ID presence: + * + *

+ */ +@InternalApi +public enum ClientType { + /** Traditional workspace client */ + WORKSPACE, + + /** Traditional account client */ + ACCOUNT, + + /** Workspace-scoped client on unified host (requires X-Databricks-Org-Id header) */ + WORKSPACE_ON_UNIFIED, + + /** Account-scoped client on unified host */ + ACCOUNT_ON_UNIFIED +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java index 687ec1dd7..c0732edb7 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java @@ -31,7 +31,8 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { } List cmd = new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); - if (config.isAccountClient()) { + if (config.getClientType() == ClientType.ACCOUNT + || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) { cmd.add("--account-id"); cmd.add(config.getAccountId()); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 572d3cb9b..eff3d5df5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -18,6 +18,23 @@ import java.util.*; import org.apache.http.HttpMessage; +/** + * Configuration for Databricks SDK clients. + * + *

This class holds all configuration needed to authenticate and connect to Databricks services, + * including support for: + * + *

+ * + *

Unified Host Support: When using a unified host, set {@code experimentalIsUnifiedHost} + * to {@code true} and optionally provide a {@code workspaceId} for workspace-scoped operations. Use + * {@link #getHostType()} and {@link #getClientType()} instead of the deprecated {@link + * #isAccountClient()} method. + */ public class DatabricksConfig { private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @@ -27,6 +44,27 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_ACCOUNT_ID") private String accountId; + /** + * Workspace ID for unified host operations. When using a unified host that supports both + * workspace and account-level operations, this field specifies which workspace context to operate + * under for workspace-level API calls. + * + *

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. When true, the host is treated as supporting + * both workspace and account-level operations through a single endpoint. + * + *

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; @@ -233,8 +271,16 @@ public synchronized Map 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 (getClientType() == ClientType.WORKSPACE_ON_UNIFIED) { + headerFactory = new UnifiedHostHeaderFactory(rawHeaderFactory, workspaceId); + } else { + headerFactory = rawHeaderFactory; + } } return headerFactory.headers(); } catch (DatabricksException e) { @@ -298,6 +344,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 +743,73 @@ 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."); } + /** + * Determines the type of host based on configuration settings and host URL. + * + *

Detection logic: + * + *

    + *
  1. If experimentalIsUnifiedHost is true → UNIFIED + *
  2. If host starts with "accounts." or "accounts-dod." → ACCOUNTS + *
  3. Otherwise → WORKSPACE + *
+ * + * @return The detected host type + */ + 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; + } + + /** + * Determines the client type based on host type and workspace ID configuration. + * + *

Client type logic: + * + *

+ * + * @return The determined client type + */ + public ClientType getClientType() { + HostType hostType = getHostType(); + switch (hostType) { + case UNIFIED: + return (workspaceId != null && !workspaceId.isEmpty()) + ? ClientType.WORKSPACE_ON_UNIFIED + : ClientType.ACCOUNT_ON_UNIFIED; + case ACCOUNTS: + return ClientType.ACCOUNT; + case WORKSPACE: + default: + return ClientType.WORKSPACE; + } + } + public OpenIDConnectEndpoints getOidcEndpoints() throws IOException { if (discoveryUrl == null) { return fetchDefaultOidcEndpoints(); @@ -705,10 +830,36 @@ private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() { return null; } + /** + * Fetches OIDC endpoints for unified hosts using the account ID. + * + *

For unified hosts, the OIDC endpoints follow the pattern: + * {host}/oidc/accounts/{accountId}/v1/{token|authorize} + * + * @param accountId The account ID to use for endpoint construction + * @return OpenIDConnectEndpoints configured for the unified host + * @throws DatabricksException if accountId is null or empty + * @throws IOException if endpoint construction fails + */ + 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); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 59ec6eca0..a97d17255 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -150,7 +150,11 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { namedIdTokenSource.idTokenSource, config.getHttpClient()) .audience(config.getTokenAudience()) - .accountId(config.isAccountClient() ? config.getAccountId() : null) + .accountId( + (config.getClientType() == ClientType.ACCOUNT + || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) + ? config.getAccountId() + : null) .scopes(config.getScopes()) .build(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java index 755c1b331..b70ffb49d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java @@ -66,7 +66,8 @@ public HeaderFactory configure(DatabricksConfig config) { Map headers = new HashMap<>(); headers.put("Authorization", String.format("Bearer %s", idToken.getTokenValue())); - if (config.isAccountClient()) { + if (config.getClientType() == ClientType.ACCOUNT + || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) { AccessToken token; try { token = finalServiceAccountCredentials.createScoped(GCP_SCOPES).refreshAccessToken(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java index c51dfd4cc..3bef2aaa0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java @@ -69,7 +69,8 @@ public HeaderFactory configure(DatabricksConfig config) { throw new DatabricksException(message, e); } - if (config.isAccountClient()) { + if (config.getClientType() == ClientType.ACCOUNT + || config.getClientType() == ClientType.ACCOUNT_ON_UNIFIED) { try { headers.put( SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java new file mode 100644 index 000000000..4bdc21648 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java @@ -0,0 +1,26 @@ +package com.databricks.sdk.core; + +import com.databricks.sdk.support.InternalApi; + +/** + * Represents the type of Databricks host being used. + * + *

This determines which APIs are available and how authentication should be handled: + * + *

+ */ +@InternalApi +public enum HostType { + /** Traditional workspace host - supports workspace-level APIs only */ + WORKSPACE, + + /** Traditional accounts host - supports account-level APIs only */ + ACCOUNTS, + + /** Unified host - supports both workspace and account APIs based on context */ + UNIFIED +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java new file mode 100644 index 000000000..c5511b50c --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java @@ -0,0 +1,39 @@ +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. + * + *

When making workspace-level API calls to a unified host, this header is required to specify + * which workspace context the operation should execute in. + */ +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 headers() { + Map headers = new HashMap<>(delegate.headers()); + headers.put("X-Databricks-Org-Id", workspaceId); + return headers; + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java new file mode 100644 index 000000000..378118c3b --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java @@ -0,0 +1,72 @@ +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-on-unified client type + assertEquals(ClientType.WORKSPACE_ON_UNIFIED, workspaceClient.config().getClientType()); + } + + @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()); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index d805de323..6d4b33e4c 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -358,4 +358,65 @@ public void testConfigFileScopes(String testName, String profile, List e List 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() { + assertEquals( + ClientType.WORKSPACE_ON_UNIFIED, + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true) + .setWorkspaceId("123456") + .getClientType()); + } + + @Test + public void testGetClientTypeAccountOnUnified() { + assertEquals( + ClientType.ACCOUNT_ON_UNIFIED, + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true) + .getClientType()); + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java new file mode 100644 index 000000000..c924acff1 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java @@ -0,0 +1,265 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.core.utils.Environment; +import java.io.IOException; +import java.util.*; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Tests for unified host support (SPOG). + * + *

Covers host type detection, client type determination, header injection, and OIDC endpoint + * resolution for unified hosts. + */ +public class UnifiedHostTest { + + // --- Host Type Detection Tests --- + + @Test + public void testHostTypeWorkspace() { + DatabricksConfig config = + new DatabricksConfig().setHost("https://adb-123456789.0.azuredatabricks.net"); + assertEquals(HostType.WORKSPACE, config.getHostType()); + } + + @Test + public void testHostTypeAccounts() { + DatabricksConfig config = + new DatabricksConfig().setHost("https://accounts.cloud.databricks.com"); + assertEquals(HostType.ACCOUNTS, config.getHostType()); + } + + @Test + public void testHostTypeAccountsDod() { + DatabricksConfig config = + new DatabricksConfig().setHost("https://accounts-dod.cloud.databricks.us"); + assertEquals(HostType.ACCOUNTS, config.getHostType()); + } + + @Test + public void testHostTypeUnifiedExplicitFlag() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true); + assertEquals(HostType.UNIFIED, config.getHostType()); + } + + @Test + public void testHostTypeUnifiedOverridesAccounts() { + // Even if host looks like accounts, explicit flag takes precedence + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setExperimentalIsUnifiedHost(true); + assertEquals(HostType.UNIFIED, config.getHostType()); + } + + @Test + public void testHostTypeNullHost() { + DatabricksConfig config = new DatabricksConfig(); + assertEquals(HostType.WORKSPACE, config.getHostType()); + } + + // --- Client Type Detection Tests --- + + private static Stream provideClientTypeTestCases() { + return Stream.of( + Arguments.of( + "Workspace host", + "https://adb-123.azuredatabricks.net", + null, + false, + ClientType.WORKSPACE), + Arguments.of( + "Account host", + "https://accounts.cloud.databricks.com", + null, + false, + ClientType.ACCOUNT), + Arguments.of( + "Unified without workspace ID", + "https://unified.databricks.com", + null, + true, + ClientType.ACCOUNT_ON_UNIFIED), + Arguments.of( + "Unified with workspace ID", + "https://unified.databricks.com", + "123456", + true, + ClientType.WORKSPACE_ON_UNIFIED)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideClientTypeTestCases") + public void testClientType( + String testName, String host, String workspaceId, boolean isUnified, ClientType expected) { + DatabricksConfig config = new DatabricksConfig().setHost(host).setWorkspaceId(workspaceId); + if (isUnified) { + config.setExperimentalIsUnifiedHost(true); + } + assertEquals(expected, config.getClientType()); + } + + // --- OIDC Endpoint Tests --- + + @Test + public void testOidcEndpointsForUnifiedHost() throws IOException { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true) + .setAccountId("test-account-123"); + + OpenIDConnectEndpoints endpoints = config.getOidcEndpoints(); + + assertEquals( + "https://unified.databricks.com/oidc/accounts/test-account-123/v1/authorize", + endpoints.getAuthorizationEndpoint()); + assertEquals( + "https://unified.databricks.com/oidc/accounts/test-account-123/v1/token", + endpoints.getTokenEndpoint()); + } + + @Test + public void testOidcEndpointsForUnifiedHostMissingAccountId() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true); + // No account ID set + + DatabricksException exception = + assertThrows(DatabricksException.class, () -> config.getOidcEndpoints()); + assertTrue(exception.getMessage().contains("account_id is required")); + } + + // --- isAccountClient() Deprecation Tests --- + + @Test + public void testIsAccountClientThrowsForUnifiedHost() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true); + + DatabricksException exception = + assertThrows(DatabricksException.class, config::isAccountClient); + assertTrue(exception.getMessage().contains("Cannot determine account client status")); + assertTrue(exception.getMessage().contains("getHostType()")); + } + + @Test + public void testIsAccountClientWorksFineForTraditionalHosts() { + assertTrue( + new DatabricksConfig().setHost("https://accounts.cloud.databricks.com").isAccountClient()); + + assertFalse( + new DatabricksConfig().setHost("https://adb-123.azuredatabricks.net").isAccountClient()); + } + + // --- Environment Variable Tests --- + + @Test + public void testUnifiedHostFromEnvironmentVariables() { + Map env = new HashMap<>(); + env.put("DATABRICKS_HOST", "https://unified.databricks.com"); + env.put("DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST", "true"); + env.put("DATABRICKS_WORKSPACE_ID", "987654321"); + env.put("DATABRICKS_ACCOUNT_ID", "account-abc"); + + DatabricksConfig config = new DatabricksConfig(); + config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name"))); + + assertEquals(HostType.UNIFIED, config.getHostType()); + assertEquals("987654321", config.getWorkspaceId()); + assertEquals("account-abc", config.getAccountId()); + assertEquals(ClientType.WORKSPACE_ON_UNIFIED, config.getClientType()); + } + + // --- UnifiedHostHeaderFactory Tests --- + + @Test + public void testUnifiedHostHeaderFactoryAddsHeader() { + Map baseHeaders = new HashMap<>(); + baseHeaders.put("Authorization", "Bearer token123"); + + HeaderFactory baseFactory = () -> baseHeaders; + UnifiedHostHeaderFactory unifiedFactory = new UnifiedHostHeaderFactory(baseFactory, "ws-456"); + + Map headers = unifiedFactory.headers(); + + assertEquals("Bearer token123", headers.get("Authorization")); + assertEquals("ws-456", headers.get("X-Databricks-Org-Id")); + } + + @Test + public void testUnifiedHostHeaderFactoryRequiresDelegate() { + assertThrows( + IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(null, "ws-123")); + } + + @Test + public void testUnifiedHostHeaderFactoryRequiresWorkspaceId() { + HeaderFactory baseFactory = () -> new HashMap<>(); + assertThrows( + IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(baseFactory, null)); + assertThrows( + IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(baseFactory, "")); + } + + // --- Header Injection Integration Tests --- + + @Test + public void testHeaderInjectionForWorkspaceOnUnified() { + String workspaceId = "123456789"; + + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true) + .setWorkspaceId(workspaceId) + .setToken("test-token"); + + Map headers = config.authenticate(); + + assertEquals("Bearer test-token", headers.get("Authorization")); + assertEquals(workspaceId, headers.get("X-Databricks-Org-Id")); + } + + @Test + public void testNoHeaderInjectionForAccountOnUnified() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://unified.databricks.com") + .setExperimentalIsUnifiedHost(true) + .setToken("test-token"); + // No workspace ID set + + Map headers = config.authenticate(); + + assertEquals("Bearer test-token", headers.get("Authorization")); + assertNull(headers.get("X-Databricks-Org-Id")); + } + + @Test + public void testNoHeaderInjectionForTraditionalWorkspace() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://adb-123.azuredatabricks.net") + .setToken("test-token"); + + Map headers = config.authenticate(); + + assertEquals("Bearer test-token", headers.get("Authorization")); + assertNull(headers.get("X-Databricks-Org-Id")); + } +} From f2d52300e7243322f275bb26ed3adc76326a3738 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 12 Jan 2026 13:57:38 +0100 Subject: [PATCH 02/13] - --- NEXT_CHANGELOG.md | 1 + .../databricks/sdk/core/DatabricksConfig.java | 20 +++++-------------- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 20b5ec452..b71cdfcb8 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -4,6 +4,7 @@ ### New Features and Improvements +* Add support for unified hosts with experimental flag. * Increase async cache stale period from 3 to 5 minutes to cover the maximum monthly downtime of a 99.99% uptime SLA. ### Bug Fixes diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index eff3d5df5..62e05eeef 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -758,13 +758,8 @@ public boolean isAccountClient() { /** * Determines the type of host based on configuration settings and host URL. * - *

Detection logic: - * - *

    - *
  1. If experimentalIsUnifiedHost is true → UNIFIED - *
  2. If host starts with "accounts." or "accounts-dod." → ACCOUNTS - *
  3. Otherwise → WORKSPACE - *
+ *

Returns UNIFIED if experimentalIsUnifiedHost is true, ACCOUNTS if the host starts with + * "accounts." or "accounts-dod.", and WORKSPACE otherwise. * * @return The detected host type */ @@ -784,14 +779,9 @@ public HostType getHostType() { /** * Determines the client type based on host type and workspace ID configuration. * - *

Client type logic: - * - *

+ *

For unified hosts, returns WORKSPACE_ON_UNIFIED if a workspace ID is set, or + * ACCOUNT_ON_UNIFIED otherwise. For traditional hosts, returns ACCOUNT or WORKSPACE based on the + * host type. * * @return The determined client type */ From 2aab1d72cb00d07cafaba999ac2f51e1e4b2567b Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 19 Jan 2026 01:45:30 +0100 Subject: [PATCH 03/13] - --- .../com/databricks/sdk/core/ClientType.java | 21 +++++-------------- .../com/databricks/sdk/core/HostType.java | 18 ++++------------ .../sdk/core/UnifiedHostHeaderFactory.java | 3 --- 3 files changed, 9 insertions(+), 33 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java index e9c603e69..ed51d8491 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java @@ -2,29 +2,18 @@ import com.databricks.sdk.support.InternalApi; -/** - * Represents the type of Databricks client being used for API operations. - * - *

This is determined by the combination of host type and workspace ID presence: - * - *

- */ +/** Represents the type of Databricks client being used for API operations. */ @InternalApi public enum ClientType { - /** Traditional workspace client */ + /** Traditional workspace client. */ WORKSPACE, - /** Traditional account client */ + /** Traditional account client. */ ACCOUNT, - /** Workspace-scoped client on unified host (requires X-Databricks-Org-Id header) */ + /** Workspace-scoped client on unified host. */ WORKSPACE_ON_UNIFIED, - /** Account-scoped client on unified host */ + /** Account-scoped client on unified host. */ ACCOUNT_ON_UNIFIED } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java index 4bdc21648..005807839 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java @@ -2,25 +2,15 @@ import com.databricks.sdk.support.InternalApi; -/** - * Represents the type of Databricks host being used. - * - *

This determines which APIs are available and how authentication should be handled: - * - *

- */ +/** Represents the type of Databricks host being used. */ @InternalApi public enum HostType { - /** Traditional workspace host - supports workspace-level APIs only */ + /** Traditional workspace host. */ WORKSPACE, - /** Traditional accounts host - supports account-level APIs only */ + /** Traditional accounts host. */ ACCOUNTS, - /** Unified host - supports both workspace and account APIs based on context */ + /** Unified host supporting both workspace and account operations. */ UNIFIED } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java index c5511b50c..2889329e5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java @@ -5,9 +5,6 @@ /** * HeaderFactory wrapper that adds X-Databricks-Org-Id header for unified host workspace operations. - * - *

When making workspace-level API calls to a unified host, this header is required to specify - * which workspace context the operation should execute in. */ class UnifiedHostHeaderFactory implements HeaderFactory { private final HeaderFactory delegate; From 278f9696f32b6cd38d156ccf66aed9fa3527ddbd Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 19 Jan 2026 01:52:22 +0100 Subject: [PATCH 04/13] - --- .../com/databricks/sdk/AccountClient.java | 10 ---- .../databricks/sdk/core/DatabricksConfig.java | 54 ++----------------- 2 files changed, 4 insertions(+), 60 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java index 13d160422..7408c0629 100755 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java @@ -1111,16 +1111,6 @@ public DatabricksConfig config() { return config; } - /** - * Creates a WorkspaceClient configured for the specified workspace. - * - *

For unified hosts, this sets the workspace ID on the config instead of changing the host. - * For traditional account hosts, this resolves the workspace deployment URL and creates a config - * with the workspace host. - * - * @param workspace The workspace to create a client for - * @return A configured WorkspaceClient for the specified workspace - */ public WorkspaceClient getWorkspaceClient(Workspace workspace) { // For unified hosts, reuse the same host and set workspace ID if (this.config.getHostType() == HostType.UNIFIED) { diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 62e05eeef..ff55f89ee 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -18,23 +18,6 @@ import java.util.*; import org.apache.http.HttpMessage; -/** - * Configuration for Databricks SDK clients. - * - *

This class holds all configuration needed to authenticate and connect to Databricks services, - * including support for: - * - *

- * - *

Unified Host Support: When using a unified host, set {@code experimentalIsUnifiedHost} - * to {@code true} and optionally provide a {@code workspaceId} for workspace-scoped operations. Use - * {@link #getHostType()} and {@link #getClientType()} instead of the deprecated {@link - * #isAccountClient()} method. - */ public class DatabricksConfig { private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @@ -45,9 +28,7 @@ public class DatabricksConfig { private String accountId; /** - * Workspace ID for unified host operations. When using a unified host that supports both - * workspace and account-level operations, this field specifies which workspace context to operate - * under for workspace-level API calls. + * Workspace ID for unified host operations. * *

Note: This API is experimental and may change or be removed in future releases * without notice. @@ -56,8 +37,7 @@ public class DatabricksConfig { private String workspaceId; /** - * Flag to explicitly mark a host as a unified host. When true, the host is treated as supporting - * both workspace and account-level operations through a single endpoint. + * 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. @@ -755,14 +735,7 @@ public boolean isAccountClient() { return host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod."); } - /** - * Determines the type of host based on configuration settings and host URL. - * - *

Returns UNIFIED if experimentalIsUnifiedHost is true, ACCOUNTS if the host starts with - * "accounts." or "accounts-dod.", and WORKSPACE otherwise. - * - * @return The detected host type - */ + /** Returns the host type based on configuration settings and host URL. */ public HostType getHostType() { if (experimentalIsUnifiedHost != null && experimentalIsUnifiedHost) { return HostType.UNIFIED; @@ -776,15 +749,7 @@ public HostType getHostType() { return HostType.WORKSPACE; } - /** - * Determines the client type based on host type and workspace ID configuration. - * - *

For unified hosts, returns WORKSPACE_ON_UNIFIED if a workspace ID is set, or - * ACCOUNT_ON_UNIFIED otherwise. For traditional hosts, returns ACCOUNT or WORKSPACE based on the - * host type. - * - * @return The determined client type - */ + /** Returns the client type based on host type and workspace ID configuration. */ public ClientType getClientType() { HostType hostType = getHostType(); switch (hostType) { @@ -820,17 +785,6 @@ private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() { return null; } - /** - * Fetches OIDC endpoints for unified hosts using the account ID. - * - *

For unified hosts, the OIDC endpoints follow the pattern: - * {host}/oidc/accounts/{accountId}/v1/{token|authorize} - * - * @param accountId The account ID to use for endpoint construction - * @return OpenIDConnectEndpoints configured for the unified host - * @throws DatabricksException if accountId is null or empty - * @throws IOException if endpoint construction fails - */ private OpenIDConnectEndpoints getUnifiedOidcEndpoints(String accountId) throws IOException { if (accountId == null || accountId.isEmpty()) { throw new DatabricksException( From 129448613115c41312cd607da14cbeb6b2c04f6a Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 19 Jan 2026 02:10:28 +0100 Subject: [PATCH 05/13] - --- .../com/databricks/sdk/core/ClientType.java | 12 +++------ .../DatabricksCliCredentialsProvider.java | 3 +-- .../databricks/sdk/core/DatabricksConfig.java | 25 ++++++++----------- .../sdk/core/DefaultCredentialsProvider.java | 5 +--- .../GoogleCredentialsCredentialsProvider.java | 3 +-- .../sdk/core/GoogleIdCredentialsProvider.java | 3 +-- .../com/databricks/sdk/AccountClientTest.java | 7 ++++-- .../sdk/core/DatabricksConfigTest.java | 6 +++-- .../databricks/sdk/core/UnifiedHostTest.java | 6 ++--- 9 files changed, 29 insertions(+), 41 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java index ed51d8491..e03956e8c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ClientType.java +++ b/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 } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java index c0732edb7..e4376f203 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java @@ -31,8 +31,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { } List 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()); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index ff55f89ee..0b8ee81f8 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/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. - * - *

Note: 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. - * - *

Note: 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. - * - *

Note: 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 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: diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index a97d17255..8d4e593f4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/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(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java index b70ffb49d..463d2bab9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java @@ -66,8 +66,7 @@ public HeaderFactory configure(DatabricksConfig config) { Map 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(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java index 3bef2aaa0..376d691c5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java +++ b/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()); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java index 378118c3b..ca20fe5a2 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java +++ b/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 diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 6d4b33e4c..eab6d991c 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/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) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java index c924acff1..6ecb9a048 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java @@ -89,13 +89,13 @@ private static Stream 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 --- From e7a4e323169639112a59b00eb7815fa7415bf482 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Tue, 20 Jan 2026 15:21:27 +0100 Subject: [PATCH 06/13] - --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 6 ++++-- .../com/databricks/sdk/core/DefaultCredentialsProvider.java | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 0b8ee81f8..7d2074b4d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -57,8 +57,10 @@ public class DatabricksConfig { private String redirectUrl; /** - * 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. + * 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; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 8d4e593f4..eabab7815 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/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(); From 4d9b21e0b369d8fddcf85b368503cf2a2c6ecd99 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Tue, 20 Jan 2026 17:55:29 +0100 Subject: [PATCH 07/13] - --- .../src/main/java/com/databricks/sdk/AccountClient.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java index 7408c0629..b1f503258 100755 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/AccountClient.java @@ -1114,9 +1114,8 @@ public DatabricksConfig config() { public WorkspaceClient getWorkspaceClient(Workspace workspace) { // For unified hosts, reuse the same host and set workspace ID if (this.config.getHostType() == HostType.UNIFIED) { - DatabricksConfig workspaceConfig = this.config.clone(); - workspaceConfig.setWorkspaceId(String.valueOf(workspace.getWorkspaceId())); - return new WorkspaceClient(workspaceConfig); + this.config.setWorkspaceId(String.valueOf(workspace.getWorkspaceId())); + return new WorkspaceClient(this.config); } // For traditional account hosts, get workspace deployment URL From 0157e222d8e708456a3fcf8a1426c415d3b91cd8 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:13:11 +0100 Subject: [PATCH 08/13] - --- .../databricks/sdk/core/DatabricksConfig.java | 24 +++++++++++++++---- .../sdk/core/UnifiedHostHeaderFactory.java | 5 ++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 7d2074b4d..fe4b71389 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -270,15 +270,31 @@ public TokenSource getTokenSource() { if (headerFactory == null) { try { 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; + } } catch (Exception e) { return new ErrorTokenSource("Failed to get token source: " + e.getMessage()); } - setAuthType(credentialsProvider.authType()); } - if (headerFactory instanceof OAuthHeaderFactory) { - return (TokenSource) headerFactory; + // For unified hosts, the underlying token source is wrapped, so extract it + HeaderFactory underlyingFactory = headerFactory; + if (headerFactory instanceof UnifiedHostHeaderFactory) { + // UnifiedHostHeaderFactory wraps an OAuthHeaderFactory, we need to get the delegate + // For token source purposes, we return the underlying OAuth token source + underlyingFactory = ((UnifiedHostHeaderFactory) headerFactory).getDelegate(); + } + + if (underlyingFactory instanceof OAuthHeaderFactory) { + return (TokenSource) underlyingFactory; } return new ErrorTokenSource( String.format("OAuth Token not supported for current auth type %s", authType)); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java index 2889329e5..cb82de867 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java @@ -33,4 +33,9 @@ public Map headers() { headers.put("X-Databricks-Org-Id", workspaceId); return headers; } + + /** Returns the underlying header factory delegate. */ + public HeaderFactory getDelegate() { + return delegate; + } } From 625417d66659b8aa39a6a36d1a4e8fbb17595755 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:18:36 +0100 Subject: [PATCH 09/13] - --- .../databricks/sdk/core/DatabricksConfig.java | 39 +++++------------- .../sdk/core/UnifiedHostHeaderFactory.java | 41 ------------------- 2 files changed, 11 insertions(+), 69 deletions(-) delete mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index fe4b71389..51d3d95a5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -247,18 +247,17 @@ public synchronized Map authenticate() throws DatabricksExceptio if (headerFactory == null) { // Calling authenticate without resolve ConfigLoader.fixHostIfNeeded(this); - HeaderFactory rawHeaderFactory = credentialsProvider.configure(this); + headerFactory = credentialsProvider.configure(this); setAuthType(credentialsProvider.authType()); + } + Map headers = new HashMap<>(headerFactory.headers()); - // 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; - } + // 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 headerFactory.headers(); + + return headers; } catch (DatabricksException e) { String msg = String.format("%s auth: %s", credentialsProvider.authType(), e.getMessage()); DatabricksException wrapperException = new DatabricksException(msg, e); @@ -270,31 +269,15 @@ public TokenSource getTokenSource() { if (headerFactory == null) { try { ConfigLoader.fixHostIfNeeded(this); - HeaderFactory rawHeaderFactory = credentialsProvider.configure(this); + headerFactory = 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; - } } catch (Exception e) { return new ErrorTokenSource("Failed to get token source: " + e.getMessage()); } } - // For unified hosts, the underlying token source is wrapped, so extract it - HeaderFactory underlyingFactory = headerFactory; - if (headerFactory instanceof UnifiedHostHeaderFactory) { - // UnifiedHostHeaderFactory wraps an OAuthHeaderFactory, we need to get the delegate - // For token source purposes, we return the underlying OAuth token source - underlyingFactory = ((UnifiedHostHeaderFactory) headerFactory).getDelegate(); - } - - if (underlyingFactory instanceof OAuthHeaderFactory) { - return (TokenSource) underlyingFactory; + if (headerFactory instanceof OAuthHeaderFactory) { + return (TokenSource) headerFactory; } return new ErrorTokenSource( String.format("OAuth Token not supported for current auth type %s", authType)); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java deleted file mode 100644 index cb82de867..000000000 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/UnifiedHostHeaderFactory.java +++ /dev/null @@ -1,41 +0,0 @@ -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 headers() { - Map headers = new HashMap<>(delegate.headers()); - headers.put("X-Databricks-Org-Id", workspaceId); - return headers; - } - - /** Returns the underlying header factory delegate. */ - public HeaderFactory getDelegate() { - return delegate; - } -} From 44e71a7c31fb3aa60c303ade6cf67b8bed53328c Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:18:56 +0100 Subject: [PATCH 10/13] - --- .../databricks/sdk/core/UnifiedHostTest.java | 33 +------------------ 1 file changed, 1 insertion(+), 32 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java index 6ecb9a048..b02ac0b34 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java @@ -185,38 +185,7 @@ public void testUnifiedHostFromEnvironmentVariables() { assertEquals(ClientType.WORKSPACE, config.getClientType()); } - // --- UnifiedHostHeaderFactory Tests --- - - @Test - public void testUnifiedHostHeaderFactoryAddsHeader() { - Map baseHeaders = new HashMap<>(); - baseHeaders.put("Authorization", "Bearer token123"); - - HeaderFactory baseFactory = () -> baseHeaders; - UnifiedHostHeaderFactory unifiedFactory = new UnifiedHostHeaderFactory(baseFactory, "ws-456"); - - Map headers = unifiedFactory.headers(); - - assertEquals("Bearer token123", headers.get("Authorization")); - assertEquals("ws-456", headers.get("X-Databricks-Org-Id")); - } - - @Test - public void testUnifiedHostHeaderFactoryRequiresDelegate() { - assertThrows( - IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(null, "ws-123")); - } - - @Test - public void testUnifiedHostHeaderFactoryRequiresWorkspaceId() { - HeaderFactory baseFactory = () -> new HashMap<>(); - assertThrows( - IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(baseFactory, null)); - assertThrows( - IllegalArgumentException.class, () -> new UnifiedHostHeaderFactory(baseFactory, "")); - } - - // --- Header Injection Integration Tests --- + // --- Header Injection Tests --- @Test public void testHeaderInjectionForWorkspaceOnUnified() { From 2b34fe7c893caafb11ed790df5d561f3a56bd499 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:23:45 +0100 Subject: [PATCH 11/13] - --- .../src/main/java/com/databricks/sdk/core/DatabricksConfig.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 51d3d95a5..b5e911b97 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -270,10 +270,10 @@ public TokenSource getTokenSource() { try { ConfigLoader.fixHostIfNeeded(this); headerFactory = credentialsProvider.configure(this); - setAuthType(credentialsProvider.authType()); } catch (Exception e) { return new ErrorTokenSource("Failed to get token source: " + e.getMessage()); } + setAuthType(credentialsProvider.authType()); } if (headerFactory instanceof OAuthHeaderFactory) { From 69cae5a60ec9ebe44542dbdf5d874ac82f12ea57 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:33:09 +0100 Subject: [PATCH 12/13] - --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index b5e911b97..dc8f5b266 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -28,8 +28,7 @@ public class DatabricksConfig { 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. + * Workspace ID for unified host operations. */ @ConfigAttribute(env = "DATABRICKS_WORKSPACE_ID") private String workspaceId; From b7a4be01468557d8e487bf79f0b76dbf46ce9080 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 21 Jan 2026 00:36:31 +0100 Subject: [PATCH 13/13] - --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index dc8f5b266..23fd990b9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -27,9 +27,7 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_ACCOUNT_ID") private String accountId; - /** - * Workspace ID for unified host operations. - */ + /** Workspace ID for unified host operations. */ @ConfigAttribute(env = "DATABRICKS_WORKSPACE_ID") private String workspaceId;