Skip to content

Commit 3afdfa8

Browse files
emmyzhou-dbemmyzhou-db
committed
Implement direct dataplane access
1 parent 78501a5 commit 3afdfa8

File tree

6 files changed

+827
-0
lines changed

6 files changed

+827
-0
lines changed

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
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 DatabricksOAuthTokenSource cpTokenSource;
18+
private final ConcurrentHashMap<TokenSourceKey, EndpointTokenSource> sourcesCache;
19+
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 DatabricksOAuthTokenSource} for control plane tokens.
66+
* @throws NullPointerException if either parameter is null
67+
*/
68+
public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource) {
69+
this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
70+
this.cpTokenSource =
71+
Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
72+
this.sourcesCache = new ConcurrentHashMap<>();
73+
}
74+
75+
/**
76+
* Retrieves a token for the specified endpoint and authorization details. It uses a cached {@link
77+
* EndpointTokenSource} if available, otherwise creates and caches a new one.
78+
*
79+
* @param endpoint The target data plane service endpoint.
80+
* @param authDetails Authorization details for the endpoint.
81+
* @return The dataplane {@link Token}.
82+
* @throws NullPointerException if either parameter is null
83+
* @throws IllegalArgumentException if either parameter is empty
84+
*/
85+
public Token getToken(String endpoint, String authDetails) {
86+
Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null");
87+
Objects.requireNonNull(authDetails, "Authorization details cannot be null");
88+
if (endpoint.isEmpty()) {
89+
throw new IllegalArgumentException("Data plane endpoint URL cannot be empty");
90+
}
91+
if (authDetails.isEmpty()) {
92+
throw new IllegalArgumentException("Authorization details cannot be empty");
93+
}
94+
TokenSourceKey key = new TokenSourceKey(endpoint, authDetails);
95+
96+
EndpointTokenSource specificSource =
97+
sourcesCache.computeIfAbsent(
98+
key, k -> new EndpointTokenSource(this.cpTokenSource, k.authDetails, this.httpClient));
99+
100+
return specificSource.getToken();
101+
}
102+
}

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

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

‎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)