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 index 20e7f883e..ba4be4c0e 100644 --- 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(DatabricksConfig.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.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() { @@ -632,6 +647,7 @@ public DatabricksEnvironment getDatabricksEnvironment() { } private DatabricksConfig clone(Set fieldsToSkip) { + fieldsToSkip.add("LOG"); DatabricksConfig newConfig = new DatabricksConfig(); for (Field f : DatabricksConfig.class.getDeclaredFields()) { if (fieldsToSkip.contains(f.getName())) { 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 index 6f4b25996..a7d30fadd 100644 --- 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 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,8 @@ 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"); 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 index e552a1427..0c899c635 100644 --- 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 =