Honor discovery URL for other flows too

databricks · Dec 4, 2024 · c5a174e · c5a174e
1 parent eed5737
commit c5a174e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
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
@@ -554,7 +554,7 @@ public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
         return oidcEndpoints;
       }
     } catch (Exception e) {
-      LOG.debug(
+      LOG.warn(
           "Failed to fetch OIDC Endpoints using discovery URL: {}. Error: {}. \nDefaulting to fetch OIDC using default endpoint.",
           discoveryUrl,
           e.getMessage(),

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java
@@ -31,6 +31,7 @@ public class OAuthClient {
   public static class Builder {
     private String host;
     private String clientId;
+    private String discoveryUrl;
     private String redirectUrl;
     private List<String> scopes;
     private String clientSecret;
@@ -53,6 +54,11 @@ public Builder withClientId(String clientId) {
       return this;
     }
 
+    public Builder withDiscoveryUrl(String discoveryUrl) {
+      this.discoveryUrl = discoveryUrl;
+      return this;
+    }
+
     public Builder withClientSecret(String clientSecret) {
       this.clientSecret = clientSecret;
       return this;
@@ -91,6 +97,7 @@ public OAuthClient(DatabricksConfig config) throws IOException {
             .withHttpClient(config.getHttpClient())
             .withClientId(config.getClientId())
             .withClientSecret(config.getClientSecret())
+            .withDiscoveryUrl(config.getDiscoveryUrl())
             .withHost(config.getHost())
             .withRedirectUrl(
                 config.getOAuthRedirectUrl() != null
@@ -106,7 +113,7 @@ private OAuthClient(Builder b) throws IOException {
     this.host = b.host;
     this.hc = b.hc;
 
-    DatabricksConfig config = new DatabricksConfig().setHost(b.host).resolve();
+    DatabricksConfig config = new DatabricksConfig().setHost(b.host).setDiscoveryUrl(b.discoveryUrl).resolve();
     OpenIDConnectEndpoints oidc = config.getOidcEndpoints();
     if (oidc == null) {
       throw new DatabricksException(b.host + " does not support OAuth");