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 43c1f5136..44d1dae64 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. @@ -723,6 +729,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; @@ -758,8 +775,15 @@ public boolean isAccountClient() { 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 configuration settings and host URL. When host metadata has been + * resolved (via /.well-known/databricks-config), the resolved host type is returned directly. + * Otherwise, the host type is inferred from URL patterns as a fallback. + */ public HostType getHostType() { + if (resolvedHostType != null) { + return resolvedHostType; + } if (experimentalIsUnifiedHost != null && experimentalIsUnifiedHost) { return HostType.UNIFIED; } @@ -889,6 +913,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"); @@ -907,8 +938,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 005807839..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 @@ -12,5 +12,25 @@ public enum HostType { ACCOUNTS, /** Unified host supporting both workspace and account operations. */ - UNIFIED + 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/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 0e86b4ef4..70081a790 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 @@ -712,6 +712,186 @@ 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()); + config.resolveHostMetadata(); + 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)) { + 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()); + config.resolveHostMetadata(); + 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()); + config.resolveHostMetadata(); + 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()); + config.resolveHostMetadata(); + 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()); + config.resolveHostMetadata(); + 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 + + "\"," + + "\"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()) + .setExperimentalIsUnifiedHost(true) + .setAccountId(DUMMY_ACCOUNT_ID); + // Note: need two fixtures — resolve() consumes first one (unified host triggers + // tryResolveHostMetadata) + // Instead, don't set unified — just test direct call + config = new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID); + config.resolve(emptyEnv()); + config.resolveHostMetadata(); + // 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()); + config.resolveHostMetadata(); + assertEquals("existing-audience", config.getTokenAudience()); + } + } + + @Test + public void testResolveHostMetadataFallsBackToAccountIdWhenNoDefaultOidcAudience() + throws IOException { + // When no default_oidc_audience, should fall back to account_id for account hosts. + // Use unified host flag so getClientType() returns ACCOUNT (no workspaceId). + 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()) + .setAccountId(DUMMY_ACCOUNT_ID) + .setExperimentalIsUnifiedHost(true); + config.resolve(emptyEnv()); + // resolve() with unified flag triggers tryResolveHostMetadata() consuming first fixture. + // Now call again with second fixture to verify account_id fallback. + assertEquals(DUMMY_ACCOUNT_ID, 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 f282b51d8..caed880cd 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 @@ -231,4 +231,71 @@ public void testNoHeaderInjectionForTraditionalWorkspace() { assertEquals("Bearer test-token", headers.get("Authorization")); assertNull(headers.get("X-Databricks-Org-Id")); } + + // --- Resolved host type from metadata tests --- + + @Test + public void testMetadataWorkspaceOverridesAccountLikeUrl() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setResolvedHostType(HostType.WORKSPACE); + assertEquals(HostType.WORKSPACE, config.getHostType()); + } + + @Test + public void testMetadataAccountOverridesWorkspaceLikeUrl() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://my-workspace.cloud.databricks.com") + .setResolvedHostType(HostType.ACCOUNTS); + assertEquals(HostType.ACCOUNTS, config.getHostType()); + } + + @Test + public void testMetadataUnifiedIsReturned() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://my-workspace.cloud.databricks.com") + .setResolvedHostType(HostType.UNIFIED); + assertEquals(HostType.UNIFIED, config.getHostType()); + } + + @Test + public void testFallsBackToUrlMatchingWhenResolvedHostTypeNull() { + DatabricksConfig config = + new DatabricksConfig().setHost("https://accounts.cloud.databricks.com"); + // resolvedHostType is null by default + assertEquals(HostType.ACCOUNTS, config.getHostType()); + } + + @Test + public void testMetadataOverridesExperimentalFlag() { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://my-workspace.cloud.databricks.com") + .setExperimentalIsUnifiedHost(true) + .setResolvedHostType(HostType.ACCOUNTS); + // Resolved host type takes priority over experimental flag + assertEquals(HostType.ACCOUNTS, config.getHostType()); + } + + @Test + public void testEndToEndResolveToGetHostType() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\"," + + "\"account_id\":\"test-account\"," + + "\"host_type\":\"unified\"}"; + 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); + config.resolve( + new Environment(new HashMap<>(), new ArrayList<>(), System.getProperty("os.name"))); + // After resolve(), tryResolveHostMetadata() should have set resolvedHostType + assertEquals(HostType.UNIFIED, config.getHostType()); + } + } }