Skip to content

Commit efcc7bd

Browse files
committed
Resolve tokenAudience from token_federation_default_oidc_audiences in host metadata
Add token_federation_default_oidc_audiences field (List<String>) to HostMetadata and resolve tokenAudience from its first element during config initialization. This takes priority over the existing accountId fallback for account hosts. Co-authored-by: Isaac
1 parent 89297b1 commit efcc7bd

4 files changed

Lines changed: 54 additions & 1 deletion

File tree

NEXT_CHANGELOG.md

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

1616
### Internal Changes
1717
* Introduced a logging abstraction (`com.databricks.sdk.core.logging`) to decouple the SDK from a specific logging backend.
18+
* 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.
1819

1920
### API Changes
2021
* 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: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -657,6 +657,42 @@ 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+
660696
// --- discoveryUrl / OIDC endpoint tests ---
661697

662698
@Test

0 commit comments

Comments
 (0)