Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions NEXT_CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
## Release v0.48.0

### New Features and Improvements
* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([423](https://github.com/databricks/databricks-sdk-java/pull/423)).
See README.md for instructions.
* [Breaking] Users running their workflows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST`
environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods.

### Bug Fixes

Expand Down
8 changes: 4 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -116,18 +116,18 @@ Depending on the Databricks authentication method, the SDK uses the following in

### Databricks native authentication

By default, the Databricks SDK for Java initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Databricks basic (username/password) authentication (`auth_type="basic"` argument).
By default, the Databricks SDK for Java initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Databricks Workload Identity Federation (WIF) authentication using OIDC (`auth_type="github-oidc"` argument).

- For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents.
- For Databricks basic authentication, you must provide `host`, `username`, and `password` _(for AWS workspace-level operations)_; or `host`, `account_id`, `username`, and `password` _(for AWS, Azure, or GCP account-level operations)_; or their environment variable or `.databrickscfg` file field equivalents.
- For Databricks OIDC authentication, you must provide the `host`, `client_id` and `token_audience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file.

| Argument | Description | Environment variable |
|--------------|-------------|-------------------|
| `host` | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` |
| `account_id` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Only has effect when `Host` is either `https://accounts.cloud.databricks.com/` _(AWS)_, `https://accounts.azuredatabricks.net/` _(Azure)_, or `https://accounts.gcp.databricks.com/` _(GCP)_. | `DATABRICKS_ACCOUNT_ID` |
| `token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` |
| `username` | _(String)_ The Databricks username part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_USERNAME` |
| `password` | _(String)_ The Databricks password part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_PASSWORD` |
| `client_id` | _(String)_ The Databricks Service Principal Application ID. | `DATABRICKS_CLIENT_ID` |
| `token_audience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. | `TOKEN_AUDIENCE` |

For example, to use Databricks token authentication:

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,13 @@ public class DatabricksConfig {

private DatabricksEnvironment databricksEnvironment;

/**
* When using Workload Identity Federation, the audience to specify when fetching an ID token from
* the ID token supplier.
*/
@ConfigAttribute(env = "TOKEN_AUDIENCE")
private String tokenAudience;

public Environment getEnv() {
return env;
}
Expand Down Expand Up @@ -512,6 +519,15 @@ public DatabricksConfig setHttpClient(HttpClient httpClient) {
return this;
}

public String getTokenAudience() {
return tokenAudience;
}

public DatabricksConfig setTokenAudience(String tokenAudience) {
this.tokenAudience = tokenAudience;
return this;
}

public boolean isAzure() {
if (azureWorkspaceResourceId != null) {
return true;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,6 @@
package com.databricks.sdk.core;

import com.databricks.sdk.core.oauth.AzureGithubOidcCredentialsProvider;
import com.databricks.sdk.core.oauth.AzureServicePrincipalCredentialsProvider;
import com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider;
import com.databricks.sdk.core.oauth.OAuthM2MServicePrincipalCredentialsProvider;
import com.databricks.sdk.core.oauth.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
Expand All @@ -18,6 +15,7 @@ public class DefaultCredentialsProvider implements CredentialsProvider {
PatCredentialsProvider.class,
BasicCredentialsProvider.class,
OAuthM2MServicePrincipalCredentialsProvider.class,
GithubOidcCredentialsProvider.class,
AzureGithubOidcCredentialsProvider.class,
AzureServicePrincipalCredentialsProvider.class,
AzureCliCredentialsProvider.class,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ private static RefreshableTokenSource tokenSourceFor(DatabricksConfig config, St
.withClientId(config.getAzureClientId())
.withClientSecret(config.getAzureClientSecret())
.withTokenUrl(tokenUrl)
.withEndpointParameters(endpointParams)
.withEndpointParametersSupplier(() -> endpointParams)
.withAuthParameterPosition(AuthParameterPosition.BODY)
.build();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.databricks.sdk.core.commons.CommonsHttpClient;
import com.databricks.sdk.core.http.HttpClient;
import java.util.*;
import java.util.function.Supplier;

/**
* An implementation of RefreshableTokenSource implementing the client_credentials OAuth grant type.
Expand All @@ -18,7 +19,11 @@ public static class Builder {
private String clientSecret;
private String tokenUrl;
private HttpClient hc = new CommonsHttpClient.Builder().withTimeoutSeconds(30).build();
private Map<String, String> endpointParams = Collections.emptyMap();

// Endpoint parameters can include tokens with expiration which
// may need to be refreshed. This supplier will be called each time
// the credentials are refreshed.
private Supplier<Map<String, String>> endpointParamsSupplier = null;
private List<String> scopes = Collections.emptyList();
private AuthParameterPosition position = AuthParameterPosition.BODY;

Expand All @@ -32,13 +37,14 @@ public Builder withClientSecret(String clientSecret) {
return this;
}

public Builder withTokenUrl(String tokenUrl) {
this.tokenUrl = tokenUrl;
public Builder withEndpointParametersSupplier(
Supplier<Map<String, String>> endpointParamsSupplier) {
this.endpointParamsSupplier = endpointParamsSupplier;
return this;
}

public Builder withEndpointParameters(Map<String, String> params) {
this.endpointParams = params;
public Builder withTokenUrl(String tokenUrl) {
this.tokenUrl = tokenUrl;
return this;
}

Expand All @@ -59,34 +65,33 @@ public Builder withHttpClient(HttpClient hc) {

public ClientCredentials build() {
Objects.requireNonNull(this.clientId, "clientId must be specified");
Objects.requireNonNull(this.clientSecret, "clientSecret must be specified");
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not required for WIF/OIDC

Objects.requireNonNull(this.tokenUrl, "tokenUrl must be specified");
return new ClientCredentials(
hc, clientId, clientSecret, tokenUrl, endpointParams, scopes, position);
hc, clientId, clientSecret, tokenUrl, endpointParamsSupplier, scopes, position);
}
}

private HttpClient hc;
private String clientId;
private String clientSecret;
private String tokenUrl;
private Map<String, String> endpointParams;
private List<String> scopes;
private AuthParameterPosition position;
private Supplier<Map<String, String>> endpointParamsSupplier;

private ClientCredentials(
HttpClient hc,
String clientId,
String clientSecret,
String tokenUrl,
Map<String, String> endpointParams,
Supplier<Map<String, String>> endpointParamsSupplier,
List<String> scopes,
AuthParameterPosition position) {
this.hc = hc;
this.clientId = clientId;
this.clientSecret = clientSecret;
this.tokenUrl = tokenUrl;
this.endpointParams = endpointParams;
this.endpointParamsSupplier = endpointParamsSupplier;
this.scopes = scopes;
this.position = position;
}
Expand All @@ -98,8 +103,8 @@ protected Token refresh() {
if (scopes != null) {
params.put("scope", String.join(" ", scopes));
}
if (endpointParams != null) {
params.putAll(endpointParams);
if (endpointParamsSupplier != null) {
params.putAll(endpointParamsSupplier.get());
}
return retrieveToken(hc, clientId, clientSecret, tokenUrl, params, new HashMap<>(), position);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
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 java.io.IOException;

public class GitHubOidcTokenSupplier {

private final ObjectMapper mapper = new ObjectMapper();
private final HttpClient httpClient;
private final String idTokenRequestUrl;
private final String idTokenRequestToken;
private final String tokenAudience;

public GitHubOidcTokenSupplier(
HttpClient httpClient,
String idTokenRequestUrl,
String idTokenRequestToken,
String tokenAudience) {
this.httpClient = httpClient;
this.idTokenRequestUrl = idTokenRequestUrl;
this.idTokenRequestToken = idTokenRequestToken;
this.tokenAudience = tokenAudience;
}

/** Checks if the required parameters are present to request a GitHub's OIDC token. */
public Boolean enabled() {
return idTokenRequestUrl != null && idTokenRequestToken != null;
}

/**
* Requests a GitHub's OIDC token.
*
* @return A GitHub OIDC token.
*/
public String getOidcToken() {
if (!enabled()) {
throw new DatabricksException("Failed to request ID token: missing required parameters");
}

String requestUrl = idTokenRequestUrl;
if (tokenAudience != null) {
requestUrl += "&audience=" + tokenAudience;
}

Request req =
new Request("GET", requestUrl).withHeader("Authorization", "Bearer " + idTokenRequestToken);

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());
}

return jsonResp.get("value").textValue();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
package com.databricks.sdk.core.oauth;

import com.databricks.sdk.core.CredentialsProvider;
import com.databricks.sdk.core.DatabricksConfig;
import com.databricks.sdk.core.DatabricksException;
import com.databricks.sdk.core.HeaderFactory;
import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
* GithubOidcCredentialsProvider uses a Token Supplier to get a GitHub OIDC JWT Token and exchanges
* it for a Databricks Token.
*/
public class GithubOidcCredentialsProvider implements CredentialsProvider {

@Override
public String authType() {
return "github-oidc";
}

@Override
public HeaderFactory configure(DatabricksConfig config) throws DatabricksException {
GitHubOidcTokenSupplier idTokenProvider =
new GitHubOidcTokenSupplier(
config.getHttpClient(),
config.getActionsIdTokenRequestUrl(),
config.getActionsIdTokenRequestToken(),
config.getTokenAudience());

if (!idTokenProvider.enabled() || config.getHost() == null || config.getClientId() == null) {
return null;
}

String endpointUrl;

try {
endpointUrl = config.getOidcEndpoints().getTokenEndpoint();
} catch (IOException e) {
throw new DatabricksException("Unable to fetch OIDC endpoint: " + e.getMessage(), e);
}

ClientCredentials clientCredentials =
new ClientCredentials.Builder()
.withHttpClient(config.getHttpClient())
.withClientId(config.getClientId())
.withTokenUrl(endpointUrl)
.withScopes(Collections.singletonList("all-apis"))
.withAuthParameterPosition(AuthParameterPosition.HEADER)
.withEndpointParametersSupplier(
() ->
new ImmutableMap.Builder<String, String>()
.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt")
.put("subject_token", idTokenProvider.getOidcToken())
.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange")
.build())
.build();

return () -> {
Map<String, String> headers = new HashMap<>();
headers.put("Authorization", "Bearer " + clientCredentials.getToken().getAccessToken());
return headers;
};
}
}
Loading
Loading