Skip to content

Commit 66d514e

Browse files
authored
Refactor SessionCredentials (#470)
## What changes are proposed in this pull request? This PR refactors the `SessionCredentials` class to achieve clearer separation of concerns by splitting its responsibilities into two distinct components: a `CredentialsProvider` implementation and a dedicated `TokenSource` implementation. This refactoring is necessary for introducing token cache wrappers in future PRs. ### Current Implementation - `SessionCredentials` both implements `CredentialsProvider` and extends `RefreshableTokenSource`. - Token refresh logic is tightly coupled with credential provider functionality, making it hard to isolate token source logic and introduce caching. - A single class is responsible for both authentication configuration and token management. ## Proposed Changes - `SessionCredentials` now implements only `CredentialsProvider` and delegates all token operations to a new `SessionCredentialsTokenSource` class. - `SessionCredentialsTokenSource` extends `RefreshableTokenSource` and is solely responsible for token refresh logic. ## Benefits - Token sources can now be easily wrapped with caching mechanisms in future updates. - `SessionCredentials` remains a public `CredentialsProvider` as intended with no breaking changes to the public API. ## How is this tested? Existing unit tests. NO_CHANGELOG=true
1 parent 6b93586 commit 66d514e

5 files changed

Lines changed: 333 additions & 190 deletions

File tree

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java

Lines changed: 125 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,131 @@ public String getClientSecret() {
155155
return clientSecret;
156156
}
157157

158+
/**
159+
* Launch a browser to collect an authorization code and exchange the code for an OAuth token.
160+
*
161+
* @return A {@code SessionCredentials} instance representing the retrieved OAuth token.
162+
* @throws IOException if the webserver cannot be started, or if the browser cannot be opened.
163+
*/
164+
public SessionCredentials launchExternalBrowser() throws IOException {
165+
Map<String, String> params = getOAuthCallbackParameters();
166+
return exchangeCallbackParameters(params);
167+
}
168+
169+
/**
170+
* Exchange callback parameters for OAuth credentials.
171+
*
172+
* @param query The callback parameters from the OAuth flow
173+
* @return A {@code SessionCredentials} instance representing the retrieved OAuth token.
174+
*/
175+
public SessionCredentials exchangeCallbackParameters(Map<String, String> query) {
176+
validateCallbackParameters(query);
177+
Token token = exchange(query.get("code"), query.get("state"));
178+
return new SessionCredentials.Builder()
179+
.withHttpClient(this.hc)
180+
.withClientId(this.clientId)
181+
.withClientSecret(this.clientSecret)
182+
.withTokenUrl(this.tokenUrl)
183+
.withToken(token)
184+
.build();
185+
}
186+
187+
/**
188+
* Launches an external browser to collect OAuth callback parameters and exchanges them for an
189+
* OAuth token.
190+
*
191+
* @return A {@code Token} instance containing the OAuth access token and related credentials
192+
* @throws IOException if the local HTTP server cannot be started, the browser cannot be opened,
193+
* or there are network issues during the token exchange
194+
* @throws DatabricksException if the OAuth callback contains an error, missing required
195+
* parameters, or if there's a state mismatch during the token exchange.
196+
*/
197+
Token getTokenFromExternalBrowser() throws IOException {
198+
Map<String, String> params = getOAuthCallbackParameters();
199+
validateCallbackParameters(params);
200+
return exchange(params.get("code"), params.get("state"));
201+
}
202+
203+
protected void desktopBrowser() throws IOException {
204+
Desktop.getDesktop().browse(URI.create(this.authUrl));
205+
}
206+
207+
/**
208+
* Handles the OAuth callback by setting up a local HTTP server, launching the browser, and
209+
* collecting the callback parameters.
210+
*
211+
* @return A map containing the callback parameters from the OAuth flow.
212+
* @throws IOException if the webserver cannot be started, or if the browser cannot be opened.
213+
*/
214+
private Map<String, String> getOAuthCallbackParameters() throws IOException {
215+
URL redirect = new URL(getRedirectUrl());
216+
if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) {
217+
throw new IllegalArgumentException(
218+
"cannot listen on "
219+
+ redirect.getHost()
220+
+ ", redirectUrl host must be one of: localhost, 127.0.0.1");
221+
}
222+
CallbackResponseHandler handler = new CallbackResponseHandler();
223+
HttpServer httpServer =
224+
HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0);
225+
httpServer.createContext("/", handler);
226+
httpServer.start();
227+
desktopBrowser();
228+
Map<String, String> params = handler.getParams();
229+
httpServer.stop(0);
230+
return params;
231+
}
232+
233+
/**
234+
* Validates the OAuth callback parameters to ensure they contain the required fields and no error
235+
* conditions.
236+
*
237+
* @param query The callback parameters to validate
238+
* @throws DatabricksException if validation fails due to error conditions or missing required
239+
* parameters
240+
*/
241+
private void validateCallbackParameters(Map<String, String> query) {
242+
if (query.containsKey("error")) {
243+
throw new DatabricksException(query.get("error") + ": " + query.get("error_description"));
244+
}
245+
if (!query.containsKey("code") || !query.containsKey("state")) {
246+
throw new DatabricksException("No code returned in callback");
247+
}
248+
}
249+
250+
/**
251+
* Exchange authorization code for OAuth token.
252+
*
253+
* @param code The authorization code from the OAuth callback
254+
* @param state The state parameter from the OAuth callback
255+
* @return A {@code Token} instance representing the OAuth token
256+
*/
257+
private Token exchange(String code, String state) {
258+
if (!this.state.equals(state)) {
259+
throw new DatabricksException(
260+
"state mismatch: original state: " + this.state + "; retrieved state: " + state);
261+
}
262+
Map<String, String> params = new HashMap<>();
263+
params.put("grant_type", "authorization_code");
264+
params.put("code", code);
265+
params.put("code_verifier", this.verifier);
266+
params.put("redirect_uri", this.redirectUrl);
267+
Map<String, String> headers = new HashMap<>();
268+
if (this.tokenUrl.contains("microsoft")) {
269+
headers.put("Origin", this.redirectUrl);
270+
}
271+
Token token =
272+
RefreshableTokenSource.retrieveToken(
273+
this.hc,
274+
this.clientId,
275+
this.clientSecret,
276+
this.tokenUrl,
277+
params,
278+
headers,
279+
AuthParameterPosition.BODY);
280+
return token;
281+
}
282+
158283
static class CallbackResponseHandler implements HttpHandler {
159284
private final Logger LOG = LoggerFactory.getLogger(getClass().getName());
160285
// Protects params
@@ -258,75 +383,4 @@ public Map<String, String> getParams() {
258383
}
259384
}
260385
}
261-
262-
/**
263-
* Launch a browser to collect an authorization code and exchange the code for an OAuth token.
264-
*
265-
* @return A {@code SessionCredentials} instance representing the retrieved OAuth token.
266-
* @throws IOException if the webserver cannot be started, or if the browser cannot be opened
267-
*/
268-
public SessionCredentials launchExternalBrowser() throws IOException {
269-
URL redirect = new URL(getRedirectUrl());
270-
if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) {
271-
throw new IllegalArgumentException(
272-
"cannot listen on "
273-
+ redirect.getHost()
274-
+ ", redirectUrl host must be one of: localhost, 127.0.0.1");
275-
}
276-
CallbackResponseHandler handler = new CallbackResponseHandler();
277-
HttpServer httpServer =
278-
HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0);
279-
httpServer.createContext("/", handler);
280-
httpServer.start();
281-
desktopBrowser();
282-
Map<String, String> params = handler.getParams();
283-
httpServer.stop(0);
284-
return exchangeCallbackParameters(params);
285-
}
286-
287-
protected void desktopBrowser() throws IOException {
288-
Desktop.getDesktop().browse(URI.create(this.authUrl));
289-
}
290-
291-
public SessionCredentials exchangeCallbackParameters(Map<String, String> query) {
292-
if (query.containsKey("error")) {
293-
throw new DatabricksException(query.get("error") + ": " + query.get("error_description"));
294-
}
295-
if (!query.containsKey("code") || !query.containsKey("state")) {
296-
throw new DatabricksException("No code returned in callback");
297-
}
298-
return exchange(query.get("code"), query.get("state"));
299-
}
300-
301-
public SessionCredentials exchange(String code, String state) {
302-
if (!this.state.equals(state)) {
303-
throw new DatabricksException(
304-
"state mismatch: original state: " + this.state + "; retrieved state: " + state);
305-
}
306-
Map<String, String> params = new HashMap<>();
307-
params.put("grant_type", "authorization_code");
308-
params.put("code", code);
309-
params.put("code_verifier", this.verifier);
310-
params.put("redirect_uri", this.redirectUrl);
311-
Map<String, String> headers = new HashMap<>();
312-
if (this.tokenUrl.contains("microsoft")) {
313-
headers.put("Origin", this.redirectUrl);
314-
}
315-
Token token =
316-
RefreshableTokenSource.retrieveToken(
317-
this.hc,
318-
this.clientId,
319-
this.clientSecret,
320-
this.tokenUrl,
321-
params,
322-
headers,
323-
AuthParameterPosition.BODY);
324-
return new SessionCredentials.Builder()
325-
.withHttpClient(this.hc)
326-
.withClientId(this.clientId)
327-
.withClientSecret(this.clientSecret)
328-
.withTokenUrl(this.tokenUrl)
329-
.withToken(token)
330-
.build();
331-
}
332386
}

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.io.IOException;
77
import java.nio.file.Path;
88
import java.util.Objects;
9+
import java.util.Optional;
910
import org.slf4j.Logger;
1011
import org.slf4j.LoggerFactory;
1112

@@ -66,39 +67,38 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
6667
LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId);
6768

6869
try {
69-
// Create SessionCredentials with the cached token and try to refresh if needed
70-
SessionCredentials cachedCreds =
71-
new SessionCredentials.Builder()
72-
.withToken(cachedToken)
73-
.withHttpClient(config.getHttpClient())
74-
.withClientId(clientId)
75-
.withClientSecret(clientSecret)
76-
.withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
77-
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
78-
.withTokenCache(tokenCache)
79-
.build();
70+
// Create SessionCredentialsTokenSource with the cached token and try to refresh if needed
71+
SessionCredentialsTokenSource tokenSource =
72+
new SessionCredentialsTokenSource(
73+
cachedToken,
74+
config.getHttpClient(),
75+
config.getOidcEndpoints().getTokenEndpoint(),
76+
clientId,
77+
clientSecret,
78+
Optional.of(config.getEffectiveOAuthRedirectUrl()),
79+
Optional.of(tokenCache));
8080

8181
LOGGER.debug("Using cached token, will immediately refresh");
82-
cachedCreds.token = cachedCreds.refresh();
83-
return cachedCreds.configure(config);
82+
tokenSource.token = tokenSource.refresh();
83+
return OAuthHeaderFactory.fromTokenSource(tokenSource);
8484
} catch (Exception e) {
8585
// If token refresh fails, log and continue to browser auth
8686
LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage());
8787
}
8888
}
8989

9090
// If no cached token or refresh failed, perform browser auth
91-
SessionCredentials credentials =
91+
SessionCredentialsTokenSource tokenSource =
9292
performBrowserAuth(config, clientId, clientSecret, tokenCache);
93-
tokenCache.save(credentials.getToken());
94-
return credentials.configure(config);
93+
tokenCache.save(tokenSource.getToken());
94+
return OAuthHeaderFactory.fromTokenSource(tokenSource);
9595
} catch (IOException | DatabricksException e) {
9696
LOGGER.error("Failed to authenticate: {}", e.getMessage());
9797
return null;
9898
}
9999
}
100100

101-
SessionCredentials performBrowserAuth(
101+
SessionCredentialsTokenSource performBrowserAuth(
102102
DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache)
103103
throws IOException {
104104
LOGGER.debug("Performing browser authentication");
@@ -113,18 +113,17 @@ SessionCredentials performBrowserAuth(
113113
.build();
114114
Consent consent = client.initiateConsent();
115115

116-
// Use the existing browser flow to get credentials
117-
SessionCredentials credentials = consent.launchExternalBrowser();
118-
119-
// Create a new SessionCredentials with the same token but with our token cache
120-
return new SessionCredentials.Builder()
121-
.withToken(credentials.getToken())
122-
.withHttpClient(config.getHttpClient())
123-
.withClientId(config.getClientId())
124-
.withClientSecret(config.getClientSecret())
125-
.withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
126-
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
127-
.withTokenCache(tokenCache)
128-
.build();
116+
// Use the existing browser flow to get credentials.
117+
Token token = consent.getTokenFromExternalBrowser();
118+
119+
// Create a SessionCredentialsTokenSource with the token from browser auth.
120+
return new SessionCredentialsTokenSource(
121+
token,
122+
config.getHttpClient(),
123+
config.getOidcEndpoints().getTokenEndpoint(),
124+
config.getClientId(),
125+
config.getClientSecret(),
126+
Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()),
127+
Optional.ofNullable(tokenCache));
129128
}
130129
}

0 commit comments

Comments
 (0)