Skip to content

Commit 79e360d

Browse files
emmyzhou-dbemmyzhou-db
committed
Add host as a field to token sources
1 parent 3afdfa8 commit 79e360d

File tree

4 files changed

+101
-36
lines changed

4 files changed

+101
-36
lines changed

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

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,9 @@
1414
*/
1515
public class DataPlaneTokenSource {
1616
private final HttpClient httpClient;
17-
private final DatabricksOAuthTokenSource cpTokenSource;
17+
private final TokenSource cpTokenSource;
18+
private final String host;
1819
private final ConcurrentHashMap<TokenSourceKey, EndpointTokenSource> sourcesCache;
19-
2020
/**
2121
* Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This
2222
* is a value object that uniquely identifies a token source configuration.
@@ -62,13 +62,19 @@ public int hashCode() {
6262
* Constructs a DataPlaneTokenSource.
6363
*
6464
* @param httpClient The {@link HttpClient} for token requests.
65-
* @param cpTokenSource The {@link DatabricksOAuthTokenSource} for control plane tokens.
65+
* @param cpTokenSource The {@link TokenSource} for control plane tokens.
66+
* @param host The host for the token exchange request.
6667
* @throws NullPointerException if either parameter is null
6768
*/
68-
public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource) {
69+
public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) {
6970
this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
7071
this.cpTokenSource =
7172
Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
73+
this.host = Objects.requireNonNull(host, "Host cannot be null");
74+
75+
if (host.isEmpty()) {
76+
throw new IllegalArgumentException("Host cannot be empty");
77+
}
7278
this.sourcesCache = new ConcurrentHashMap<>();
7379
}
7480

@@ -85,17 +91,22 @@ public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cp
8591
public Token getToken(String endpoint, String authDetails) {
8692
Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null");
8793
Objects.requireNonNull(authDetails, "Authorization details cannot be null");
94+
8895
if (endpoint.isEmpty()) {
8996
throw new IllegalArgumentException("Data plane endpoint URL cannot be empty");
9097
}
9198
if (authDetails.isEmpty()) {
9299
throw new IllegalArgumentException("Authorization details cannot be empty");
93100
}
101+
94102
TokenSourceKey key = new TokenSourceKey(endpoint, authDetails);
95103

96104
EndpointTokenSource specificSource =
97105
sourcesCache.computeIfAbsent(
98-
key, k -> new EndpointTokenSource(this.cpTokenSource, k.authDetails, this.httpClient));
106+
key,
107+
k ->
108+
new EndpointTokenSource(
109+
this.cpTokenSource, k.authDetails, this.httpClient, this.host));
99110

100111
return specificSource.getToken();
101112
}

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

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,7 @@
1111

1212
/**
1313
* 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.
14+
* token. It utilizes an underlying {@link TokenSource} to obtain the initial control plane token.
1615
*/
1716
public class EndpointTokenSource extends RefreshableTokenSource {
1817
private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class);
@@ -22,29 +21,35 @@ public class EndpointTokenSource extends RefreshableTokenSource {
2221
private static final String ASSERTION_PARAM = "assertion";
2322
private static final String TOKEN_ENDPOINT = "/oidc/v1/token";
2423

25-
private final DatabricksOAuthTokenSource cpTokenSource;
24+
private final TokenSource cpTokenSource;
2625
private final String authDetails;
2726
private final HttpClient httpClient;
27+
private final String host;
2828

2929
/**
3030
* Constructs a new EndpointTokenSource.
3131
*
32-
* @param cpTokenSource The {@link DatabricksOAuthTokenSource} used to obtain the control plane
33-
* token.
32+
* @param cpTokenSource The {@link TokenSource} used to obtain the control plane token.
3433
* @param authDetails The authorization details required for the token exchange.
3534
* @param httpClient The {@link HttpClient} used to make the token exchange request.
36-
* @throws IllegalArgumentException if authDetails is empty.
35+
* @param host The host for the token exchange request.
36+
* @throws IllegalArgumentException if authDetails is empty or host is empty.
3737
* @throws NullPointerException if any of the parameters are null.
3838
*/
3939
public EndpointTokenSource(
40-
DatabricksOAuthTokenSource cpTokenSource, String authDetails, HttpClient httpClient) {
40+
TokenSource cpTokenSource, String authDetails, HttpClient httpClient, String host) {
4141
this.cpTokenSource =
4242
Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
4343
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+
4447
if (authDetails.isEmpty()) {
4548
throw new IllegalArgumentException("Authorization details cannot be empty");
4649
}
47-
this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
50+
if (host.isEmpty()) {
51+
throw new IllegalArgumentException("Host cannot be empty");
52+
}
4853
}
4954

5055
/**
@@ -64,15 +69,15 @@ public EndpointTokenSource(
6469
@Override
6570
protected Token refresh() {
6671
Token cpToken = cpTokenSource.getToken();
67-
6872
Map<String, String> params = new HashMap<>();
6973
params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE);
7074
params.put(AUTHORIZATION_DETAILS_PARAM, authDetails);
7175
params.put(ASSERTION_PARAM, cpToken.getAccessToken());
7276

7377
OAuthResponse oauthResponse;
7478
try {
75-
oauthResponse = TokenEndpointClient.requestToken(this.httpClient, TOKEN_ENDPOINT, params);
79+
oauthResponse =
80+
TokenEndpointClient.requestToken(this.httpClient, this.host + TOKEN_ENDPOINT, params);
7681
} catch (DatabricksException | IllegalArgumentException | NullPointerException e) {
7782
LOG.error(
7883
"Failed to exchange control plane token for dataplane token at endpoint {}: {}",

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public class DataPlaneTokenSourceTest {
2323
private static final String TEST_TOKEN_TYPE = "Bearer";
2424
private static final String TEST_REFRESH_TOKEN = "refresh-token";
2525
private static final int TEST_EXPIRES_IN = 3600;
26+
private static final String TEST_HOST = "https://test.databricks.com";
2627

2728
private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Exception {
2829
// Mock DatabricksOAuthTokenSource for control plane token
@@ -31,7 +32,6 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
3132
DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class);
3233
when(mockCpTokenSource.getToken()).thenReturn(cpToken);
3334

34-
// --- Mock HttpClient for different scenarios ---
3535
// Success JSON for endpoint1/auth1
3636
String successJson1 =
3737
"{"
@@ -56,7 +56,6 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
5656
when(mockSuccessClient2.execute(any()))
5757
.thenReturn(new Response(successJson2, 200, "OK", new URL(TEST_ENDPOINT_2)));
5858

59-
// Error response JSON
6059
String errorJson =
6160
"{" + "\"error\":\"invalid_request\"," + "\"error_description\":\"Bad request\"" + "}";
6261
HttpClient mockErrorClient = mock(HttpClient.class);
@@ -67,12 +66,6 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
6766
HttpClient mockIOExceptionClient = mock(HttpClient.class);
6867
when(mockIOExceptionClient.execute(any())).thenThrow(new IOException("Network error"));
6968

70-
// For null cpTokenSource
71-
DatabricksOAuthTokenSource nullCpTokenSource = null;
72-
73-
// For null httpClient
74-
HttpClient nullHttpClient = null;
75-
7669
// For null/empty endpoint or authDetails
7770
return Stream.of(
7871
Arguments.of(
@@ -81,6 +74,7 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
8174
TEST_AUTH_DETAILS_1,
8275
mockSuccessClient1,
8376
mockCpTokenSource,
77+
TEST_HOST,
8478
new Token(
8579
"dp-access-token1",
8680
TEST_TOKEN_TYPE,
@@ -94,6 +88,7 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
9488
TEST_AUTH_DETAILS_2,
9589
mockSuccessClient2,
9690
mockCpTokenSource,
91+
TEST_HOST,
9792
new Token(
9893
"dp-access-token2",
9994
TEST_TOKEN_TYPE,
@@ -106,6 +101,7 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
106101
TEST_AUTH_DETAILS_1,
107102
mockErrorClient,
108103
mockCpTokenSource,
104+
TEST_HOST,
109105
null,
110106
com.databricks.sdk.core.DatabricksException.class),
111107
Arguments.of(
@@ -114,22 +110,25 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
114110
TEST_AUTH_DETAILS_1,
115111
mockIOExceptionClient,
116112
mockCpTokenSource,
113+
TEST_HOST,
117114
null,
118115
com.databricks.sdk.core.DatabricksException.class),
119116
Arguments.of(
120117
"Null cpTokenSource",
121118
TEST_ENDPOINT_1,
122119
TEST_AUTH_DETAILS_1,
123120
mockSuccessClient1,
124-
nullCpTokenSource,
121+
null,
122+
TEST_HOST,
125123
null,
126124
NullPointerException.class),
127125
Arguments.of(
128126
"Null httpClient",
129127
TEST_ENDPOINT_1,
130128
TEST_AUTH_DETAILS_1,
131-
nullHttpClient,
129+
null,
132130
mockCpTokenSource,
131+
TEST_HOST,
133132
null,
134133
NullPointerException.class),
135134
Arguments.of(
@@ -138,6 +137,7 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
138137
TEST_AUTH_DETAILS_1,
139138
mockSuccessClient1,
140139
mockCpTokenSource,
140+
TEST_HOST,
141141
null,
142142
NullPointerException.class),
143143
Arguments.of(
@@ -146,8 +146,27 @@ private static Stream<Arguments> provideDataPlaneTokenScenarios() throws Excepti
146146
null,
147147
mockSuccessClient1,
148148
mockCpTokenSource,
149+
TEST_HOST,
150+
null,
151+
NullPointerException.class),
152+
Arguments.of(
153+
"Null host",
154+
TEST_ENDPOINT_1,
155+
TEST_AUTH_DETAILS_1,
156+
mockSuccessClient1,
157+
mockCpTokenSource,
158+
null,
159+
null,
160+
NullPointerException.class),
161+
Arguments.of(
162+
"Empty host",
163+
TEST_ENDPOINT_1,
164+
TEST_AUTH_DETAILS_1,
165+
mockSuccessClient1,
166+
mockCpTokenSource,
167+
"",
149168
null,
150-
NullPointerException.class));
169+
IllegalArgumentException.class));
151170
}
152171

153172
@ParameterizedTest(name = "{0}")
@@ -158,17 +177,18 @@ void testDataPlaneTokenSource(
158177
String authDetails,
159178
HttpClient httpClient,
160179
DatabricksOAuthTokenSource cpTokenSource,
180+
String host,
161181
Token expectedToken,
162182
Class<? extends Exception> expectedException) {
163183
if (expectedException != null) {
164184
assertThrows(
165185
expectedException,
166186
() -> {
167-
DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource);
187+
DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host);
168188
source.getToken(endpoint, authDetails);
169189
});
170190
} else {
171-
DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource);
191+
DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host);
172192
Token token = source.getToken(endpoint, authDetails);
173193
assertNotNull(token);
174194
assertEquals(expectedToken.getAccessToken(), token.getAccessToken());

0 commit comments

Comments
 (0)