Wire HostMetadata discovery into DatabricksConfig

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

@@ -827,6 +827,45 @@ HostMetadata getHostMetadata() throws IOException {
     }
   }
 
+  /**
+   * [Experimental] Populate missing config fields from the host's /.well-known/databricks-config
+   * discovery endpoint.
+   *
+   * <p>Fills in {@code accountId}, {@code workspaceId}, and {@code discoveryUrl} (derived from
+   * {@code oidc_endpoint}, with any {@code {account_id}} placeholder substituted) if not already
+   * set.
+   *
+   * <p><b>Note:</b> This API is experimental and may change or be removed in future releases
+   * without notice.
+   *
+   * @throws DatabricksException if {@code accountId} cannot be resolved or {@code oidc_endpoint} is
+   *     missing from the host metadata.
+   */
+  void resolveHostMetadata() throws IOException {
+    if (host == null) {
+      return;
+    }
+    HostMetadata meta = getHostMetadata();
+    if (accountId == null && meta.getAccountId() != null) {
+      accountId = meta.getAccountId();
+    }
+    if (accountId == null) {
+      throw new DatabricksException(
+          "account_id is not configured and could not be resolved from host metadata");
+    }
+    if (workspaceId == null && meta.getWorkspaceId() != null) {
+      workspaceId = meta.getWorkspaceId();
+    }
+    if (discoveryUrl == null) {
+      if (meta.getOidcEndpoint() != null && !meta.getOidcEndpoint().isEmpty()) {
+        discoveryUrl = meta.getOidcEndpoint().replace("{account_id}", accountId);
+      } else {
+        throw new DatabricksException(
+            "discovery_url is not configured and could not be resolved from host metadata");
+      }
+    }
+  }
+
   private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
     try {
       Request request = new Request("GET", discoveryUrl);

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

@@ -480,6 +480,105 @@ public void testGetHostMetadataRaisesOnHttpError() throws IOException {
     }
   }
 
+  // --- resolveHostMetadata tests ---
+
+  @Test
+  public void testResolveHostMetadataWorkspacePopulatesAllFields() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"workspace_id\":\""
+            + DUMMY_WORKSPACE_ID
+            + "\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(DUMMY_ACCOUNT_ID, config.getAccountId());
+      assertEquals(DUMMY_WORKSPACE_ID, config.getWorkspaceId());
+      assertEquals("https://ws.databricks.com/oidc", config.getDiscoveryUrl());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataAccountSubstitutesAccountId() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{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);
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(
+          "https://acc.databricks.com/oidc/accounts/" + DUMMY_ACCOUNT_ID, config.getDiscoveryUrl());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDoesNotOverwriteExistingFields() throws IOException {
+    String existingAccountId = "existing-account-id";
+    String existingWorkspaceId = "existing-workspace-id";
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\"other-account\","
+            + "\"workspace_id\":\"other-ws\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig()
+              .setHost(server.getUrl())
+              .setAccountId(existingAccountId)
+              .setWorkspaceId(existingWorkspaceId);
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals(existingAccountId, config.getAccountId());
+      assertEquals(existingWorkspaceId, config.getWorkspaceId());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataRaisesWhenAccountIdUnresolvable() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{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());
+      DatabricksException ex =
+          assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
+      assertTrue(ex.getMessage().contains("account_id is not configured"));
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataRaisesWhenOidcEndpointMissing() throws IOException {
+    String response = "{\"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());
+      DatabricksException ex =
+          assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
+      assertTrue(ex.getMessage().contains("discovery_url is not configured"));
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataRaisesOnHttpError() throws IOException {
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", "{}", 500)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      DatabricksException ex =
+          assertThrows(DatabricksException.class, () -> config.resolveHostMetadata());
+      assertTrue(ex.getMessage().contains("Failed to fetch host metadata"));
+    }
+  }
+
   // --- discoveryUrl / OIDC endpoint tests ---
 
   @Test