Skip to content

Commit f1240ce

Browse files
committed
Resolve TokenAudience from host metadata for account hosts
Port of Go SDK #1543. When resolveHostMetadata() runs on an account host and tokenAudience is not already set, automatically sets it to the accountId. This enables OIDC token exchange to work correctly for account-level operations without explicit TOKEN_AUDIENCE config. Co-authored-by: Isaac
1 parent c37f21f commit f1240ce

File tree

2 files changed

+52
-0
lines changed

2 files changed

+52
-0
lines changed

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -884,6 +884,10 @@ void resolveHostMetadata() throws IOException {
884884
"discovery_url is not configured and could not be resolved from host metadata");
885885
}
886886
}
887+
// For account hosts, use the accountId as the token audience if not already set.
888+
if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) {
889+
tokenAudience = accountId;
890+
}
887891
}
888892

889893
private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,54 @@ public void testResolveHostMetadataRaisesOnHttpError() throws IOException {
599599
}
600600
}
601601

602+
@Test
603+
public void testResolveHostMetadataSetsTokenAudienceForAccountHost() throws IOException {
604+
// For a unified host with no workspaceId (ACCOUNT client type), resolveHostMetadata should
605+
// set tokenAudience to accountId when not already configured.
606+
String response =
607+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
608+
+ "\"account_id\":\""
609+
+ DUMMY_ACCOUNT_ID
610+
+ "\"}";
611+
try (FixtureServer server =
612+
new FixtureServer()
613+
.with("GET", "/.well-known/databricks-config", response, 200)
614+
.with("GET", "/.well-known/databricks-config", response, 200)) {
615+
DatabricksConfig config =
616+
new DatabricksConfig()
617+
.setHost(server.getUrl())
618+
.setExperimentalIsUnifiedHost(true)
619+
.setAccountId(DUMMY_ACCOUNT_ID);
620+
config.resolve(emptyEnv());
621+
// Client type should be ACCOUNT (unified host, no workspaceId)
622+
assertEquals(ClientType.ACCOUNT, config.getClientType());
623+
config.resolveHostMetadata();
624+
assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience());
625+
}
626+
}
627+
628+
@Test
629+
public void testResolveHostMetadataDoesNotOverwriteTokenAudience() throws IOException {
630+
String response =
631+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
632+
+ "\"account_id\":\""
633+
+ DUMMY_ACCOUNT_ID
634+
+ "\"}";
635+
try (FixtureServer server =
636+
new FixtureServer()
637+
.with("GET", "/.well-known/databricks-config", response, 200)
638+
.with("GET", "/.well-known/databricks-config", response, 200)) {
639+
DatabricksConfig config =
640+
new DatabricksConfig()
641+
.setHost(server.getUrl())
642+
.setAccountId(DUMMY_ACCOUNT_ID)
643+
.setTokenAudience("custom-audience");
644+
config.resolve(emptyEnv());
645+
config.resolveHostMetadata();
646+
assertEquals("custom-audience", config.getTokenAudience());
647+
}
648+
}
649+
602650
// --- discoveryUrl / OIDC endpoint tests ---
603651

604652
@Test

0 commit comments

Comments
 (0)