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..35e93fd28 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; @@ -889,6 +906,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"); 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..dfbf91fa5 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,9 @@ public class HostMetadata { @JsonProperty("cloud") private String cloud; + @JsonProperty("host_type") + private String hostType; + public HostMetadata() {} public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) { @@ -53,4 +56,8 @@ public String getWorkspaceId() { public String getCloud() { return cloud; } + + public String getHostType() { + return hostType; + } } 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..b2a83392f 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,94 @@ 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()); + } + } + // --- 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")); + } +}