diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 574bacef4..220a222fd 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -884,6 +884,10 @@ void resolveHostMetadata() throws IOException { "discovery_url is not configured and could not be resolved from host metadata"); } } + // For account hosts, use the accountId as the token audience if not already set. + if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) { + tokenAudience = accountId; + } } private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index b7f8dab02..c6762652b 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -599,6 +599,54 @@ public void testResolveHostMetadataRaisesOnHttpError() throws IOException { } } + @Test + public void testResolveHostMetadataSetsTokenAudienceForAccountHost() throws IOException { + // For a unified host with no workspaceId (ACCOUNT client type), resolveHostMetadata should + // set tokenAudience to accountId when not already configured. + String response = + "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"}"; + try (FixtureServer server = + new FixtureServer() + .with("GET", "/.well-known/databricks-config", response, 200) + .with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = + new DatabricksConfig() + .setHost(server.getUrl()) + .setExperimentalIsUnifiedHost(true) + .setAccountId(DUMMY_ACCOUNT_ID); + config.resolve(emptyEnv()); + // Client type should be ACCOUNT (unified host, no workspaceId) + assertEquals(ClientType.ACCOUNT, config.getClientType()); + config.resolveHostMetadata(); + assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience()); + } + } + + @Test + public void testResolveHostMetadataDoesNotOverwriteTokenAudience() throws IOException { + String response = + "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"," + + "\"account_id\":\"" + + DUMMY_ACCOUNT_ID + + "\"}"; + try (FixtureServer server = + new FixtureServer() + .with("GET", "/.well-known/databricks-config", response, 200) + .with("GET", "/.well-known/databricks-config", response, 200)) { + DatabricksConfig config = + new DatabricksConfig() + .setHost(server.getUrl()) + .setAccountId(DUMMY_ACCOUNT_ID) + .setTokenAudience("custom-audience"); + config.resolve(emptyEnv()); + config.resolveHostMetadata(); + assertEquals("custom-audience", config.getTokenAudience()); + } + } + // --- discoveryUrl / OIDC endpoint tests --- @Test