Skip to content

Commit cc5965f

Browse files
Resolve TokenAudience from host metadata for account hosts (#714)
## 🥞 Stacked PR Use this [link](https://github.com/databricks/databricks-sdk-java/pull/714/files) to review incremental changes. - [**hectorcast-db/stack/port-5-token-audience-from-metadata**](#714) [[Files changed](https://github.com/databricks/databricks-sdk-java/pull/714/files)] - [hectorcast-db/stack/port-6-gcp-sa-nonblocking](#718) [[Files changed](https://github.com/databricks/databricks-sdk-java/pull/718/files/2dd4a6da83dd5de9f279c0b2bfe37d3abf7a74a8..ed4ef1be20407d8797dfc2dc71528a059167cead)] - [hectorcast-db/stack/port-7-integration-test-metadata](#719) [[Files changed](https://github.com/databricks/databricks-sdk-java/pull/719/files/ed4ef1be20407d8797dfc2dc71528a059167cead..206961d539aca2b5eb89c25154d2c4f41d958c64)] - [hectorcast-db/stack/port-8-remove-unified-flag](#720) [[Files changed](https://github.com/databricks/databricks-sdk-java/pull/720/files/206961d539aca2b5eb89c25154d2c4f41d958c64..a380c97d6fa383c573f02d4c4784a38d520c27d2)] --------- ## Summary Port of Go SDK [#1543](databricks/databricks-sdk-go#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` configuration. **Changes:** - `DatabricksConfig.resolveHostMetadata()`: sets `tokenAudience = accountId` for ACCOUNT clients when not already configured - Tests: `testResolveHostMetadataSetsTokenAudienceForAccountHost`, `testResolveHostMetadataDoesNotOverwriteTokenAudience` `NO_CHANGELOG=true` ## Test plan - [x] `DatabricksConfigTest`: token audience resolution tests pass
1 parent 31fd142 commit cc5965f

File tree

2 files changed

+50
-0
lines changed

2 files changed

+50
-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
@@ -907,6 +907,10 @@ void resolveHostMetadata() throws IOException {
907907
discoveryUrl = oidcUri.resolve(".well-known/oauth-authorization-server").toString();
908908
LOG.debug("Resolved discovery_url from host metadata: \"{}\"", discoveryUrl);
909909
}
910+
// For account hosts, use the accountId as the token audience if not already set.
911+
if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) {
912+
tokenAudience = accountId;
913+
}
910914
}
911915

912916
private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,6 +586,52 @@ public void testResolveHostMetadataRaisesOnHttpError() throws IOException {
586586
}
587587
}
588588

589+
@Test
590+
public void testResolveHostMetadataSetsTokenAudienceForAccountHost() throws IOException {
591+
// For a unified host with no workspaceId (ACCOUNT client type), resolveHostMetadata should
592+
// set tokenAudience to accountId when not already configured.
593+
String response =
594+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
595+
+ "\"account_id\":\""
596+
+ DUMMY_ACCOUNT_ID
597+
+ "\"}";
598+
try (FixtureServer server =
599+
new FixtureServer()
600+
.with("GET", "/.well-known/databricks-config", response, 200)
601+
.with("GET", "/.well-known/databricks-config", response, 200)) {
602+
DatabricksConfig config =
603+
new DatabricksConfig()
604+
.setHost(server.getUrl())
605+
.setExperimentalIsUnifiedHost(true)
606+
.setAccountId(DUMMY_ACCOUNT_ID);
607+
config.resolve(emptyEnv());
608+
// Client type should be ACCOUNT (unified host, no workspaceId)
609+
assertEquals(ClientType.ACCOUNT, config.getClientType());
610+
config.resolveHostMetadata();
611+
assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience());
612+
}
613+
}
614+
615+
@Test
616+
public void testResolveHostMetadataDoesNotOverwriteTokenAudience() throws IOException {
617+
String response =
618+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
619+
+ "\"account_id\":\""
620+
+ DUMMY_ACCOUNT_ID
621+
+ "\"}";
622+
try (FixtureServer server =
623+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
624+
DatabricksConfig config =
625+
new DatabricksConfig()
626+
.setHost(server.getUrl())
627+
.setAccountId(DUMMY_ACCOUNT_ID)
628+
.setTokenAudience("custom-audience");
629+
config.resolve(emptyEnv());
630+
config.resolveHostMetadata();
631+
assertEquals("custom-audience", config.getTokenAudience());
632+
}
633+
}
634+
589635
// --- tryResolveHostMetadata (config init) tests ---
590636

591637
@Test

0 commit comments

Comments
 (0)