Honor discovery url for flows other than M2M

Author: samikshya-db

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

@@ -17,8 +17,11 @@
 import java.time.Duration;
 import java.util.*;
 import org.apache.http.HttpMessage;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class DatabricksConfig {
+  private static final Logger LOG = LoggerFactory.getLogger(DatabricksConfig.class);
   private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider();
 
   @ConfigAttribute(env = "DATABRICKS_HOST")
@@ -647,7 +650,19 @@ public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
     if (discoveryUrl == null) {
       return fetchDefaultOidcEndpoints();
     }
-    return fetchOidcEndpointsFromDiscovery();
+    try {
+      OpenIDConnectEndpoints oidcEndpoints = fetchOidcEndpointsFromDiscovery();
+      if (oidcEndpoints != null) {
+        return oidcEndpoints;
+      }
+    } catch (Exception e) {
+      LOG.warn(
+              "Failed to fetch OIDC Endpoints using discovery URL: {}. Error: {}. \nDefaulting to fetch OIDC using default endpoint.",
+              discoveryUrl,
+              e.getMessage(),
+              e);
+    }
+    return fetchDefaultOidcEndpoints();
   }
 
   private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
@@ -737,6 +752,7 @@ public DatabricksEnvironment getDatabricksEnvironment() {
   }
 
   private DatabricksConfig clone(Set<String> fieldsToSkip) {
+    fieldsToSkip.add("LOG");
     DatabricksConfig newConfig = new DatabricksConfig();
     for (Field f : DatabricksConfig.class.getDeclaredFields()) {
       if (fieldsToSkip.contains(f.getName())) {

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

@@ -43,6 +43,7 @@ public static class Builder {
     private HttpClient hc;
     private String accountId;
     private Optional<Duration> browserTimeout = Optional.empty();
+    private String discoveryUrl;
 
     public Builder() {}
 
@@ -89,6 +90,11 @@ public Builder withBrowserTimeout(Duration browserTimeout) {
       this.browserTimeout = Optional.of(browserTimeout);
       return this;
     }
+
+    public Builder withDiscoveryUrl(String discoveryUrl) {
+      this.discoveryUrl = discoveryUrl;
+      return this;
+    }
   }
 
   private final String clientId;
@@ -112,7 +118,7 @@ private OAuthClient(Builder b) throws IOException {
     this.hc = b.hc;
 
     DatabricksConfig config =
-        new DatabricksConfig().setHost(b.host).setAccountId(b.accountId).resolve();
+        new DatabricksConfig().setHost(b.host).setAccountId(b.accountId).setDiscoveryUrl(b.discoveryUrl).resolve();
     OpenIDConnectEndpoints oidc = config.getOidcEndpoints();
     if (oidc == null) {
       throw new DatabricksException(b.host + " does not support OAuth");