From eed57372f6433bee37d8928cb236b1306e389f82 Mon Sep 17 00:00:00 2001 From: Samikshya Chand Date: Mon, 2 Dec 2024 15:31:29 +0530 Subject: [PATCH 1/4] Fallback for discovery URL --- .../databricks/sdk/core/DatabricksConfig.java | 17 ++++++++++- .../sdk/core/DatabricksConfigTest.java | 29 +++++++++++++++++++ 2 files changed, 45 insertions(+), 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 index 20e7f883e..b4f32b5a1 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(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 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 = From c5a174efa72a6431d5de9c5a82142fd39c92d437 Mon Sep 17 00:00:00 2001 From: Samikshya Chand Date: Wed, 4 Dec 2024 16:49:01 +0530 Subject: [PATCH 2/4] Honor discovery URL for other flows too --- .../java/com/databricks/sdk/core/DatabricksConfig.java | 2 +- .../java/com/databricks/sdk/core/oauth/OAuthClient.java | 9 ++++++++- 2 files changed, 9 insertions(+), 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 index b4f32b5a1..1f7106659 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 @@ -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 index 6f4b25996..bc09ec4d9 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,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"); From 55145dbb2ac15e7aeceae6aa5170f766a36c8d2c Mon Sep 17 00:00:00 2001 From: Samikshya Chand Date: Wed, 4 Dec 2024 18:21:31 +0530 Subject: [PATCH 3/4] Exclude log in clone method --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 3 ++- 1 file changed, 2 insertions(+), 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 index 1f7106659..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 @@ -17,7 +17,7 @@ import org.slf4j.LoggerFactory; public class DatabricksConfig { - private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); + private static final Logger LOG = LoggerFactory.getLogger(DatabricksConfig.class); private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider(); @ConfigAttribute(env = "DATABRICKS_HOST") @@ -647,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())) { From 3db82e4661612e15c423a56f99d485864b1ae057 Mon Sep 17 00:00:00 2001 From: Samikshya Chand Date: Wed, 4 Dec 2024 22:37:50 +0530 Subject: [PATCH 4/4] fmt --- .../main/java/com/databricks/sdk/core/oauth/OAuthClient.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 bc09ec4d9..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 @@ -113,7 +113,8 @@ private OAuthClient(Builder b) throws IOException { this.host = b.host; this.hc = b.hc; - DatabricksConfig config = new DatabricksConfig().setHost(b.host).setDiscoveryUrl(b.discoveryUrl).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");