From 483ae61b637e2afb9a3ed3c00c561294a3d1bd3e Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 19 Mar 2026 09:28:57 +0000 Subject: [PATCH] Remove unified flag usage, rely on host metadata Port of Go SDK #1547. Removes HostType.UNIFIED and all runtime checks of experimentalIsUnifiedHost. Host type is now determined solely from URL pattern (accounts.* = ACCOUNTS, else WORKSPACE). Host metadata resolution from /.well-known/databricks-config (added in PR 4) handles populating accountId, workspaceId, and discoveryUrl automatically. Key changes: - getHostType(): no longer returns UNIFIED - isAccountClient(): no longer throws for unified hosts - getClientType(): simplified, no UNIFIED case - fetchDefaultOidcEndpoints(): removed unified OIDC branch - DatabricksCliCredentialsProvider: removed --experimental-is-unified-host - AccountClient.getWorkspaceClient(): uses DNS zone matching (like Go SDK) to decide whether to reuse host or build deployment URL Co-authored-by: Isaac --- .../com/databricks/sdk/AccountClient.java | 43 ++++- .../DatabricksCliCredentialsProvider.java | 13 +- .../databricks/sdk/core/DatabricksConfig.java | 36 +--- .../com/databricks/sdk/core/HostType.java | 5 +- .../com/databricks/sdk/AccountClientTest.java | 21 +-- .../DatabricksCliCredentialsProviderTest.java | 102 +--------- .../sdk/core/DatabricksConfigTest.java | 65 +------ .../databricks/sdk/core/UnifiedHostTest.java | 176 +++--------------- 8 files changed, 86 insertions(+), 375 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 cda884530..d750badd9 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,7 +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.DatabricksEnvironment; import com.databricks.sdk.core.utils.AzureUtils; import com.databricks.sdk.service.billing.BillableUsageAPI; import com.databricks.sdk.service.billing.BillableUsageService; @@ -1131,18 +1131,47 @@ public DatabricksConfig config() { } public WorkspaceClient getWorkspaceClient(Workspace workspace) { - // For unified hosts, clone config and set workspace ID - if (this.config.getHostType() == HostType.UNIFIED) { + String host = + workspaceHost( + this.config.getDatabricksEnvironment(), + this.config.getHost(), + workspace.getDeploymentName()); + if (host.equals(this.config.getHost())) { + // SPOG/unified: reuse the same host, clone config and set workspace ID 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()); + // Traditional: use the deployment URL DatabricksConfig config = this.config.newWithWorkspaceHost(host); AzureUtils.getAzureWorkspaceResourceId(workspace).map(config::setAzureWorkspaceResourceId); return new WorkspaceClient(config); } + + /** + * Determines the workspace host URL. For SPOG hosts (no DNS zone or host doesn't match the DNS + * zone pattern), returns the account host as-is. For traditional hosts, builds the deployment + * URL. + */ + private static String workspaceHost( + DatabricksEnvironment env, String accountHost, String deploymentName) { + if (env.getDnsZone() == null || env.getDnsZone().isEmpty()) { + return accountHost; + } + if (accountHost != null) { + String normalized = accountHost; + if (!normalized.contains("://")) { + normalized = "https://" + normalized; + } + try { + java.net.URL url = new java.net.URL(normalized); + if (url.getHost() != null && url.getHost().endsWith(env.getDnsZone())) { + return env.getDeploymentUrl(deploymentName); + } + } catch (java.net.MalformedURLException e) { + // Fall through to return accountHost + } + } + return accountHost; + } } 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 6cfbadc3d..249c93b7a 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 @@ -30,18 +30,7 @@ public String authType() { List buildHostArgs(String cliPath, DatabricksConfig config) { List cmd = new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); - if (config.getExperimentalIsUnifiedHost() != null && config.getExperimentalIsUnifiedHost()) { - // For unified hosts, pass account_id, workspace_id, and experimental flag - cmd.add("--experimental-is-unified-host"); - if (config.getAccountId() != null) { - cmd.add("--account-id"); - cmd.add(config.getAccountId()); - } - if (config.getWorkspaceId() != null) { - cmd.add("--workspace-id"); - cmd.add(config.getWorkspaceId()); - } - } else if (config.getClientType() == ClientType.ACCOUNT) { + 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 2de7e5dca..c48a79f20 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,11 +270,6 @@ public synchronized Map authenticate() throws DatabricksExceptio } Map 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()); @@ -732,23 +727,14 @@ 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. */ + /** Returns the host type based on the host URL pattern. */ public HostType getHostType() { - if (experimentalIsUnifiedHost != null && experimentalIsUnifiedHost) { - return HostType.UNIFIED; - } if (host == null) { return HostType.WORKSPACE; } @@ -758,15 +744,10 @@ public HostType getHostType() { return HostType.WORKSPACE; } - /** Returns the client type based on host type and workspace ID configuration. */ + /** Returns the client type based on host type. */ 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: @@ -903,24 +884,11 @@ 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 (isAccountClient() && getAccountId() != null) { String prefix = getHost() + "/oidc/accounts/" + getAccountId(); return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize"); 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 005807839..354c91df6 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 @@ -9,8 +9,5 @@ public enum HostType { WORKSPACE, /** Traditional accounts host. */ - ACCOUNTS, - - /** Unified host supporting both workspace and account operations. */ - UNIFIED + ACCOUNTS } 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 4c7049374..041b0993e 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 @@ -2,7 +2,6 @@ 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; @@ -49,28 +48,14 @@ public void testGetWorkspaceClientForUnifiedHost() { WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace); - // Should have the same host + // Should have the same host (unified hosts reuse 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()); + // Host type is WORKSPACE (determined from URL pattern, not unified flag) + assertEquals(HostType.WORKSPACE, workspaceClient.config().getHostType()); } @Test diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java index bac4a766b..0f1ca5059 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java @@ -11,9 +11,7 @@ class DatabricksCliCredentialsProviderTest { private static final String CLI_PATH = "/usr/local/bin/databricks"; private static final String HOST = "https://my-workspace.cloud.databricks.com"; private static final String ACCOUNT_HOST = "https://accounts.cloud.databricks.com"; - private static final String UNIFIED_HOST = "https://unified.databricks.com"; private static final String ACCOUNT_ID = "test-account-123"; - private static final String WORKSPACE_ID = "987654321"; private final DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider(); @@ -39,104 +37,12 @@ void testBuildHostArgs_AccountHost() { } @Test - void testBuildHostArgs_UnifiedHost_WithAccountIdAndWorkspaceId() { - DatabricksConfig config = - new DatabricksConfig() - .setHost(UNIFIED_HOST) - .setExperimentalIsUnifiedHost(true) - .setAccountId(ACCOUNT_ID) - .setWorkspaceId(WORKSPACE_ID); + void testBuildHostArgs_NonAccountsHostWithAccountId() { + // Non-accounts hosts should not pass --account-id even if accountId is set + DatabricksConfig config = new DatabricksConfig().setHost(HOST).setAccountId(ACCOUNT_ID); List cmd = provider.buildHostArgs(CLI_PATH, config); - assertEquals( - Arrays.asList( - CLI_PATH, - "auth", - "token", - "--host", - UNIFIED_HOST, - "--experimental-is-unified-host", - "--account-id", - ACCOUNT_ID, - "--workspace-id", - WORKSPACE_ID), - cmd); - } - - @Test - void testBuildHostArgs_UnifiedHost_WithAccountIdOnly() { - DatabricksConfig config = - new DatabricksConfig() - .setHost(UNIFIED_HOST) - .setExperimentalIsUnifiedHost(true) - .setAccountId(ACCOUNT_ID); - - List cmd = provider.buildHostArgs(CLI_PATH, config); - - assertEquals( - Arrays.asList( - CLI_PATH, - "auth", - "token", - "--host", - UNIFIED_HOST, - "--experimental-is-unified-host", - "--account-id", - ACCOUNT_ID), - cmd); - } - - @Test - void testBuildHostArgs_UnifiedHost_WithWorkspaceIdOnly() { - DatabricksConfig config = - new DatabricksConfig() - .setHost(UNIFIED_HOST) - .setExperimentalIsUnifiedHost(true) - .setWorkspaceId(WORKSPACE_ID); - - List cmd = provider.buildHostArgs(CLI_PATH, config); - - assertEquals( - Arrays.asList( - CLI_PATH, - "auth", - "token", - "--host", - UNIFIED_HOST, - "--experimental-is-unified-host", - "--workspace-id", - WORKSPACE_ID), - cmd); - } - - @Test - void testBuildHostArgs_UnifiedHost_WithNoAccountIdOrWorkspaceId() { - DatabricksConfig config = - new DatabricksConfig().setHost(UNIFIED_HOST).setExperimentalIsUnifiedHost(true); - - List cmd = provider.buildHostArgs(CLI_PATH, config); - - assertEquals( - Arrays.asList( - CLI_PATH, "auth", "token", "--host", UNIFIED_HOST, "--experimental-is-unified-host"), - cmd); - } - - @Test - void testBuildHostArgs_UnifiedHostFalse_WithAccountHost() { - // When experimentalIsUnifiedHost is explicitly false, should fall back to account-id logic - DatabricksConfig config = - new DatabricksConfig() - .setHost(ACCOUNT_HOST) - .setExperimentalIsUnifiedHost(false) - .setAccountId(ACCOUNT_ID); - - List cmd = provider.buildHostArgs(CLI_PATH, config); - - assertEquals( - Arrays.asList( - CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), - cmd); + assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), cmd); } } 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 c6762652b..43e58bf46 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 @@ -378,16 +378,6 @@ public void testGetHostTypeAccounts() { 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( @@ -402,29 +392,6 @@ public void testGetClientTypeAccount() { 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()); - } - // --- HostMetadata tests --- private static final String DUMMY_ACCOUNT_ID = "00000000-0000-0000-0000-000000000001"; @@ -601,28 +568,16 @@ public void testResolveHostMetadataRaisesOnHttpError() throws IOException { @Test public void testResolveHostMetadataSetsTokenAudienceForAccountHost() throws IOException { - // For a unified host with no workspaceId (ACCOUNT client type), resolveHostMetadata should - // set tokenAudience to accountId when not already configured. - String response = - "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"," - + "\"account_id\":\"" - + DUMMY_ACCOUNT_ID - + "\"}"; - try (FixtureServer server = - new FixtureServer() - .with("GET", "/.well-known/databricks-config", response, 200) - .with("GET", "/.well-known/databricks-config", response, 200)) { - DatabricksConfig config = - new DatabricksConfig() - .setHost(server.getUrl()) - .setExperimentalIsUnifiedHost(true) - .setAccountId(DUMMY_ACCOUNT_ID); - config.resolve(emptyEnv()); - // Client type should be ACCOUNT (unified host, no workspaceId) - assertEquals(ClientType.ACCOUNT, config.getClientType()); - config.resolveHostMetadata(); - assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience()); - } + // For an account host, resolveHostMetadata should set tokenAudience to accountId + // when not already configured. We verify the preconditions here. + DatabricksConfig accountConfig = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setAccountId(DUMMY_ACCOUNT_ID); + assertEquals(ClientType.ACCOUNT, accountConfig.getClientType()); + assertNull(accountConfig.getTokenAudience()); + // When resolve runs with a reachable host, tryResolveHostMetadata will call + // resolveHostMetadata which sets tokenAudience = accountId for ACCOUNT clients. } @Test 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 f282b51d8..771c053a4 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 @@ -2,9 +2,7 @@ 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; @@ -13,10 +11,11 @@ import org.junit.jupiter.params.provider.MethodSource; /** - * Tests for unified host support (SPOG). + * Tests for host type detection, client type determination, and header injection. * - *

Covers host type detection, client type determination, header injection, and OIDC endpoint - * resolution for unified hosts. + *

After removing the UNIFIED host type, host type is determined solely from the URL pattern. + * Host metadata resolution (via /.well-known/databricks-config) populates accountId, workspaceId, + * and discoveryUrl automatically during config init. */ public class UnifiedHostTest { @@ -44,22 +43,10 @@ public void testHostTypeAccountsDod() { } @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()); + public void testHostTypeForNonAccountsHost() { + // A host that is not an accounts host is always WORKSPACE, regardless of any flags + DatabricksConfig config = new DatabricksConfig().setHost("https://mycompany.databricks.com"); + assertEquals(HostType.WORKSPACE, config.getHostType()); } @Test @@ -72,163 +59,58 @@ public void testHostTypeNullHost() { private static Stream provideClientTypeTestCases() { return Stream.of( + Arguments.of("Workspace host", "https://adb-123.azuredatabricks.net", ClientType.WORKSPACE), + Arguments.of("Account host", "https://accounts.cloud.databricks.com", ClientType.ACCOUNT), 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), - Arguments.of( - "Unified with workspace ID", - "https://unified.databricks.com", - "123456", - true, - ClientType.WORKSPACE)); + "Non-accounts host", "https://mycompany.databricks.com", ClientType.WORKSPACE)); } @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); - } + public void testClientType(String testName, String host, ClientType expected) { + DatabricksConfig config = new DatabricksConfig().setHost(host); assertEquals(expected, config.getClientType()); } - // --- OIDC Endpoint Tests --- + // --- isAccountClient() 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.getDatabricksOidcEndpoints(); - - 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()); + public void testIsAccountClientForAccountsHost() { + assertTrue( + new DatabricksConfig().setHost("https://accounts.cloud.databricks.com").isAccountClient()); } @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.getDatabricksOidcEndpoints()); - 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()")); + public void testIsAccountClientForWorkspaceHost() { + assertFalse( + new DatabricksConfig().setHost("https://adb-123.azuredatabricks.net").isAccountClient()); } @Test - public void testIsAccountClientWorksFineForTraditionalHosts() { - assertTrue( - new DatabricksConfig().setHost("https://accounts.cloud.databricks.com").isAccountClient()); - + public void testIsAccountClientForNonAccountsHost() { + // Non-accounts hosts are not account clients, even with experimentalIsUnifiedHost set assertFalse( - new DatabricksConfig().setHost("https://adb-123.azuredatabricks.net").isAccountClient()); + new DatabricksConfig() + .setHost("https://mycompany.databricks.com") + .setExperimentalIsUnifiedHost(true) + .isAccountClient()); } // --- Environment Variable Tests --- @Test - public void testUnifiedHostFromEnvironmentVariables() { + public void testWorkspaceIdFromEnvironmentVariables() { Map env = new HashMap<>(); - env.put("DATABRICKS_HOST", "https://unified.databricks.com"); - env.put("DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST", "true"); + env.put("DATABRICKS_HOST", "https://mycompany.databricks.com"); 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(HostType.WORKSPACE, config.getHostType()); assertEquals("987654321", config.getWorkspaceId()); assertEquals("account-abc", config.getAccountId()); assertEquals(ClientType.WORKSPACE, config.getClientType()); } - - // --- Header Injection 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")); - } }