[Internal] Resolve TokenAudience from token_federation_default_oidc_audiences in host metadata

NEXT_CHANGELOG.md

@@ -15,6 +15,7 @@
 
 ### Internal Changes
 * Introduced a logging abstraction (`com.databricks.sdk.core.logging`) to decouple the SDK from a specific logging backend.
+* 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.
 
 ### API Changes
 * 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

@@ -884,8 +884,16 @@ void resolveHostMetadata() throws IOException {
       discoveryUrl = oidcUri.resolve(".well-known/oauth-authorization-server").toString();
       LOG.debug("Resolved discovery_url from host metadata: \"{}\"", discoveryUrl);
     }
-    // For account hosts, use the accountId as the token audience if not already set.
+    List<String> audiences = meta.getTokenFederationDefaultOidcAudiences();
+    if (tokenAudience == null && audiences != null && !audiences.isEmpty()) {
+      LOG.debug(
+          "Resolved token_audience from host metadata token_federation_default_oidc_audiences: \"{}\"",
+          audiences.get(0));
+      tokenAudience = audiences.get(0);
+    }
+    // Fallback: for account hosts, use the accountId as the token audience if not already set.
     if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) {
+      LOG.debug("Setting token_audience to account_id for account host: \"{}\"", accountId);
       tokenAudience = accountId;
     }
   }

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

@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.List;
 
 /**
  * [Experimental] Parsed response from the /.well-known/databricks-config discovery endpoint.
@@ -23,6 +24,9 @@ public class HostMetadata {
   @JsonProperty("cloud")
   private String cloud;
 
+  @JsonProperty("token_federation_default_oidc_audiences")
+  private List<String> tokenFederationDefaultOidcAudiences;
+
   public HostMetadata() {}
 
   public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) {
@@ -53,4 +57,8 @@ public String getWorkspaceId() {
   public String getCloud() {
     return cloud;
   }
+
+  public List<String> getTokenFederationDefaultOidcAudiences() {
+    return tokenFederationDefaultOidcAudiences;
+  }
 }

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

@@ -657,6 +657,149 @@ public void testEnsureResolvedHostMetadataMissingAccountIdWithPlaceholderNonFata
     }
   }
 
+  // --- resolveHostMetadata token_federation_default_oidc_audiences tests ---
+
+  @Test
+  public void testResolveHostMetadataSetsTokenAudienceFromOidcAudiences() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[\"https://ws.databricks.com/oidc/v1/token\"]}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      assertEquals("https://ws.databricks.com/oidc/v1/token", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDoesNotOverrideExistingTokenAudienceWithOidcAudiences()
+      throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[\"metadata-audience\"]}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig().setHost(server.getUrl()).setTokenAudience("existing-audience");
+      config.resolve(emptyEnv());
+      assertEquals("existing-audience", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataOidcAudiencesPriorityOverAccountIdFallback()
+      throws IOException {
+    // token_federation_default_oidc_audiences should take priority over the account_id fallback.
+    // The audiences check runs before the fallback, so once tokenAudience is set from audiences,
+    // the fallback's null check prevents it from overriding.
+    String response =
+        "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[\"custom-audience\"]}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID);
+      config.resolve(emptyEnv());
+      // Should use first element of token_federation_default_oidc_audiences, NOT account_id
+      assertEquals("custom-audience", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataNoAudiencesFieldLeavesTokenAudienceNull() throws IOException {
+    // When token_federation_default_oidc_audiences is absent, tokenAudience stays null
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\"}";
+    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.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataEmptyAudiencesListLeavesTokenAudienceNull()
+      throws IOException {
+    // When token_federation_default_oidc_audiences is an empty array, tokenAudience stays null
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[]}";
+    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.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataEmptyStringAudienceSetsTokenAudience() throws IOException {
+    // When first element is empty string, tokenAudience is set to empty string
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[\"\"]}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      assertEquals("", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataMultipleAudiencesPicksFirst() throws IOException {
+    // When multiple audiences, the first element is used
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[\"first-audience\",\"second-audience\",\"third-audience\"]}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      assertEquals("first-audience", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataNullFirstAudienceLeavesTokenAudienceNull() throws IOException {
+    // When first element is null, tokenAudience stays null
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"token_federation_default_oidc_audiences\":[null,\"second-audience\"]}";
+    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.getTokenAudience());
+    }
+  }
+
   // --- discoveryUrl / OIDC endpoint tests ---
 
   @Test