Skip to content

Commit a8187d6

Browse files
emmyzhou-dbemmyzhou-db
authored
Add DataPlaneTokenSource and EndpointTokenSource (#449)
## What changes are proposed in this pull request? # Direct Dataplane Access for Databricks Java SDK This PR adds the required classes for direct dataplane access in the Databricks Java SDK. ### `DataPlaneTokenSource` - Manages and caches `EndpointTokenSource` instances. - Ensures a single `EndpointTokenSource` per unique endpoint and set of authorization details. - Provides a thread-safe cache to avoid redundant token source creation. ### `EndpointTokenSource` - Handles the OAuth token exchange process. - Exchanges a control plane token for a dataplane token. - Caches the dataplane token and manages its refresh and expiry. ### `TokenEndpointClient` - Utility for making HTTP requests to OAuth token endpoints. - Handles request formatting, response parsing, and error handling. ## How is this tested? - Unit tests have been added for all the aforementioned classes. NO_CHANGELOG=true
1 parent 3d852c5 commit a8187d6

File tree

6 files changed

+950
-0
lines changed

6 files changed

+950
-0
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.http.HttpClient;
4+
import java.util.Objects;
5+
import java.util.concurrent.ConcurrentHashMap;
6+
7+
/**
8+
* Manages and provides Databricks data plane tokens. This class is responsible for acquiring and
9+
* caching OAuth tokens that are specific to a particular Databricks data plane service endpoint and
10+
* a set of authorization details. It utilizes a {@link DatabricksOAuthTokenSource} for obtaining
11+
* control plane tokens, which may then be exchanged or used to authorize requests for data plane
12+
* tokens. Cached {@link EndpointTokenSource} instances are used to efficiently reuse tokens for
13+
* repeated requests to the same endpoint with the same authorization context.
14+
*/
15+
public class DataPlaneTokenSource {
16+
private final HttpClient httpClient;
17+
private final TokenSource cpTokenSource;
18+
private final String host;
19+
private final ConcurrentHashMap<TokenSourceKey, EndpointTokenSource> sourcesCache;
20+
/**
21+
* Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This
22+
* is a value object that uniquely identifies a token source configuration.
23+
*/
24+
private static final class TokenSourceKey {
25+
/** The target service endpoint URL. */
26+
private final String endpoint;
27+
28+
/** Specific authorization details for the endpoint. */
29+
private final String authDetails;
30+
31+
/**
32+
* Constructs a TokenSourceKey.
33+
*
34+
* @param endpoint The target service endpoint URL.
35+
* @param authDetails Specific authorization details.
36+
*/
37+
public TokenSourceKey(String endpoint, String authDetails) {
38+
this.endpoint = endpoint;
39+
this.authDetails = authDetails;
40+
}
41+
42+
@Override
43+
public boolean equals(Object o) {
44+
if (this == o) {
45+
return true;
46+
}
47+
if (o == null || getClass() != o.getClass()) {
48+
return false;
49+
}
50+
TokenSourceKey that = (TokenSourceKey) o;
51+
return Objects.equals(endpoint, that.endpoint)
52+
&& Objects.equals(authDetails, that.authDetails);
53+
}
54+
55+
@Override
56+
public int hashCode() {
57+
return Objects.hash(endpoint, authDetails);
58+
}
59+
}
60+
61+
/**
62+
* Constructs a DataPlaneTokenSource.
63+
*
64+
* @param httpClient The {@link HttpClient} for token requests.
65+
* @param cpTokenSource The {@link TokenSource} for control plane tokens.
66+
* @param host The host for the token exchange request.
67+
* @throws NullPointerException if any parameter is null.
68+
* @throws IllegalArgumentException if the host is empty.
69+
*/
70+
public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) {
71+
this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
72+
this.cpTokenSource =
73+
Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
74+
this.host = Objects.requireNonNull(host, "Host cannot be null");
75+
76+
if (host.isEmpty()) {
77+
throw new IllegalArgumentException("Host cannot be empty");
78+
}
79+
this.sourcesCache = new ConcurrentHashMap<>();
80+
}
81+
82+
/**
83+
* Retrieves a token for the specified endpoint and authorization details. It uses a cached {@link
84+
* EndpointTokenSource} if available, otherwise creates and caches a new one.
85+
*
86+
* @param endpoint The target data plane service endpoint.
87+
* @param authDetails Authorization details for the endpoint.
88+
* @return The dataplane {@link Token}.
89+
* @throws NullPointerException if either parameter is null.
90+
* @throws IllegalArgumentException if either parameter is empty.
91+
* @throws DatabricksException if the token request fails.
92+
*/
93+
public Token getToken(String endpoint, String authDetails) {
94+
Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null");
95+
Objects.requireNonNull(authDetails, "Authorization details cannot be null");
96+
97+
if (endpoint.isEmpty()) {
98+
throw new IllegalArgumentException("Data plane endpoint URL cannot be empty");
99+
}
100+
if (authDetails.isEmpty()) {
101+
throw new IllegalArgumentException("Authorization details cannot be empty");
102+
}
103+
104+
TokenSourceKey key = new TokenSourceKey(endpoint, authDetails);
105+
106+
EndpointTokenSource specificSource =
107+
sourcesCache.computeIfAbsent(
108+
key,
109+
k ->
110+
new EndpointTokenSource(
111+
this.cpTokenSource, k.authDetails, this.httpClient, this.host));
112+
113+
return specificSource.getToken();
114+
}
115+
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksException;
4+
import com.databricks.sdk.core.http.HttpClient;
5+
import java.time.LocalDateTime;
6+
import java.util.HashMap;
7+
import java.util.Map;
8+
import java.util.Objects;
9+
import org.slf4j.Logger;
10+
import org.slf4j.LoggerFactory;
11+
12+
/**
13+
* Represents a token source that exchanges a control plane token for an endpoint-specific dataplane
14+
* token. It utilizes an underlying {@link TokenSource} to obtain the initial control plane token.
15+
*/
16+
public class EndpointTokenSource extends RefreshableTokenSource {
17+
private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class);
18+
private static final String JWT_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer";
19+
private static final String GRANT_TYPE_PARAM = "grant_type";
20+
private static final String AUTHORIZATION_DETAILS_PARAM = "authorization_details";
21+
private static final String ASSERTION_PARAM = "assertion";
22+
private static final String TOKEN_ENDPOINT = "/oidc/v1/token";
23+
24+
private final TokenSource cpTokenSource;
25+
private final String authDetails;
26+
private final HttpClient httpClient;
27+
private final String host;
28+
29+
/**
30+
* Constructs a new EndpointTokenSource.
31+
*
32+
* @param cpTokenSource The {@link TokenSource} used to obtain the control plane token.
33+
* @param authDetails The authorization details required for the token exchange.
34+
* @param httpClient The {@link HttpClient} used to make the token exchange request.
35+
* @param host The host for the token exchange request.
36+
* @throws IllegalArgumentException if authDetails is empty or host is empty.
37+
* @throws NullPointerException if any of the parameters are null.
38+
*/
39+
public EndpointTokenSource(
40+
TokenSource cpTokenSource, String authDetails, HttpClient httpClient, String host) {
41+
this.cpTokenSource =
42+
Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
43+
this.authDetails = Objects.requireNonNull(authDetails, "Authorization details cannot be null");
44+
this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
45+
this.host = Objects.requireNonNull(host, "Host cannot be null");
46+
47+
if (authDetails.isEmpty()) {
48+
throw new IllegalArgumentException("Authorization details cannot be empty");
49+
}
50+
if (host.isEmpty()) {
51+
throw new IllegalArgumentException("Host cannot be empty");
52+
}
53+
}
54+
55+
/**
56+
* Fetches an endpoint-specific dataplane token by exchanging a control plane token.
57+
*
58+
* <p>This method first obtains a control plane token from the configured {@code cpTokenSource}.
59+
* It then uses this token as an assertion along with the provided {@code authDetails} to request
60+
* a new, more scoped dataplane token from the Databricks OAuth token endpoint ({@value
61+
* #TOKEN_ENDPOINT}).
62+
*
63+
* @return A new {@link Token} containing the exchanged dataplane access token, its type, any
64+
* accompanying refresh token, and its expiry time.
65+
* @throws DatabricksException if the token exchange with the OAuth endpoint fails.
66+
* @throws IllegalArgumentException if the token endpoint url is empty.
67+
* @throws NullPointerException if any of the parameters are null.
68+
*/
69+
@Override
70+
protected Token refresh() {
71+
Token cpToken = cpTokenSource.getToken();
72+
Map<String, String> params = new HashMap<>();
73+
params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE);
74+
params.put(AUTHORIZATION_DETAILS_PARAM, authDetails);
75+
params.put(ASSERTION_PARAM, cpToken.getAccessToken());
76+
77+
OAuthResponse oauthResponse;
78+
try {
79+
oauthResponse =
80+
TokenEndpointClient.requestToken(this.httpClient, this.host + TOKEN_ENDPOINT, params);
81+
} catch (DatabricksException | IllegalArgumentException | NullPointerException e) {
82+
LOG.error(
83+
"Failed to exchange control plane token for dataplane token at endpoint {}: {}",
84+
TOKEN_ENDPOINT,
85+
e.getMessage(),
86+
e);
87+
throw e;
88+
}
89+
90+
LocalDateTime expiry = LocalDateTime.now().plusSeconds(oauthResponse.getExpiresIn());
91+
return new Token(
92+
oauthResponse.getAccessToken(),
93+
oauthResponse.getTokenType(),
94+
oauthResponse.getRefreshToken(),
95+
expiry);
96+
}
97+
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksException;
4+
import com.databricks.sdk.core.http.FormRequest;
5+
import com.databricks.sdk.core.http.HttpClient;
6+
import com.databricks.sdk.core.http.Response;
7+
import com.fasterxml.jackson.databind.ObjectMapper;
8+
import java.io.IOException;
9+
import java.util.Map;
10+
import java.util.Objects;
11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
13+
14+
/**
15+
* Client for interacting with an OAuth token endpoint.
16+
*
17+
* <p>This class provides a method to request an OAuth token from a specified token endpoint URL
18+
* using the provided HTTP client and request parameters. It handles the HTTP request and parses the
19+
* JSON response into an {@link OAuthResponse} object.
20+
*/
21+
public final class TokenEndpointClient {
22+
private static final Logger LOG = LoggerFactory.getLogger(TokenEndpointClient.class);
23+
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
24+
25+
private TokenEndpointClient() {}
26+
27+
/**
28+
* Requests an OAuth token from the specified token endpoint.
29+
*
30+
* @param httpClient The {@link HttpClient} to use for making the request.
31+
* @param tokenEndpointUrl The URL of the token endpoint.
32+
* @param params A map of parameters to include in the token request.
33+
* @return An {@link OAuthResponse} containing the token information.
34+
* @throws DatabricksException if an error occurs during the token request or response parsing.
35+
* @throws IllegalArgumentException if the token endpoint URL is empty.
36+
* @throws NullPointerException if any of the parameters are null.
37+
*/
38+
public static OAuthResponse requestToken(
39+
HttpClient httpClient, String tokenEndpointUrl, Map<String, String> params)
40+
throws DatabricksException {
41+
Objects.requireNonNull(httpClient, "HttpClient cannot be null");
42+
Objects.requireNonNull(params, "Request parameters map cannot be null");
43+
Objects.requireNonNull(tokenEndpointUrl, "Token endpoint URL cannot be null");
44+
45+
if (tokenEndpointUrl.isEmpty()) {
46+
throw new IllegalArgumentException("Token endpoint URL cannot be empty");
47+
}
48+
49+
Response rawResponse;
50+
try {
51+
LOG.debug("Requesting token from endpoint: {}", tokenEndpointUrl);
52+
rawResponse = httpClient.execute(new FormRequest(tokenEndpointUrl, params));
53+
} catch (IOException e) {
54+
LOG.error("Failed to request token from {}: {}", tokenEndpointUrl, e.getMessage(), e);
55+
throw new DatabricksException(
56+
String.format("Failed to request token from %s: %s", tokenEndpointUrl, e.getMessage()),
57+
e);
58+
}
59+
60+
OAuthResponse response;
61+
try {
62+
response = OBJECT_MAPPER.readValue(rawResponse.getBody(), OAuthResponse.class);
63+
} catch (IOException e) {
64+
LOG.error(
65+
"Failed to parse OAuth response from token endpoint {}: {}",
66+
tokenEndpointUrl,
67+
e.getMessage(),
68+
e);
69+
throw new DatabricksException(
70+
String.format(
71+
"Failed to parse OAuth response from token endpoint %s: %s",
72+
tokenEndpointUrl, e.getMessage()),
73+
e);
74+
}
75+
76+
if (response.getErrorCode() != null) {
77+
String errorSummary =
78+
response.getErrorSummary() != null ? response.getErrorSummary() : "No summary provided.";
79+
LOG.error(
80+
"Token request to {} failed with error: {} - {}",
81+
tokenEndpointUrl,
82+
response.getErrorCode(),
83+
errorSummary);
84+
throw new DatabricksException(
85+
String.format(
86+
"Token request failed with error: %s - %s", response.getErrorCode(), errorSummary));
87+
}
88+
LOG.debug("Successfully obtained token response from {}", tokenEndpointUrl);
89+
return response;
90+
}
91+
}

0 commit comments

Comments
 (0)