Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -53,4 +56,8 @@ public String getWorkspaceId() {
public String getCloud() {
return cloud;
}

public String getHostType() {
return hostType;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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"));
}
}
Loading