Merge remote-tracking branch 'origin/main' into default-profile

Author: simonfaltum

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() {

databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleCredentialsCredentialsProvider.java

@@ -66,17 +66,12 @@ public HeaderFactory configure(DatabricksConfig config) {
       Map<String, String> headers = new HashMap<>();
       headers.put("Authorization", String.format("Bearer %s", idToken.getTokenValue()));
 
-      if (config.getClientType() == ClientType.ACCOUNT) {
-        AccessToken token;
-        try {
-          token = finalServiceAccountCredentials.createScoped(GCP_SCOPES).refreshAccessToken();
-        } catch (IOException e) {
-          String message =
-              "Failed to refresh access token from Google service account credentials.";
-          LOG.error(message + e);
-          throw new DatabricksException(message, e);
-        }
+      try {
+        AccessToken token =
+            finalServiceAccountCredentials.createScoped(GCP_SCOPES).refreshAccessToken();
         headers.put(SA_ACCESS_TOKEN_HEADER, token.getTokenValue());
+      } catch (IOException e) {
+        LOG.warn("Failed to refresh GCP SA access token, skipping header: {}", e.getMessage());
       }
 
       return headers;

databricks-sdk-java/src/main/java/com/databricks/sdk/core/GoogleIdCredentialsProvider.java

@@ -69,15 +69,11 @@ public HeaderFactory configure(DatabricksConfig config) {
         throw new DatabricksException(message, e);
       }
 
-      if (config.getClientType() == ClientType.ACCOUNT) {
-        try {
-          headers.put(
-              SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue());
-        } catch (IOException e) {
-          String message = "Failed to refresh access token from scoped id token credentials.";
-          LOG.error(message + e);
-          throw new DatabricksException(message, e);
-        }
+      try {
+        headers.put(
+            SA_ACCESS_TOKEN_HEADER, gcpScopedCredentials.refreshAccessToken().getTokenValue());
+      } catch (IOException e) {
+        LOG.warn("Failed to refresh GCP SA access token, skipping header: {}", e.getMessage());
       }
 
       return headers;

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

databricks-sdk-java/src/test/java/com/databricks/sdk/integration/HostMetadataIT.java

@@ -0,0 +1,42 @@
+package com.databricks.sdk.integration;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.integration.framework.EnvContext;
+import com.databricks.sdk.integration.framework.EnvOrSkip;
+import com.databricks.sdk.integration.framework.EnvTest;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Integration test for host metadata resolution via /.well-known/databricks-config. Port of Go SDK
+ * #1546.
+ */
+@ExtendWith(EnvTest.class)
+@EnvContext("ucws")
+public class HostMetadataIT {
+  private static final Logger LOG = LoggerFactory.getLogger(HostMetadataIT.class);
+
+  @Test
+  void testResolvePopulatesFieldsFromMetadata(@EnvOrSkip("DATABRICKS_HOST") String host) {
+    DatabricksConfig config =
+        new DatabricksConfig().setHost(host).setExperimentalIsUnifiedHost(true);
+    config.resolve();
+
+    LOG.info(
+        "Resolved metadata for {}: accountId={}, workspaceId={}, discoveryUrl={}",
+        host,
+        config.getAccountId(),
+        config.getWorkspaceId(),
+        config.getDiscoveryUrl());
+
+    // After resolve(), host metadata should have been resolved
+    assertNotNull(config.getDiscoveryUrl(), "Expected discoveryUrl to be populated after resolve");
+    assertFalse(config.getDiscoveryUrl().isEmpty(), "Expected non-empty discoveryUrl");
+    assertNotNull(config.getAccountId(), "Expected accountId to be populated after resolve");
+    assertFalse(config.getAccountId().isEmpty(), "Expected non-empty accountId");
+  }
+}