|
2 | 2 |
|
3 | 3 | import com.databricks.sdk.core.oauth.*; |
4 | 4 | import java.util.ArrayList; |
5 | | -import java.util.Arrays; |
6 | 5 | import java.util.List; |
7 | 6 | import org.slf4j.Logger; |
8 | 7 | import org.slf4j.LoggerFactory; |
9 | 8 |
|
| 9 | +/** |
| 10 | + * The DefaultCredentialsProvider is the primary authentication handler for the Databricks SDK. It |
| 11 | + * implements a chain of responsibility pattern to manage multiple authentication methods, including |
| 12 | + * Personal Access Tokens (PAT), OAuth, Azure, Google, and OpenID Connect (OIDC). The provider |
| 13 | + * attempts each authentication method in sequence until a valid credential is obtained. |
| 14 | + */ |
10 | 15 | public class DefaultCredentialsProvider implements CredentialsProvider { |
11 | 16 | private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class); |
12 | 17 |
|
13 | | - private static final List<Class<?>> providerClasses = |
14 | | - Arrays.asList( |
15 | | - PatCredentialsProvider.class, |
16 | | - BasicCredentialsProvider.class, |
17 | | - OAuthM2MServicePrincipalCredentialsProvider.class, |
18 | | - GithubOidcCredentialsProvider.class, |
19 | | - AzureGithubOidcCredentialsProvider.class, |
20 | | - AzureServicePrincipalCredentialsProvider.class, |
21 | | - AzureCliCredentialsProvider.class, |
22 | | - ExternalBrowserCredentialsProvider.class, |
23 | | - DatabricksCliCredentialsProvider.class, |
24 | | - NotebookNativeCredentialsProvider.class, |
25 | | - GoogleCredentialsCredentialsProvider.class, |
26 | | - GoogleIdCredentialsProvider.class); |
27 | | - |
28 | | - private final List<CredentialsProvider> providers; |
| 18 | + /* List of credential providers that will be tried in sequence */ |
| 19 | + private List<CredentialsProvider> providers = new ArrayList<>(); |
29 | 20 |
|
| 21 | + /* The currently selected authentication type */ |
30 | 22 | private String authType = "default"; |
31 | 23 |
|
32 | | - public String authType() { |
33 | | - return authType; |
34 | | - } |
| 24 | + /** |
| 25 | + * Internal class to associate an ID token source with a name for identification purposes. Used |
| 26 | + * primarily for OIDC (OpenID Connect) authentication flows. |
| 27 | + */ |
| 28 | + private static class NamedIDTokenSource { |
| 29 | + private final String name; |
| 30 | + private final IDTokenSource idTokenSource; |
35 | 31 |
|
36 | | - public DefaultCredentialsProvider() { |
37 | | - providers = new ArrayList<>(); |
38 | | - for (Class<?> clazz : providerClasses) { |
39 | | - try { |
40 | | - providers.add((CredentialsProvider) clazz.newInstance()); |
41 | | - } catch (NoClassDefFoundError | InstantiationException | IllegalAccessException e) { |
42 | | - LOG.warn( |
43 | | - "Failed to instantiate credentials provider: " |
44 | | - + clazz.getName() |
45 | | - + ", skipping. Cause: " |
46 | | - + e.getClass().getCanonicalName() |
47 | | - + ": " |
48 | | - + e.getMessage()); |
49 | | - } |
| 32 | + public NamedIDTokenSource(String name, IDTokenSource idTokenSource) { |
| 33 | + this.name = name; |
| 34 | + this.idTokenSource = idTokenSource; |
| 35 | + } |
| 36 | + |
| 37 | + public String getName() { |
| 38 | + return name; |
50 | 39 | } |
| 40 | + |
| 41 | + public IDTokenSource getIdTokenSource() { |
| 42 | + return idTokenSource; |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + public DefaultCredentialsProvider() {} |
| 47 | + |
| 48 | + /** |
| 49 | + * Returns the current authentication type being used |
| 50 | + * |
| 51 | + * @return String representing the authentication type |
| 52 | + */ |
| 53 | + public String authType() { |
| 54 | + return authType; |
51 | 55 | } |
52 | 56 |
|
| 57 | + /** |
| 58 | + * Configures the credentials provider with the given Databricks configuration. This method tries |
| 59 | + * each available credential provider in sequence until one succeeds. |
| 60 | + * |
| 61 | + * @param config The Databricks configuration containing authentication details |
| 62 | + * @return HeaderFactory for making authenticated requests |
| 63 | + * @throws DatabricksException if no valid credentials can be configured |
| 64 | + */ |
53 | 65 | @Override |
54 | 66 | public synchronized HeaderFactory configure(DatabricksConfig config) { |
| 67 | + addDefaultCredentialsProviders(config); |
55 | 68 | for (CredentialsProvider provider : providers) { |
56 | 69 | if (config.getAuthType() != null |
57 | 70 | && !config.getAuthType().isEmpty() |
@@ -80,4 +93,77 @@ public synchronized HeaderFactory configure(DatabricksConfig config) { |
80 | 93 | + authFlowUrl |
81 | 94 | + " to configure credentials for your preferred authentication method"); |
82 | 95 | } |
| 96 | + |
| 97 | + /** |
| 98 | + * Adds OpenID Connect (OIDC) based credential providers to the list of available providers. |
| 99 | + * |
| 100 | + * @param config The Databricks configuration containing OIDC settings |
| 101 | + */ |
| 102 | + private void addOIDCCredentialsProviders(DatabricksConfig config) { |
| 103 | + // TODO: refactor the code so that the IdTokenSources are created within the |
| 104 | + // configure call of their corresponding CredentialsProvider. This will allow |
| 105 | + // us to simplify the code by validating IdTokenSources when they are created. |
| 106 | + OpenIDConnectEndpoints endpoints = null; |
| 107 | + try { |
| 108 | + endpoints = config.getOidcEndpoints(); |
| 109 | + } catch (Exception e) { |
| 110 | + LOG.warn("Failed to get OpenID Connect endpoints", e); |
| 111 | + } |
| 112 | + |
| 113 | + List<NamedIDTokenSource> namedIdTokenSources = new ArrayList<>(); |
| 114 | + namedIdTokenSources.add( |
| 115 | + new NamedIDTokenSource( |
| 116 | + "github-oidc", |
| 117 | + new GithubIDTokenSource( |
| 118 | + config.getActionsIdTokenRequestUrl(), |
| 119 | + config.getActionsIdTokenRequestToken(), |
| 120 | + config.getHttpClient()))); |
| 121 | + // Add new IDTokenSources and ID providers here. Example: |
| 122 | + // namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...))); |
| 123 | + |
| 124 | + // Configure OAuth token sources for each ID token source |
| 125 | + for (NamedIDTokenSource namedIdTokenSource : namedIdTokenSources) { |
| 126 | + DatabricksOAuthTokenSource oauthTokenSource = |
| 127 | + new DatabricksOAuthTokenSource.Builder( |
| 128 | + config.getClientId(), |
| 129 | + config.getHost(), |
| 130 | + endpoints, |
| 131 | + namedIdTokenSource.getIdTokenSource(), |
| 132 | + config.getHttpClient()) |
| 133 | + .audience(config.getTokenAudience()) |
| 134 | + .accountId(config.isAccountClient() ? config.getAccountId() : null) |
| 135 | + .build(); |
| 136 | + |
| 137 | + providers.add( |
| 138 | + new TokenSourceCredentialsProvider(oauthTokenSource, namedIdTokenSource.getName())); |
| 139 | + } |
| 140 | + } |
| 141 | + |
| 142 | + /** |
| 143 | + * Initializes all available credential providers in the preferred order. The order of providers |
| 144 | + * determines the authentication fallback sequence. |
| 145 | + * |
| 146 | + * @param config The Databricks configuration to use for provider initialization |
| 147 | + */ |
| 148 | + private synchronized void addDefaultCredentialsProviders(DatabricksConfig config) { |
| 149 | + if (!providers.isEmpty()) { |
| 150 | + return; |
| 151 | + } |
| 152 | + |
| 153 | + providers.add(new PatCredentialsProvider()); |
| 154 | + providers.add(new BasicCredentialsProvider()); |
| 155 | + providers.add(new OAuthM2MServicePrincipalCredentialsProvider()); |
| 156 | + |
| 157 | + // Add OIDC-based providers |
| 158 | + addOIDCCredentialsProviders(config); |
| 159 | + |
| 160 | + providers.add(new AzureGithubOidcCredentialsProvider()); |
| 161 | + providers.add(new AzureServicePrincipalCredentialsProvider()); |
| 162 | + providers.add(new AzureCliCredentialsProvider()); |
| 163 | + providers.add(new ExternalBrowserCredentialsProvider()); |
| 164 | + providers.add(new DatabricksCliCredentialsProvider()); |
| 165 | + providers.add(new NotebookNativeCredentialsProvider()); |
| 166 | + providers.add(new GoogleCredentialsCredentialsProvider()); |
| 167 | + providers.add(new GoogleIdCredentialsProvider()); |
| 168 | + } |
83 | 169 | } |
0 commit comments