Skip to content
Open
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 @@ -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;
Expand Down Expand Up @@ -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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
}
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 @@ -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 {
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -657,6 +657,92 @@ 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());
}
}

// --- 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