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 8cf3d9597..1f7259255 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 @@ -162,6 +162,12 @@ public class DatabricksConfig { private DatabricksEnvironment databricksEnvironment; + /** + * The host type resolved from the /.well-known/databricks-config discovery endpoint. When set, + * this takes priority over URL-based host type detection in {@link #getHostType()}. + */ + private HostType resolvedHostType; + /** * When using Workload Identity Federation, the audience to specify when fetching an ID token from * the ID token supplier. @@ -714,6 +720,17 @@ public DatabricksConfig setDisableOauthRefreshToken(boolean disable) { return this; } + /** Returns the host type resolved from host metadata, or {@code null} if not yet resolved. */ + HostType getResolvedHostType() { + return resolvedHostType; + } + + /** Sets the resolved host type. Package-private for testing. */ + DatabricksConfig setResolvedHostType(HostType resolvedHostType) { + this.resolvedHostType = resolvedHostType; + return this; + } + public boolean isAzure() { if (azureWorkspaceResourceId != null) { return true; @@ -866,6 +883,13 @@ void resolveHostMetadata() throws IOException { LOG.debug("Resolved workspace_id from host metadata: \"{}\"", meta.getWorkspaceId()); workspaceId = meta.getWorkspaceId(); } + if (resolvedHostType == null && meta.getHostType() != null) { + HostType ht = HostType.fromApiValue(meta.getHostType()); + if (ht != null) { + LOG.debug("Resolved host_type from host metadata: \"{}\"", ht); + resolvedHostType = ht; + } + } if (discoveryUrl == null) { if (meta.getOidcEndpoint() == null || meta.getOidcEndpoint().isEmpty()) { LOG.warn("Host metadata missing oidc_endpoint; skipping discovery URL resolution"); @@ -884,8 +908,17 @@ void resolveHostMetadata() throws IOException { discoveryUrl = oidcUri.resolve(".well-known/oauth-authorization-server").toString(); LOG.debug("Resolved discovery_url from host metadata: \"{}\"", discoveryUrl); } - // For account hosts, use the accountId as the token audience if not already set. + if (tokenAudience == null + && meta.getDefaultOidcAudience() != null + && !meta.getDefaultOidcAudience().isEmpty()) { + LOG.debug( + "Resolved token_audience from host metadata default_oidc_audience: \"{}\"", + meta.getDefaultOidcAudience()); + tokenAudience = meta.getDefaultOidcAudience(); + } + // Fallback: for account hosts, use the accountId as the token audience if not already set. if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) { + LOG.debug("Setting token_audience to account_id for account host: \"{}\"", accountId); tokenAudience = accountId; } } 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 354c91df6..5976ae345 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,5 +9,28 @@ public enum HostType { WORKSPACE, /** Traditional accounts host. */ - ACCOUNTS + ACCOUNTS, + + /** Unified host supporting both workspace and account operations. */ + UNIFIED; + + /** + * Converts an API-level host type string (e.g. "workspace", "account", "unified") to the + * corresponding enum value. Returns {@code null} for unknown or empty values. + */ + public static HostType fromApiValue(String value) { + if (value == null || value.isEmpty()) { + return null; + } + switch (value.toLowerCase()) { + case "workspace": + return WORKSPACE; + case "account": + return ACCOUNTS; + case "unified": + return UNIFIED; + default: + return null; + } + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java index 3962f75d9..0fe5a15ef 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java @@ -23,6 +23,12 @@ public class HostMetadata { @JsonProperty("cloud") private String cloud; + @JsonProperty("host_type") + private String hostType; + + @JsonProperty("default_oidc_audience") + private String defaultOidcAudience; + public HostMetadata() {} public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) { @@ -53,4 +59,12 @@ public String getWorkspaceId() { public String getCloud() { return cloud; } + + public String getHostType() { + return hostType; + } + + public String getDefaultOidcAudience() { + return defaultOidcAudience; + } } 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 f529d3525..5ff637d65 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 @@ -3,8 +3,10 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.FixtureServer; import com.databricks.sdk.core.HostType; import com.databricks.sdk.service.provisioning.Workspace; +import java.io.IOException; import org.junit.jupiter.api.Test; public class AccountClientTest { @@ -31,30 +33,33 @@ public void testGetWorkspaceClientForTraditionalAccount() { } @Test - public void testGetWorkspaceClientForUnifiedHost() { - String unifiedHost = "https://unified.databricks.com"; - DatabricksConfig accountConfig = - new DatabricksConfig() - .setHost(unifiedHost) - .setAccountId("test-account") - .setToken("test-token"); + public void testGetWorkspaceClientForUnifiedHost() throws IOException { + try (FixtureServer server = new FixtureServer()) { + String unifiedHost = server.getUrl(); + DatabricksConfig accountConfig = + new DatabricksConfig() + .setHost(unifiedHost) + .setAccountId("test-account") + .setToken("test-token"); - AccountClient accountClient = new AccountClient(accountConfig); + AccountClient accountClient = new AccountClient(accountConfig); - Workspace workspace = new Workspace(); - workspace.setWorkspaceId(123456L); - workspace.setDeploymentName("test-workspace"); + Workspace workspace = new Workspace(); + workspace.setWorkspaceId(123456L); + workspace.setDeploymentName("test-workspace"); - WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace); + WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace); - // Should have the same host (unified hosts reuse the same host) - assertEquals(unifiedHost, workspaceClient.config().getHost()); + // Should have the same host (non-matching DNS zone means SPOG path) + assertEquals(unifiedHost, workspaceClient.config().getHost()); - // Should have workspace ID set - assertEquals("123456", workspaceClient.config().getWorkspaceId()); + // Should have workspace ID set + assertEquals("123456", workspaceClient.config().getWorkspaceId()); - // Host type is WORKSPACE (determined from URL pattern, not unified flag) - assertEquals(HostType.WORKSPACE, workspaceClient.config().getHostType()); + // Host type is WORKSPACE (no resolved host type from metadata, URL doesn't match accounts + // pattern) + assertEquals(HostType.WORKSPACE, 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 1ba934c5a..a7517cec5 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 @@ -657,6 +657,149 @@ public void testEnsureResolvedHostMetadataMissingAccountIdWithPlaceholderNonFata } } + // --- resolveHostMetadata host type tests --- + + @Test + public void testResolveHostMetadataPopulatesResolvedHostType() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"workspace\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl()); + config.resolve(emptyEnv()); + assertEquals(HostType.WORKSPACE, config.getResolvedHostType()); + } + } + + @Test + public void testResolveHostMetadataDoesNotOverwriteExistingHostType() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"workspace\"}"; + 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()); + config.resolve(emptyEnv()); + config.setResolvedHostType(HostType.UNIFIED); + config.resolveHostMetadata(); + assertEquals(HostType.UNIFIED, config.getResolvedHostType()); + } + } + + @Test + public void testResolveHostMetadataUnknownHostTypeIgnored() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"unknown_value\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl()); + config.resolve(emptyEnv()); + assertNull(config.getResolvedHostType()); + } + } + + @Test + public void testResolveHostMetadataHostTypeAccount() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"account\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl()); + config.resolve(emptyEnv()); + assertEquals(HostType.ACCOUNTS, config.getResolvedHostType()); + } + } + + @Test + public void testResolveHostMetadataHostTypeUnified() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"unified\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl()); + config.resolve(emptyEnv()); + assertEquals(HostType.UNIFIED, config.getResolvedHostType()); + } + } + + // --- resolveHostMetadata default_oidc_audience tests --- + + @Test + public void testResolveHostMetadataSetsTokenAudienceFromDefaultOidcAudience() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"default_oidc_audience\":\"https://ws.databricks.com/oidc/v1/token\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl()); + config.resolve(emptyEnv()); + assertEquals("https://ws.databricks.com/oidc/v1/token", config.getTokenAudience()); + } + } + + @Test + public void testResolveHostMetadataDefaultOidcAudiencePriorityOverAccountIdFallback() + throws IOException { + // default_oidc_audience should take priority over the account_id fallback for account hosts + String response = + "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"host_type\":\"account\"," + + "\"default_oidc_audience\":\"custom-audience\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = + new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID); + config.resolve(emptyEnv()); + // Should use default_oidc_audience, NOT account_id + assertEquals("custom-audience", config.getTokenAudience()); + } + } + + @Test + public void testResolveHostMetadataDoesNotOverrideExistingTokenAudienceWithOidcAudience() + throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"," + + "\"default_oidc_audience\":\"metadata-audience\"}"; + try (FixtureServer server = + new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = + new DatabricksConfig().setHost(server.getUrl()).setTokenAudience("existing-audience"); + config.resolve(emptyEnv()); + assertEquals("existing-audience", config.getTokenAudience()); + } + } + // --- discoveryUrl / OIDC endpoint tests --- @Test diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/HostTypeTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/HostTypeTest.java new file mode 100644 index 000000000..63cd9b100 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/HostTypeTest.java @@ -0,0 +1,46 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; + +public class HostTypeTest { + + @Test + public void testFromApiValueWorkspace() { + assertEquals(HostType.WORKSPACE, HostType.fromApiValue("workspace")); + } + + @Test + public void testFromApiValueAccount() { + assertEquals(HostType.ACCOUNTS, HostType.fromApiValue("account")); + } + + @Test + public void testFromApiValueUnified() { + assertEquals(HostType.UNIFIED, HostType.fromApiValue("unified")); + } + + @Test + public void testFromApiValueCaseInsensitive() { + assertEquals(HostType.WORKSPACE, HostType.fromApiValue("WORKSPACE")); + assertEquals(HostType.ACCOUNTS, HostType.fromApiValue("Account")); + assertEquals(HostType.UNIFIED, HostType.fromApiValue("UNIFIED")); + } + + @Test + public void testFromApiValueNull() { + assertNull(HostType.fromApiValue(null)); + } + + @Test + public void testFromApiValueEmpty() { + assertNull(HostType.fromApiValue("")); + } + + @Test + public void testFromApiValueUnknown() { + assertNull(HostType.fromApiValue("unknown")); + assertNull(HostType.fromApiValue("something_else")); + } +} 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 3d25b9276..0de127859 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 @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.*; 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; @@ -96,18 +97,20 @@ public void testIsAccountClientForNonAccountsHost() { // --- Environment Variable Tests --- @Test - public void testWorkspaceIdFromEnvironmentVariables() { - Map env = new HashMap<>(); - 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.WORKSPACE, config.getHostType()); - assertEquals("987654321", config.getWorkspaceId()); - assertEquals("account-abc", config.getAccountId()); - assertEquals(ClientType.WORKSPACE, config.getClientType()); + public void testWorkspaceIdFromEnvironmentVariables() throws IOException { + try (FixtureServer server = new FixtureServer()) { + Map env = new HashMap<>(); + env.put("DATABRICKS_HOST", server.getUrl()); + 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.WORKSPACE, config.getHostType()); + assertEquals("987654321", config.getWorkspaceId()); + assertEquals("account-abc", config.getAccountId()); + assertEquals(ClientType.WORKSPACE, config.getClientType()); + } } }