Fallback for discovery URL

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

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
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
@@ -13,8 +13,11 @@
 import java.lang.reflect.Field;
 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(DefaultCredentialsProvider.class);
   private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider();
 
   @ConfigAttribute(env = "DATABRICKS_HOST")
@@ -545,7 +548,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.debug(
+          "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() {

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
@@ -159,6 +159,35 @@ public void testDiscoveryEndpoint() throws IOException {
     }
   }
 
+  @Test
+  public void testDiscoveryEndpointFetchThrowsError() throws IOException {
+    String discoveryUrlSuffix = "/test.discovery.url";
+    String OIDCResponse =
+        "{\n"
+            + "  \"authorization_endpoint\": \"https://test.auth.endpoint/oidc/v1/authorize\",\n"
+            + "  \"token_endpoint\": \"https://test.auth.endpoint/oidc/v1/token\"\n"
+            + "}";
+
+    try (FixtureServer server =
+        new FixtureServer()
+            .with("GET", discoveryUrlSuffix, "", 400)
+            .with("GET", "/oidc/.well-known/oauth-authorization-server", OIDCResponse, 200)) {
+
+      String discoveryUrl = server.getUrl() + discoveryUrlSuffix;
+
+      OpenIDConnectEndpoints oidcEndpoints =
+          new DatabricksConfig()
+              .setHost(server.getUrl())
+              .setDiscoveryUrl(discoveryUrl)
+              .setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build())
+              .getOidcEndpoints();
+
+      assertEquals(
+          oidcEndpoints.getAuthorizationEndpoint(), "https://test.auth.endpoint/oidc/v1/authorize");
+      assertEquals(oidcEndpoints.getTokenEndpoint(), "https://test.auth.endpoint/oidc/v1/token");
+    }
+  }
+
   @Test
   public void testNewWithWorkspaceHost() {
     DatabricksConfig config =