From 5203e57b56a0308f512179b51d478af3a3d02d9e Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 7 May 2025 15:56:56 +0000 Subject: [PATCH 01/13] Add GithubIDTokenSource and TokenSourceCredentialsProvider --- .../sdk/core/DefaultCredentialsProvider.java | 109 +++++++++---- .../sdk/core/oauth/GithubIDTokenSource.java | 89 ++++++++++ .../oauth/TokenSourceCredentialsProvider.java | 45 +++++ .../sdk/DatabricksAuthManualTest.java | 11 +- .../databricks/sdk/DatabricksAuthTest.java | 74 +++++---- .../core/oauth/GithubIDTokenSourceTest.java | 154 ++++++++++++++++++ 6 files changed, 410 insertions(+), 72 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index b8f4d7867..0571c01f0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -2,7 +2,6 @@ import com.databricks.sdk.core.oauth.*; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -10,48 +9,37 @@ public class DefaultCredentialsProvider implements CredentialsProvider { private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); - private static final List> providerClasses = - Arrays.asList( - PatCredentialsProvider.class, - BasicCredentialsProvider.class, - OAuthM2MServicePrincipalCredentialsProvider.class, - GithubOidcCredentialsProvider.class, - AzureGithubOidcCredentialsProvider.class, - AzureServicePrincipalCredentialsProvider.class, - AzureCliCredentialsProvider.class, - ExternalBrowserCredentialsProvider.class, - DatabricksCliCredentialsProvider.class, - NotebookNativeCredentialsProvider.class, - GoogleCredentialsCredentialsProvider.class, - GoogleIdCredentialsProvider.class); - - private final List providers; + private List providers = new ArrayList<>(); private String authType = "default"; - public String authType() { - return authType; - } + private static class NamedIDTokenSource { + private final String name; + private final IDTokenSource idTokenSource; - public DefaultCredentialsProvider() { - providers = new ArrayList<>(); - for (Class clazz : providerClasses) { - try { - providers.add((CredentialsProvider) clazz.newInstance()); - } catch (NoClassDefFoundError | InstantiationException | IllegalAccessException e) { - LOG.warn( - "Failed to instantiate credentials provider: " - + clazz.getName() - + ", skipping. Cause: " - + e.getClass().getCanonicalName() - + ": " - + e.getMessage()); - } + public NamedIDTokenSource(String name, IDTokenSource idTokenSource) { + this.name = name; + this.idTokenSource = idTokenSource; + } + + public String getName() { + return name; + } + + public IDTokenSource getIdTokenSource() { + return idTokenSource; } } + public DefaultCredentialsProvider() {} + + public String authType() { + return authType; + } + @Override public synchronized HeaderFactory configure(DatabricksConfig config) { + addDefaultCredentialsProviders(config); for (CredentialsProvider provider : providers) { if (config.getAuthType() != null && !config.getAuthType().isEmpty() @@ -80,4 +68,57 @@ public synchronized HeaderFactory configure(DatabricksConfig config) { + authFlowUrl + " to configure credentials for your preferred authentication method"); } + + private void addOIDCCredentialsProviders(DatabricksConfig config) { + OpenIDConnectEndpoints endpoints = null; + try { + endpoints = config.getOidcEndpoints(); + } catch (Exception e) { + LOG.warn("Failed to get OpenID Connect endpoints", e); + } + + List namedIdTokenSources = new ArrayList<>(); + namedIdTokenSources.add( + new NamedIDTokenSource( + "github-oidc", + new GithubIDTokenSource( + config.getActionsIdTokenRequestUrl(), + config.getActionsIdTokenRequestToken(), + config.getHttpClient()))); + // Add new IDTokenSources and ID providers here. Example: + // namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...))); + + for (NamedIDTokenSource namedIdTokenSource : namedIdTokenSources) { + DatabricksOAuthTokenSource oauthTokenSource = + new DatabricksOAuthTokenSource.Builder( + config.getClientId(), + config.getHost(), + endpoints, + namedIdTokenSource.getIdTokenSource(), + config.getHttpClient()) + .audience(config.getTokenAudience()) + .accountId(config.isAccountClient() ? config.getAccountId() : null) + .build(); + + providers.add( + new TokenSourceCredentialsProvider(oauthTokenSource, namedIdTokenSource.getName())); + } + } + + private void addDefaultCredentialsProviders(DatabricksConfig config) { + providers.add(new PatCredentialsProvider()); + providers.add(new BasicCredentialsProvider()); + providers.add(new OAuthM2MServicePrincipalCredentialsProvider()); + + addOIDCCredentialsProviders(config); + + providers.add(new AzureGithubOidcCredentialsProvider()); + providers.add(new AzureServicePrincipalCredentialsProvider()); + providers.add(new AzureCliCredentialsProvider()); + providers.add(new ExternalBrowserCredentialsProvider()); + providers.add(new DatabricksCliCredentialsProvider()); + providers.add(new NotebookNativeCredentialsProvider()); + providers.add(new GoogleCredentialsCredentialsProvider()); + providers.add(new GoogleIdCredentialsProvider()); + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java new file mode 100644 index 000000000..362719bda --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java @@ -0,0 +1,89 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.base.Strings; +import java.io.IOException; + +/** GithubIDTokenSource retrieves JWT Tokens from GitHub Actions. */ +public class GithubIDTokenSource implements IDTokenSource { + private final String actionsIDTokenRequestURL; + private final String actionsIDTokenRequestToken; + private final HttpClient httpClient; + private final ObjectMapper mapper = new ObjectMapper(); + + /** + * Constructs a new GithubIDTokenSource. + * + * @param actionsIDTokenRequestURL The URL to request the ID token from GitHub Actions. + * @param actionsIDTokenRequestToken The token used to authenticate the request. + * @param httpClient The HTTP client to use for making requests. + */ + public GithubIDTokenSource( + String actionsIDTokenRequestURL, String actionsIDTokenRequestToken, HttpClient httpClient) { + this.actionsIDTokenRequestURL = actionsIDTokenRequestURL; + this.actionsIDTokenRequestToken = actionsIDTokenRequestToken; + this.httpClient = httpClient; + } + + @Override + public IDToken getIDToken(String audience) { + if (Strings.isNullOrEmpty(actionsIDTokenRequestURL)) { + throw new DatabricksException("Missing ActionsIDTokenRequestURL"); + } + if (Strings.isNullOrEmpty(actionsIDTokenRequestToken)) { + throw new DatabricksException("Missing ActionsIDTokenRequestToken"); + } + if (httpClient == null) { + throw new DatabricksException("HttpClient cannot be null"); + } + + String requestUrl = actionsIDTokenRequestURL; + if (!Strings.isNullOrEmpty(audience)) { + requestUrl = String.format("%s&audience=%s", requestUrl, audience); + } + + Request req = + new Request("GET", requestUrl) + .withHeader("Authorization", "Bearer " + actionsIDTokenRequestToken); + + Response resp; + try { + resp = httpClient.execute(req); + } catch (IOException e) { + throw new DatabricksException( + "Failed to request ID token from " + requestUrl + ": " + e.getMessage(), e); + } + + if (resp.getStatusCode() != 200) { + throw new DatabricksException( + "Failed to request ID token: status code " + + resp.getStatusCode() + + ", response body: " + + resp.getBody().toString()); + } + + ObjectNode jsonResp; + try { + jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class); + } catch (IOException e) { + throw new DatabricksException( + "Failed to request ID token: corrupted token: " + e.getMessage()); + } + + if (!jsonResp.has("value")) { + throw new DatabricksException("ID token response missing 'value' field"); + } + + String tokenValue = jsonResp.get("value").textValue(); + if (Strings.isNullOrEmpty(tokenValue)) { + throw new DatabricksException("Received empty ID token from GitHub Actions"); + } + + return new IDToken(tokenValue); + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java new file mode 100644 index 000000000..53cfc23d4 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java @@ -0,0 +1,45 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.CredentialsProvider; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.HeaderFactory; +import java.util.HashMap; +import java.util.Map; + +/** Base class for token-based credentials providers. */ +public class TokenSourceCredentialsProvider implements CredentialsProvider { + private final TokenSource tokenSource; + private final String authType; + + /** + * Creates a new TokenSourceCredentialsProvider with the specified token source and auth type. + * + * @param tokenSource The token source to use for token exchange + * @param authType The authentication type string + */ + public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType) { + this.tokenSource = tokenSource; + this.authType = authType; + } + + @Override + public HeaderFactory configure(DatabricksConfig config) { + try { + // Validate that we can get a token before returning the HeaderFactory + String accessToken = tokenSource.getToken().getAccessToken(); + + return () -> { + Map headers = new HashMap<>(); + headers.put("Authorization", "Bearer " + accessToken); + return headers; + }; + } catch (Exception e) { + return null; + } + } + + @Override + public String authType() { + return authType; + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java index c467d5f4b..9b870458a 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java @@ -2,12 +2,17 @@ import com.databricks.sdk.core.ConfigResolving; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DummyHttpClient; import com.databricks.sdk.core.utils.TestOSUtils; import java.util.Map; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class DatabricksAuthManualTest implements ConfigResolving { + private DatabricksConfig createConfigWithMockClient() { + return new DatabricksConfig().setHttpClient(new DummyHttpClient()); + } + @Test void azureCliWorkspaceHeaderPresent() { StaticEnv env = @@ -18,7 +23,7 @@ void azureCliWorkspaceHeaderPresent() { String azureWorkspaceResourceId = "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123"; DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setAuthType("azure-cli") .setHost("https://x") .setAzureWorkspaceResourceId(azureWorkspaceResourceId); @@ -38,7 +43,7 @@ void azureCliUserWithManagementAccess() { String azureWorkspaceResourceId = "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123"; DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setAuthType("azure-cli") .setHost("https://x") .setAzureWorkspaceResourceId(azureWorkspaceResourceId); @@ -58,7 +63,7 @@ void azureCliUserNoManagementAccess() { String azureWorkspaceResourceId = "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123"; DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setAuthType("azure-cli") .setHost("https://x") .setAzureWorkspaceResourceId(azureWorkspaceResourceId); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java index 710c44611..d742b1e98 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java @@ -7,6 +7,7 @@ import com.databricks.sdk.core.ConfigResolving; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DummyHttpClient; import com.databricks.sdk.core.utils.GitHubUtils; import com.databricks.sdk.core.utils.TestOSUtils; import java.io.File; @@ -17,18 +18,20 @@ public class DatabricksAuthTest implements GitHubUtils, ConfigResolving { private String errorMessageBase = "default auth: cannot configure default credentials, please check https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication to configure credentials for your preferred authentication method"; + private DatabricksConfig createConfigWithMockClient() { + return new DatabricksConfig().setHttpClient(new DummyHttpClient()); + } + public DatabricksAuthTest() { setPermissionOnTestAz(); } @Test public void testTestConfigNoParams() { - raises( errorMessageBase, () -> { - DatabricksConfig config = new DatabricksConfig(); - + DatabricksConfig config = createConfigWithMockClient(); config.authenticate(); }); } @@ -40,7 +43,7 @@ public void testTestConfigHostEnv() { raises( errorMessageBase + ". Config: host=https://x. Env: DATABRICKS_HOST", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -53,7 +56,7 @@ public void testTestConfigTokenEnv() { raises( errorMessageBase + ". Config: token=***. Env: DATABRICKS_TOKEN", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -63,7 +66,7 @@ public void testTestConfigTokenEnv() { public void testTestConfigHostTokenEnv() { // Set environment variables StaticEnv env = new StaticEnv().with("DATABRICKS_HOST", "x").with("DATABRICKS_TOKEN", "x"); - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -75,7 +78,7 @@ public void testTestConfigHostTokenEnv() { public void testTestConfigHostParamTokenEnv() { // Set environment variables StaticEnv env = new StaticEnv().with("DATABRICKS_TOKEN", "x"); - DatabricksConfig config = new DatabricksConfig().setHost("https://x"); + DatabricksConfig config = createConfigWithMockClient().setHost("https://x"); resolveConfig(config, env); config.authenticate(); @@ -92,7 +95,7 @@ public void testTestConfigUserPasswordEnv() { errorMessageBase + ". Config: username=x, password=***. Env: DATABRICKS_USERNAME, DATABRICKS_PASSWORD", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -108,7 +111,7 @@ public void testTestConfigBasicAuth() { .with("DATABRICKS_HOST", "x") .with("DATABRICKS_PASSWORD", "x") .with("DATABRICKS_USERNAME", "x"); - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -124,7 +127,7 @@ public void testTestConfigAttributePrecedence() { .with("DATABRICKS_HOST", "x") .with("DATABRICKS_PASSWORD", "x") .with("DATABRICKS_USERNAME", "x"); - DatabricksConfig config = new DatabricksConfig().setHost("y"); + DatabricksConfig config = createConfigWithMockClient().setHost("y"); resolveConfig(config, env); config.authenticate(); @@ -136,7 +139,7 @@ public void testTestConfigAttributePrecedence() { public void testTestConfigBasicAuthMix() { // Set environment variables StaticEnv env = new StaticEnv().with("DATABRICKS_PASSWORD", "x"); - DatabricksConfig config = new DatabricksConfig().setHost("y").setUsername("x"); + DatabricksConfig config = createConfigWithMockClient().setHost("y").setUsername("x"); resolveConfig(config, env); config.authenticate(); @@ -146,8 +149,8 @@ public void testTestConfigBasicAuthMix() { @Test public void testTestConfigBasicAuthAttrs() { - - DatabricksConfig config = new DatabricksConfig().setHost("y").setUsername("x").setPassword("x"); + DatabricksConfig config = + createConfigWithMockClient().setHost("y").setUsername("x").setPassword("x"); config.authenticate(); @@ -167,7 +170,7 @@ public void testTestConfigConflictingEnvs() { raises( "validate: more than one authorization method configured: basic and pat. Config: host=x, token=***, username=x, password=***. Env: DATABRICKS_HOST, DATABRICKS_TOKEN, DATABRICKS_USERNAME, DATABRICKS_PASSWORD", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -182,7 +185,7 @@ public void testTestConfigConflictingEnvsAuthType() { .with("DATABRICKS_PASSWORD", "x") .with("DATABRICKS_TOKEN", "x") .with("DATABRICKS_USERNAME", "x"); - DatabricksConfig config = new DatabricksConfig().setAuthType("basic"); + DatabricksConfig config = createConfigWithMockClient().setAuthType("basic"); resolveConfig(config, env); config.authenticate(); @@ -197,7 +200,7 @@ public void testTestConfigConfigFile() { raises( errorMessageBase + ". Config: config_file=x. Env: DATABRICKS_CONFIG_FILE", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -210,7 +213,7 @@ public void testTestConfigConfigFileSkipDefaultProfileIfHostSpecified() { raises( errorMessageBase + ". Config: host=https://x", () -> { - DatabricksConfig config = new DatabricksConfig().setHost("x"); + DatabricksConfig config = createConfigWithMockClient().setHost("x"); resolveConfig(config, env); config.authenticate(); }); @@ -223,7 +226,7 @@ public void testTestConfigConfigFileWithEmptyDefaultProfileSelectDefault() { raises( errorMessageBase, () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -236,7 +239,7 @@ public void testTestConfigConfigFileWithEmptyDefaultProfileSelectAbc() { new StaticEnv() .with("DATABRICKS_CONFIG_PROFILE", "abc") .with("HOME", TestOSUtils.resource("/testdata/empty_default")); - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -248,7 +251,7 @@ public void testTestConfigConfigFileWithEmptyDefaultProfileSelectAbc() { public void testTestConfigPatFromDatabricksCfg() { // Set environment variables StaticEnv env = new StaticEnv().with("HOME", TestOSUtils.resource("/testdata")); - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -263,7 +266,7 @@ public void testTestConfigPatFromDatabricksCfgDotProfile() { new StaticEnv() .with("DATABRICKS_CONFIG_PROFILE", "pat.with.dot") .with("HOME", TestOSUtils.resource("/testdata")); - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); @@ -281,7 +284,7 @@ public void testTestConfigPatFromDatabricksCfgNohostProfile() { raises( errorMessageBase + ". Config: token=***, profile=nohost. Env: DATABRICKS_CONFIG_PROFILE", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -299,7 +302,7 @@ public void testTestConfigConfigProfileAndToken() { errorMessageBase + ". Config: token=***, profile=nohost. Env: DATABRICKS_TOKEN, DATABRICKS_CONFIG_PROFILE", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -316,7 +319,7 @@ public void testTestConfigConfigProfileAndPassword() { raises( "validate: more than one authorization method configured: basic and pat. Config: token=***, username=x, profile=nohost. Env: DATABRICKS_USERNAME, DATABRICKS_CONFIG_PROFILE", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -324,9 +327,10 @@ public void testTestConfigConfigProfileAndPassword() { @Test public void testTestConfigAzurePat() { - DatabricksConfig config = - new DatabricksConfig().setHost("https://adb-xxx.y.azuredatabricks.net/").setToken("y"); + createConfigWithMockClient() + .setHost("https://adb-xxx.y.azuredatabricks.net/") + .setToken("y"); config.authenticate(); @@ -343,7 +347,7 @@ public void testTestConfigAzureCliHost() { .with("HOME", TestOSUtils.resource("/testdata/azure")) .with("PATH", "testdata:/bin"); DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setHost("https://adb-123.4.azuredatabricks.net") .setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); @@ -366,7 +370,7 @@ public void testTestConfigAzureCliHostFail() { "default auth: azure-cli: cannot get access token: This is just a failing script.\n. Config: azure_workspace_resource_id=/sub/rg/ws", () -> { DatabricksConfig config = - new DatabricksConfig().setAzureWorkspaceResourceId("/sub/rg/ws"); + createConfigWithMockClient().setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); config.authenticate(); }); @@ -383,7 +387,7 @@ public void testTestConfigAzureCliHostAzNotInstalled() { errorMessageBase + ". Config: azure_workspace_resource_id=/sub/rg/ws", () -> { DatabricksConfig config = - new DatabricksConfig().setAzureWorkspaceResourceId("/sub/rg/ws"); + createConfigWithMockClient().setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); config.authenticate(); }); @@ -400,7 +404,7 @@ public void testTestConfigAzureCliHostPatConflictWithConfigFilePresentWithoutDef "validate: more than one authorization method configured: azure and pat. Config: token=***, azure_workspace_resource_id=/sub/rg/ws", () -> { DatabricksConfig config = - new DatabricksConfig().setToken("x").setAzureWorkspaceResourceId("/sub/rg/ws"); + createConfigWithMockClient().setToken("x").setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); config.authenticate(); }); @@ -414,7 +418,7 @@ public void testTestConfigAzureCliHostAndResourceId() { .with("HOME", TestOSUtils.resource("/testdata")) .with("PATH", "testdata:/bin"); DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setHost("https://adb-123.4.azuredatabricks.net") .setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); @@ -434,7 +438,7 @@ public void testTestConfigAzureCliHostAndResourceIDConfigurationPrecedence() { .with("HOME", TestOSUtils.resource("/testdata/azure")) .with("PATH", "testdata:/bin"); DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setHost("https://adb-123.4.azuredatabricks.net") .setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); @@ -457,7 +461,7 @@ public void testTestConfigAzureAndPasswordConflict() { "validate: more than one authorization method configured: azure and basic. Config: host=https://adb-123.4.azuredatabricks.net, username=x, azure_workspace_resource_id=/sub/rg/ws. Env: DATABRICKS_USERNAME", () -> { DatabricksConfig config = - new DatabricksConfig() + createConfigWithMockClient() .setHost("https://adb-123.4.azuredatabricks.net") .setAzureWorkspaceResourceId("/sub/rg/ws"); resolveConfig(config, env); @@ -475,7 +479,7 @@ public void testTestConfigCorruptConfig() { raises( "resolve: testdata/corrupt/.databrickscfg has no DEFAULT profile configured. Config: profile=DEFAULT. Env: DATABRICKS_CONFIG_PROFILE", () -> { - DatabricksConfig config = new DatabricksConfig(); + DatabricksConfig config = createConfigWithMockClient(); resolveConfig(config, env); config.authenticate(); }); @@ -490,7 +494,7 @@ public void testTestConfigAuthTypeFromEnv() { .with("DATABRICKS_PASSWORD", "password") .with("DATABRICKS_TOKEN", "token") .with("DATABRICKS_USERNAME", "user"); - DatabricksConfig config = new DatabricksConfig().setHost("x"); + DatabricksConfig config = createConfigWithMockClient().setHost("x"); resolveConfig(config, env); config.authenticate(); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java new file mode 100644 index 000000000..e53aaefa8 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java @@ -0,0 +1,154 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.IOException; +import java.net.MalformedURLException; +import java.net.URL; +import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +public class GithubIDTokenSourceTest { + private static final String TEST_REQUEST_URL = "https://github.com/token"; + private static final String TEST_REQUEST_TOKEN = "test-request-token"; + private static final String TEST_ID_TOKEN = "test-id-token"; + private static final String TEST_AUDIENCE = "test-audience"; + + @Mock private static HttpClient mockHttpClient; + + private GithubIDTokenSource tokenSource; + private ObjectMapper mapper; + + @BeforeEach + void setUp() throws IOException { + MockitoAnnotations.openMocks(this); + mapper = new ObjectMapper(); + tokenSource = new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, mockHttpClient); + } + + @Test + void testSuccessfulTokenRetrieval() throws IOException { + // Prepare mock response + ObjectNode responseJson = mapper.createObjectNode(); + responseJson.put("value", TEST_ID_TOKEN); + Response mockResponse = makeResponse(responseJson.toString(), 200); + when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse); + + // Test token retrieval + IDToken token = tokenSource.getIDToken(TEST_AUDIENCE); + + assertNotNull(token); + assertEquals(TEST_ID_TOKEN, token.getValue()); + + // Verify the request was made with correct parameters + verify(mockHttpClient) + .execute( + argThat( + request -> { + return request.getMethod().equals("GET") + && request.getUrl().startsWith(TEST_REQUEST_URL) + && request.getUrl().contains("audience=" + TEST_AUDIENCE) + && request + .getHeaders() + .get("Authorization") + .equals("Bearer " + TEST_REQUEST_TOKEN); + })); + } + + @Test + void testSuccessfulTokenRetrievalWithoutAudience() throws IOException { + // Prepare mock response + ObjectNode responseJson = mapper.createObjectNode(); + responseJson.put("value", TEST_ID_TOKEN); + Response mockResponse = makeResponse(responseJson.toString(), 200); + when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse); + + // Test token retrieval without audience + IDToken token = tokenSource.getIDToken(""); + + assertNotNull(token); + assertEquals(TEST_ID_TOKEN, token.getValue()); + + // Verify the request was made with correct parameters + verify(mockHttpClient) + .execute( + argThat( + request -> { + return request.getMethod().equals("GET") + && request.getUrl().equals(TEST_REQUEST_URL) + && request + .getHeaders() + .get("Authorization") + .equals("Bearer " + TEST_REQUEST_TOKEN); + })); + } + + private static Stream provideInvalidConstructorParameters() { + return Stream.of( + Arguments.of("Missing Request URL", null, TEST_REQUEST_TOKEN, mockHttpClient), + Arguments.of("Missing Request Token", TEST_REQUEST_URL, null, mockHttpClient), + Arguments.of("Null HttpClient", TEST_REQUEST_URL, TEST_REQUEST_TOKEN, null)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideInvalidConstructorParameters") + void testInvalidConstructorParameters( + String testName, String requestUrl, String requestToken, HttpClient httpClient) { + GithubIDTokenSource invalidSource = + new GithubIDTokenSource(requestUrl, requestToken, httpClient); + assertThrows(DatabricksException.class, () -> invalidSource.getIDToken(TEST_AUDIENCE)); + } + + private static Stream provideHttpErrorScenarios() throws IOException { + HttpClient httpClientError = mock(HttpClient.class); + when(httpClientError.execute(any(Request.class))).thenThrow(new IOException("Network error")); + + HttpClient nonSuccessClient = mock(HttpClient.class); + when(nonSuccessClient.execute(any(Request.class))) + .thenReturn(makeResponse("Error response", 400)); + + HttpClient invalidJsonClient = mock(HttpClient.class); + when(invalidJsonClient.execute(any(Request.class))) + .thenReturn(makeResponse("Invalid json", 200)); + + HttpClient missingTokenClient = mock(HttpClient.class); + when(missingTokenClient.execute(any(Request.class))).thenReturn(makeResponse("{}", 200)); + + HttpClient emptyTokenClient = mock(HttpClient.class); + when(emptyTokenClient.execute(any(Request.class))) + .thenReturn(makeResponse("{\"value\":\"\"}", 200)); + + return Stream.of( + Arguments.of("HTTP Client Error", httpClientError), + Arguments.of("Non-Success Status Code", nonSuccessClient), + Arguments.of("Invalid JSON Response", invalidJsonClient), + Arguments.of("Missing Token Value", missingTokenClient), + Arguments.of("Empty Token Value", emptyTokenClient)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideHttpErrorScenarios") + void testHttpErrorScenarios(String testName, HttpClient httpClient) { + GithubIDTokenSource source = + new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, httpClient); + assertThrows(DatabricksException.class, () -> source.getIDToken(TEST_AUDIENCE)); + } + + private static Response makeResponse(String body, int status) throws MalformedURLException { + return new Response(body, status, "status", new URL("https://databricks.com/")); + } +} From 69a0215950eb9f19fd6a45b01bea968016a77ad0 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 8 May 2025 10:58:14 +0000 Subject: [PATCH 02/13] Add example code --- .../example/GithubOIDCAuthExample.java | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java diff --git a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java new file mode 100644 index 000000000..56e3647a8 --- /dev/null +++ b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java @@ -0,0 +1,42 @@ +package com.databricks.sdk.examples; + +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.service.workspace.WorkspaceClient; + +/** + * Example demonstrating how to use GitHub OIDC authentication with Databricks. + * + * This example assumes you have: + * 1. A Databricks workspace + * 2. A service principal configured with GitHub OIDC federation + * 3. The following environment variables set: + * - DATABRICKS_HOST: Your Databricks workspace URL + * - DATABRICKS_CLIENT_ID: Your service principal's client ID + * - ACTIONS_ID_TOKEN_REQUEST_URL: GitHub Actions token request URL + * - ACTIONS_ID_TOKEN_REQUEST_TOKEN: GitHub Actions token request token + */ +public class GithubOIDCAuthExample { + public static void main(String[] args) { + // Create Databricks configuration with GitHub OIDC authentication + DatabricksConfig config = new DatabricksConfig() + .setAuthType("github-oidc"); + + // Create a workspace client + WorkspaceClient workspace = new WorkspaceClient(config); + + try { + // Test the authentication by getting the current user + var currentUser = workspace.currentUser().me(); + System.out.println("Successfully authenticated as: " + currentUser.getUserName()); + + // You can add more API calls here to test other functionality + // For example: + // var clusters = workspace.clusters().list(); + // System.out.println("Available clusters: " + clusters); + + } catch (Exception e) { + System.err.println("Authentication failed: " + e.getMessage()); + e.printStackTrace(); + } + } +} \ No newline at end of file From e0789e66314691a6dd26b6173ddb896aaa247dd3 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 8 May 2025 11:43:28 +0000 Subject: [PATCH 03/13] Fix example --- .../main/java/com/databricks/example/GithubOIDCAuthExample.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java index 56e3647a8..e3a8d745c 100644 --- a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java +++ b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java @@ -1,7 +1,7 @@ package com.databricks.sdk.examples; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.service.workspace.WorkspaceClient; +import com.databricks.sdk.WorkspaceClient; /** * Example demonstrating how to use GitHub OIDC authentication with Databricks. From 15798980ca20e29d9d66663f9644dec9d83d1a0f Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 8 May 2025 12:03:47 +0000 Subject: [PATCH 04/13] Fix code example --- .../java/com/databricks/example/GithubOIDCAuthExample.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java index e3a8d745c..f6ae4fdb2 100644 --- a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java +++ b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java @@ -2,6 +2,7 @@ import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.WorkspaceClient; +import com.databricks.sdk.service.iam.User; /** * Example demonstrating how to use GitHub OIDC authentication with Databricks. @@ -25,8 +26,8 @@ public static void main(String[] args) { WorkspaceClient workspace = new WorkspaceClient(config); try { - // Test the authentication by getting the current user - var currentUser = workspace.currentUser().me(); + // Test the authentication by getting the current User + User currentUser = workspace.currentUser().me(); System.out.println("Successfully authenticated as: " + currentUser.getUserName()); // You can add more API calls here to test other functionality From 2bc27a9d264d68b39c0901eeacffc67418a250a3 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 8 May 2025 15:40:09 +0000 Subject: [PATCH 05/13] Working Github OIDC code example --- .../example/GithubOIDCAuthExample.java | 41 ++++++++----------- 1 file changed, 17 insertions(+), 24 deletions(-) diff --git a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java index f6ae4fdb2..cb86e9cb0 100644 --- a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java +++ b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java @@ -1,43 +1,36 @@ -package com.databricks.sdk.examples; +package com.databricks.example; +import com.databricks.sdk.AccountClient; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.WorkspaceClient; import com.databricks.sdk.service.iam.User; +import com.databricks.sdk.service.iam.ListAccountGroupsRequest; +import com.databricks.sdk.service.iam.Group; /** * Example demonstrating how to use GitHub OIDC authentication with Databricks. - * - * This example assumes you have: - * 1. A Databricks workspace - * 2. A service principal configured with GitHub OIDC federation - * 3. The following environment variables set: - * - DATABRICKS_HOST: Your Databricks workspace URL - * - DATABRICKS_CLIENT_ID: Your service principal's client ID - * - ACTIONS_ID_TOKEN_REQUEST_URL: GitHub Actions token request URL - * - ACTIONS_ID_TOKEN_REQUEST_TOKEN: GitHub Actions token request token */ public class GithubOIDCAuthExample { + public static void main(String[] args) { // Create Databricks configuration with GitHub OIDC authentication DatabricksConfig config = new DatabricksConfig() - .setAuthType("github-oidc"); + .setAuthType("github-oidc") + .setHost("https://accounts.cloud.databricks.com/") + .setAccountId("968367da-7edd-44f7-9dea-3e0b20b0ec97") + .setClientId("dd3b5d58-5a6e-45e3-a3ae-81d8f89fc6fd"); - // Create a workspace client - WorkspaceClient workspace = new WorkspaceClient(config); + AccountClient account = new AccountClient(config); try { - // Test the authentication by getting the current User - User currentUser = workspace.currentUser().me(); - System.out.println("Successfully authenticated as: " + currentUser.getUserName()); - - // You can add more API calls here to test other functionality - // For example: - // var clusters = workspace.clusters().list(); - // System.out.println("Available clusters: " + clusters); - + System.out.println("\nListing account groups:"); + Iterable groups = account.groups().list(new ListAccountGroupsRequest()); + for (Group group : groups) { + System.out.println("- Group: " + group.getDisplayName() + " (ID: " + group.getId() + ")"); + } + } catch (Exception e) { System.err.println("Authentication failed: " + e.getMessage()); e.printStackTrace(); } } -} \ No newline at end of file +} \ No newline at end of file From 6f492bfc1dd24d09dea371925c3534ad47c3d552 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 9 May 2025 13:27:32 +0000 Subject: [PATCH 06/13] Add unit tests and comments --- .../sdk/core/DefaultCredentialsProvider.java | 38 ++++++++++ .../sdk/core/oauth/GithubIDTokenSource.java | 23 +++++- .../oauth/TokenSourceCredentialsProvider.java | 28 +++++++- .../TokenSourceCredentialsProviderTest.java | 72 +++++++++++++++++++ .../example/GithubOIDCAuthExample.java | 14 ++-- 5 files changed, 167 insertions(+), 8 deletions(-) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 0571c01f0..1d9607c76 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -6,13 +6,25 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * The DefaultCredentialsProvider is the primary authentication handler for the Databricks SDK. It + * implements a chain of responsibility pattern to manage multiple authentication methods, including + * Personal Access Tokens (PAT), OAuth, Azure, Google, and OpenID Connect (OIDC). The provider + * attempts each authentication method in sequence until a valid credential is obtained. + */ public class DefaultCredentialsProvider implements CredentialsProvider { private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); + // List of credential providers that will be tried in sequence private List providers = new ArrayList<>(); + // The currently selected authentication type private String authType = "default"; + /** + * Internal class to associate an ID token source with a name for identification purposes. Used + * primarily for OIDC (OpenID Connect) authentication flows. + */ private static class NamedIDTokenSource { private final String name; private final IDTokenSource idTokenSource; @@ -33,10 +45,23 @@ public IDTokenSource getIdTokenSource() { public DefaultCredentialsProvider() {} + /** + * Returns the current authentication type being used + * + * @return String representing the authentication type + */ public String authType() { return authType; } + /** + * Configures the credentials provider with the given Databricks configuration. This method tries + * each available credential provider in sequence until one succeeds. + * + * @param config The Databricks configuration containing authentication details + * @return HeaderFactory for making authenticated requests + * @throws DatabricksException if no valid credentials can be configured + */ @Override public synchronized HeaderFactory configure(DatabricksConfig config) { addDefaultCredentialsProviders(config); @@ -69,6 +94,11 @@ public synchronized HeaderFactory configure(DatabricksConfig config) { + " to configure credentials for your preferred authentication method"); } + /** + * Adds OpenID Connect (OIDC) based credential providers to the list of available providers. + * + * @param config The Databricks configuration containing OIDC settings + */ private void addOIDCCredentialsProviders(DatabricksConfig config) { OpenIDConnectEndpoints endpoints = null; try { @@ -88,6 +118,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { // Add new IDTokenSources and ID providers here. Example: // namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...))); + // Configure OAuth token sources for each ID token source for (NamedIDTokenSource namedIdTokenSource : namedIdTokenSources) { DatabricksOAuthTokenSource oauthTokenSource = new DatabricksOAuthTokenSource.Builder( @@ -105,11 +136,18 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { } } + /** + * Initializes all available credential providers in the preferred order. The order of providers + * determines the authentication fallback sequence. + * + * @param config The Databricks configuration to use for provider initialization + */ private void addDefaultCredentialsProviders(DatabricksConfig config) { providers.add(new PatCredentialsProvider()); providers.add(new BasicCredentialsProvider()); providers.add(new OAuthM2MServicePrincipalCredentialsProvider()); + // Add OIDC-based providers addOIDCCredentialsProviders(config); providers.add(new AzureGithubOidcCredentialsProvider()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java index 362719bda..e9d4e26c4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java @@ -9,11 +9,19 @@ import com.google.common.base.Strings; import java.io.IOException; -/** GithubIDTokenSource retrieves JWT Tokens from GitHub Actions. */ +/** + * GithubIDTokenSource retrieves JWT Tokens from GitHub Actions. This class implements the + * IDTokenSource interface and provides a method for obtaining ID tokens specifically from GitHub + * Actions environment. + */ public class GithubIDTokenSource implements IDTokenSource { + // URL endpoint for requesting ID tokens from GitHub Actions private final String actionsIDTokenRequestURL; + // Authentication token required to request ID tokens from GitHub Actions private final String actionsIDTokenRequestToken; + // HTTP client for making requests to GitHub Actions private final HttpClient httpClient; + // JSON mapper for parsing response data private final ObjectMapper mapper = new ObjectMapper(); /** @@ -30,8 +38,18 @@ public GithubIDTokenSource( this.httpClient = httpClient; } + /** + * Retrieves an ID token from GitHub Actions. This method makes an authenticated request to GitHub + * Actions to obtain a JWT token that later can be exchanged for a Databricks access token. + * + * @param audience Optional audience claim for the token. If provided, it will be included in the + * token request to GitHub Actions. + * @return An IDToken object containing the JWT token value + * @throws DatabricksException if the token request fails or if required configuration is missing + */ @Override public IDToken getIDToken(String audience) { + // Validate required configuration if (Strings.isNullOrEmpty(actionsIDTokenRequestURL)) { throw new DatabricksException("Missing ActionsIDTokenRequestURL"); } @@ -59,6 +77,7 @@ public IDToken getIDToken(String audience) { "Failed to request ID token from " + requestUrl + ": " + e.getMessage(), e); } + // Validate response status code if (resp.getStatusCode() != 200) { throw new DatabricksException( "Failed to request ID token: status code " @@ -67,6 +86,7 @@ public IDToken getIDToken(String audience) { + resp.getBody().toString()); } + // Parse the JSON response ObjectNode jsonResp; try { jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class); @@ -75,6 +95,7 @@ public IDToken getIDToken(String audience) { "Failed to request ID token: corrupted token: " + e.getMessage()); } + // Validate response structure and token value if (!jsonResp.has("value")) { throw new DatabricksException("ID token response missing 'value' field"); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java index 53cfc23d4..1428b4a5f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java @@ -6,7 +6,14 @@ import java.util.HashMap; import java.util.Map; -/** Base class for token-based credentials providers. */ +/** + * A credentials provider that uses a TokenSource to obtain and manage authentication tokens. This + * class serves as a base implementation for token-based authentication, handling the conversion of + * tokens into HTTP authorization headers. + * + *

The provider validates token availability during configuration and creates. appropriate + * authorization headers for API requests. + */ public class TokenSourceCredentialsProvider implements CredentialsProvider { private final TokenSource tokenSource; private final String authType; @@ -14,14 +21,23 @@ public class TokenSourceCredentialsProvider implements CredentialsProvider { /** * Creates a new TokenSourceCredentialsProvider with the specified token source and auth type. * - * @param tokenSource The token source to use for token exchange - * @param authType The authentication type string + * @param tokenSource The token source responsible for token acquisition and management. + * @param authType The authentication type identifier. */ public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType) { this.tokenSource = tokenSource; this.authType = authType; } + /** + * Configures the credentials provider and creates a HeaderFactory for generating authentication + * headers. This method validates token availability by attempting to obtain an access token + * before returning the HeaderFactory. + * + * @param config The Databricks configuration object. + * @return A HeaderFactory that generates "Bearer" token authorization headers, or null if token + * acquisition fails. + */ @Override public HeaderFactory configure(DatabricksConfig config) { try { @@ -38,6 +54,12 @@ public HeaderFactory configure(DatabricksConfig config) { } } + /** + * Returns the authentication type identifier for this credentials provider. This is used to + * identify the authentication method being used. + * + * @return The authentication type string + */ @Override public String authType() { return authType; diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java new file mode 100644 index 000000000..0b7661409 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java @@ -0,0 +1,72 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.HeaderFactory; +import java.time.LocalDateTime; +import java.util.Map; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** Tests for TokenSourceCredentialsProvider */ +class TokenSourceCredentialsProviderTest { + private TokenSourceCredentialsProvider provider; + private static final String TEST_AUTH_TYPE = "test-auth-type"; + private static final String TEST_ACCESS_TOKEN_VALUE = "test-access-token"; + private static final Token TEST_TOKEN = + new Token(TEST_ACCESS_TOKEN_VALUE, "Bearer", LocalDateTime.now().plusHours(1)); + + /** Tests token retrieval scenarios */ + @ParameterizedTest(name = "{0}") + @MethodSource("provideTokenScenarios") + void testTokenScenarios( + String testName, + TokenSource mockTokenSource, + String expectedAuthHeader, + Exception expectedException) { + provider = new TokenSourceCredentialsProvider(mockTokenSource, TEST_AUTH_TYPE); + DatabricksConfig config = new DatabricksConfig(); + HeaderFactory headerFactory = provider.configure(config); + + if (expectedException != null) { + assertNull(headerFactory); + } else { + assertNotNull(headerFactory); + Map headers = headerFactory.headers(); + assertEquals(expectedAuthHeader, headers.get("Authorization")); + } + verify(mockTokenSource).getToken(); + assertEquals(TEST_AUTH_TYPE, provider.authType()); + } + + /** Provides test scenarios */ + private static Stream provideTokenScenarios() { + // Mock behaviour of successful token retrieval + TokenSource mockSuccessTokenSource = mock(TokenSource.class); + when(mockSuccessTokenSource.getToken()).thenReturn(TEST_TOKEN); + + // Mock behaviour of failing token retrieval + TokenSource mockFailureTokenSource = mock(TokenSource.class); + when(mockFailureTokenSource.getToken()) + .thenThrow(new DatabricksException("Token retrieval failed")); + + return Stream.of( + // Success case + Arguments.of( + "Successful token retrieval", + mockSuccessTokenSource, + "Bearer " + TEST_ACCESS_TOKEN_VALUE, + null), + // Failure case + Arguments.of( + "Failed token retrieval", + mockFailureTokenSource, + null, + new DatabricksException("Token retrieval failed"))); + } +} diff --git a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java index cb86e9cb0..478d4a21a 100644 --- a/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java +++ b/examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java @@ -8,20 +8,26 @@ /** * Example demonstrating how to use GitHub OIDC authentication with Databricks. + * + * IMPORTANT: This example only works when running within GitHub Actions. */ public class GithubOIDCAuthExample { public static void main(String[] args) { // Create Databricks configuration with GitHub OIDC authentication + // Note: This configuration assumes the code is running in GitHub Actions DatabricksConfig config = new DatabricksConfig() - .setAuthType("github-oidc") - .setHost("https://accounts.cloud.databricks.com/") - .setAccountId("968367da-7edd-44f7-9dea-3e0b20b0ec97") - .setClientId("dd3b5d58-5a6e-45e3-a3ae-81d8f89fc6fd"); + .setAuthType("github-oidc") // Specifies GitHub OIDC as the authentication method + .setHost("") // Databricks account URL + .setAccountId("") // Your Databricks account ID + .setClientId(""); // Service Principal ID + // Initialize the Account client with the OIDC configuration AccountClient account = new AccountClient(config); try { + // Example: List all groups in the Databricks account + // This demonstrates that the OIDC authentication is working System.out.println("\nListing account groups:"); Iterable groups = account.groups().list(new ListAccountGroupsRequest()); for (Group group : groups) { From 71a8d3d7330b989f49ee8f3741b267d6963b021a Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 12 May 2025 13:40:52 +0000 Subject: [PATCH 07/13] Minor fix --- .../sdk/core/DefaultCredentialsProvider.java | 7 ++++--- .../sdk/core/oauth/GithubIDTokenSource.java | 18 +++++++++--------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 1d9607c76..a09cf35f1 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -15,10 +15,10 @@ public class DefaultCredentialsProvider implements CredentialsProvider { private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); - // List of credential providers that will be tried in sequence - private List providers = new ArrayList<>(); + /* List of credential providers that will be tried in sequence */ + private List providers; - // The currently selected authentication type + /* The currently selected authentication type */ private String authType = "default"; /** @@ -143,6 +143,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { * @param config The Databricks configuration to use for provider initialization */ private void addDefaultCredentialsProviders(DatabricksConfig config) { + providers = new ArrayList<>(); providers.add(new PatCredentialsProvider()); providers.add(new BasicCredentialsProvider()); providers.add(new OAuthM2MServicePrincipalCredentialsProvider()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java index e9d4e26c4..fe5414e9d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java @@ -15,14 +15,14 @@ * Actions environment. */ public class GithubIDTokenSource implements IDTokenSource { - // URL endpoint for requesting ID tokens from GitHub Actions + /* URL endpoint for requesting ID tokens from GitHub Actions */ private final String actionsIDTokenRequestURL; - // Authentication token required to request ID tokens from GitHub Actions + /* Authentication token required to request ID tokens from GitHub Actions */ private final String actionsIDTokenRequestToken; - // HTTP client for making requests to GitHub Actions + /* HTTP client for making requests to GitHub Actions */ private final HttpClient httpClient; - // JSON mapper for parsing response data - private final ObjectMapper mapper = new ObjectMapper(); + /* JSON mapper for parsing response data */ + private static final ObjectMapper mapper = new ObjectMapper(); /** * Constructs a new GithubIDTokenSource. @@ -100,11 +100,11 @@ public IDToken getIDToken(String audience) { throw new DatabricksException("ID token response missing 'value' field"); } - String tokenValue = jsonResp.get("value").textValue(); - if (Strings.isNullOrEmpty(tokenValue)) { + try { + String tokenValue = jsonResp.get("value").textValue(); + return new IDToken(tokenValue); + } catch (IllegalArgumentException e) { throw new DatabricksException("Received empty ID token from GitHub Actions"); } - - return new IDToken(tokenValue); } } From 01a0bc246d3db5a8ade75fc29557a86b3b76f541 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 12 May 2025 14:26:06 +0000 Subject: [PATCH 08/13] Update unit tests --- .../databricks/sdk/DatabricksAuthTest.java | 8 +- .../core/oauth/GithubIDTokenSourceTest.java | 248 +++++++++--------- 2 files changed, 136 insertions(+), 120 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java index d742b1e98..295a8f9d9 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java @@ -4,22 +4,26 @@ package com.databricks.sdk; import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.mock; import com.databricks.sdk.core.ConfigResolving; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.DummyHttpClient; +import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.utils.GitHubUtils; import com.databricks.sdk.core.utils.TestOSUtils; import java.io.File; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; +@ExtendWith(MockitoExtension.class) public class DatabricksAuthTest implements GitHubUtils, ConfigResolving { private String errorMessageBase = "default auth: cannot configure default credentials, please check https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication to configure credentials for your preferred authentication method"; private DatabricksConfig createConfigWithMockClient() { - return new DatabricksConfig().setHttpClient(new DummyHttpClient()); + return new DatabricksConfig().setHttpClient(mock(HttpClient.class)); } public DatabricksAuthTest() { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java index e53aaefa8..dabe34754 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/GithubIDTokenSourceTest.java @@ -8,19 +8,13 @@ import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ObjectNode; import java.io.IOException; import java.net.MalformedURLException; import java.net.URL; import java.util.stream.Stream; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; public class GithubIDTokenSourceTest { private static final String TEST_REQUEST_URL = "https://github.com/token"; @@ -28,124 +22,142 @@ public class GithubIDTokenSourceTest { private static final String TEST_ID_TOKEN = "test-id-token"; private static final String TEST_AUDIENCE = "test-audience"; - @Mock private static HttpClient mockHttpClient; - - private GithubIDTokenSource tokenSource; - private ObjectMapper mapper; - - @BeforeEach - void setUp() throws IOException { - MockitoAnnotations.openMocks(this); - mapper = new ObjectMapper(); - tokenSource = new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, mockHttpClient); - } - - @Test - void testSuccessfulTokenRetrieval() throws IOException { - // Prepare mock response - ObjectNode responseJson = mapper.createObjectNode(); - responseJson.put("value", TEST_ID_TOKEN); - Response mockResponse = makeResponse(responseJson.toString(), 200); - when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse); - - // Test token retrieval - IDToken token = tokenSource.getIDToken(TEST_AUDIENCE); - - assertNotNull(token); - assertEquals(TEST_ID_TOKEN, token.getValue()); - - // Verify the request was made with correct parameters - verify(mockHttpClient) - .execute( - argThat( - request -> { - return request.getMethod().equals("GET") - && request.getUrl().startsWith(TEST_REQUEST_URL) - && request.getUrl().contains("audience=" + TEST_AUDIENCE) - && request - .getHeaders() - .get("Authorization") - .equals("Bearer " + TEST_REQUEST_TOKEN); - })); - } - - @Test - void testSuccessfulTokenRetrievalWithoutAudience() throws IOException { - // Prepare mock response - ObjectNode responseJson = mapper.createObjectNode(); - responseJson.put("value", TEST_ID_TOKEN); - Response mockResponse = makeResponse(responseJson.toString(), 200); - when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse); - - // Test token retrieval without audience - IDToken token = tokenSource.getIDToken(""); - - assertNotNull(token); - assertEquals(TEST_ID_TOKEN, token.getValue()); - - // Verify the request was made with correct parameters - verify(mockHttpClient) - .execute( - argThat( - request -> { - return request.getMethod().equals("GET") - && request.getUrl().equals(TEST_REQUEST_URL) - && request - .getHeaders() - .get("Authorization") - .equals("Bearer " + TEST_REQUEST_TOKEN); - })); + private static HttpClient createHttpMock( + String responseBody, int statusCode, IOException exception) throws IOException { + HttpClient client = mock(HttpClient.class); + if (exception != null) { + when(client.execute(any(Request.class))).thenThrow(exception); + } else { + when(client.execute(any(Request.class))).thenReturn(makeResponse(responseBody, statusCode)); + } + return client; } - private static Stream provideInvalidConstructorParameters() { - return Stream.of( - Arguments.of("Missing Request URL", null, TEST_REQUEST_TOKEN, mockHttpClient), - Arguments.of("Missing Request Token", TEST_REQUEST_URL, null, mockHttpClient), - Arguments.of("Null HttpClient", TEST_REQUEST_URL, TEST_REQUEST_TOKEN, null)); - } - - @ParameterizedTest(name = "{0}") - @MethodSource("provideInvalidConstructorParameters") - void testInvalidConstructorParameters( - String testName, String requestUrl, String requestToken, HttpClient httpClient) { - GithubIDTokenSource invalidSource = - new GithubIDTokenSource(requestUrl, requestToken, httpClient); - assertThrows(DatabricksException.class, () -> invalidSource.getIDToken(TEST_AUDIENCE)); - } - - private static Stream provideHttpErrorScenarios() throws IOException { - HttpClient httpClientError = mock(HttpClient.class); - when(httpClientError.execute(any(Request.class))).thenThrow(new IOException("Network error")); - - HttpClient nonSuccessClient = mock(HttpClient.class); - when(nonSuccessClient.execute(any(Request.class))) - .thenReturn(makeResponse("Error response", 400)); - - HttpClient invalidJsonClient = mock(HttpClient.class); - when(invalidJsonClient.execute(any(Request.class))) - .thenReturn(makeResponse("Invalid json", 200)); - - HttpClient missingTokenClient = mock(HttpClient.class); - when(missingTokenClient.execute(any(Request.class))).thenReturn(makeResponse("{}", 200)); - - HttpClient emptyTokenClient = mock(HttpClient.class); - when(emptyTokenClient.execute(any(Request.class))) - .thenReturn(makeResponse("{\"value\":\"\"}", 200)); - + private static Stream provideAllTestScenarios() throws IOException { return Stream.of( - Arguments.of("HTTP Client Error", httpClientError), - Arguments.of("Non-Success Status Code", nonSuccessClient), - Arguments.of("Invalid JSON Response", invalidJsonClient), - Arguments.of("Missing Token Value", missingTokenClient), - Arguments.of("Empty Token Value", emptyTokenClient)); + // Successful token retrieval with audience + Arguments.of( + "Successful token retrieval with audience", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null), + TEST_AUDIENCE, + (java.util.function.Predicate) + request -> + request.getMethod().equals("GET") + && request.getUrl().startsWith(TEST_REQUEST_URL) + && request.getUrl().contains("audience=" + TEST_AUDIENCE) + && request + .getHeaders() + .get("Authorization") + .equals("Bearer " + TEST_REQUEST_TOKEN), + null), + // Successful token retrieval without audience + Arguments.of( + "Successful token retrieval without audience", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null), + "", + (java.util.function.Predicate) + request -> + request.getMethod().equals("GET") + && request.getUrl().equals(TEST_REQUEST_URL) + && request + .getHeaders() + .get("Authorization") + .equals("Bearer " + TEST_REQUEST_TOKEN), + null), + // Invalid constructor parameters + Arguments.of( + "Missing Request URL", + null, + TEST_REQUEST_TOKEN, + createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Missing Request Token", + TEST_REQUEST_URL, + null, + createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Null HttpClient", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + null, + TEST_AUDIENCE, + null, + DatabricksException.class), + // HTTP error scenarios + Arguments.of( + "HTTP Client Error", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock(null, 0, new IOException("Network error")), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Non-Success Status Code", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{\"error\":\"Bad Request\"}", 400, null), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Invalid JSON Response", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{invalid json}", 200, null), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Missing Token Value", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{\"other\":\"field\"}", 200, null), + TEST_AUDIENCE, + null, + DatabricksException.class), + Arguments.of( + "Empty Token Value", + TEST_REQUEST_URL, + TEST_REQUEST_TOKEN, + createHttpMock("{\"value\":\"\"}", 200, null), + TEST_AUDIENCE, + null, + DatabricksException.class)); } @ParameterizedTest(name = "{0}") - @MethodSource("provideHttpErrorScenarios") - void testHttpErrorScenarios(String testName, HttpClient httpClient) { - GithubIDTokenSource source = - new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, httpClient); - assertThrows(DatabricksException.class, () -> source.getIDToken(TEST_AUDIENCE)); + @MethodSource("provideAllTestScenarios") + void testAllScenarios( + String testName, + String requestUrl, + String requestToken, + HttpClient httpClient, + String audience, + java.util.function.Predicate requestValidator, + Class expectedException) + throws IOException { + + GithubIDTokenSource tokenSource = new GithubIDTokenSource(requestUrl, requestToken, httpClient); + + if (expectedException != null) { + assertThrows(expectedException, () -> tokenSource.getIDToken(audience)); + } else { + IDToken token = tokenSource.getIDToken(audience); + assertNotNull(token); + assertEquals(TEST_ID_TOKEN, token.getValue()); + verify(httpClient).execute(argThat(request -> requestValidator.test(request))); + } } private static Response makeResponse(String body, int status) throws MalformedURLException { From b527b0cc69b7b923f086e14ec5108dfaa7599035 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 12 May 2025 18:56:11 +0200 Subject: [PATCH 09/13] Update databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java Co-authored-by: Renaud Hartert --- .../com/databricks/sdk/core/DefaultCredentialsProvider.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index a09cf35f1..d6835f544 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -100,6 +100,9 @@ public synchronized HeaderFactory configure(DatabricksConfig config) { * @param config The Databricks configuration containing OIDC settings */ private void addOIDCCredentialsProviders(DatabricksConfig config) { + // TODO: refactor the code so that the IdTokenSources are created within the + // configure call of their corresponding CredentialsProvider. This will allow + // us to simplify the code by validating IdTokenSources when they are created. OpenIDConnectEndpoints endpoints = null; try { endpoints = config.getOidcEndpoints(); From 7384c0b0c71d0ff1e906d38d00aa7ab260169851 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 12 May 2025 16:57:17 +0000 Subject: [PATCH 10/13] Add TODO --- .../com/databricks/sdk/core/DefaultCredentialsProvider.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index d6835f544..5602665e9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -100,9 +100,9 @@ public synchronized HeaderFactory configure(DatabricksConfig config) { * @param config The Databricks configuration containing OIDC settings */ private void addOIDCCredentialsProviders(DatabricksConfig config) { - // TODO: refactor the code so that the IdTokenSources are created within the + // TODO: refactor the code so that the IdTokenSources are created within the // configure call of their corresponding CredentialsProvider. This will allow - // us to simplify the code by validating IdTokenSources when they are created. + // us to simplify the code by validating IdTokenSources when they are created. OpenIDConnectEndpoints endpoints = null; try { endpoints = config.getOidcEndpoints(); From 5bf9b3732c23bb1d28651260c64bbb7a44fc305f Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 12 May 2025 17:25:37 +0000 Subject: [PATCH 11/13] Small fix --- .../databricks/sdk/core/DefaultCredentialsProvider.java | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 5602665e9..c3fc3b1e4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -16,7 +16,7 @@ public class DefaultCredentialsProvider implements CredentialsProvider { private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); /* List of credential providers that will be tried in sequence */ - private List providers; + private List providers = new ArrayList<>(); /* The currently selected authentication type */ private String authType = "default"; @@ -145,8 +145,11 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { * * @param config The Databricks configuration to use for provider initialization */ - private void addDefaultCredentialsProviders(DatabricksConfig config) { - providers = new ArrayList<>(); + private synchronized void addDefaultCredentialsProviders(DatabricksConfig config) { + if (!providers.isEmpty()) { + return; + } + providers.add(new PatCredentialsProvider()); providers.add(new BasicCredentialsProvider()); providers.add(new OAuthM2MServicePrincipalCredentialsProvider()); From e3ca4afa42d561afb5e1826a10b1d2ebc99fd448 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 13 May 2025 08:39:59 +0000 Subject: [PATCH 12/13] Add Mock HttpClient to tests --- .../sdk/DatabricksAuthManualTest.java | 17 +++++++++++------ .../com/databricks/sdk/DatabricksAuthTest.java | 10 +++++----- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java index 9b870458a..e33ecd54f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java @@ -1,16 +1,21 @@ package com.databricks.sdk; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.mock; + import com.databricks.sdk.core.ConfigResolving; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.DummyHttpClient; +import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.utils.TestOSUtils; import java.util.Map; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class DatabricksAuthManualTest implements ConfigResolving { + private DatabricksConfig createConfigWithMockClient() { - return new DatabricksConfig().setHttpClient(new DummyHttpClient()); + HttpClient mockClient = mock(HttpClient.class); + return new DatabricksConfig().setHttpClient(mockClient); } @Test @@ -29,7 +34,7 @@ void azureCliWorkspaceHeaderPresent() { .setAzureWorkspaceResourceId(azureWorkspaceResourceId); resolveConfig(config, env); Map headers = config.authenticate(); - Assertions.assertEquals( + assertEquals( azureWorkspaceResourceId, headers.get("X-Databricks-Azure-Workspace-Resource-Id")); } @@ -49,7 +54,7 @@ void azureCliUserWithManagementAccess() { .setAzureWorkspaceResourceId(azureWorkspaceResourceId); resolveConfig(config, env); Map headers = config.authenticate(); - Assertions.assertEquals("...", headers.get("X-Databricks-Azure-SP-Management-Token")); + assertEquals("...", headers.get("X-Databricks-Azure-SP-Management-Token")); } @Test @@ -69,6 +74,6 @@ void azureCliUserNoManagementAccess() { .setAzureWorkspaceResourceId(azureWorkspaceResourceId); resolveConfig(config, env); Map headers = config.authenticate(); - Assertions.assertNull(headers.get("X-Databricks-Azure-SP-Management-Token")); + assertNull(headers.get("X-Databricks-Azure-SP-Management-Token")); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java index 295a8f9d9..31a1cdd62 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java @@ -3,7 +3,9 @@ package com.databricks.sdk; -import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.Mockito.mock; import com.databricks.sdk.core.ConfigResolving; @@ -13,17 +15,15 @@ import com.databricks.sdk.core.utils.TestOSUtils; import java.io.File; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.junit.jupiter.MockitoExtension; -@ExtendWith(MockitoExtension.class) public class DatabricksAuthTest implements GitHubUtils, ConfigResolving { private String errorMessageBase = "default auth: cannot configure default credentials, please check https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication to configure credentials for your preferred authentication method"; private DatabricksConfig createConfigWithMockClient() { - return new DatabricksConfig().setHttpClient(mock(HttpClient.class)); + HttpClient mockClient = mock(HttpClient.class); + return new DatabricksConfig().setHttpClient(mockClient); } public DatabricksAuthTest() { From e1816201ba9f669d8ae086362c7adcb40cd88b53 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 13 May 2025 09:10:14 +0000 Subject: [PATCH 13/13] Fix formatting --- .../test/java/com/databricks/sdk/DatabricksAuthManualTest.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java index e33ecd54f..288975445 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java @@ -34,8 +34,7 @@ void azureCliWorkspaceHeaderPresent() { .setAzureWorkspaceResourceId(azureWorkspaceResourceId); resolveConfig(config, env); Map headers = config.authenticate(); - assertEquals( - azureWorkspaceResourceId, headers.get("X-Databricks-Azure-Workspace-Resource-Id")); + assertEquals(azureWorkspaceResourceId, headers.get("X-Databricks-Azure-Workspace-Resource-Id")); } @Test