From 40d986e2a595a5e951ba101d98c980dffd5fc83c Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 19 Jun 2025 16:27:56 +0000 Subject: [PATCH 1/6] First version --- .../databricks/sdk/core/oauth/Consent.java | 196 +++++++++++------- .../ExternalBrowserCredentialsProvider.java | 46 ++-- .../sdk/core/oauth/SessionCredentials.java | 67 ++---- .../oauth/SessionCredentialsTokenSource.java | 151 ++++++++++++++ ...xternalBrowserCredentialsProviderTest.java | 66 +++--- 5 files changed, 354 insertions(+), 172 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 77045df97..c66ac8d5f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -155,6 +155,131 @@ public String getClientSecret() { return clientSecret; } + /** + * Launch a browser to collect an authorization code and exchange the code for an OAuth token. + * + * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened + */ + public SessionCredentials launchExternalBrowser() throws IOException { + Map params = getOAuthCallbackParameters(); + return exchangeCallbackParameters(params); + } + + /** + * Exchange callback parameters for OAuth credentials. + * + * @param query The callback parameters from the OAuth flow + * @return A {@code SessionCredentials} instance representing the retrieved OAuth token + */ + public SessionCredentials exchangeCallbackParameters(Map query) { + validateCallbackParameters(query); + Token token = exchange(query.get("code"), query.get("state")); + return new SessionCredentials.Builder() + .withHttpClient(this.hc) + .withClientId(this.clientId) + .withClientSecret(this.clientSecret) + .withTokenUrl(this.tokenUrl) + .withToken(token) + .build(); + } + + /** + * Launches an external browser to collect OAuth callback parameters and exchanges them for an + * OAuth token. + * + * @return A {@code Token} instance containing the OAuth access token and related credentials + * @throws IOException if the local HTTP server cannot be started, the browser cannot be opened, + * or there are network issues during the token exchange + * @throws DatabricksException if the OAuth callback contains an error, missing required + * parameters, or if there's a state mismatch during the token exchange. + */ + Token getTokenFromExternalBrowser() throws IOException { + Map params = getOAuthCallbackParameters(); + validateCallbackParameters(params); + return exchange(params.get("code"), params.get("state")); + } + + protected void desktopBrowser() throws IOException { + Desktop.getDesktop().browse(URI.create(this.authUrl)); + } + + /** + * Handles the OAuth callback by setting up a local HTTP server, launching the browser, and + * collecting the callback parameters. + * + * @return A map containing the callback parameters from the OAuth flow + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened + */ + private Map getOAuthCallbackParameters() throws IOException { + URL redirect = new URL(getRedirectUrl()); + if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) { + throw new IllegalArgumentException( + "cannot listen on " + + redirect.getHost() + + ", redirectUrl host must be one of: localhost, 127.0.0.1"); + } + CallbackResponseHandler handler = new CallbackResponseHandler(); + HttpServer httpServer = + HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0); + httpServer.createContext("/", handler); + httpServer.start(); + desktopBrowser(); + Map params = handler.getParams(); + httpServer.stop(0); + return params; + } + + /** + * Validates the OAuth callback parameters to ensure they contain the required fields and no error + * conditions. + * + * @param query The callback parameters to validate + * @throws DatabricksException if validation fails due to error conditions or missing required + * parameters + */ + private void validateCallbackParameters(Map query) { + if (query.containsKey("error")) { + throw new DatabricksException(query.get("error") + ": " + query.get("error_description")); + } + if (!query.containsKey("code") || !query.containsKey("state")) { + throw new DatabricksException("No code returned in callback"); + } + } + + /** + * Exchange authorization code for OAuth token. + * + * @param code The authorization code from the OAuth callback + * @param state The state parameter from the OAuth callback + * @return A {@code Token} instance representing the OAuth token + */ + private Token exchange(String code, String state) { + if (!this.state.equals(state)) { + throw new DatabricksException( + "state mismatch: original state: " + this.state + "; retrieved state: " + state); + } + Map params = new HashMap<>(); + params.put("grant_type", "authorization_code"); + params.put("code", code); + params.put("code_verifier", this.verifier); + params.put("redirect_uri", this.redirectUrl); + Map headers = new HashMap<>(); + if (this.tokenUrl.contains("microsoft")) { + headers.put("Origin", this.redirectUrl); + } + Token token = + RefreshableTokenSource.retrieveToken( + this.hc, + this.clientId, + this.clientSecret, + this.tokenUrl, + params, + headers, + AuthParameterPosition.BODY); + return token; + } + static class CallbackResponseHandler implements HttpHandler { private final Logger LOG = LoggerFactory.getLogger(getClass().getName()); // Protects params @@ -258,75 +383,4 @@ public Map getParams() { } } } - - /** - * Launch a browser to collect an authorization code and exchange the code for an OAuth token. - * - * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. - * @throws IOException if the webserver cannot be started, or if the browser cannot be opened - */ - public SessionCredentials launchExternalBrowser() throws IOException { - URL redirect = new URL(getRedirectUrl()); - if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) { - throw new IllegalArgumentException( - "cannot listen on " - + redirect.getHost() - + ", redirectUrl host must be one of: localhost, 127.0.0.1"); - } - CallbackResponseHandler handler = new CallbackResponseHandler(); - HttpServer httpServer = - HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0); - httpServer.createContext("/", handler); - httpServer.start(); - desktopBrowser(); - Map params = handler.getParams(); - httpServer.stop(0); - return exchangeCallbackParameters(params); - } - - protected void desktopBrowser() throws IOException { - Desktop.getDesktop().browse(URI.create(this.authUrl)); - } - - public SessionCredentials exchangeCallbackParameters(Map query) { - if (query.containsKey("error")) { - throw new DatabricksException(query.get("error") + ": " + query.get("error_description")); - } - if (!query.containsKey("code") || !query.containsKey("state")) { - throw new DatabricksException("No code returned in callback"); - } - return exchange(query.get("code"), query.get("state")); - } - - public SessionCredentials exchange(String code, String state) { - if (!this.state.equals(state)) { - throw new DatabricksException( - "state mismatch: original state: " + this.state + "; retrieved state: " + state); - } - Map params = new HashMap<>(); - params.put("grant_type", "authorization_code"); - params.put("code", code); - params.put("code_verifier", this.verifier); - params.put("redirect_uri", this.redirectUrl); - Map headers = new HashMap<>(); - if (this.tokenUrl.contains("microsoft")) { - headers.put("Origin", this.redirectUrl); - } - Token token = - RefreshableTokenSource.retrieveToken( - this.hc, - this.clientId, - this.clientSecret, - this.tokenUrl, - params, - headers, - AuthParameterPosition.BODY); - return new SessionCredentials.Builder() - .withHttpClient(this.hc) - .withClientId(this.clientId) - .withClientSecret(this.clientSecret) - .withTokenUrl(this.tokenUrl) - .withToken(token) - .build(); - } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index 7bae60022..a698e7586 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java @@ -66,21 +66,21 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId); try { - // Create SessionCredentials with the cached token and try to refresh if needed - SessionCredentials cachedCreds = - new SessionCredentials.Builder() - .withToken(cachedToken) - .withHttpClient(config.getHttpClient()) - .withClientId(clientId) - .withClientSecret(clientSecret) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) + // Create SessionCredentialsTokenSource with the cached token and try to refresh if needed + SessionCredentialsTokenSource tokenSource = + new SessionCredentialsTokenSource.Builder( + cachedToken, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + clientId, + clientSecret) .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .withTokenCache(tokenCache) .build(); LOGGER.debug("Using cached token, will immediately refresh"); - cachedCreds.token = cachedCreds.refresh(); - return cachedCreds.configure(config); + tokenSource.token = tokenSource.refresh(); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); @@ -88,17 +88,17 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } // If no cached token or refresh failed, perform browser auth - SessionCredentials credentials = + SessionCredentialsTokenSource tokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); - tokenCache.save(credentials.getToken()); - return credentials.configure(config); + tokenCache.save(tokenSource.getToken()); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; } } - SessionCredentials performBrowserAuth( + SessionCredentialsTokenSource performBrowserAuth( DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache) throws IOException { LOGGER.debug("Performing browser authentication"); @@ -113,16 +113,16 @@ SessionCredentials performBrowserAuth( .build(); Consent consent = client.initiateConsent(); - // Use the existing browser flow to get credentials - SessionCredentials credentials = consent.launchExternalBrowser(); + // Use the existing browser flow to get credentials. + Token token = consent.getTokenFromExternalBrowser(); - // Create a new SessionCredentials with the same token but with our token cache - return new SessionCredentials.Builder() - .withToken(credentials.getToken()) - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) + // Create a SessionCredentialsTokenSource with the token from browser auth. + return new SessionCredentialsTokenSource.Builder( + token, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + config.getClientId(), + config.getClientSecret()) .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .withTokenCache(tokenCache) .build(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index 4d2d512e3..2f20032f4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -2,11 +2,8 @@ import com.databricks.sdk.core.CredentialsProvider; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.HttpClient; import java.io.Serializable; -import java.util.HashMap; -import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -17,11 +14,21 @@ * requests to an API, and a long-lived refresh token, which can be used to fetch new access tokens. * Calling refresh() uses the refresh token to retrieve a new access token to authenticate to APIs. */ -public class SessionCredentials extends RefreshableTokenSource - implements CredentialsProvider, Serializable { +public class SessionCredentials implements CredentialsProvider, Serializable { private static final long serialVersionUID = 3083941540130596650L; private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentials.class); + private final SessionCredentialsTokenSource tokenSource; + + private SessionCredentials(Builder b) { + this.tokenSource = + new SessionCredentialsTokenSource.Builder( + b.token, b.hc, b.tokenUrl, b.clientId, b.clientSecret) + .withRedirectUrl(b.redirectUrl) + .withTokenCache(b.tokenCache) + .build(); + } + @Override public String authType() { return "oauth-u2m"; @@ -29,7 +36,7 @@ public String authType() { @Override public OAuthHeaderFactory configure(DatabricksConfig config) { - return OAuthHeaderFactory.fromTokenSource(this); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } static class Builder { @@ -80,52 +87,4 @@ public SessionCredentials build() { return new SessionCredentials(this); } } - - private final HttpClient hc; - private final String tokenUrl; - private final String redirectUrl; - private final String clientId; - private final String clientSecret; - private final TokenCache tokenCache; - - private SessionCredentials(Builder b) { - super(b.token); - this.hc = b.hc; - this.tokenUrl = b.tokenUrl; - this.redirectUrl = b.redirectUrl; - this.clientId = b.clientId; - this.clientSecret = b.clientSecret; - this.tokenCache = b.tokenCache; - } - - @Override - protected Token refresh() { - if (this.token == null) { - throw new DatabricksException("oauth2: token is not set"); - } - String refreshToken = this.token.getRefreshToken(); - if (refreshToken == null) { - throw new DatabricksException("oauth2: token expired and refresh token is not set"); - } - - Map params = new HashMap<>(); - params.put("grant_type", "refresh_token"); - params.put("refresh_token", refreshToken); - Map headers = new HashMap<>(); - if (tokenUrl.contains("microsoft")) { - // Tokens issued for the 'Single-Page Application' client-type may only be redeemed via - // cross-origin requests - headers.put("Origin", redirectUrl); - } - Token newToken = - retrieveToken( - hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); - - // Save the refreshed token directly to cache - if (tokenCache != null) { - tokenCache.save(newToken); - LOGGER.debug("Saved refreshed token to cache"); - } - return newToken; - } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java new file mode 100644 index 000000000..65b585cb7 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java @@ -0,0 +1,151 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import java.util.HashMap; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * TokenSource that handles OAuth token refresh for SessionCredentials. + * + *

Implements the refresh_token OAuth grant type with optional token caching. + * + * @see RefreshableTokenSource + * @see Token + * @see TokenCache + */ +public class SessionCredentialsTokenSource extends RefreshableTokenSource { + private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentialsTokenSource.class); + + // HTTP client for making token refresh requests + private final HttpClient hc; + // OAuth token endpoint URL for refresh requests + private final String tokenUrl; + // OAuth redirect URL, used for Microsoft OAuth endpoints + private final String redirectUrl; + // OAuth client ID for authentication + private final String clientId; + // OAuth client secret for authentication + private final String clientSecret; + // Optional token cache for persisting refreshed tokens + private final TokenCache tokenCache; + + /** + * Constructs a new SessionCredentialsTokenSource with the given builder. + * + * @param builder The builder containing all configuration parameters + */ + private SessionCredentialsTokenSource(Builder builder) { + super(builder.token); + this.hc = builder.hc; + this.tokenUrl = builder.tokenUrl; + this.redirectUrl = builder.redirectUrl; + this.clientId = builder.clientId; + this.clientSecret = builder.clientSecret; + this.tokenCache = builder.tokenCache; + } + + /** + * Refreshes the OAuth token using the refresh_token grant type. + * + *

This method attempts to refresh the current token using the refresh token. If successful, + * the new token is automatically saved to the token cache if one is configured. For Microsoft + * OAuth endpoints, it includes the Origin header. + * + * @return A new {@link Token} with updated access and refresh tokens + * @throws DatabricksException if the token is not set, refresh token is missing, or the refresh + * request fails + */ + @Override + protected Token refresh() { + if (this.token == null) { + throw new DatabricksException("oauth2: token is not set"); + } + String refreshToken = this.token.getRefreshToken(); + if (refreshToken == null) { + throw new DatabricksException("oauth2: token expired and refresh token is not set"); + } + + Map params = new HashMap<>(); + params.put("grant_type", "refresh_token"); + params.put("refresh_token", refreshToken); + Map headers = new HashMap<>(); + if (tokenUrl.contains("microsoft")) { + // Tokens issued for the 'Single-Page Application' client-type may only be redeemed via + // cross-origin requests + headers.put("Origin", redirectUrl); + } + Token newToken = + retrieveToken( + hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); + + // Save the refreshed token directly to cache + if (tokenCache != null) { + tokenCache.save(newToken); + LOGGER.debug("Saved refreshed token to cache"); + } + return newToken; + } + + /** Builder for creating SessionCredentialsTokenSource instances. */ + public static class Builder { + private final Token token; + private final HttpClient hc; + private final String tokenUrl; + private final String clientId; + private final String clientSecret; + private String redirectUrl; + private TokenCache tokenCache; + + /** + * Creates a new Builder for SessionCredentialsTokenSource with required parameters. + * + * @param token The initial token to use + * @param hc The HTTP client for making token refresh requests + * @param tokenUrl The OAuth token endpoint URL + * @param clientId The OAuth client ID + * @param clientSecret The OAuth client secret + */ + public Builder( + Token token, HttpClient hc, String tokenUrl, String clientId, String clientSecret) { + this.token = token; + this.hc = hc; + this.tokenUrl = tokenUrl; + this.clientId = clientId; + this.clientSecret = clientSecret; + } + + /** + * Sets the redirect URL (optional). + * + * @param redirectUrl The OAuth redirect URL + * @return This Builder instance for method chaining + */ + public Builder withRedirectUrl(String redirectUrl) { + this.redirectUrl = redirectUrl; + return this; + } + + /** + * Sets the token cache (optional). + * + * @param tokenCache The token cache for persisting refreshed tokens + * @return This Builder instance for method chaining + */ + public Builder withTokenCache(TokenCache tokenCache) { + this.tokenCache = tokenCache; + return this; + } + + /** + * Builds the SessionCredentialsTokenSource instance. + * + * @return A new SessionCredentialsTokenSource instance + */ + public SessionCredentialsTokenSource build() { + return new SessionCredentialsTokenSource(this); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 932690bd7..3b438da80 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -129,7 +129,7 @@ void openIDConnectEndPointsTestAccounts() throws IOException { void exchangeTest() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); // Mock because it's a POST Request to http client @@ -166,16 +166,28 @@ void exchangeTest() throws IOException { Map queryCreds = new HashMap<>(); queryCreds.put("code", "testCode"); queryCreds.put("state", "testState"); + + // Verify that SessionCredentials is created successfully with the exchange SessionCredentials creds = testConsent.exchangeCallbackParameters(queryCreds); - assertEquals("accessTokenFromServer", creds.token.getAccessToken()); - assertEquals("refreshTokenFromServer", creds.token.getRefreshToken()); + assertNotNull(creds); + + // Create a minimal config for testing the configure method + DatabricksConfig testConfig = new DatabricksConfig(); + + // Configure the SessionCredentials to get the OAuthHeaderFactory + OAuthHeaderFactory headerFactory = creds.configure(testConfig); + assertNotNull(headerFactory); + + // Verify the headers are correctly formatted + Map headers = headerFactory.headers(); + assertEquals("tokenTypeFromServer accessTokenFromServer", headers.get("Authorization")); } @Test void clientCredentials() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); @@ -195,24 +207,23 @@ void clientCredentials() throws IOException { void sessionCredentials() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); - SessionCredentials sessionCredentials = - new SessionCredentials.Builder() - .withHttpClient(hc) - .withClientId("testClientId") - .withClientSecret("abc") - .withTokenUrl("https://tokenUrl") - .withToken( + SessionCredentialsTokenSource sessionCredentialsTokenSource = + new SessionCredentialsTokenSource.Builder( new Token( "originalAccessToken", "originalTokenType", "originalRefreshToken", - Instant.MAX)) + Instant.MAX), + hc, + "https://tokenUrl", + "testClientId", + "abc") .build(); - Token token = sessionCredentials.refresh(); + Token token = sessionCredentialsTokenSource.refresh(); // We check that we are actually getting the token from server response (that is defined // above) rather than what was given while creating session credentials @@ -408,11 +419,13 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentials browserAuthCreds = - new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") + SessionCredentialsTokenSource browserAuthCreds = + new SessionCredentialsTokenSource.Builder( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret") .build(); // Create config with failing HTTP client and mock token cache @@ -459,6 +472,9 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { @Test void cacheWithInvalidTokensTest() throws IOException { + // Create mock HTTP client + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + // Create completely invalid token (no refresh token) Instant pastTime = Instant.now().minusSeconds(3600); Token invalidToken = new Token("expired_access_token", "Bearer", null, pastTime); @@ -475,11 +491,13 @@ void cacheWithInvalidTokensTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentials browserAuthCreds = - new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") + SessionCredentialsTokenSource browserAuthCreds = + new SessionCredentialsTokenSource.Builder( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret") .build(); // Create simple config From 2b79852bbb4a50ef33c7a54c94634424efaf1ec1 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 20 Jun 2025 10:08:30 +0000 Subject: [PATCH 2/6] Fix bug in test --- ...xternalBrowserCredentialsProviderTest.java | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 3b438da80..420f5eb17 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -129,11 +129,13 @@ void openIDConnectEndPointsTestAccounts() throws IOException { void exchangeTest() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); // Mock because it's a POST Request to http client - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); Consent testConsent = new Consent.Builder() @@ -187,9 +189,11 @@ void exchangeTest() throws IOException { void clientCredentials() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); ClientCredentials clientCredentials = new ClientCredentials.Builder() @@ -207,9 +211,11 @@ void clientCredentials() throws IOException { void sessionCredentials() throws IOException { HttpClient hc = Mockito.mock(HttpClient.class); String response = - "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"3600\", \"refresh_token\": \"refreshTokenFromServer\"}"; + "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); SessionCredentialsTokenSource sessionCredentialsTokenSource = new SessionCredentialsTokenSource.Builder( @@ -237,10 +243,10 @@ void sessionCredentials() throws IOException { void cacheWithValidTokenTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String refreshResponse = + String response = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doReturn(new Response(refreshResponse, url)) + Mockito.doAnswer(invocation -> new Response(response, url)) .when(mockHttpClient) .execute(any(Request.class)); @@ -317,10 +323,10 @@ void cacheWithValidTokenTest() throws IOException { void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String refreshResponse = + String response = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doReturn(new Response(refreshResponse, url)) + Mockito.doAnswer(invocation -> new Response(response, url)) .when(mockHttpClient) .execute(any(Request.class)); From abc6ded620da6f1169f25ccc4f75d284dec87ebd Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 20 Jun 2025 11:33:24 +0000 Subject: [PATCH 3/6] Update names --- .../oauth/ExternalBrowserCredentialsProviderTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 420f5eb17..7479b259e 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -243,10 +243,10 @@ void sessionCredentials() throws IOException { void cacheWithValidTokenTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String response = + String refreshResponse = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doAnswer(invocation -> new Response(response, url)) + Mockito.doAnswer(invocation -> new Response(refreshResponse, url)) .when(mockHttpClient) .execute(any(Request.class)); @@ -323,10 +323,10 @@ void cacheWithValidTokenTest() throws IOException { void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String response = + String refreshResponse = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doAnswer(invocation -> new Response(response, url)) + Mockito.doAnswer(invocation -> new Response(refreshResponse, url)) .when(mockHttpClient) .execute(any(Request.class)); From 063fbbd19f22b933e594e389284a4d0e43dc8808 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 20 Jun 2025 11:40:12 +0000 Subject: [PATCH 4/6] Better names --- .../oauth/ExternalBrowserCredentialsProviderTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 7479b259e..c7c3897b6 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -425,7 +425,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentialsTokenSource browserAuthCreds = + SessionCredentialsTokenSource browserAuthTokenSource = new SessionCredentialsTokenSource.Builder( browserAuthToken, mockHttpClient, @@ -450,7 +450,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); @@ -497,7 +497,7 @@ void cacheWithInvalidTokensTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentialsTokenSource browserAuthCreds = + SessionCredentialsTokenSource browserAuthTokenSource = new SessionCredentialsTokenSource.Builder( browserAuthToken, mockHttpClient, @@ -516,7 +516,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); From 3c62ad7669d1314368f4fe755291459f338fcff0 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 20 Jun 2025 11:43:38 +0000 Subject: [PATCH 5/6] Clean up --- .../sdk/core/oauth/SessionCredentialsTokenSource.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java index 65b585cb7..67434cbec 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java @@ -11,10 +11,6 @@ * TokenSource that handles OAuth token refresh for SessionCredentials. * *

Implements the refresh_token OAuth grant type with optional token caching. - * - * @see RefreshableTokenSource - * @see Token - * @see TokenCache */ public class SessionCredentialsTokenSource extends RefreshableTokenSource { private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentialsTokenSource.class); From 4b9a4037a3d3424dd2be7ed83a07e1ed897b3c82 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 25 Jun 2025 13:16:25 +0000 Subject: [PATCH 6/6] Address comments --- .../databricks/sdk/core/oauth/Consent.java | 8 +- .../ExternalBrowserCredentialsProvider.java | 35 +++--- .../sdk/core/oauth/SessionCredentials.java | 14 ++- .../oauth/SessionCredentialsTokenSource.java | 113 ++++++------------ ...xternalBrowserCredentialsProviderTest.java | 51 ++++---- 5 files changed, 90 insertions(+), 131 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index c66ac8d5f..57c97a2ff 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -159,7 +159,7 @@ public String getClientSecret() { * Launch a browser to collect an authorization code and exchange the code for an OAuth token. * * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. - * @throws IOException if the webserver cannot be started, or if the browser cannot be opened + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened. */ public SessionCredentials launchExternalBrowser() throws IOException { Map params = getOAuthCallbackParameters(); @@ -170,7 +170,7 @@ public SessionCredentials launchExternalBrowser() throws IOException { * Exchange callback parameters for OAuth credentials. * * @param query The callback parameters from the OAuth flow - * @return A {@code SessionCredentials} instance representing the retrieved OAuth token + * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. */ public SessionCredentials exchangeCallbackParameters(Map query) { validateCallbackParameters(query); @@ -208,8 +208,8 @@ protected void desktopBrowser() throws IOException { * Handles the OAuth callback by setting up a local HTTP server, launching the browser, and * collecting the callback parameters. * - * @return A map containing the callback parameters from the OAuth flow - * @throws IOException if the webserver cannot be started, or if the browser cannot be opened + * @return A map containing the callback parameters from the OAuth flow. + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened. */ private Map getOAuthCallbackParameters() throws IOException { URL redirect = new URL(getRedirectUrl()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index a698e7586..af67daeba 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java @@ -6,6 +6,7 @@ import java.io.IOException; import java.nio.file.Path; import java.util.Objects; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -68,15 +69,14 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { try { // Create SessionCredentialsTokenSource with the cached token and try to refresh if needed SessionCredentialsTokenSource tokenSource = - new SessionCredentialsTokenSource.Builder( - cachedToken, - config.getHttpClient(), - config.getOidcEndpoints().getTokenEndpoint(), - clientId, - clientSecret) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withTokenCache(tokenCache) - .build(); + new SessionCredentialsTokenSource( + cachedToken, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + clientId, + clientSecret, + Optional.of(config.getEffectiveOAuthRedirectUrl()), + Optional.of(tokenCache)); LOGGER.debug("Using cached token, will immediately refresh"); tokenSource.token = tokenSource.refresh(); @@ -117,14 +117,13 @@ SessionCredentialsTokenSource performBrowserAuth( Token token = consent.getTokenFromExternalBrowser(); // Create a SessionCredentialsTokenSource with the token from browser auth. - return new SessionCredentialsTokenSource.Builder( - token, - config.getHttpClient(), - config.getOidcEndpoints().getTokenEndpoint(), - config.getClientId(), - config.getClientSecret()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withTokenCache(tokenCache) - .build(); + return new SessionCredentialsTokenSource( + token, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + config.getClientId(), + config.getClientSecret(), + Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()), + Optional.ofNullable(tokenCache)); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index 2f20032f4..5a05b8751 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -4,6 +4,7 @@ import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.http.HttpClient; import java.io.Serializable; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -22,11 +23,14 @@ public class SessionCredentials implements CredentialsProvider, Serializable { private SessionCredentials(Builder b) { this.tokenSource = - new SessionCredentialsTokenSource.Builder( - b.token, b.hc, b.tokenUrl, b.clientId, b.clientSecret) - .withRedirectUrl(b.redirectUrl) - .withTokenCache(b.tokenCache) - .build(); + new SessionCredentialsTokenSource( + b.token, + b.hc, + b.tokenUrl, + b.clientId, + b.clientSecret, + Optional.ofNullable(b.redirectUrl), + Optional.ofNullable(b.tokenCache)); } @Override diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java index 67434cbec..8c71f2eb3 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java @@ -4,6 +4,7 @@ import com.databricks.sdk.core.http.HttpClient; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -20,27 +21,40 @@ public class SessionCredentialsTokenSource extends RefreshableTokenSource { // OAuth token endpoint URL for refresh requests private final String tokenUrl; // OAuth redirect URL, used for Microsoft OAuth endpoints - private final String redirectUrl; + private final Optional redirectUrl; // OAuth client ID for authentication private final String clientId; // OAuth client secret for authentication private final String clientSecret; // Optional token cache for persisting refreshed tokens - private final TokenCache tokenCache; + private final Optional tokenCache; /** - * Constructs a new SessionCredentialsTokenSource with the given builder. + * Constructs a new SessionCredentialsTokenSource. * - * @param builder The builder containing all configuration parameters + * @param token The initial token to use + * @param hc The HTTP client for making token refresh requests + * @param tokenUrl The OAuth token endpoint URL + * @param clientId The OAuth client ID + * @param clientSecret The OAuth client secret + * @param redirectUrl The OAuth redirect URL (optional) + * @param tokenCache The token cache for persisting refreshed tokens (optional) */ - private SessionCredentialsTokenSource(Builder builder) { - super(builder.token); - this.hc = builder.hc; - this.tokenUrl = builder.tokenUrl; - this.redirectUrl = builder.redirectUrl; - this.clientId = builder.clientId; - this.clientSecret = builder.clientSecret; - this.tokenCache = builder.tokenCache; + public SessionCredentialsTokenSource( + Token token, + HttpClient hc, + String tokenUrl, + String clientId, + String clientSecret, + Optional redirectUrl, + Optional tokenCache) { + super(token); + this.hc = hc; + this.tokenUrl = tokenUrl; + this.clientId = clientId; + this.clientSecret = clientSecret; + this.redirectUrl = redirectUrl; + this.tokenCache = tokenCache; } /** @@ -50,9 +64,9 @@ private SessionCredentialsTokenSource(Builder builder) { * the new token is automatically saved to the token cache if one is configured. For Microsoft * OAuth endpoints, it includes the Origin header. * - * @return A new {@link Token} with updated access and refresh tokens + * @return A new {@link Token} with updated access and refresh tokens. * @throws DatabricksException if the token is not set, refresh token is missing, or the refresh - * request fails + * request fails. */ @Override protected Token refresh() { @@ -71,77 +85,18 @@ protected Token refresh() { if (tokenUrl.contains("microsoft")) { // Tokens issued for the 'Single-Page Application' client-type may only be redeemed via // cross-origin requests - headers.put("Origin", redirectUrl); + redirectUrl.ifPresent(url -> headers.put("Origin", url)); } Token newToken = retrieveToken( hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); // Save the refreshed token directly to cache - if (tokenCache != null) { - tokenCache.save(newToken); - LOGGER.debug("Saved refreshed token to cache"); - } + tokenCache.ifPresent( + cache -> { + cache.save(newToken); + LOGGER.debug("Saved refreshed token to cache"); + }); return newToken; } - - /** Builder for creating SessionCredentialsTokenSource instances. */ - public static class Builder { - private final Token token; - private final HttpClient hc; - private final String tokenUrl; - private final String clientId; - private final String clientSecret; - private String redirectUrl; - private TokenCache tokenCache; - - /** - * Creates a new Builder for SessionCredentialsTokenSource with required parameters. - * - * @param token The initial token to use - * @param hc The HTTP client for making token refresh requests - * @param tokenUrl The OAuth token endpoint URL - * @param clientId The OAuth client ID - * @param clientSecret The OAuth client secret - */ - public Builder( - Token token, HttpClient hc, String tokenUrl, String clientId, String clientSecret) { - this.token = token; - this.hc = hc; - this.tokenUrl = tokenUrl; - this.clientId = clientId; - this.clientSecret = clientSecret; - } - - /** - * Sets the redirect URL (optional). - * - * @param redirectUrl The OAuth redirect URL - * @return This Builder instance for method chaining - */ - public Builder withRedirectUrl(String redirectUrl) { - this.redirectUrl = redirectUrl; - return this; - } - - /** - * Sets the token cache (optional). - * - * @param tokenCache The token cache for persisting refreshed tokens - * @return This Builder instance for method chaining - */ - public Builder withTokenCache(TokenCache tokenCache) { - this.tokenCache = tokenCache; - return this; - } - - /** - * Builds the SessionCredentialsTokenSource instance. - * - * @return A new SessionCredentialsTokenSource instance - */ - public SessionCredentialsTokenSource build() { - return new SessionCredentialsTokenSource(this); - } - } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index c7c3897b6..b7e237ddd 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -17,6 +17,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; @@ -218,17 +219,15 @@ void sessionCredentials() throws IOException { .execute(any(Request.class)); SessionCredentialsTokenSource sessionCredentialsTokenSource = - new SessionCredentialsTokenSource.Builder( - new Token( - "originalAccessToken", - "originalTokenType", - "originalRefreshToken", - Instant.MAX), - hc, - "https://tokenUrl", - "testClientId", - "abc") - .build(); + new SessionCredentialsTokenSource( + new Token( + "originalAccessToken", "originalTokenType", "originalRefreshToken", Instant.MAX), + hc, + "https://tokenUrl", + "testClientId", + "abc", + Optional.empty(), + Optional.empty()); Token token = sessionCredentialsTokenSource.refresh(); // We check that we are actually getting the token from server response (that is defined @@ -426,13 +425,14 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { Instant.now().plusSeconds(3600)); SessionCredentialsTokenSource browserAuthTokenSource = - new SessionCredentialsTokenSource.Builder( - browserAuthToken, - mockHttpClient, - "https://test-token-url", - "test-client-id", - "test-client-secret") - .build(); + new SessionCredentialsTokenSource( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret", + Optional.empty(), + Optional.empty()); // Create config with failing HTTP client and mock token cache DatabricksConfig config = @@ -498,13 +498,14 @@ void cacheWithInvalidTokensTest() throws IOException { Instant.now().plusSeconds(3600)); SessionCredentialsTokenSource browserAuthTokenSource = - new SessionCredentialsTokenSource.Builder( - browserAuthToken, - mockHttpClient, - "https://test-token-url", - "test-client-id", - "test-client-secret") - .build(); + new SessionCredentialsTokenSource( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret", + Optional.empty(), + Optional.empty()); // Create simple config DatabricksConfig config =