Skip to content

Commit 7175953

Browse files
tanmay-dbtanmay-db
authored
[Internal] Resolve TokenAudience from token_federation_default_oidc_audiences in host metadata (#759)
## Summary Resolves `tokenAudience` automatically from the `token_federation_default_oidc_audiences` field returned by the `/.well-known/databricks-config` host metadata endpoint, removing the need for manual audience configuration when using OIDC-based authentication. ## Why Today, when using Workload Identity Federation or other OIDC-based credential providers, the `tokenAudience` must either be explicitly configured by the user or falls back to `accountId` for account-level hosts. The host metadata endpoint now returns a `token_federation_default_oidc_audiences` field containing the recommended audience values. Without this change, users must manually configure `tokenAudience` even though the server already advertises the correct value — adding unnecessary friction to OIDC auth setup. This PR reads the new field during config initialization so that the SDK automatically picks up the correct audience from host metadata, with user-configured values still taking priority. ## What changed ### Interface changes - **`HostMetadata.getTokenFederationDefaultOidcAudiences()`** — New getter returning `List<String>` of OIDC audiences from host metadata. ### Behavioral changes - `tokenAudience` resolution now follows a three-tier priority chain: 1. User-configured `tokenAudience` (highest priority, unchanged) 2. First element of `token_federation_default_oidc_audiences` from host metadata (**new**) 3. `accountId` for account hosts (fallback, unchanged) ### Internal changes - `HostMetadata`: Added `token_federation_default_oidc_audiences` field (`List<String>`) with `@JsonProperty` annotation - `DatabricksConfig.resolveHostMetadata()`: Added audience resolution logic before the existing `accountId` fallback - `NEXT_CHANGELOG.md`: Added internal changelog entry ## How is this tested? - All integration tests passed (manually triggered) - Three new unit tests in `DatabricksConfigTest.java`: - `testResolveHostMetadataSetsTokenAudienceFromOidcAudiences` — verifies audience is resolved from metadata - `testResolveHostMetadataDoesNotOverrideExistingTokenAudienceWithOidcAudiences` — verifies user-configured audience takes priority - `testResolveHostMetadataOidcAudiencesPriorityOverAccountIdFallback` — verifies metadata audience takes priority over accountId fallback
1 parent 70ba8a8 commit 7175953

File tree

4 files changed

+161
-1
lines changed

4 files changed

+161
-1
lines changed

‎NEXT_CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
### Internal Changes
1818
* Introduced a logging abstraction (`com.databricks.sdk.core.logging`) to decouple the SDK from a specific logging backend.
19+
* Added `token_federation_default_oidc_audiences` resolution from host metadata. The SDK now sets `tokenAudience` from the first element of this field during config initialization, with fallback to `accountId` for account hosts.
1920

2021
### API Changes
2122
* Add `createCatalog()`, `createSyncedTable()`, `deleteCatalog()`, `deleteSyncedTable()`, `getCatalog()` and `getSyncedTable()` methods for `workspaceClient.postgres()` service.

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -884,8 +884,16 @@ void resolveHostMetadata() throws IOException {
884884
discoveryUrl = oidcUri.resolve(".well-known/oauth-authorization-server").toString();
885885
LOG.debug("Resolved discovery_url from host metadata: \"{}\"", discoveryUrl);
886886
}
887-
// For account hosts, use the accountId as the token audience if not already set.
887+
List<String> audiences = meta.getTokenFederationDefaultOidcAudiences();
888+
if (tokenAudience == null && audiences != null && !audiences.isEmpty()) {
889+
LOG.debug(
890+
"Resolved token_audience from host metadata token_federation_default_oidc_audiences: \"{}\"",
891+
audiences.get(0));
892+
tokenAudience = audiences.get(0);
893+
}
894+
// Fallback: for account hosts, use the accountId as the token audience if not already set.
888895
if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) {
896+
LOG.debug("Setting token_audience to account_id for account host: \"{}\"", accountId);
889897
tokenAudience = accountId;
890898
}
891899
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
44
import com.fasterxml.jackson.annotation.JsonProperty;
5+
import java.util.List;
56

67
/**
78
* [Experimental] Parsed response from the /.well-known/databricks-config discovery endpoint.
@@ -23,6 +24,9 @@ public class HostMetadata {
2324
@JsonProperty("cloud")
2425
private String cloud;
2526

27+
@JsonProperty("token_federation_default_oidc_audiences")
28+
private List<String> tokenFederationDefaultOidcAudiences;
29+
2630
public HostMetadata() {}
2731

2832
public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) {
@@ -53,4 +57,8 @@ public String getWorkspaceId() {
5357
public String getCloud() {
5458
return cloud;
5559
}
60+
61+
public List<String> getTokenFederationDefaultOidcAudiences() {
62+
return tokenFederationDefaultOidcAudiences;
63+
}
5664
}

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -657,6 +657,149 @@ public void testEnsureResolvedHostMetadataMissingAccountIdWithPlaceholderNonFata
657657
}
658658
}
659659

660+
// --- resolveHostMetadata token_federation_default_oidc_audiences tests ---
661+
662+
@Test
663+
public void testResolveHostMetadataSetsTokenAudienceFromOidcAudiences() throws IOException {
664+
String response =
665+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
666+
+ "\"account_id\":\""
667+
+ DUMMY_ACCOUNT_ID
668+
+ "\","
669+
+ "\"token_federation_default_oidc_audiences\":[\"https://ws.databricks.com/oidc/v1/token\"]}";
670+
try (FixtureServer server =
671+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
672+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
673+
config.resolve(emptyEnv());
674+
assertEquals("https://ws.databricks.com/oidc/v1/token", config.getTokenAudience());
675+
}
676+
}
677+
678+
@Test
679+
public void testResolveHostMetadataDoesNotOverrideExistingTokenAudienceWithOidcAudiences()
680+
throws IOException {
681+
String response =
682+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
683+
+ "\"account_id\":\""
684+
+ DUMMY_ACCOUNT_ID
685+
+ "\","
686+
+ "\"token_federation_default_oidc_audiences\":[\"metadata-audience\"]}";
687+
try (FixtureServer server =
688+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
689+
DatabricksConfig config =
690+
new DatabricksConfig().setHost(server.getUrl()).setTokenAudience("existing-audience");
691+
config.resolve(emptyEnv());
692+
assertEquals("existing-audience", config.getTokenAudience());
693+
}
694+
}
695+
696+
@Test
697+
public void testResolveHostMetadataOidcAudiencesPriorityOverAccountIdFallback()
698+
throws IOException {
699+
// token_federation_default_oidc_audiences should take priority over the account_id fallback.
700+
// The audiences check runs before the fallback, so once tokenAudience is set from audiences,
701+
// the fallback's null check prevents it from overriding.
702+
String response =
703+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
704+
+ "\"account_id\":\""
705+
+ DUMMY_ACCOUNT_ID
706+
+ "\","
707+
+ "\"token_federation_default_oidc_audiences\":[\"custom-audience\"]}";
708+
try (FixtureServer server =
709+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
710+
DatabricksConfig config =
711+
new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID);
712+
config.resolve(emptyEnv());
713+
// Should use first element of token_federation_default_oidc_audiences, NOT account_id
714+
assertEquals("custom-audience", config.getTokenAudience());
715+
}
716+
}
717+
718+
@Test
719+
public void testResolveHostMetadataNoAudiencesFieldLeavesTokenAudienceNull() throws IOException {
720+
// When token_federation_default_oidc_audiences is absent, tokenAudience stays null
721+
String response =
722+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
723+
+ "\"account_id\":\""
724+
+ DUMMY_ACCOUNT_ID
725+
+ "\"}";
726+
try (FixtureServer server =
727+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
728+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
729+
config.resolve(emptyEnv());
730+
assertNull(config.getTokenAudience());
731+
}
732+
}
733+
734+
@Test
735+
public void testResolveHostMetadataEmptyAudiencesListLeavesTokenAudienceNull()
736+
throws IOException {
737+
// When token_federation_default_oidc_audiences is an empty array, tokenAudience stays null
738+
String response =
739+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
740+
+ "\"account_id\":\""
741+
+ DUMMY_ACCOUNT_ID
742+
+ "\","
743+
+ "\"token_federation_default_oidc_audiences\":[]}";
744+
try (FixtureServer server =
745+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
746+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
747+
config.resolve(emptyEnv());
748+
assertNull(config.getTokenAudience());
749+
}
750+
}
751+
752+
@Test
753+
public void testResolveHostMetadataEmptyStringAudienceSetsTokenAudience() throws IOException {
754+
// When first element is empty string, tokenAudience is set to empty string
755+
String response =
756+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
757+
+ "\"account_id\":\""
758+
+ DUMMY_ACCOUNT_ID
759+
+ "\","
760+
+ "\"token_federation_default_oidc_audiences\":[\"\"]}";
761+
try (FixtureServer server =
762+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
763+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
764+
config.resolve(emptyEnv());
765+
assertEquals("", config.getTokenAudience());
766+
}
767+
}
768+
769+
@Test
770+
public void testResolveHostMetadataMultipleAudiencesPicksFirst() throws IOException {
771+
// When multiple audiences, the first element is used
772+
String response =
773+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
774+
+ "\"account_id\":\""
775+
+ DUMMY_ACCOUNT_ID
776+
+ "\","
777+
+ "\"token_federation_default_oidc_audiences\":[\"first-audience\",\"second-audience\",\"third-audience\"]}";
778+
try (FixtureServer server =
779+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
780+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
781+
config.resolve(emptyEnv());
782+
assertEquals("first-audience", config.getTokenAudience());
783+
}
784+
}
785+
786+
@Test
787+
public void testResolveHostMetadataNullFirstAudienceLeavesTokenAudienceNull() throws IOException {
788+
// When first element is null, tokenAudience stays null
789+
String response =
790+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
791+
+ "\"account_id\":\""
792+
+ DUMMY_ACCOUNT_ID
793+
+ "\","
794+
+ "\"token_federation_default_oidc_audiences\":[null,\"second-audience\"]}";
795+
try (FixtureServer server =
796+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
797+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
798+
config.resolve(emptyEnv());
799+
assertNull(config.getTokenAudience());
800+
}
801+
}
802+
660803
// --- discoveryUrl / OIDC endpoint tests ---
661804

662805
@Test

0 commit comments

Comments
 (0)