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 b2b0d75e4..43c1f5136 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 @@ -907,6 +907,10 @@ 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. + 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 c6e149d76..0e86b4ef4 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 @@ -586,6 +586,52 @@ 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)) { + DatabricksConfig config = + new DatabricksConfig() + .setHost(server.getUrl()) + .setAccountId(DUMMY_ACCOUNT_ID) + .setTokenAudience("custom-audience"); + config.resolve(emptyEnv()); + config.resolveHostMetadata(); + assertEquals("custom-audience", config.getTokenAudience()); + } + } + // --- tryResolveHostMetadata (config init) tests --- @Test