From eb0fad072590652353d4728fccbc4bbc0988aecf Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 10 Apr 2025 14:44:55 +0530 Subject: [PATCH 01/13] Adding token cache to u2m oauth --- databricks-sdk-java/pom.xml | 5 + .../databricks/sdk/core/DatabricksConfig.java | 71 ++++++ .../ExternalBrowserCredentialsProvider.java | 53 ++++- .../sdk/core/oauth/OAuthClient.java | 5 +- .../sdk/core/oauth/SessionCredentials.java | 20 +- .../com/databricks/sdk/core/oauth/Token.java | 8 +- .../databricks/sdk/core/oauth/TokenCache.java | 219 ++++++++++++++++++ .../databricks/sdk/core/utils/SerDeUtils.java | 2 + ...xternalBrowserCredentialsProviderTest.java | 153 ++++++++++++ .../sdk/core/oauth/TokenCacheTest.java | 155 +++++++++++++ 10 files changed, 679 insertions(+), 12 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java diff --git a/databricks-sdk-java/pom.xml b/databricks-sdk-java/pom.xml index 717e3aa3c..8e5596679 100644 --- a/databricks-sdk-java/pom.xml +++ b/databricks-sdk-java/pom.xml @@ -97,5 +97,10 @@ google-auth-library-oauth2-http 1.20.0 + + com.fasterxml.jackson.datatype + jackson-datatype-jsr310 + ${jackson.version} + diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 2943ce82e..542761048 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -5,6 +5,8 @@ import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.core.oauth.Token; +import com.databricks.sdk.core.oauth.TokenCache; import com.databricks.sdk.core.utils.Cloud; import com.databricks.sdk.core.utils.Environment; import com.fasterxml.jackson.databind.ObjectMapper; @@ -38,6 +40,13 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_REDIRECT_URL", auth = "oauth") private String redirectUrl; + /** + * The passphrase used to encrypt the OAuth token cache. + * This is optional and only used with token caching. + */ + @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", auth = "oauth", sensitive = true) + private String oAuthTokenCachePassphrase; + /** * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints. * @@ -141,6 +150,9 @@ public class DatabricksConfig { private DatabricksEnvironment databricksEnvironment; + // Lazily initialized OAuth token cache + private transient TokenCache tokenCache; + public Environment getEnv() { return env; } @@ -669,4 +681,63 @@ public DatabricksConfig newWithWorkspaceHost(String host) { "headerFactory")); return clone(fieldsToSkip).setHost(host); } + + public String getOAuthPassphrase() { + return oAuthTokenCachePassphrase; + } + + public DatabricksConfig setOAuthPassphrase(String oAuthPassphrase) { + this.oAuthTokenCachePassphrase = oAuthPassphrase; + return this; + } + + /** + * Gets the default OAuth redirect URL. + * If one is not provided explicitly, uses http://localhost:8080/callback + * + * @return The OAuth redirect URL to use + */ + public String getEffectiveOAuthRedirectUrl() { + return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; + } + + /** + * Gets the TokenCache instance for the current configuration. + * Creates it if it doesn't exist yet. + * + * @return A TokenCache instance for the current host, client ID, and scopes + */ + public synchronized TokenCache getTokenCache() { + if (tokenCache == null && getHost() != null && getClientId() != null) { + tokenCache = new TokenCache(getHost(), getClientId(), getScopes(), getOAuthPassphrase()); + } + return tokenCache; + } + + /** + * Loads a cached token if one exists for the current configuration. + * + * @return The cached token, or null if no valid token is cached + */ + public Token loadCachedToken() { + TokenCache cache = getTokenCache(); + if (cache == null) { + return null; + } + return cache.load(); + } + + /** + * Saves a token to the cache for the current configuration. + * + * @param token The token to cache + * @throws IOException If there is an error saving the token + */ + public void saveTokenToCache(Token token) throws IOException { + TokenCache cache = getTokenCache(); + if (cache == null) { + return; + } + cache.save(token); + } } 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 f054f06ed..cc29e8deb 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 @@ -5,12 +5,16 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.HeaderFactory; import java.io.IOException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a - * browser for the user to authorize the application. + * browser for the user to authorize the application. When cache support is enabled with + * {@link DatabricksConfig#setOAuthPassphrase}, tokens will be cached to avoid repeated authentication. */ public class ExternalBrowserCredentialsProvider implements CredentialsProvider { + private static final Logger LOGGER = LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class); @Override public String authType() { @@ -19,16 +23,53 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { - if (config.getHost() == null || config.getAuthType() != "external-browser") { + if (config.getHost() == null || config.getClientId() == null || !config.getAuthType().equals("external-browser")) { return null; } + try { - OAuthClient client = new OAuthClient(config); - Consent consent = client.initiateConsent(); - SessionCredentials creds = consent.launchExternalBrowser(); - return creds.configure(config); + // Get the token cache from config + TokenCache tokenCache = config.getTokenCache(); + + // First try to use the cached token if available + Token cachedToken = tokenCache.load(); + if (cachedToken != null && cachedToken.getRefreshToken() != null) { + LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId()); + + try { + // Create SessionCredentials with the cached token and try to refresh if needed + SessionCredentials cachedCreds = new SessionCredentials.Builder() + .withToken(cachedToken) + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .build(); + + LOGGER.debug("Using cached token (will refresh if needed)"); + return cachedCreds.configure(config); + } 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()); + } + } + + // If no cached token or refresh failed, perform browser auth + SessionCredentials credentials = performBrowserAuth(config); + tokenCache.save(credentials.getToken()); + LOGGER.debug("Saved browser auth token to cache"); + return credentials.configure(config); } catch (IOException | DatabricksException e) { + LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; } } + + SessionCredentials performBrowserAuth(DatabricksConfig config) throws IOException { + LOGGER.debug("Performing browser authentication"); + OAuthClient client = new OAuthClient(config); + Consent consent = client.initiateConsent(); + return consent.launchExternalBrowser(); + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index 6f4b25996..632fc831c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java @@ -92,10 +92,7 @@ public OAuthClient(DatabricksConfig config) throws IOException { .withClientId(config.getClientId()) .withClientSecret(config.getClientSecret()) .withHost(config.getHost()) - .withRedirectUrl( - config.getOAuthRedirectUrl() != null - ? config.getOAuthRedirectUrl() - : "http://localhost:8080/callback") + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .withScopes(config.getScopes())); } 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 aee173ea4..bc0035637 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 @@ -9,6 +9,8 @@ import java.util.HashMap; import java.util.Map; import org.apache.http.HttpHeaders; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An implementation of RefreshableTokenSource implementing the refresh_token OAuth grant type. @@ -20,6 +22,7 @@ public class SessionCredentials extends RefreshableTokenSource implements CredentialsProvider, Serializable { private static final long serialVersionUID = 3083941540130596650L; + private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentials.class); @Override public String authType() { @@ -28,6 +31,8 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { + this.tokenCache = config.getTokenCache(); + return () -> { Map headers = new HashMap<>(); headers.put( @@ -84,6 +89,7 @@ public SessionCredentials build() { private final String redirectUrl; private final String clientId; private final String clientSecret; + private transient TokenCache tokenCache; private SessionCredentials(Builder b) { super(b.token); @@ -113,7 +119,19 @@ protected Token refresh() { // cross-origin requests headers.put("Origin", redirectUrl); } - return retrieveToken( + Token newToken = retrieveToken( hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); + + // Save the refreshed token directly to cache + if (tokenCache != null) { + try { + tokenCache.save(newToken); + LOGGER.debug("Saved refreshed token to cache"); + } catch (Exception e) { + LOGGER.warn("Failed to save token to cache: {}", e.getMessage()); + } + } + + return newToken; } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index 153abe066..92a04a1a5 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -2,6 +2,7 @@ import com.databricks.sdk.core.utils.ClockSupplier; import com.databricks.sdk.core.utils.SystemClockSupplier; +import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; @@ -37,7 +38,12 @@ public Token( } /** Constructor for refreshable tokens. */ - public Token(String accessToken, String tokenType, String refreshToken, LocalDateTime expiry) { + @JsonCreator + public Token( + @JsonProperty("accessToken") String accessToken, + @JsonProperty("tokenType") String tokenType, + @JsonProperty("refreshToken") String refreshToken, + @JsonProperty("expiry") LocalDateTime expiry) { this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier()); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java new file mode 100644 index 000000000..e15a74ec3 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java @@ -0,0 +1,219 @@ +package com.databricks.sdk.core.oauth; + +import java.io.*; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.security.spec.KeySpec; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; +import java.util.Objects; +import javax.crypto.Cipher; +import javax.crypto.SecretKey; +import javax.crypto.SecretKeyFactory; +import javax.crypto.spec.PBEKeySpec; +import javax.crypto.spec.SecretKeySpec; +import com.databricks.sdk.core.utils.SerDeUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; + +/** + * TokenCache stores OAuth tokens on disk to avoid repeated authentication. + * It generates a unique cache filename based on the host, client ID, and scopes. + * If a passphrase is provided, the token data is encrypted for added security. + */ +public class TokenCache { + // Base path for token cache files, aligned with Python implementation + private static final String BASE_PATH = ".config/databricks-sdk-java/oauth"; + + // Encryption constants + private static final String ALGORITHM = "AES"; + private static final String SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256"; + private static final byte[] SALT = "DatabricksTokenCache".getBytes(); // Fixed salt for simplicity + private static final int ITERATION_COUNT = 65536; + private static final int KEY_LENGTH = 256; + + private final String host; + private final String clientId; + private final List scopes; + private final Path cacheFile; + private final String passphrase; + private final ObjectMapper mapper; + + /** + * Constructs a new TokenCache instance for OAuth token caching + * + * @param host The Databricks host URL + * @param clientId The OAuth client ID + * @param scopes The OAuth scopes requested (optional) + * @param passphrase The passphrase used to encrypt/decrypt the token cache (optional) + */ + public TokenCache( + String host, + String clientId, + List scopes, + String passphrase) { + this.host = Objects.requireNonNull(host, "host must be defined"); + this.clientId = Objects.requireNonNull(clientId, "clientId must be defined"); + this.scopes = scopes != null ? scopes : new ArrayList<>(); + this.passphrase = passphrase; // Can be null or empty, encryption will be skipped in that case + this.mapper = SerDeUtils.createMapper(); + + this.cacheFile = getFilename(); + } + + /** + * Returns the path to the cache file for the current configuration. + * The filename is based on a hash of the host, client ID, and scopes. + * + * @return The path to the token cache file + */ + @VisibleForTesting + Path getFilename() { + try { + // Create SHA-256 hash of host, client_id, and scopes + MessageDigest hash = MessageDigest.getInstance("SHA-256"); + for (String chunk : new String[] { + this.host, + this.clientId, + String.join(",", this.scopes) + }) { + hash.update(chunk.getBytes(StandardCharsets.UTF_8)); + } + + // Convert hash bytes to hexadecimal string + StringBuilder hexString = new StringBuilder(); + for (byte b : hash.digest()) { + String hex = Integer.toHexString(0xff & b); + if (hex.length() == 1) { + hexString.append('0'); + } + hexString.append(hex); + } + + String userHome = System.getProperty("user.home"); + Path basePath = Paths.get(userHome, BASE_PATH); + return basePath.resolve(hexString.toString()); + } catch (NoSuchAlgorithmException e) { + throw new RuntimeException("Failed to create hash for token cache filename", e); + } + } + + /** + * Determines if encryption should be used based on whether a valid passphrase is provided + * @return true if encryption should be used, false otherwise + */ + private boolean shouldUseEncryption() { + return passphrase != null && !passphrase.isEmpty(); + } + + /** + * Saves a Token to the cache file, encrypting if a passphrase is provided + * + * @param token The Token to save + * @throws IOException If an error occurs writing to the file + */ + public void save(Token token) throws IOException { + try { + Files.createDirectories(cacheFile.getParent()); + + byte[] dataToWrite; + + // Serialize token to JSON + String json = mapper.writeValueAsString(token); + dataToWrite = json.getBytes(StandardCharsets.UTF_8); + + // Encrypt data if a passphrase is provided + if (shouldUseEncryption()) { + dataToWrite = encrypt(dataToWrite); + } + + Files.write(cacheFile, dataToWrite); + // Set file permissions to be readable only by the owner (equivalent to 0600) + cacheFile.toFile().setReadable(false, false); + cacheFile.toFile().setReadable(true, true); + cacheFile.toFile().setWritable(false, false); + cacheFile.toFile().setWritable(true, true); + } catch (Exception e) { + throw new IOException("Failed to save token cache: " + e.getMessage(), e); + } + } + + /** + * Loads a Token from the cache file, decrypting if a passphrase was provided + * + * @return The Token from the cache or null if the cache file doesn't exist or is invalid + */ + public Token load() { + try { + if (!Files.exists(cacheFile)) { + return null; + } + + byte[] fileContent = Files.readAllBytes(cacheFile); + byte[] decodedContent; + + if (shouldUseEncryption()) { + try { + decodedContent = decrypt(fileContent); + } catch (Exception e) { + // If decryption fails, it might be because the file was saved without encryption + // or the passphrase is incorrect + return null; + } + } else { + decodedContent = fileContent; + } + + // Deserialize token from JSON + String json = new String(decodedContent, StandardCharsets.UTF_8); + return mapper.readValue(json, Token.class); + } catch (Exception e) { + // If there's any issue loading the token, return null + // to allow a fresh token to be obtained + return null; + } + } + + /** + * Generates a secret key from the passphrase using PBKDF2 with HMAC-SHA256. + * + * @return A SecretKey generated from the passphrase + * @throws Exception If an error occurs generating the key + */ + private SecretKey generateSecretKey() throws Exception { + SecretKeyFactory factory = SecretKeyFactory.getInstance(SECRET_KEY_ALGORITHM); + KeySpec spec = new PBEKeySpec(passphrase.toCharArray(), SALT, ITERATION_COUNT, KEY_LENGTH); + return new SecretKeySpec(factory.generateSecret(spec).getEncoded(), ALGORITHM); + } + + /** + * Encrypts the given data using AES encryption with a key derived from the passphrase. + * + * @param data The data to encrypt + * @return The encrypted data + * @throws Exception If an error occurs during encryption + */ + private byte[] encrypt(byte[] data) throws Exception { + Cipher cipher = Cipher.getInstance(ALGORITHM); + cipher.init(Cipher.ENCRYPT_MODE, generateSecretKey()); + return Base64.getEncoder().encode(cipher.doFinal(data)); + } + + /** + * Decrypts the given encrypted data using AES decryption with a key derived from the passphrase. + * + * @param encryptedData The encrypted data, Base64 encoded + * @return The decrypted data + * @throws Exception If an error occurs during decryption + */ + private byte[] decrypt(byte[] encryptedData) throws Exception { + Cipher cipher = Cipher.getInstance(ALGORITHM); + cipher.init(Cipher.DECRYPT_MODE, generateSecretKey()); + return cipher.doFinal(Base64.getDecoder().decode(encryptedData)); + } +} \ No newline at end of file diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SerDeUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SerDeUtils.java index 916eafd23..9730b135e 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SerDeUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SerDeUtils.java @@ -4,12 +4,14 @@ import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; /** Utilities for serialization and deserialization in the Databricks Java SDK. */ public class SerDeUtils { public static ObjectMapper createMapper() { ObjectMapper mapper = new ObjectMapper(); mapper + .registerModule(new JavaTimeModule()) .configure(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, false) .configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false) .configure(DeserializationFeature.ACCEPT_EMPTY_STRING_AS_NULL_OBJECT, true) 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 1551045a7..2aed9cd93 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 @@ -6,6 +6,7 @@ import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.FixtureServer; +import com.databricks.sdk.core.HeaderFactory; import com.databricks.sdk.core.commons.CommonsHttpClient; import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; @@ -201,4 +202,156 @@ void sessionCredentials() throws IOException { assertEquals("accessTokenFromServer", token.getAccessToken()); assertEquals("refreshTokenFromServer", token.getRefreshToken()); } + + // Token caching tests + + @Test + void cacheWithValidTokenTest() throws IOException { + // Create a valid token (not expired) + LocalDateTime futureTime = LocalDateTime.now().plusHours(1); + Token validToken = new Token("valid_access_token", "Bearer", "valid_refresh_token", futureTime); + + // Create mock token cache that returns the valid token + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.doReturn(validToken).when(mockTokenCache).load(); + + // Create config with the mock token cache + DatabricksConfig config = new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id"); + + // Spy on the config to inject the mock token cache + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + OpenIDConnectEndpoints oidcEndpoints = Mockito.mock(OpenIDConnectEndpoints.class); + Mockito.doReturn(oidcEndpoints).when(spyConfig).getOidcEndpoints(); + Mockito.when(oidcEndpoints.getTokenEndpoint()).thenReturn("https://test.databricks.com/v1/token"); + + // Create our provider and mock the browser auth method + ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + + // Configure provider + HeaderFactory headerFactory = provider.configure(spyConfig); + + // Verify token was loaded from cache + Mockito.verify(mockTokenCache, Mockito.times(1)).load(); + + // Verify performBrowserAuth was NOT called since we had a valid cached token + Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); + + // Verify headers contain the cached token + Map headers = headerFactory.headers(); + assertEquals("Bearer valid_access_token", headers.get("Authorization")); + } + + @Test + void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { + // Create mock HTTP client for token refresh + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + 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.doReturn(new Response(refreshResponse, url)).when(mockHttpClient).execute(any(Request.class)); + + // Create an expired token with valid refresh token + LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Token expiredToken = new Token("expired_access_token", "Bearer", "valid_refresh_token", pastTime); + + // Create mock token cache that returns the expired token + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); + + // Create config with HTTP client and mock token cache + DatabricksConfig config = new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); + + // We need to provide OIDC endpoints for token refresh + OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( + "https://test.databricks.com/token", + "https://test.databricks.com/authorize"); + + // Create our provider and mock the browser auth method + ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + + // Spy on the config to inject the mock token cache and endpoints + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); + + // Configure provider + HeaderFactory headerFactory = provider.configure(spyConfig); + + // Verify headers contain the refreshed token, not the browser auth token or expired token + Map headers = headerFactory.headers(); + assertEquals("Bearer refreshed_access_token", headers.get("Authorization")); + + // Verify token was loaded from cache + Mockito.verify(mockTokenCache, Mockito.times(1)).load(); + + // Verify HTTP call was made to refresh the token + Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); + + // Verify performBrowserAuth was NOT called since refresh succeeded + Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); + + // Verify token was saved back to cache + Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + } + + @Test + void cacheWithInvalidTokensTest() throws IOException { + // Create completely invalid token (no refresh token) + LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Token invalidToken = new Token("expired_access_token", "Bearer", null, pastTime); + + // Create mock token cache that returns the invalid token + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.when(mockTokenCache.load()).thenReturn(invalidToken); + + // Setup browser auth result (should be used as fallback) + Token browserAuthToken = new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); + + // Create simple config + DatabricksConfig config = new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id"); + + // Create our provider and mock the browser auth method + ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); + + // Spy on the config to inject the mock token cache + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + + // Configure provider + HeaderFactory headerFactory = provider.configure(spyConfig); + // Verify headers contain the browser auth token (fallback) + Map headers = headerFactory.headers(); + assertEquals("Bearer browser_access_token", headers.get("Authorization")); + + // Verify token was loaded from cache + Mockito.verify(mockTokenCache, Mockito.times(1)).load(); + + // Verify performBrowserAuth was called since we had an invalid token + Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + + // Verify token was saved after browser auth (for the new token) + Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java new file mode 100644 index 000000000..f1db091ff --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java @@ -0,0 +1,155 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksConfig; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.LocalDateTime; +import java.util.Arrays; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.spy; + +public class TokenCacheTest { + private static final String TEST_HOST = "https://test-host.cloud.databricks.com"; + private static final String TEST_CLIENT_ID = "test-client-id"; + private static final List TEST_SCOPES = Arrays.asList("offline_access", "clusters", "sql"); + private static final String TEST_PASS = "test-passphrase"; + private Path cacheFile; + private TokenCache tokenCache; + + @BeforeEach + void setUp() { + tokenCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS); + cacheFile = tokenCache.getFilename(); + } + + @AfterEach + void tearDown() throws IOException { + Files.deleteIfExists(cacheFile); + } + + @Test + void testEmptyCache() throws IOException { + assertNull(tokenCache.load()); + } + + @Test + void testSaveAndLoadToken() throws IOException { + LocalDateTime expiry = LocalDateTime.now().plusHours(1); + Token token = new Token("access-token", "Bearer", "refresh-token", expiry); + + tokenCache.save(token); + + Token loadedToken = tokenCache.load(); + assertNotNull(loadedToken, "Loaded token should not be null"); + assertEquals("access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("refresh-token", loadedToken.getRefreshToken()); + // No direct way to compare expiry times exactly, but we can compare if they're both in the future + assertFalse(loadedToken.isExpired(), "Token should not be expired"); + } + + @Test + void testTokenExpiry() throws IOException { + LocalDateTime now = LocalDateTime.now(); + LocalDateTime expiry = now.plusMinutes(30); + Token token = new Token("access-token", "Bearer", "refresh-token", expiry); + + tokenCache.save(token); + + Token loadedToken = tokenCache.load(); + assertNotNull(loadedToken, "Loaded token should not be null"); + assertFalse(loadedToken.isExpired(), "Token should not be expired"); + + // Create a new token with past expiry + Token expiredToken = new Token("access-token", "Bearer", "refresh-token", now.minusHours(1)); + assertTrue(expiredToken.isExpired(), "Token should be expired"); + } + + @Test + void testInvalidPassphrase() { + assertThrows(NullPointerException.class, () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS)); + } + + @Test + void testOverwriteToken() throws IOException { + Token token1 = new Token("token1", "Bearer", "refresh1", LocalDateTime.now().plusHours(1)); + Token token2 = new Token("token2", "Bearer", "refresh2", LocalDateTime.now().plusHours(2)); + + tokenCache.save(token1); + tokenCache.save(token2); + + Token loadedToken = tokenCache.load(); + assertNotNull(loadedToken, "Loaded token should not be null"); + assertEquals("token2", loadedToken.getAccessToken()); + assertEquals("refresh2", loadedToken.getRefreshToken()); + } + + @Test + void testTokenCacheSaveLoad(@TempDir Path tempDir) throws IOException { + // Create a token cache that uses a temp directory + TokenCache cache = spy(new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS)); + + // Create test token + Token testToken = new Token( + "test-access-token", + "Bearer", + "test-refresh-token", + LocalDateTime.now().plusHours(1)); + + // Save to cache + cache.save(testToken); + + // Load from cache + Token loadedToken = cache.load(); + assertNotNull(loadedToken, "Should load token from cache"); + + // Verify token details + assertEquals("test-access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("test-refresh-token", loadedToken.getRefreshToken()); + } + + @Test + void testDatabricksConfigTokenCache() throws IOException { + // Create a DatabricksConfig with the test values + DatabricksConfig config = new DatabricksConfig() + .setHost(TEST_HOST) + .setClientId(TEST_CLIENT_ID) + .setScopes(TEST_SCOPES) + .setOAuthPassphrase(TEST_PASS); + + // Get TokenCache from config + TokenCache cache = config.getTokenCache(); + assertNotNull(cache, "TokenCache from config should not be null"); + + // Create test token + Token testToken = new Token( + "config-access-token", + "Bearer", + "config-refresh-token", + LocalDateTime.now().plusHours(1)); + + // Save to cache directly + cache.save(testToken); + + // Load from cache directly + Token loadedToken = cache.load(); + assertNotNull(loadedToken, "Should load token from cache directly"); + + // Verify token details + assertEquals("config-access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("config-refresh-token", loadedToken.getRefreshToken()); + + // Clean up + Files.deleteIfExists(cache.getFilename()); + } +} From e8080b37f7e1bb15e31f6b189fa8b72701d83413 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 10 Apr 2025 16:28:48 +0530 Subject: [PATCH 02/13] add a way to enable/disable token cache --- .../databricks/sdk/core/DatabricksConfig.java | 64 +++--- .../ExternalBrowserCredentialsProvider.java | 10 +- .../databricks/sdk/core/oauth/TokenCache.java | 34 +++- ...xternalBrowserCredentialsProviderTest.java | 191 ++++++++++++++++-- .../sdk/core/oauth/TokenCacheTest.java | 66 +++++- 5 files changed, 300 insertions(+), 65 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 542761048..e3cef53d7 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -47,6 +47,13 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", auth = "oauth", sensitive = true) private String oAuthTokenCachePassphrase; + /** + * Controls whether OAuth token caching is enabled. + * When set to false, tokens will not be cached or loaded from cache. + */ + @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_ENABLED", auth = "oauth") + private Boolean isTokenCacheEnabled; + /** * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints. * @@ -682,15 +689,36 @@ public DatabricksConfig newWithWorkspaceHost(String host) { return clone(fieldsToSkip).setHost(host); } - public String getOAuthPassphrase() { + public String getOAuthTokenCachePassphrase() { return oAuthTokenCachePassphrase; } - public DatabricksConfig setOAuthPassphrase(String oAuthPassphrase) { + public DatabricksConfig setOAuthTokenCachePassphrase(String oAuthPassphrase) { this.oAuthTokenCachePassphrase = oAuthPassphrase; return this; } + /** + * Gets whether OAuth token caching is enabled. + * Default is true. + * + * @return true if token caching is enabled, false otherwise + */ + public boolean isTokenCacheEnabled() { + return isTokenCacheEnabled == null || isTokenCacheEnabled; + } + + /** + * Sets whether OAuth token caching is enabled. + * + * @param enabled true to enable token caching, false to disable + * @return this config instance + */ + public DatabricksConfig setTokenCacheEnabled(boolean enabled) { + this.isTokenCacheEnabled = enabled; + return this; + } + /** * Gets the default OAuth redirect URL. * If one is not provided explicitly, uses http://localhost:8080/callback @@ -704,40 +732,14 @@ public String getEffectiveOAuthRedirectUrl() { /** * Gets the TokenCache instance for the current configuration. * Creates it if it doesn't exist yet. + * When token caching is disabled, the TokenCache will be created but operations will be no-ops. * * @return A TokenCache instance for the current host, client ID, and scopes */ public synchronized TokenCache getTokenCache() { - if (tokenCache == null && getHost() != null && getClientId() != null) { - tokenCache = new TokenCache(getHost(), getClientId(), getScopes(), getOAuthPassphrase()); + if (tokenCache == null) { + tokenCache = new TokenCache(getHost(), getClientId(), getScopes(), getOAuthTokenCachePassphrase(), isTokenCacheEnabled()); } return tokenCache; } - - /** - * Loads a cached token if one exists for the current configuration. - * - * @return The cached token, or null if no valid token is cached - */ - public Token loadCachedToken() { - TokenCache cache = getTokenCache(); - if (cache == null) { - return null; - } - return cache.load(); - } - - /** - * Saves a token to the cache for the current configuration. - * - * @param token The token to cache - * @throws IOException If there is an error saving the token - */ - public void saveTokenToCache(Token token) throws IOException { - TokenCache cache = getTokenCache(); - if (cache == null) { - return; - } - cache.save(token); - } } 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 cc29e8deb..b89a2b821 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 @@ -11,7 +11,8 @@ /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a * browser for the user to authorize the application. When cache support is enabled with - * {@link DatabricksConfig#setOAuthPassphrase}, tokens will be cached to avoid repeated authentication. + * {@link DatabricksConfig#setOAuthTokenCachePassphrase} and {@link DatabricksConfig#setTokenCacheEnabled(boolean)}, + * tokens will be cached to avoid repeated authentication. */ public class ExternalBrowserCredentialsProvider implements CredentialsProvider { private static final Logger LOGGER = LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class); @@ -31,7 +32,7 @@ public HeaderFactory configure(DatabricksConfig config) { // Get the token cache from config TokenCache tokenCache = config.getTokenCache(); - // First try to use the cached token if available + // First try to use the cached token if available (will return null if disabled) Token cachedToken = tokenCache.load(); if (cachedToken != null && cachedToken.getRefreshToken() != null) { LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId()); @@ -47,7 +48,9 @@ public HeaderFactory configure(DatabricksConfig config) { .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .build(); - LOGGER.debug("Using cached token (will refresh if needed)"); + LOGGER.debug("Using cached token, will immediately refresh"); + cachedCreds.token = cachedCreds.refresh(); + tokenCache.save(cachedToken); return cachedCreds.configure(config); } catch (Exception e) { // If token refresh fails, log and continue to browser auth @@ -58,7 +61,6 @@ public HeaderFactory configure(DatabricksConfig config) { // If no cached token or refresh failed, perform browser auth SessionCredentials credentials = performBrowserAuth(config); tokenCache.save(credentials.getToken()); - LOGGER.debug("Saved browser auth token to cache"); return credentials.configure(config); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java index e15a74ec3..45090e865 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java @@ -20,13 +20,18 @@ import com.databricks.sdk.core.utils.SerDeUtils; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * TokenCache stores OAuth tokens on disk to avoid repeated authentication. * It generates a unique cache filename based on the host, client ID, and scopes. * If a passphrase is provided, the token data is encrypted for added security. + * Cache operations can be disabled by setting isEnabled to false. */ public class TokenCache { + private static final Logger LOGGER = LoggerFactory.getLogger(TokenCache.class); + // Base path for token cache files, aligned with Python implementation private static final String BASE_PATH = ".config/databricks-sdk-java/oauth"; @@ -43,6 +48,7 @@ public class TokenCache { private final Path cacheFile; private final String passphrase; private final ObjectMapper mapper; + private final boolean isEnabled; /** * Constructs a new TokenCache instance for OAuth token caching @@ -51,17 +57,20 @@ public class TokenCache { * @param clientId The OAuth client ID * @param scopes The OAuth scopes requested (optional) * @param passphrase The passphrase used to encrypt/decrypt the token cache (optional) + * @param isEnabled Whether token caching is enabled */ public TokenCache( String host, String clientId, List scopes, - String passphrase) { + String passphrase, + boolean isEnabled) { this.host = Objects.requireNonNull(host, "host must be defined"); this.clientId = Objects.requireNonNull(clientId, "clientId must be defined"); this.scopes = scopes != null ? scopes : new ArrayList<>(); this.passphrase = passphrase; // Can be null or empty, encryption will be skipped in that case this.mapper = SerDeUtils.createMapper(); + this.isEnabled = isEnabled; this.cacheFile = getFilename(); } @@ -113,11 +122,17 @@ private boolean shouldUseEncryption() { /** * Saves a Token to the cache file, encrypting if a passphrase is provided + * Does nothing if caching is disabled * * @param token The Token to save * @throws IOException If an error occurs writing to the file */ public void save(Token token) throws IOException { + if (!isEnabled) { + LOGGER.debug("Token caching is disabled, skipping save operation"); + return; + } + try { Files.createDirectories(cacheFile.getParent()); @@ -138,6 +153,8 @@ public void save(Token token) throws IOException { cacheFile.toFile().setReadable(true, true); cacheFile.toFile().setWritable(false, false); cacheFile.toFile().setWritable(true, true); + + LOGGER.debug("Successfully saved token to cache: {}", cacheFile); } catch (Exception e) { throw new IOException("Failed to save token cache: " + e.getMessage(), e); } @@ -145,12 +162,19 @@ public void save(Token token) throws IOException { /** * Loads a Token from the cache file, decrypting if a passphrase was provided + * Returns null if caching is disabled * - * @return The Token from the cache or null if the cache file doesn't exist or is invalid + * @return The Token from the cache or null if the cache file doesn't exist, is invalid, or caching is disabled */ public Token load() { + if (!isEnabled) { + LOGGER.debug("Token caching is disabled, skipping load operation"); + return null; + } + try { if (!Files.exists(cacheFile)) { + LOGGER.debug("No token cache file found at: {}", cacheFile); return null; } @@ -163,6 +187,7 @@ public Token load() { } catch (Exception e) { // If decryption fails, it might be because the file was saved without encryption // or the passphrase is incorrect + LOGGER.debug("Failed to decrypt token cache: {}", e.getMessage()); return null; } } else { @@ -171,10 +196,13 @@ public Token load() { // Deserialize token from JSON String json = new String(decodedContent, StandardCharsets.UTF_8); - return mapper.readValue(json, Token.class); + Token token = mapper.readValue(json, Token.class); + LOGGER.debug("Successfully loaded token from cache: {}", cacheFile); + return token; } catch (Exception e) { // If there's any issue loading the token, return null // to allow a fresh token to be obtained + LOGGER.debug("Failed to load token from cache: {}", e.getMessage()); return null; } } 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 2aed9cd93..cd176c6d3 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 @@ -19,6 +19,7 @@ import java.util.Map; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import java.nio.file.Files; public class ExternalBrowserCredentialsProviderTest { @Test @@ -207,42 +208,59 @@ void sessionCredentials() throws IOException { @Test void cacheWithValidTokenTest() throws IOException { - // Create a valid token (not expired) + // Create mock HTTP client for token refresh + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + 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.doReturn(new Response(refreshResponse, url)).when(mockHttpClient).execute(any(Request.class)); + + // Create an valid token with valid refresh token LocalDateTime futureTime = LocalDateTime.now().plusHours(1); Token validToken = new Token("valid_access_token", "Bearer", "valid_refresh_token", futureTime); - + // Create mock token cache that returns the valid token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.doReturn(validToken).when(mockTokenCache).load(); + Mockito.when(mockTokenCache.load()).thenReturn(validToken); - // Create config with the mock token cache + // Create config with HTTP client and mock token cache DatabricksConfig config = new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id"); + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); + + // We need to provide OIDC endpoints for token refresh + OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( + "https://test.databricks.com/token", + "https://test.databricks.com/authorize"); - // Spy on the config to inject the mock token cache - DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); - OpenIDConnectEndpoints oidcEndpoints = Mockito.mock(OpenIDConnectEndpoints.class); - Mockito.doReturn(oidcEndpoints).when(spyConfig).getOidcEndpoints(); - Mockito.when(oidcEndpoints.getTokenEndpoint()).thenReturn("https://test.databricks.com/v1/token"); - // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + // Spy on the config to inject the mock token cache and endpoints + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); + // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); - + + // Verify headers contain the refreshed token even though the cached token is valid + Map headers = headerFactory.headers(); + assertEquals("Bearer refreshed_access_token", headers.get("Authorization")); + // Verify token was loaded from cache Mockito.verify(mockTokenCache, Mockito.times(1)).load(); - - // Verify performBrowserAuth was NOT called since we had a valid cached token + + // Verify HTTP call was made to refresh the token + Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); + + // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); - - // Verify headers contain the cached token - Map headers = headerFactory.headers(); - assertEquals("Bearer valid_access_token", headers.get("Authorization")); + + // Verify token was saved back to cache + Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } @Test @@ -301,6 +319,71 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } + + @Test + void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { + // Create HTTP client that fails when refreshing token + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + Mockito.when(mockHttpClient.execute(any(Request.class))).thenThrow(new IOException("Failed to refresh token")); + + // Create an expired token with invalid refresh token + LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Token expiredToken = new Token("expired_access_token", "Bearer", "invalid_refresh_token", pastTime); + + // Create mock token cache that returns the expired token + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); + + // Setup browser auth result (should be used as fallback) + Token browserAuthToken = new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); + + // Create config with failing HTTP client and mock token cache + DatabricksConfig config = new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); + + // We need to provide OIDC endpoints for token refresh attempt + OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( + "https://test.databricks.com/token", + "https://test.databricks.com/authorize"); + + // Create our provider and mock the browser auth method + ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); + + // Spy on the config to inject the mock token cache and endpoints + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); + + // Configure provider + HeaderFactory headerFactory = provider.configure(spyConfig); + + // Verify headers contain the browser auth token (fallback) + Map headers = headerFactory.headers(); + assertEquals("Bearer browser_access_token", headers.get("Authorization")); + + // Verify token was loaded from cache + Mockito.verify(mockTokenCache, Mockito.times(1)).load(); + + // Verify performBrowserAuth was called since refresh failed + Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + + // Verify token was saved after browser auth (for the new token) + Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + } @Test void cacheWithInvalidTokensTest() throws IOException { @@ -344,7 +427,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify headers contain the browser auth token (fallback) Map headers = headerFactory.headers(); assertEquals("Bearer browser_access_token", headers.get("Authorization")); - + // Verify token was loaded from cache Mockito.verify(mockTokenCache, Mockito.times(1)).load(); @@ -354,4 +437,68 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } + + @Test + void disabledTokenCacheTest() throws IOException { + // Create mock HTTP client for token operations + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + + // Setup browser auth result + Token browserAuthToken = new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); + + // Create config with caching explicitly disabled + DatabricksConfig config = new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient) + .setTokenCacheEnabled(false); + + // We need to provide OIDC endpoints + OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( + "https://test.databricks.com/token", + "https://test.databricks.com/authorize"); + + // Create our provider and mock the browser auth method + ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); + + // Create a real token cache with caching disabled + TokenCache tokenCache = new TokenCache( + "https://test.databricks.com", + "test-client-id", + Arrays.asList("offline_access"), + "test-passphrase", + false); + + // Spy on the config to inject the token cache and endpoints + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.doReturn(tokenCache).when(spyConfig).getTokenCache(); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); + + // Configure provider + HeaderFactory headerFactory = provider.configure(spyConfig); + + // Verify headers contain the browser auth token + Map headers = headerFactory.headers(); + assertEquals("Bearer browser_access_token", headers.get("Authorization")); + + // Verify performBrowserAuth was called immediately (no attempt to use cache) + Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + + // The save operation should be called but be a no-op + // We can't easily verify the no-op behavior in this test, but we can + // verify that the cache file doesn't exist + assertFalse(Files.exists(tokenCache.getFilename()), "Cache file should not be created when disabled"); + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java index f1db091ff..9fa3c1312 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java @@ -24,9 +24,9 @@ public class TokenCacheTest { private Path cacheFile; private TokenCache tokenCache; - @BeforeEach + @BeforeEach void setUp() { - tokenCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS); + tokenCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true); cacheFile = tokenCache.getFilename(); } @@ -75,7 +75,7 @@ void testTokenExpiry() throws IOException { @Test void testInvalidPassphrase() { - assertThrows(NullPointerException.class, () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS)); + assertThrows(NullPointerException.class, () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); } @Test @@ -95,7 +95,7 @@ void testOverwriteToken() throws IOException { @Test void testTokenCacheSaveLoad(@TempDir Path tempDir) throws IOException { // Create a token cache that uses a temp directory - TokenCache cache = spy(new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS)); + TokenCache cache = spy(new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); // Create test token Token testToken = new Token( @@ -124,7 +124,8 @@ void testDatabricksConfigTokenCache() throws IOException { .setHost(TEST_HOST) .setClientId(TEST_CLIENT_ID) .setScopes(TEST_SCOPES) - .setOAuthPassphrase(TEST_PASS); + .setOAuthTokenCachePassphrase(TEST_PASS) + .setTokenCacheEnabled(true); // Get TokenCache from config TokenCache cache = config.getTokenCache(); @@ -152,4 +153,59 @@ void testDatabricksConfigTokenCache() throws IOException { // Clean up Files.deleteIfExists(cache.getFilename()); } + + @Test + void testDisabledTokenCache() throws IOException { + // Create a disabled token cache + TokenCache disabledCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, false); + + // Create test token + Token testToken = new Token( + "disabled-token", + "Bearer", + "disabled-refresh", + LocalDateTime.now().plusHours(1)); + + // Try to save to cache - should be a no-op + disabledCache.save(testToken); + + // Load from cache - should return null + Token loadedToken = disabledCache.load(); + assertNull(loadedToken, "Disabled cache should not load a token"); + + // Verify the file wasn't created + assertFalse(Files.exists(disabledCache.getFilename()), "Cache file should not be created when disabled"); + } + + @Test + void testDatabricksConfigDisabledTokenCache() throws IOException { + // Create a DatabricksConfig with caching disabled + DatabricksConfig config = new DatabricksConfig() + .setHost(TEST_HOST) + .setClientId(TEST_CLIENT_ID) + .setScopes(TEST_SCOPES) + .setOAuthTokenCachePassphrase(TEST_PASS) + .setTokenCacheEnabled(false); + + // Get TokenCache from config + TokenCache cache = config.getTokenCache(); + assertNotNull(cache, "TokenCache should still be created even when disabled"); + + // Create test token + Token testToken = new Token( + "disabled-config-token", + "Bearer", + "disabled-config-refresh", + LocalDateTime.now().plusHours(1)); + + // Try to save to cache - should be a no-op + cache.save(testToken); + + // Load from cache - should return null + Token loadedToken = cache.load(); + assertNull(loadedToken, "Disabled cache should not load a token"); + + // Verify the file wasn't created + assertFalse(Files.exists(cache.getFilename()), "Cache file should not be created when disabled"); + } } From b57f6e2ad79b9b1504ae4da4e0c58a914ecb5ce9 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 10 Apr 2025 21:31:07 +0530 Subject: [PATCH 03/13] fmt --- .../databricks/sdk/core/DatabricksConfig.java | 42 +-- .../ExternalBrowserCredentialsProvider.java | 39 +-- .../sdk/core/oauth/SessionCredentials.java | 7 +- .../com/databricks/sdk/core/oauth/Token.java | 6 +- .../databricks/sdk/core/oauth/TokenCache.java | 78 +++-- ...xternalBrowserCredentialsProviderTest.java | 269 ++++++++++-------- .../sdk/core/oauth/TokenCacheTest.java | 146 +++++----- 7 files changed, 315 insertions(+), 272 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index e3cef53d7..304e6e04b 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -5,7 +5,6 @@ import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; -import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.oauth.TokenCache; import com.databricks.sdk.core.utils.Cloud; import com.databricks.sdk.core.utils.Environment; @@ -41,15 +40,18 @@ public class DatabricksConfig { private String redirectUrl; /** - * The passphrase used to encrypt the OAuth token cache. - * This is optional and only used with token caching. + * The passphrase used to encrypt the OAuth token cache. This is optional and only used with token + * caching. */ - @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", auth = "oauth", sensitive = true) + @ConfigAttribute( + env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", + auth = "oauth", + sensitive = true) private String oAuthTokenCachePassphrase; /** - * Controls whether OAuth token caching is enabled. - * When set to false, tokens will not be cached or loaded from cache. + * Controls whether OAuth token caching is enabled. When set to false, tokens will not be cached + * or loaded from cache. */ @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_ENABLED", auth = "oauth") private Boolean isTokenCacheEnabled; @@ -699,9 +701,8 @@ public DatabricksConfig setOAuthTokenCachePassphrase(String oAuthPassphrase) { } /** - * Gets whether OAuth token caching is enabled. - * Default is true. - * + * Gets whether OAuth token caching is enabled. Default is true. + * * @return true if token caching is enabled, false otherwise */ public boolean isTokenCacheEnabled() { @@ -710,7 +711,7 @@ public boolean isTokenCacheEnabled() { /** * Sets whether OAuth token caching is enabled. - * + * * @param enabled true to enable token caching, false to disable * @return this config instance */ @@ -720,25 +721,30 @@ public DatabricksConfig setTokenCacheEnabled(boolean enabled) { } /** - * Gets the default OAuth redirect URL. - * If one is not provided explicitly, uses http://localhost:8080/callback - * + * Gets the default OAuth redirect URL. If one is not provided explicitly, uses + * http://localhost:8080/callback + * * @return The OAuth redirect URL to use */ public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } - + /** - * Gets the TokenCache instance for the current configuration. - * Creates it if it doesn't exist yet. + * Gets the TokenCache instance for the current configuration. Creates it if it doesn't exist yet. * When token caching is disabled, the TokenCache will be created but operations will be no-ops. - * + * * @return A TokenCache instance for the current host, client ID, and scopes */ public synchronized TokenCache getTokenCache() { if (tokenCache == null) { - tokenCache = new TokenCache(getHost(), getClientId(), getScopes(), getOAuthTokenCachePassphrase(), isTokenCacheEnabled()); + tokenCache = + new TokenCache( + getHost(), + getClientId(), + getScopes(), + getOAuthTokenCachePassphrase(), + isTokenCacheEnabled()); } return tokenCache; } 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 b89a2b821..ae751f394 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 @@ -10,12 +10,14 @@ /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a - * browser for the user to authorize the application. When cache support is enabled with - * {@link DatabricksConfig#setOAuthTokenCachePassphrase} and {@link DatabricksConfig#setTokenCacheEnabled(boolean)}, - * tokens will be cached to avoid repeated authentication. + * browser for the user to authorize the application. When cache support is enabled with {@link + * DatabricksConfig#setOAuthTokenCachePassphrase} and {@link + * DatabricksConfig#setTokenCacheEnabled(boolean)}, tokens will be cached to avoid repeated + * authentication. */ public class ExternalBrowserCredentialsProvider implements CredentialsProvider { - private static final Logger LOGGER = LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class); + private static final Logger LOGGER = + LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class); @Override public String authType() { @@ -24,10 +26,12 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { - if (config.getHost() == null || config.getClientId() == null || !config.getAuthType().equals("external-browser")) { + if (config.getHost() == null + || config.getClientId() == null + || !config.getAuthType().equals("external-browser")) { return null; } - + try { // Get the token cache from config TokenCache tokenCache = config.getTokenCache(); @@ -36,18 +40,19 @@ public HeaderFactory configure(DatabricksConfig config) { Token cachedToken = tokenCache.load(); if (cachedToken != null && cachedToken.getRefreshToken() != null) { LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId()); - + try { // Create SessionCredentials with the cached token and try to refresh if needed - SessionCredentials cachedCreds = new SessionCredentials.Builder() - .withToken(cachedToken) - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .build(); - + SessionCredentials cachedCreds = + new SessionCredentials.Builder() + .withToken(cachedToken) + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .build(); + LOGGER.debug("Using cached token, will immediately refresh"); cachedCreds.token = cachedCreds.refresh(); tokenCache.save(cachedToken); @@ -67,7 +72,7 @@ public HeaderFactory configure(DatabricksConfig config) { return null; } } - + SessionCredentials performBrowserAuth(DatabricksConfig config) throws IOException { LOGGER.debug("Performing browser authentication"); OAuthClient client = new OAuthClient(config); 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 bc0035637..4a2f0e512 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 @@ -119,8 +119,9 @@ protected Token refresh() { // cross-origin requests headers.put("Origin", redirectUrl); } - Token newToken = retrieveToken( - hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); + Token newToken = + retrieveToken( + hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); // Save the refreshed token directly to cache if (tokenCache != null) { @@ -131,7 +132,7 @@ protected Token refresh() { LOGGER.warn("Failed to save token to cache: {}", e.getMessage()); } } - + return newToken; } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index 92a04a1a5..f0fd72f68 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -40,9 +40,9 @@ public Token( /** Constructor for refreshable tokens. */ @JsonCreator public Token( - @JsonProperty("accessToken") String accessToken, - @JsonProperty("tokenType") String tokenType, - @JsonProperty("refreshToken") String refreshToken, + @JsonProperty("accessToken") String accessToken, + @JsonProperty("tokenType") String tokenType, + @JsonProperty("refreshToken") String refreshToken, @JsonProperty("expiry") LocalDateTime expiry) { this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier()); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java index 45090e865..2ad1356cc 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java @@ -1,5 +1,8 @@ package com.databricks.sdk.core.oauth; +import com.databricks.sdk.core.utils.SerDeUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; import java.io.*; import java.nio.charset.StandardCharsets; import java.nio.file.Files; @@ -17,24 +20,21 @@ import javax.crypto.SecretKeyFactory; import javax.crypto.spec.PBEKeySpec; import javax.crypto.spec.SecretKeySpec; -import com.databricks.sdk.core.utils.SerDeUtils; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** - * TokenCache stores OAuth tokens on disk to avoid repeated authentication. - * It generates a unique cache filename based on the host, client ID, and scopes. - * If a passphrase is provided, the token data is encrypted for added security. - * Cache operations can be disabled by setting isEnabled to false. + * TokenCache stores OAuth tokens on disk to avoid repeated authentication. It generates a unique + * cache filename based on the host, client ID, and scopes. If a passphrase is provided, the token + * data is encrypted for added security. Cache operations can be disabled by setting isEnabled to + * false. */ public class TokenCache { private static final Logger LOGGER = LoggerFactory.getLogger(TokenCache.class); - + // Base path for token cache files, aligned with Python implementation private static final String BASE_PATH = ".config/databricks-sdk-java/oauth"; - + // Encryption constants private static final String ALGORITHM = "AES"; private static final String SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256"; @@ -60,25 +60,21 @@ public class TokenCache { * @param isEnabled Whether token caching is enabled */ public TokenCache( - String host, - String clientId, - List scopes, - String passphrase, - boolean isEnabled) { + String host, String clientId, List scopes, String passphrase, boolean isEnabled) { this.host = Objects.requireNonNull(host, "host must be defined"); this.clientId = Objects.requireNonNull(clientId, "clientId must be defined"); this.scopes = scopes != null ? scopes : new ArrayList<>(); this.passphrase = passphrase; // Can be null or empty, encryption will be skipped in that case this.mapper = SerDeUtils.createMapper(); this.isEnabled = isEnabled; - + this.cacheFile = getFilename(); } /** - * Returns the path to the cache file for the current configuration. - * The filename is based on a hash of the host, client ID, and scopes. - * + * Returns the path to the cache file for the current configuration. The filename is based on a + * hash of the host, client ID, and scopes. + * * @return The path to the token cache file */ @VisibleForTesting @@ -86,14 +82,10 @@ Path getFilename() { try { // Create SHA-256 hash of host, client_id, and scopes MessageDigest hash = MessageDigest.getInstance("SHA-256"); - for (String chunk : new String[] { - this.host, - this.clientId, - String.join(",", this.scopes) - }) { + for (String chunk : new String[] {this.host, this.clientId, String.join(",", this.scopes)}) { hash.update(chunk.getBytes(StandardCharsets.UTF_8)); } - + // Convert hash bytes to hexadecimal string StringBuilder hexString = new StringBuilder(); for (byte b : hash.digest()) { @@ -103,7 +95,7 @@ Path getFilename() { } hexString.append(hex); } - + String userHome = System.getProperty("user.home"); Path basePath = Paths.get(userHome, BASE_PATH); return basePath.resolve(hexString.toString()); @@ -114,6 +106,7 @@ Path getFilename() { /** * Determines if encryption should be used based on whether a valid passphrase is provided + * * @return true if encryption should be used, false otherwise */ private boolean shouldUseEncryption() { @@ -121,8 +114,8 @@ private boolean shouldUseEncryption() { } /** - * Saves a Token to the cache file, encrypting if a passphrase is provided - * Does nothing if caching is disabled + * Saves a Token to the cache file, encrypting if a passphrase is provided Does nothing if caching + * is disabled * * @param token The Token to save * @throws IOException If an error occurs writing to the file @@ -132,28 +125,28 @@ public void save(Token token) throws IOException { LOGGER.debug("Token caching is disabled, skipping save operation"); return; } - + try { Files.createDirectories(cacheFile.getParent()); - + byte[] dataToWrite; - + // Serialize token to JSON String json = mapper.writeValueAsString(token); dataToWrite = json.getBytes(StandardCharsets.UTF_8); - + // Encrypt data if a passphrase is provided if (shouldUseEncryption()) { dataToWrite = encrypt(dataToWrite); } - + Files.write(cacheFile, dataToWrite); // Set file permissions to be readable only by the owner (equivalent to 0600) cacheFile.toFile().setReadable(false, false); cacheFile.toFile().setReadable(true, true); cacheFile.toFile().setWritable(false, false); cacheFile.toFile().setWritable(true, true); - + LOGGER.debug("Successfully saved token to cache: {}", cacheFile); } catch (Exception e) { throw new IOException("Failed to save token cache: " + e.getMessage(), e); @@ -161,26 +154,27 @@ public void save(Token token) throws IOException { } /** - * Loads a Token from the cache file, decrypting if a passphrase was provided - * Returns null if caching is disabled + * Loads a Token from the cache file, decrypting if a passphrase was provided Returns null if + * caching is disabled * - * @return The Token from the cache or null if the cache file doesn't exist, is invalid, or caching is disabled + * @return The Token from the cache or null if the cache file doesn't exist, is invalid, or + * caching is disabled */ public Token load() { if (!isEnabled) { LOGGER.debug("Token caching is disabled, skipping load operation"); return null; } - + try { if (!Files.exists(cacheFile)) { LOGGER.debug("No token cache file found at: {}", cacheFile); return null; } - + byte[] fileContent = Files.readAllBytes(cacheFile); byte[] decodedContent; - + if (shouldUseEncryption()) { try { decodedContent = decrypt(fileContent); @@ -193,7 +187,7 @@ public Token load() { } else { decodedContent = fileContent; } - + // Deserialize token from JSON String json = new String(decodedContent, StandardCharsets.UTF_8); Token token = mapper.readValue(json, Token.class); @@ -206,7 +200,7 @@ public Token load() { return null; } } - + /** * Generates a secret key from the passphrase using PBKDF2 with HMAC-SHA256. * @@ -244,4 +238,4 @@ private byte[] decrypt(byte[] encryptedData) throws Exception { cipher.init(Cipher.DECRYPT_MODE, generateSecretKey()); return cipher.doFinal(Base64.getDecoder().decode(encryptedData)); } -} \ No newline at end of file +} 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 cd176c6d3..34205c2dc 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 @@ -13,13 +13,13 @@ import com.databricks.sdk.core.http.Response; import java.io.IOException; import java.net.URL; +import java.nio.file.Files; import java.time.LocalDateTime; import java.util.Arrays; import java.util.HashMap; import java.util.Map; import org.junit.jupiter.api.Test; import org.mockito.Mockito; -import java.nio.file.Files; public class ExternalBrowserCredentialsProviderTest { @Test @@ -211,9 +211,11 @@ void cacheWithValidTokenTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); String refreshResponse = - "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; + "{\"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)).when(mockHttpClient).execute(any(Request.class)); + Mockito.doReturn(new Response(refreshResponse, url)) + .when(mockHttpClient) + .execute(any(Request.class)); // Create an valid token with valid refresh token LocalDateTime futureTime = LocalDateTime.now().plusHours(1); @@ -224,19 +226,21 @@ void cacheWithValidTokenTest() throws IOException { Mockito.when(mockTokenCache.load()).thenReturn(validToken); // Create config with HTTP client and mock token cache - DatabricksConfig config = new DatabricksConfig() + DatabricksConfig config = + new DatabricksConfig() .setAuthType("external-browser") .setHost("https://test.databricks.com") .setClientId("test-client-id") .setHttpClient(mockHttpClient); // We need to provide OIDC endpoints for token refresh - OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", - "https://test.databricks.com/authorize"); + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/token", "https://test.databricks.com/authorize"); // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider()); // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); @@ -262,44 +266,49 @@ void cacheWithValidTokenTest() throws IOException { // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } - + @Test void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Create mock HTTP client for token refresh HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String refreshResponse = + 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.doReturn(new Response(refreshResponse, url)).when(mockHttpClient).execute(any(Request.class)); - + Mockito.doReturn(new Response(refreshResponse, url)) + .when(mockHttpClient) + .execute(any(Request.class)); + // Create an expired token with valid refresh token LocalDateTime pastTime = LocalDateTime.now().minusHours(1); - Token expiredToken = new Token("expired_access_token", "Bearer", "valid_refresh_token", pastTime); - + Token expiredToken = + new Token("expired_access_token", "Bearer", "valid_refresh_token", pastTime); + // Create mock token cache that returns the expired token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); - + // Create config with HTTP client and mock token cache - DatabricksConfig config = new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id") - .setHttpClient(mockHttpClient); - + DatabricksConfig config = + new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); + // We need to provide OIDC endpoints for token refresh - OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", - "https://test.databricks.com/authorize"); - + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/token", "https://test.databricks.com/authorize"); + // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider()); // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); - + // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); @@ -309,13 +318,13 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Verify token was loaded from cache Mockito.verify(mockTokenCache, Mockito.times(1)).load(); - + // Verify HTTP call was made to refresh the token Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); - + // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); - + // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } @@ -324,44 +333,52 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create HTTP client that fails when refreshing token HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - Mockito.when(mockHttpClient.execute(any(Request.class))).thenThrow(new IOException("Failed to refresh token")); + Mockito.when(mockHttpClient.execute(any(Request.class))) + .thenThrow(new IOException("Failed to refresh token")); // Create an expired token with invalid refresh token LocalDateTime pastTime = LocalDateTime.now().minusHours(1); - Token expiredToken = new Token("expired_access_token", "Bearer", "invalid_refresh_token", pastTime); + Token expiredToken = + new Token("expired_access_token", "Bearer", "invalid_refresh_token", pastTime); // Create mock token cache that returns the expired token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); // Setup browser auth result (should be used as fallback) - Token browserAuthToken = new Token( - "browser_access_token", - "Bearer", - "browser_refresh_token", - LocalDateTime.now().plusHours(1)); - - SessionCredentials browserAuthCreds = new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); + Token browserAuthToken = + new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = + new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); // Create config with failing HTTP client and mock token cache - DatabricksConfig config = new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id") - .setHttpClient(mockHttpClient); + DatabricksConfig config = + new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); // We need to provide OIDC endpoints for token refresh attempt - OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", - "https://test.databricks.com/authorize"); + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/token", "https://test.databricks.com/authorize"); // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); - Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds) + .when(provider) + .performBrowserAuth(any(DatabricksConfig.class)); // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); @@ -384,44 +401,50 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } - + @Test void cacheWithInvalidTokensTest() throws IOException { // Create completely invalid token (no refresh token) LocalDateTime pastTime = LocalDateTime.now().minusHours(1); Token invalidToken = new Token("expired_access_token", "Bearer", null, pastTime); - + // Create mock token cache that returns the invalid token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); Mockito.when(mockTokenCache.load()).thenReturn(invalidToken); - + // Setup browser auth result (should be used as fallback) - Token browserAuthToken = new Token( - "browser_access_token", - "Bearer", - "browser_refresh_token", - LocalDateTime.now().plusHours(1)); - - SessionCredentials browserAuthCreds = new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); - + Token browserAuthToken = + new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = + new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); + // Create simple config - DatabricksConfig config = new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id"); - + DatabricksConfig config = + new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id"); + // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); - Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); - + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds) + .when(provider) + .performBrowserAuth(any(DatabricksConfig.class)); + // Spy on the config to inject the mock token cache DatabricksConfig spyConfig = Mockito.spy(config); Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); - + // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); // Verify headers contain the browser auth token (fallback) @@ -430,75 +453,83 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify token was loaded from cache Mockito.verify(mockTokenCache, Mockito.times(1)).load(); - + // Verify performBrowserAuth was called since we had an invalid token Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); - + // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } - + @Test void disabledTokenCacheTest() throws IOException { // Create mock HTTP client for token operations HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - + // Setup browser auth result - Token browserAuthToken = new Token( - "browser_access_token", - "Bearer", - "browser_refresh_token", - LocalDateTime.now().plusHours(1)); - - SessionCredentials browserAuthCreds = new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); - + Token browserAuthToken = + new Token( + "browser_access_token", + "Bearer", + "browser_refresh_token", + LocalDateTime.now().plusHours(1)); + + SessionCredentials browserAuthCreds = + new SessionCredentials.Builder() + .withToken(browserAuthToken) + .withClientId("test-client-id") + .withTokenUrl("https://test-token-url") + .build(); + // Create config with caching explicitly disabled - DatabricksConfig config = new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id") - .setHttpClient(mockHttpClient) - .setTokenCacheEnabled(false); - + DatabricksConfig config = + new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient) + .setTokenCacheEnabled(false); + // We need to provide OIDC endpoints - OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", - "https://test.databricks.com/authorize"); - + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/token", "https://test.databricks.com/authorize"); + // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider()); - Mockito.doReturn(browserAuthCreds).when(provider).performBrowserAuth(any(DatabricksConfig.class)); - + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.doReturn(browserAuthCreds) + .when(provider) + .performBrowserAuth(any(DatabricksConfig.class)); + // Create a real token cache with caching disabled - TokenCache tokenCache = new TokenCache( - "https://test.databricks.com", - "test-client-id", - Arrays.asList("offline_access"), - "test-passphrase", - false); - + TokenCache tokenCache = + new TokenCache( + "https://test.databricks.com", + "test-client-id", + Arrays.asList("offline_access"), + "test-passphrase", + false); + // Spy on the config to inject the token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); Mockito.doReturn(tokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); - + // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); - + // Verify headers contain the browser auth token Map headers = headerFactory.headers(); assertEquals("Bearer browser_access_token", headers.get("Authorization")); - + // Verify performBrowserAuth was called immediately (no attempt to use cache) Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); - + // The save operation should be called but be a no-op // We can't easily verify the no-op behavior in this test, but we can // verify that the cache file doesn't exist - assertFalse(Files.exists(tokenCache.getFilename()), "Cache file should not be created when disabled"); + assertFalse( + Files.exists(tokenCache.getFilename()), "Cache file should not be created when disabled"); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java index 9fa3c1312..923913aaf 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java @@ -1,25 +1,25 @@ package com.databricks.sdk.core.oauth; -import com.databricks.sdk.core.DatabricksConfig; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.io.TempDir; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.spy; +import com.databricks.sdk.core.DatabricksConfig; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.time.LocalDateTime; import java.util.Arrays; import java.util.List; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.spy; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; public class TokenCacheTest { private static final String TEST_HOST = "https://test-host.cloud.databricks.com"; private static final String TEST_CLIENT_ID = "test-client-id"; - private static final List TEST_SCOPES = Arrays.asList("offline_access", "clusters", "sql"); + private static final List TEST_SCOPES = + Arrays.asList("offline_access", "clusters", "sql"); private static final String TEST_PASS = "test-passphrase"; private Path cacheFile; private TokenCache tokenCache; @@ -44,15 +44,16 @@ void testEmptyCache() throws IOException { void testSaveAndLoadToken() throws IOException { LocalDateTime expiry = LocalDateTime.now().plusHours(1); Token token = new Token("access-token", "Bearer", "refresh-token", expiry); - + tokenCache.save(token); - + Token loadedToken = tokenCache.load(); assertNotNull(loadedToken, "Loaded token should not be null"); assertEquals("access-token", loadedToken.getAccessToken()); assertEquals("Bearer", loadedToken.getTokenType()); assertEquals("refresh-token", loadedToken.getRefreshToken()); - // No direct way to compare expiry times exactly, but we can compare if they're both in the future + // No direct way to compare expiry times exactly, but we can compare if they're both in the + // future assertFalse(loadedToken.isExpired(), "Token should not be expired"); } @@ -61,9 +62,9 @@ void testTokenExpiry() throws IOException { LocalDateTime now = LocalDateTime.now(); LocalDateTime expiry = now.plusMinutes(30); Token token = new Token("access-token", "Bearer", "refresh-token", expiry); - + tokenCache.save(token); - + Token loadedToken = tokenCache.load(); assertNotNull(loadedToken, "Loaded token should not be null"); assertFalse(loadedToken.isExpired(), "Token should not be expired"); @@ -75,7 +76,9 @@ void testTokenExpiry() throws IOException { @Test void testInvalidPassphrase() { - assertThrows(NullPointerException.class, () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); + assertThrows( + NullPointerException.class, + () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); } @Test @@ -98,114 +101,117 @@ void testTokenCacheSaveLoad(@TempDir Path tempDir) throws IOException { TokenCache cache = spy(new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); // Create test token - Token testToken = new Token( - "test-access-token", - "Bearer", - "test-refresh-token", - LocalDateTime.now().plusHours(1)); - + Token testToken = + new Token( + "test-access-token", "Bearer", "test-refresh-token", LocalDateTime.now().plusHours(1)); + // Save to cache cache.save(testToken); // Load from cache Token loadedToken = cache.load(); assertNotNull(loadedToken, "Should load token from cache"); - + // Verify token details assertEquals("test-access-token", loadedToken.getAccessToken()); assertEquals("Bearer", loadedToken.getTokenType()); assertEquals("test-refresh-token", loadedToken.getRefreshToken()); } - + @Test void testDatabricksConfigTokenCache() throws IOException { // Create a DatabricksConfig with the test values - DatabricksConfig config = new DatabricksConfig() - .setHost(TEST_HOST) - .setClientId(TEST_CLIENT_ID) - .setScopes(TEST_SCOPES) - .setOAuthTokenCachePassphrase(TEST_PASS) - .setTokenCacheEnabled(true); - + DatabricksConfig config = + new DatabricksConfig() + .setHost(TEST_HOST) + .setClientId(TEST_CLIENT_ID) + .setScopes(TEST_SCOPES) + .setOAuthTokenCachePassphrase(TEST_PASS) + .setTokenCacheEnabled(true); + // Get TokenCache from config TokenCache cache = config.getTokenCache(); assertNotNull(cache, "TokenCache from config should not be null"); - + // Create test token - Token testToken = new Token( - "config-access-token", - "Bearer", - "config-refresh-token", - LocalDateTime.now().plusHours(1)); - + Token testToken = + new Token( + "config-access-token", + "Bearer", + "config-refresh-token", + LocalDateTime.now().plusHours(1)); + // Save to cache directly cache.save(testToken); - + // Load from cache directly Token loadedToken = cache.load(); assertNotNull(loadedToken, "Should load token from cache directly"); - + // Verify token details assertEquals("config-access-token", loadedToken.getAccessToken()); assertEquals("Bearer", loadedToken.getTokenType()); assertEquals("config-refresh-token", loadedToken.getRefreshToken()); - + // Clean up Files.deleteIfExists(cache.getFilename()); } - + @Test void testDisabledTokenCache() throws IOException { // Create a disabled token cache - TokenCache disabledCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, false); - + TokenCache disabledCache = + new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, false); + // Create test token - Token testToken = new Token( - "disabled-token", - "Bearer", - "disabled-refresh", - LocalDateTime.now().plusHours(1)); - + Token testToken = + new Token("disabled-token", "Bearer", "disabled-refresh", LocalDateTime.now().plusHours(1)); + // Try to save to cache - should be a no-op disabledCache.save(testToken); - + // Load from cache - should return null Token loadedToken = disabledCache.load(); assertNull(loadedToken, "Disabled cache should not load a token"); - + // Verify the file wasn't created - assertFalse(Files.exists(disabledCache.getFilename()), "Cache file should not be created when disabled"); + assertFalse( + Files.exists(disabledCache.getFilename()), + "Cache file should not be created when disabled"); } - + @Test void testDatabricksConfigDisabledTokenCache() throws IOException { // Create a DatabricksConfig with caching disabled - DatabricksConfig config = new DatabricksConfig() - .setHost(TEST_HOST) - .setClientId(TEST_CLIENT_ID) - .setScopes(TEST_SCOPES) - .setOAuthTokenCachePassphrase(TEST_PASS) - .setTokenCacheEnabled(false); - + DatabricksConfig config = + new DatabricksConfig() + .setHost(TEST_HOST) + .setClientId(TEST_CLIENT_ID) + .setScopes(TEST_SCOPES) + .setOAuthTokenCachePassphrase(TEST_PASS) + .setTokenCacheEnabled(false); + // Get TokenCache from config TokenCache cache = config.getTokenCache(); assertNotNull(cache, "TokenCache should still be created even when disabled"); - + // Create test token - Token testToken = new Token( - "disabled-config-token", - "Bearer", - "disabled-config-refresh", - LocalDateTime.now().plusHours(1)); - + Token testToken = + new Token( + "disabled-config-token", + "Bearer", + "disabled-config-refresh", + LocalDateTime.now().plusHours(1)); + // Try to save to cache - should be a no-op cache.save(testToken); - + // Load from cache - should return null Token loadedToken = cache.load(); assertNull(loadedToken, "Disabled cache should not load a token"); - + // Verify the file wasn't created - assertFalse(Files.exists(cache.getFilename()), "Cache file should not be created when disabled"); + assertFalse( + Files.exists(cache.getFilename()), "Cache file should not be created when disabled"); } } From 6f1b69d1ddcf1404c29e2b226206baa437dfa8ba Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 10 Apr 2025 21:33:58 +0530 Subject: [PATCH 04/13] add changelog --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 50df105cb..f9f9f9787 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,6 +3,7 @@ ## Release v0.44.0 ### New Features and Improvements +Added `TokenCache` to `ExternalBrowserCredentialsProvider` to reduce number of authentications needed for U2M OAuth ### Bug Fixes From d8c366e03bec14d8bb966de55c2f1a36ef0853ed Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Fri, 11 Apr 2025 18:31:34 +0530 Subject: [PATCH 05/13] update implementation based on discussion --- NEXT_CHANGELOG.md | 4 +- databricks-sdk-java/pom.xml | 1 + .../databricks/sdk/core/DatabricksConfig.java | 71 ++---- .../ExternalBrowserCredentialsProvider.java | 6 +- .../sdk/core/oauth/FileTokenCache.java | 75 ++++++ .../sdk/core/oauth/SessionCredentials.java | 9 +- .../databricks/sdk/core/oauth/TokenCache.java | 236 +----------------- .../sdk/core/oauth/TokenCacheUtils.java | 49 ++++ .../sdk/core/DatabricksConfigTest.java | 43 ++++ ...xternalBrowserCredentialsProviderTest.java | 49 ++-- .../sdk/core/oauth/FileTokenCacheTest.java | 128 ++++++++++ .../sdk/core/oauth/TokenCacheTest.java | 217 ---------------- 12 files changed, 350 insertions(+), 538 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java delete mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index f9f9f9787..782458501 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,8 +3,8 @@ ## Release v0.44.0 ### New Features and Improvements -Added `TokenCache` to `ExternalBrowserCredentialsProvider` to reduce number of authentications needed for U2M OAuth - + * Added `TokenCache` to `ExternalBrowserCredentialsProvider` to reduce number of authentications needed for U2M OAuth. + ### Bug Fixes ### Documentation diff --git a/databricks-sdk-java/pom.xml b/databricks-sdk-java/pom.xml index 8e5596679..f2df8e656 100644 --- a/databricks-sdk-java/pom.xml +++ b/databricks-sdk-java/pom.xml @@ -97,6 +97,7 @@ google-auth-library-oauth2-http 1.20.0 + com.fasterxml.jackson.datatype jackson-datatype-jsr310 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 304e6e04b..a065b44b2 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -4,14 +4,17 @@ import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; +import com.databricks.sdk.core.oauth.FileTokenCache; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; import com.databricks.sdk.core.oauth.TokenCache; +import com.databricks.sdk.core.oauth.TokenCacheUtils; import com.databricks.sdk.core.utils.Cloud; import com.databricks.sdk.core.utils.Environment; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.File; import java.io.IOException; import java.lang.reflect.Field; +import java.nio.file.Path; import java.util.*; import org.apache.http.HttpMessage; @@ -39,23 +42,6 @@ public class DatabricksConfig { @ConfigAttribute(env = "DATABRICKS_REDIRECT_URL", auth = "oauth") private String redirectUrl; - /** - * The passphrase used to encrypt the OAuth token cache. This is optional and only used with token - * caching. - */ - @ConfigAttribute( - env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", - auth = "oauth", - sensitive = true) - private String oAuthTokenCachePassphrase; - - /** - * Controls whether OAuth token caching is enabled. When set to false, tokens will not be cached - * or loaded from cache. - */ - @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_ENABLED", auth = "oauth") - private Boolean isTokenCacheEnabled; - /** * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints. * @@ -395,13 +381,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** @deprecated Use {@link #getAzureUseMsi()} instead. */ + /** + * @deprecated Use {@link #getAzureUseMsi()} instead. + */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ + /** + * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. + */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; @@ -691,32 +681,14 @@ public DatabricksConfig newWithWorkspaceHost(String host) { return clone(fieldsToSkip).setHost(host); } - public String getOAuthTokenCachePassphrase() { - return oAuthTokenCachePassphrase; - } - - public DatabricksConfig setOAuthTokenCachePassphrase(String oAuthPassphrase) { - this.oAuthTokenCachePassphrase = oAuthPassphrase; - return this; - } - - /** - * Gets whether OAuth token caching is enabled. Default is true. - * - * @return true if token caching is enabled, false otherwise - */ - public boolean isTokenCacheEnabled() { - return isTokenCacheEnabled == null || isTokenCacheEnabled; - } - /** - * Sets whether OAuth token caching is enabled. + * Sets a custom TokenCache implementation. * - * @param enabled true to enable token caching, false to disable + * @param tokenCache the TokenCache implementation to use * @return this config instance */ - public DatabricksConfig setTokenCacheEnabled(boolean enabled) { - this.isTokenCacheEnabled = enabled; + public DatabricksConfig setTokenCache(TokenCache tokenCache) { + this.tokenCache = tokenCache; return this; } @@ -731,20 +703,17 @@ public String getEffectiveOAuthRedirectUrl() { } /** - * Gets the TokenCache instance for the current configuration. Creates it if it doesn't exist yet. - * When token caching is disabled, the TokenCache will be created but operations will be no-ops. + * Gets the TokenCache instance for the current configuration. + * + *

If a custom TokenCache has been set, it will be returned. Otherwise, a SimpleFileTokenCache + * will be created based on the configuration properties. * - * @return A TokenCache instance for the current host, client ID, and scopes + * @return A TokenCache instance */ public synchronized TokenCache getTokenCache() { if (tokenCache == null) { - tokenCache = - new TokenCache( - getHost(), - getClientId(), - getScopes(), - getOAuthTokenCachePassphrase(), - isTokenCacheEnabled()); + Path cachePath = TokenCacheUtils.getCacheFilePath(getHost(), getClientId(), getScopes()); + tokenCache = new FileTokenCache(cachePath); } return tokenCache; } 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 ae751f394..40bb16910 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 @@ -10,10 +10,8 @@ /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a - * browser for the user to authorize the application. When cache support is enabled with {@link - * DatabricksConfig#setOAuthTokenCachePassphrase} and {@link - * DatabricksConfig#setTokenCacheEnabled(boolean)}, tokens will be cached to avoid repeated - * authentication. + * browser for the user to authorize the application. Uses the cache from {@link + * DatabricksConfig#getTokenCache()}, tokens will be cached to avoid repeated authentication. */ public class ExternalBrowserCredentialsProvider implements CredentialsProvider { private static final Logger LOGGER = diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java new file mode 100644 index 000000000..fc64dd273 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java @@ -0,0 +1,75 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.utils.SerDeUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Objects; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** A TokenCache implementation that stores tokens as plain files. */ +public class FileTokenCache implements TokenCache { + private static final Logger LOGGER = LoggerFactory.getLogger(FileTokenCache.class); + + private final Path cacheFile; + private final ObjectMapper mapper; + + /** + * Constructs a new SimpleFileTokenCache instance. + * + * @param cacheFilePath The path where the token cache will be stored + */ + public FileTokenCache(Path cacheFilePath) { + Objects.requireNonNull(cacheFilePath, "cacheFilePath must be defined"); + + this.cacheFile = cacheFilePath; + this.mapper = SerDeUtils.createMapper(); + } + + @Override + public void save(Token token) { + try { + Files.createDirectories(cacheFile.getParent()); + + // Serialize token to JSON + String json = mapper.writeValueAsString(token); + byte[] dataToWrite = json.getBytes(StandardCharsets.UTF_8); + + Files.write(cacheFile, dataToWrite); + // Set file permissions to be readable only by the owner (equivalent to 0600) + cacheFile.toFile().setReadable(false, false); + cacheFile.toFile().setReadable(true, true); + cacheFile.toFile().setWritable(false, false); + cacheFile.toFile().setWritable(true, true); + + LOGGER.debug("Successfully saved token to cache: {}", cacheFile); + } catch (Exception e) { + LOGGER.warn("Failed to save token to cache: {}", cacheFile, e); + } + } + + @Override + public Token load() { + try { + if (!Files.exists(cacheFile)) { + LOGGER.debug("No token cache file found at: {}", cacheFile); + return null; + } + + byte[] fileContent = Files.readAllBytes(cacheFile); + + // Deserialize token from JSON + String json = new String(fileContent, StandardCharsets.UTF_8); + Token token = mapper.readValue(json, Token.class); + LOGGER.debug("Successfully loaded token from cache: {}", cacheFile); + return token; + } catch (Exception e) { + // If there's any issue loading the token, return null + // to allow a fresh token to be obtained + LOGGER.warn("Failed to load token from cache: {}", e.getMessage()); + return null; + } + } +} 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 4a2f0e512..56ec588e6 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 @@ -125,14 +125,9 @@ protected Token refresh() { // Save the refreshed token directly to cache if (tokenCache != null) { - try { - tokenCache.save(newToken); - LOGGER.debug("Saved refreshed token to cache"); - } catch (Exception e) { - LOGGER.warn("Failed to save token to cache: {}", e.getMessage()); - } + 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/TokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java index 2ad1356cc..ed8fb533b 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCache.java @@ -1,241 +1,21 @@ package com.databricks.sdk.core.oauth; -import com.databricks.sdk.core.utils.SerDeUtils; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.annotations.VisibleForTesting; -import java.io.*; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.security.MessageDigest; -import java.security.NoSuchAlgorithmException; -import java.security.spec.KeySpec; -import java.util.ArrayList; -import java.util.Base64; -import java.util.List; -import java.util.Objects; -import javax.crypto.Cipher; -import javax.crypto.SecretKey; -import javax.crypto.SecretKeyFactory; -import javax.crypto.spec.PBEKeySpec; -import javax.crypto.spec.SecretKeySpec; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - /** - * TokenCache stores OAuth tokens on disk to avoid repeated authentication. It generates a unique - * cache filename based on the host, client ID, and scopes. If a passphrase is provided, the token - * data is encrypted for added security. Cache operations can be disabled by setting isEnabled to - * false. + * TokenCache interface for storing and retrieving OAuth tokens. Implementations can use different + * storage mechanisms and security approaches. */ -public class TokenCache { - private static final Logger LOGGER = LoggerFactory.getLogger(TokenCache.class); - - // Base path for token cache files, aligned with Python implementation - private static final String BASE_PATH = ".config/databricks-sdk-java/oauth"; - - // Encryption constants - private static final String ALGORITHM = "AES"; - private static final String SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256"; - private static final byte[] SALT = "DatabricksTokenCache".getBytes(); // Fixed salt for simplicity - private static final int ITERATION_COUNT = 65536; - private static final int KEY_LENGTH = 256; - - private final String host; - private final String clientId; - private final List scopes; - private final Path cacheFile; - private final String passphrase; - private final ObjectMapper mapper; - private final boolean isEnabled; - +public interface TokenCache { /** - * Constructs a new TokenCache instance for OAuth token caching - * - * @param host The Databricks host URL - * @param clientId The OAuth client ID - * @param scopes The OAuth scopes requested (optional) - * @param passphrase The passphrase used to encrypt/decrypt the token cache (optional) - * @param isEnabled Whether token caching is enabled - */ - public TokenCache( - String host, String clientId, List scopes, String passphrase, boolean isEnabled) { - this.host = Objects.requireNonNull(host, "host must be defined"); - this.clientId = Objects.requireNonNull(clientId, "clientId must be defined"); - this.scopes = scopes != null ? scopes : new ArrayList<>(); - this.passphrase = passphrase; // Can be null or empty, encryption will be skipped in that case - this.mapper = SerDeUtils.createMapper(); - this.isEnabled = isEnabled; - - this.cacheFile = getFilename(); - } - - /** - * Returns the path to the cache file for the current configuration. The filename is based on a - * hash of the host, client ID, and scopes. - * - * @return The path to the token cache file - */ - @VisibleForTesting - Path getFilename() { - try { - // Create SHA-256 hash of host, client_id, and scopes - MessageDigest hash = MessageDigest.getInstance("SHA-256"); - for (String chunk : new String[] {this.host, this.clientId, String.join(",", this.scopes)}) { - hash.update(chunk.getBytes(StandardCharsets.UTF_8)); - } - - // Convert hash bytes to hexadecimal string - StringBuilder hexString = new StringBuilder(); - for (byte b : hash.digest()) { - String hex = Integer.toHexString(0xff & b); - if (hex.length() == 1) { - hexString.append('0'); - } - hexString.append(hex); - } - - String userHome = System.getProperty("user.home"); - Path basePath = Paths.get(userHome, BASE_PATH); - return basePath.resolve(hexString.toString()); - } catch (NoSuchAlgorithmException e) { - throw new RuntimeException("Failed to create hash for token cache filename", e); - } - } - - /** - * Determines if encryption should be used based on whether a valid passphrase is provided - * - * @return true if encryption should be used, false otherwise - */ - private boolean shouldUseEncryption() { - return passphrase != null && !passphrase.isEmpty(); - } - - /** - * Saves a Token to the cache file, encrypting if a passphrase is provided Does nothing if caching - * is disabled + * Saves a Token to the cache. * * @param token The Token to save - * @throws IOException If an error occurs writing to the file - */ - public void save(Token token) throws IOException { - if (!isEnabled) { - LOGGER.debug("Token caching is disabled, skipping save operation"); - return; - } - - try { - Files.createDirectories(cacheFile.getParent()); - - byte[] dataToWrite; - - // Serialize token to JSON - String json = mapper.writeValueAsString(token); - dataToWrite = json.getBytes(StandardCharsets.UTF_8); - - // Encrypt data if a passphrase is provided - if (shouldUseEncryption()) { - dataToWrite = encrypt(dataToWrite); - } - - Files.write(cacheFile, dataToWrite); - // Set file permissions to be readable only by the owner (equivalent to 0600) - cacheFile.toFile().setReadable(false, false); - cacheFile.toFile().setReadable(true, true); - cacheFile.toFile().setWritable(false, false); - cacheFile.toFile().setWritable(true, true); - - LOGGER.debug("Successfully saved token to cache: {}", cacheFile); - } catch (Exception e) { - throw new IOException("Failed to save token cache: " + e.getMessage(), e); - } - } - - /** - * Loads a Token from the cache file, decrypting if a passphrase was provided Returns null if - * caching is disabled - * - * @return The Token from the cache or null if the cache file doesn't exist, is invalid, or - * caching is disabled - */ - public Token load() { - if (!isEnabled) { - LOGGER.debug("Token caching is disabled, skipping load operation"); - return null; - } - - try { - if (!Files.exists(cacheFile)) { - LOGGER.debug("No token cache file found at: {}", cacheFile); - return null; - } - - byte[] fileContent = Files.readAllBytes(cacheFile); - byte[] decodedContent; - - if (shouldUseEncryption()) { - try { - decodedContent = decrypt(fileContent); - } catch (Exception e) { - // If decryption fails, it might be because the file was saved without encryption - // or the passphrase is incorrect - LOGGER.debug("Failed to decrypt token cache: {}", e.getMessage()); - return null; - } - } else { - decodedContent = fileContent; - } - - // Deserialize token from JSON - String json = new String(decodedContent, StandardCharsets.UTF_8); - Token token = mapper.readValue(json, Token.class); - LOGGER.debug("Successfully loaded token from cache: {}", cacheFile); - return token; - } catch (Exception e) { - // If there's any issue loading the token, return null - // to allow a fresh token to be obtained - LOGGER.debug("Failed to load token from cache: {}", e.getMessage()); - return null; - } - } - - /** - * Generates a secret key from the passphrase using PBKDF2 with HMAC-SHA256. - * - * @return A SecretKey generated from the passphrase - * @throws Exception If an error occurs generating the key - */ - private SecretKey generateSecretKey() throws Exception { - SecretKeyFactory factory = SecretKeyFactory.getInstance(SECRET_KEY_ALGORITHM); - KeySpec spec = new PBEKeySpec(passphrase.toCharArray(), SALT, ITERATION_COUNT, KEY_LENGTH); - return new SecretKeySpec(factory.generateSecret(spec).getEncoded(), ALGORITHM); - } - - /** - * Encrypts the given data using AES encryption with a key derived from the passphrase. - * - * @param data The data to encrypt - * @return The encrypted data - * @throws Exception If an error occurs during encryption */ - private byte[] encrypt(byte[] data) throws Exception { - Cipher cipher = Cipher.getInstance(ALGORITHM); - cipher.init(Cipher.ENCRYPT_MODE, generateSecretKey()); - return Base64.getEncoder().encode(cipher.doFinal(data)); - } + void save(Token token); /** - * Decrypts the given encrypted data using AES decryption with a key derived from the passphrase. + * Loads a Token from the cache. * - * @param encryptedData The encrypted data, Base64 encoded - * @return The decrypted data - * @throws Exception If an error occurs during decryption + * @return The Token from the cache or null if the cache doesn't exist or is invalid */ - private byte[] decrypt(byte[] encryptedData) throws Exception { - Cipher cipher = Cipher.getInstance(ALGORITHM); - cipher.init(Cipher.DECRYPT_MODE, generateSecretKey()); - return cipher.doFinal(Base64.getDecoder().decode(encryptedData)); - } + Token load(); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java new file mode 100644 index 000000000..aaaeca4c8 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenCacheUtils.java @@ -0,0 +1,49 @@ +package com.databricks.sdk.core.oauth; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.List; + +/** Utility methods for TokenCache implementations. */ +public class TokenCacheUtils { + // Base path for token cache files + private static final String BASE_PATH = ".config/databricks-sdk-java/oauth"; + + /** + * Returns the path to the cache file for the given configuration. The filename is based on a hash + * of the host, client ID, and scopes. + * + * @param host The Databricks host URL + * @param clientId The OAuth client ID + * @param scopes The OAuth scopes requested + * @return The path to the token cache file + */ + public static Path getCacheFilePath(String host, String clientId, List scopes) { + try { + // Create SHA-256 hash of host, client_id, and scopes + MessageDigest hash = MessageDigest.getInstance("SHA-256"); + for (String chunk : new String[] {host, clientId, String.join(",", scopes)}) { + hash.update(chunk.getBytes(StandardCharsets.UTF_8)); + } + + // Convert hash bytes to hexadecimal string + StringBuilder hexString = new StringBuilder(); + for (byte b : hash.digest()) { + String hex = Integer.toHexString(0xff & b); + if (hex.length() == 1) { + hexString.append('0'); + } + hexString.append(hex); + } + + String userHome = System.getProperty("user.home"); + Path basePath = Paths.get(userHome, BASE_PATH); + return basePath.resolve(hexString.toString()); + } catch (NoSuchAlgorithmException e) { + throw new RuntimeException("Failed to create hash for token cache filename", e); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index e552a1427..51edac500 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -3,14 +3,18 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.commons.CommonsHttpClient; +import com.databricks.sdk.core.oauth.FileTokenCache; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.core.oauth.TokenCache; import com.databricks.sdk.core.utils.Environment; import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; public class DatabricksConfigTest { @Test @@ -195,4 +199,43 @@ public void testClone() { assert newWorkspaceConfig.getClientId().equals("my-client-id"); assert newWorkspaceConfig.getClientSecret().equals("my-client-secret"); } + + @Test + public void testGetTokenCache() { + // Test that getTokenCache returns a FileTokenCache by default + String testHost = "https://test-host.cloud.databricks.com"; + String testClientId = "test-client-id"; + List testScopes = Arrays.asList("offline_access", "clusters", "sql"); + + DatabricksConfig config = + new DatabricksConfig().setHost(testHost).setClientId(testClientId).setScopes(testScopes); + + TokenCache cache = config.getTokenCache(); + + assertNotNull(cache, "TokenCache from config should not be null"); + assertInstanceOf( + FileTokenCache.class, cache, "Default cache should be a FileTokenCache instance"); + } + + @Test + public void testSetTokenCache() { + // Test that a custom token cache can be set and is returned by getTokenCache + String testHost = "https://test-host.cloud.databricks.com"; + String testClientId = "test-client-id"; + List testScopes = Arrays.asList("offline_access", "clusters", "sql"); + + // Create a mock token cache + TokenCache mockCache = Mockito.mock(TokenCache.class); + + DatabricksConfig config = + new DatabricksConfig() + .setHost(testHost) + .setClientId(testClientId) + .setScopes(testScopes) + .setTokenCache(mockCache); + + // Verify the custom cache is returned + TokenCache returnedCache = config.getTokenCache(); + assertSame(mockCache, returnedCache, "Should return the custom token cache"); + } } 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 34205c2dc..f1718a298 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 @@ -13,7 +13,6 @@ import com.databricks.sdk.core.http.Response; import java.io.IOException; import java.net.URL; -import java.nio.file.Files; import java.time.LocalDateTime; import java.util.Arrays; import java.util.HashMap; @@ -223,7 +222,7 @@ void cacheWithValidTokenTest() throws IOException { // Create mock token cache that returns the valid token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.when(mockTokenCache.load()).thenReturn(validToken); + Mockito.doReturn(validToken).when(mockTokenCache).load(); // Create config with HTTP client and mock token cache DatabricksConfig config = @@ -244,7 +243,7 @@ void cacheWithValidTokenTest() throws IOException { // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -285,7 +284,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Create mock token cache that returns the expired token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); + Mockito.doReturn(expiredToken).when(mockTokenCache).load(); // Create config with HTTP client and mock token cache DatabricksConfig config = @@ -306,7 +305,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -333,8 +332,9 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create HTTP client that fails when refreshing token HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - Mockito.when(mockHttpClient.execute(any(Request.class))) - .thenThrow(new IOException("Failed to refresh token")); + Mockito.doThrow(new IOException("Failed to refresh token")) + .when(mockHttpClient) + .execute(any(Request.class)); // Create an expired token with invalid refresh token LocalDateTime pastTime = LocalDateTime.now().minusHours(1); @@ -343,7 +343,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create mock token cache that returns the expired token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.when(mockTokenCache.load()).thenReturn(expiredToken); + Mockito.doReturn(expiredToken).when(mockTokenCache).load(); // Setup browser auth result (should be used as fallback) Token browserAuthToken = @@ -382,7 +382,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -410,7 +410,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Create mock token cache that returns the invalid token TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.when(mockTokenCache.load()).thenReturn(invalidToken); + Mockito.doReturn(invalidToken).when(mockTokenCache).load(); // Setup browser auth result (should be used as fallback) Token browserAuthToken = @@ -443,7 +443,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Spy on the config to inject the mock token cache DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.when(spyConfig.getTokenCache()).thenReturn(mockTokenCache); + Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); @@ -481,14 +481,13 @@ void disabledTokenCacheTest() throws IOException { .withTokenUrl("https://test-token-url") .build(); - // Create config with caching explicitly disabled + // Create config with browser auth type DatabricksConfig config = new DatabricksConfig() .setAuthType("external-browser") .setHost("https://test.databricks.com") .setClientId("test-client-id") - .setHttpClient(mockHttpClient) - .setTokenCacheEnabled(false); + .setHttpClient(mockHttpClient); // We need to provide OIDC endpoints OpenIDConnectEndpoints endpoints = @@ -502,18 +501,13 @@ void disabledTokenCacheTest() throws IOException { .when(provider) .performBrowserAuth(any(DatabricksConfig.class)); - // Create a real token cache with caching disabled - TokenCache tokenCache = - new TokenCache( - "https://test.databricks.com", - "test-client-id", - Arrays.asList("offline_access"), - "test-passphrase", - false); + // Create a mock token cache that simulates being disabled by returning null on load + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.doReturn(null).when(mockTokenCache).load(); - // Spy on the config to inject the token cache and endpoints + // Spy on the config to inject the mock token cache and endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(tokenCache).when(spyConfig).getTokenCache(); + Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -526,10 +520,7 @@ void disabledTokenCacheTest() throws IOException { // Verify performBrowserAuth was called immediately (no attempt to use cache) Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); - // The save operation should be called but be a no-op - // We can't easily verify the no-op behavior in this test, but we can - // verify that the cache file doesn't exist - assertFalse( - Files.exists(tokenCache.getFilename()), "Cache file should not be created when disabled"); + // Verify token was saved to cache + Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java new file mode 100644 index 000000000..ede6cfd11 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -0,0 +1,128 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.LocalDateTime; +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** Tests for the FileTokenCache implementation of TokenCache. */ +public class FileTokenCacheTest { + private static final String TEST_HOST = "https://test-host.cloud.databricks.com"; + private static final String TEST_CLIENT_ID = "test-client-id"; + private static final List TEST_SCOPES = + Arrays.asList("offline_access", "clusters", "sql"); + private Path cacheFile; + private FileTokenCache tokenCache; + + @BeforeEach + void setUp() { + cacheFile = TokenCacheUtils.getCacheFilePath(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES); + tokenCache = new FileTokenCache(cacheFile); + } + + @AfterEach + void tearDown() throws IOException { + Files.deleteIfExists(cacheFile); + } + + @Test + void testEmptyCache() { + // When no cache file exists + assertNull(tokenCache.load(), "Loading from non-existent cache should return null"); + } + + @Test + void testSaveAndLoadToken() { + // Given a token + LocalDateTime expiry = LocalDateTime.now().plusHours(1); + Token token = new Token("access-token", "Bearer", "refresh-token", expiry); + + // When saving and loading the token + tokenCache.save(token); + Token loadedToken = tokenCache.load(); + + // Then the loaded token should match the original + assertNotNull(loadedToken, "Loaded token should not be null"); + assertEquals("access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("refresh-token", loadedToken.getRefreshToken()); + assertFalse(loadedToken.isExpired(), "Token should not be expired"); + } + + @Test + void testTokenExpiry() { + // Create an expired token + LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Token expiredToken = new Token("access-token", "Bearer", "refresh-token", pastTime); + + // Verify it's marked as expired + assertTrue(expiredToken.isExpired(), "Token should be expired"); + + // Create a valid token + LocalDateTime futureTime = LocalDateTime.now().plusMinutes(30); + Token validToken = new Token("access-token", "Bearer", "refresh-token", futureTime); + + // Verify it's not marked as expired + assertFalse(validToken.isExpired(), "Token should not be expired"); + } + + @Test + void testNullPathRejection() { + // FileTokenCache should reject null path + assertThrows( + NullPointerException.class, + () -> new FileTokenCache(null), + "Should throw NullPointerException for null path"); + } + + @Test + void testOverwriteToken() { + // Given two tokens saved in sequence + Token token1 = new Token("token1", "Bearer", "refresh1", LocalDateTime.now().plusHours(1)); + Token token2 = new Token("token2", "Bearer", "refresh2", LocalDateTime.now().plusHours(2)); + + tokenCache.save(token1); + tokenCache.save(token2); + + // When loading from cache + Token loadedToken = tokenCache.load(); + + // Then the second token should be loaded + assertNotNull(loadedToken, "Loaded token should not be null"); + assertEquals("token2", loadedToken.getAccessToken()); + assertEquals("refresh2", loadedToken.getRefreshToken()); + } + + @Test + void testWithCustomPath(@TempDir Path tempDir) { + // Given a token cache with a custom path + Path tempPath = tempDir.resolve("custom-token-cache"); + FileTokenCache cache = new FileTokenCache(tempPath); + + // And a token + Token testToken = + new Token( + "test-access-token", "Bearer", "test-refresh-token", LocalDateTime.now().plusHours(1)); + + // When saving and loading + cache.save(testToken); + Token loadedToken = cache.load(); + + // Then the token should be loaded from the custom path + assertNotNull(loadedToken, "Should load token from custom cache path"); + assertEquals("test-access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("test-refresh-token", loadedToken.getRefreshToken()); + + // Verify the file exists + assertTrue(Files.exists(tempPath), "Cache file should exist at custom path"); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java deleted file mode 100644 index 923913aaf..000000000 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java +++ /dev/null @@ -1,217 +0,0 @@ -package com.databricks.sdk.core.oauth; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.spy; - -import com.databricks.sdk.core.DatabricksConfig; -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.time.LocalDateTime; -import java.util.Arrays; -import java.util.List; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.io.TempDir; - -public class TokenCacheTest { - private static final String TEST_HOST = "https://test-host.cloud.databricks.com"; - private static final String TEST_CLIENT_ID = "test-client-id"; - private static final List TEST_SCOPES = - Arrays.asList("offline_access", "clusters", "sql"); - private static final String TEST_PASS = "test-passphrase"; - private Path cacheFile; - private TokenCache tokenCache; - - @BeforeEach - void setUp() { - tokenCache = new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true); - cacheFile = tokenCache.getFilename(); - } - - @AfterEach - void tearDown() throws IOException { - Files.deleteIfExists(cacheFile); - } - - @Test - void testEmptyCache() throws IOException { - assertNull(tokenCache.load()); - } - - @Test - void testSaveAndLoadToken() throws IOException { - LocalDateTime expiry = LocalDateTime.now().plusHours(1); - Token token = new Token("access-token", "Bearer", "refresh-token", expiry); - - tokenCache.save(token); - - Token loadedToken = tokenCache.load(); - assertNotNull(loadedToken, "Loaded token should not be null"); - assertEquals("access-token", loadedToken.getAccessToken()); - assertEquals("Bearer", loadedToken.getTokenType()); - assertEquals("refresh-token", loadedToken.getRefreshToken()); - // No direct way to compare expiry times exactly, but we can compare if they're both in the - // future - assertFalse(loadedToken.isExpired(), "Token should not be expired"); - } - - @Test - void testTokenExpiry() throws IOException { - LocalDateTime now = LocalDateTime.now(); - LocalDateTime expiry = now.plusMinutes(30); - Token token = new Token("access-token", "Bearer", "refresh-token", expiry); - - tokenCache.save(token); - - Token loadedToken = tokenCache.load(); - assertNotNull(loadedToken, "Loaded token should not be null"); - assertFalse(loadedToken.isExpired(), "Token should not be expired"); - - // Create a new token with past expiry - Token expiredToken = new Token("access-token", "Bearer", "refresh-token", now.minusHours(1)); - assertTrue(expiredToken.isExpired(), "Token should be expired"); - } - - @Test - void testInvalidPassphrase() { - assertThrows( - NullPointerException.class, - () -> new TokenCache(null, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); - } - - @Test - void testOverwriteToken() throws IOException { - Token token1 = new Token("token1", "Bearer", "refresh1", LocalDateTime.now().plusHours(1)); - Token token2 = new Token("token2", "Bearer", "refresh2", LocalDateTime.now().plusHours(2)); - - tokenCache.save(token1); - tokenCache.save(token2); - - Token loadedToken = tokenCache.load(); - assertNotNull(loadedToken, "Loaded token should not be null"); - assertEquals("token2", loadedToken.getAccessToken()); - assertEquals("refresh2", loadedToken.getRefreshToken()); - } - - @Test - void testTokenCacheSaveLoad(@TempDir Path tempDir) throws IOException { - // Create a token cache that uses a temp directory - TokenCache cache = spy(new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, true)); - - // Create test token - Token testToken = - new Token( - "test-access-token", "Bearer", "test-refresh-token", LocalDateTime.now().plusHours(1)); - - // Save to cache - cache.save(testToken); - - // Load from cache - Token loadedToken = cache.load(); - assertNotNull(loadedToken, "Should load token from cache"); - - // Verify token details - assertEquals("test-access-token", loadedToken.getAccessToken()); - assertEquals("Bearer", loadedToken.getTokenType()); - assertEquals("test-refresh-token", loadedToken.getRefreshToken()); - } - - @Test - void testDatabricksConfigTokenCache() throws IOException { - // Create a DatabricksConfig with the test values - DatabricksConfig config = - new DatabricksConfig() - .setHost(TEST_HOST) - .setClientId(TEST_CLIENT_ID) - .setScopes(TEST_SCOPES) - .setOAuthTokenCachePassphrase(TEST_PASS) - .setTokenCacheEnabled(true); - - // Get TokenCache from config - TokenCache cache = config.getTokenCache(); - assertNotNull(cache, "TokenCache from config should not be null"); - - // Create test token - Token testToken = - new Token( - "config-access-token", - "Bearer", - "config-refresh-token", - LocalDateTime.now().plusHours(1)); - - // Save to cache directly - cache.save(testToken); - - // Load from cache directly - Token loadedToken = cache.load(); - assertNotNull(loadedToken, "Should load token from cache directly"); - - // Verify token details - assertEquals("config-access-token", loadedToken.getAccessToken()); - assertEquals("Bearer", loadedToken.getTokenType()); - assertEquals("config-refresh-token", loadedToken.getRefreshToken()); - - // Clean up - Files.deleteIfExists(cache.getFilename()); - } - - @Test - void testDisabledTokenCache() throws IOException { - // Create a disabled token cache - TokenCache disabledCache = - new TokenCache(TEST_HOST, TEST_CLIENT_ID, TEST_SCOPES, TEST_PASS, false); - - // Create test token - Token testToken = - new Token("disabled-token", "Bearer", "disabled-refresh", LocalDateTime.now().plusHours(1)); - - // Try to save to cache - should be a no-op - disabledCache.save(testToken); - - // Load from cache - should return null - Token loadedToken = disabledCache.load(); - assertNull(loadedToken, "Disabled cache should not load a token"); - - // Verify the file wasn't created - assertFalse( - Files.exists(disabledCache.getFilename()), - "Cache file should not be created when disabled"); - } - - @Test - void testDatabricksConfigDisabledTokenCache() throws IOException { - // Create a DatabricksConfig with caching disabled - DatabricksConfig config = - new DatabricksConfig() - .setHost(TEST_HOST) - .setClientId(TEST_CLIENT_ID) - .setScopes(TEST_SCOPES) - .setOAuthTokenCachePassphrase(TEST_PASS) - .setTokenCacheEnabled(false); - - // Get TokenCache from config - TokenCache cache = config.getTokenCache(); - assertNotNull(cache, "TokenCache should still be created even when disabled"); - - // Create test token - Token testToken = - new Token( - "disabled-config-token", - "Bearer", - "disabled-config-refresh", - LocalDateTime.now().plusHours(1)); - - // Try to save to cache - should be a no-op - cache.save(testToken); - - // Load from cache - should return null - Token loadedToken = cache.load(); - assertNull(loadedToken, "Disabled cache should not load a token"); - - // Verify the file wasn't created - assertFalse( - Files.exists(cache.getFilename()), "Cache file should not be created when disabled"); - } -} From 3dad7709dd52823bedc6fea71d668c9d0bbf0b33 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Fri, 11 Apr 2025 18:43:53 +0530 Subject: [PATCH 06/13] fmt --- .../java/com/databricks/sdk/core/DatabricksConfig.java | 8 ++------ .../com/databricks/sdk/core/oauth/FileTokenCache.java | 10 ++++++---- .../com/databricks/sdk/core/DatabricksConfigTest.java | 2 +- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index a065b44b2..a23258b05 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -381,17 +381,13 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** - * @deprecated Use {@link #getAzureUseMsi()} instead. - */ + /** @deprecated Use {@link #getAzureUseMsi()} instead. */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** - * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. - */ + /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java index fc64dd273..ff62ca835 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileTokenCache.java @@ -2,6 +2,7 @@ import com.databricks.sdk.core.utils.SerDeUtils; import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.File; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; @@ -39,10 +40,11 @@ public void save(Token token) { Files.write(cacheFile, dataToWrite); // Set file permissions to be readable only by the owner (equivalent to 0600) - cacheFile.toFile().setReadable(false, false); - cacheFile.toFile().setReadable(true, true); - cacheFile.toFile().setWritable(false, false); - cacheFile.toFile().setWritable(true, true); + File file = cacheFile.toFile(); + file.setReadable(false, false); + file.setReadable(true, true); + file.setWritable(false, false); + file.setWritable(true, true); LOGGER.debug("Successfully saved token to cache: {}", cacheFile); } catch (Exception e) { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 51edac500..43de984c2 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -223,7 +223,7 @@ public void testSetTokenCache() { String testHost = "https://test-host.cloud.databricks.com"; String testClientId = "test-client-id"; List testScopes = Arrays.asList("offline_access", "clusters", "sql"); - + // Create a mock token cache TokenCache mockCache = Mockito.mock(TokenCache.class); From c9fc0088ada206cb740df375893866aef651f9ae Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Fri, 11 Apr 2025 19:11:34 +0530 Subject: [PATCH 07/13] fix --- .../ExternalBrowserCredentialsProvider.java | 2 +- ...xternalBrowserCredentialsProviderTest.java | 31 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) 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 40bb16910..b01f36ef5 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 @@ -53,7 +53,7 @@ public HeaderFactory configure(DatabricksConfig config) { LOGGER.debug("Using cached token, will immediately refresh"); cachedCreds.token = cachedCreds.refresh(); - tokenCache.save(cachedToken); + tokenCache.save(cachedCreds.getToken()); return cachedCreds.configure(config); } catch (Exception e) { // If token refresh fails, log and continue to browser auth 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 f1718a298..ff71a8661 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 @@ -18,6 +18,7 @@ import java.util.HashMap; import java.util.Map; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; public class ExternalBrowserCredentialsProviderTest { @@ -264,6 +265,21 @@ void cacheWithValidTokenTest() throws IOException { // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + + // Capture the token that was saved to cache to verify it's the refreshed token + ArgumentCaptor tokenCaptor = ArgumentCaptor.forClass(Token.class); + Mockito.verify(mockTokenCache).save(tokenCaptor.capture()); + Token savedToken = tokenCaptor.getValue(); + + // Verify the saved token contains the refreshed values from the HTTP response + assertEquals( + "refreshed_access_token", + savedToken.getAccessToken(), + "Should save refreshed access token to cache"); + assertEquals( + "new_refresh_token", + savedToken.getRefreshToken(), + "Should save new refresh token to cache"); } @Test @@ -326,6 +342,21 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + + // Capture the token that was saved to cache to verify it's the refreshed token + ArgumentCaptor tokenCaptor = ArgumentCaptor.forClass(Token.class); + Mockito.verify(mockTokenCache).save(tokenCaptor.capture()); + Token savedToken = tokenCaptor.getValue(); + + // Verify the saved token contains the refreshed values from the HTTP response + assertEquals( + "refreshed_access_token", + savedToken.getAccessToken(), + "Should save refreshed access token to cache"); + assertEquals( + "new_refresh_token", + savedToken.getRefreshToken(), + "Should save new refresh token to cache"); } @Test From c4433deaa2742e8a15d562505ab167b4d86c3b7a Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 16:21:10 +0530 Subject: [PATCH 08/13] address comments --- .../databricks/sdk/core/DatabricksConfig.java | 34 ---------- .../ExternalBrowserCredentialsProvider.java | 60 ++++++++++++++---- .../sdk/core/oauth/SessionCredentials.java | 11 +++- .../sdk/core/DatabricksConfigTest.java | 43 ------------- ...xternalBrowserCredentialsProviderTest.java | 63 +++++++++---------- 5 files changed, 86 insertions(+), 125 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index a23258b05..fa89f5041 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -4,17 +4,13 @@ import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; -import com.databricks.sdk.core.oauth.FileTokenCache; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; -import com.databricks.sdk.core.oauth.TokenCache; -import com.databricks.sdk.core.oauth.TokenCacheUtils; import com.databricks.sdk.core.utils.Cloud; import com.databricks.sdk.core.utils.Environment; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.File; import java.io.IOException; import java.lang.reflect.Field; -import java.nio.file.Path; import java.util.*; import org.apache.http.HttpMessage; @@ -145,9 +141,6 @@ public class DatabricksConfig { private DatabricksEnvironment databricksEnvironment; - // Lazily initialized OAuth token cache - private transient TokenCache tokenCache; - public Environment getEnv() { return env; } @@ -677,17 +670,6 @@ public DatabricksConfig newWithWorkspaceHost(String host) { return clone(fieldsToSkip).setHost(host); } - /** - * Sets a custom TokenCache implementation. - * - * @param tokenCache the TokenCache implementation to use - * @return this config instance - */ - public DatabricksConfig setTokenCache(TokenCache tokenCache) { - this.tokenCache = tokenCache; - return this; - } - /** * Gets the default OAuth redirect URL. If one is not provided explicitly, uses * http://localhost:8080/callback @@ -697,20 +679,4 @@ public DatabricksConfig setTokenCache(TokenCache tokenCache) { public String getEffectiveOAuthRedirectUrl() { return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback"; } - - /** - * Gets the TokenCache instance for the current configuration. - * - *

If a custom TokenCache has been set, it will be returned. Otherwise, a SimpleFileTokenCache - * will be created based on the configuration properties. - * - * @return A TokenCache instance - */ - public synchronized TokenCache getTokenCache() { - if (tokenCache == null) { - Path cachePath = TokenCacheUtils.getCacheFilePath(getHost(), getClientId(), getScopes()); - tokenCache = new FileTokenCache(cachePath); - } - return tokenCache; - } } 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 b01f36ef5..fb7795b03 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 @@ -5,18 +5,38 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.HeaderFactory; import java.io.IOException; +import java.nio.file.Path; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a - * browser for the user to authorize the application. Uses the cache from {@link - * DatabricksConfig#getTokenCache()}, tokens will be cached to avoid repeated authentication. + * browser for the user to authorize the application. Uses a specified TokenCache or creates a + * default one if none is provided. */ public class ExternalBrowserCredentialsProvider implements CredentialsProvider { private static final Logger LOGGER = LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class); + private TokenCache tokenCache; + + /** + * Creates a new ExternalBrowserCredentialsProvider with the specified TokenCache. + * + * @param tokenCache the TokenCache to use for caching tokens + */ + public ExternalBrowserCredentialsProvider(TokenCache tokenCache) { + this.tokenCache = tokenCache; + } + + /** + * Creates a new ExternalBrowserCredentialsProvider with a default TokenCache. A FileTokenCache + * will be created when credentials are configured. + */ + public ExternalBrowserCredentialsProvider() { + this(null); + } + @Override public String authType() { return "external-browser"; @@ -24,15 +44,17 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { - if (config.getHost() == null - || config.getClientId() == null - || !config.getAuthType().equals("external-browser")) { + if (config.getHost() == null || config.getAuthType() != "external-browser") { return null; } - try { - // Get the token cache from config - TokenCache tokenCache = config.getTokenCache(); + if (tokenCache == null) { + // Create a default FileTokenCache based on config + Path cachePath = + TokenCacheUtils.getCacheFilePath( + config.getHost(), config.getClientId(), config.getScopes()); + tokenCache = new FileTokenCache(cachePath); + } // First try to use the cached token if available (will return null if disabled) Token cachedToken = tokenCache.load(); @@ -49,11 +71,11 @@ public HeaderFactory configure(DatabricksConfig config) { .withClientSecret(config.getClientSecret()) .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withTokenCache(tokenCache) .build(); LOGGER.debug("Using cached token, will immediately refresh"); cachedCreds.token = cachedCreds.refresh(); - tokenCache.save(cachedCreds.getToken()); return cachedCreds.configure(config); } catch (Exception e) { // If token refresh fails, log and continue to browser auth @@ -62,7 +84,7 @@ public HeaderFactory configure(DatabricksConfig config) { } // If no cached token or refresh failed, perform browser auth - SessionCredentials credentials = performBrowserAuth(config); + SessionCredentials credentials = performBrowserAuth(config, tokenCache); tokenCache.save(credentials.getToken()); return credentials.configure(config); } catch (IOException | DatabricksException e) { @@ -71,10 +93,24 @@ public HeaderFactory configure(DatabricksConfig config) { } } - SessionCredentials performBrowserAuth(DatabricksConfig config) throws IOException { + SessionCredentials performBrowserAuth(DatabricksConfig config, TokenCache tokenCache) + throws IOException { LOGGER.debug("Performing browser authentication"); OAuthClient client = new OAuthClient(config); Consent consent = client.initiateConsent(); - return consent.launchExternalBrowser(); + + // Use the existing browser flow to get credentials + SessionCredentials credentials = consent.launchExternalBrowser(); + + // 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()) + .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 56ec588e6..9114b6d6c 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 @@ -31,8 +31,6 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { - this.tokenCache = config.getTokenCache(); - return () -> { Map headers = new HashMap<>(); headers.put( @@ -48,6 +46,7 @@ static class Builder { private String redirectUrl; private String clientId; private String clientSecret; + private TokenCache tokenCache; public Builder withHttpClient(HttpClient hc) { this.hc = hc; @@ -79,6 +78,11 @@ public Builder withClientSecret(String clientSecret) { return this; } + public Builder withTokenCache(TokenCache tokenCache) { + this.tokenCache = tokenCache; + return this; + } + public SessionCredentials build() { return new SessionCredentials(this); } @@ -89,7 +93,7 @@ public SessionCredentials build() { private final String redirectUrl; private final String clientId; private final String clientSecret; - private transient TokenCache tokenCache; + private final TokenCache tokenCache; private SessionCredentials(Builder b) { super(b.token); @@ -98,6 +102,7 @@ private SessionCredentials(Builder b) { this.redirectUrl = b.redirectUrl; this.clientId = b.clientId; this.clientSecret = b.clientSecret; + this.tokenCache = b.tokenCache; } @Override diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 43de984c2..e552a1427 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -3,18 +3,14 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.commons.CommonsHttpClient; -import com.databricks.sdk.core.oauth.FileTokenCache; import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; -import com.databricks.sdk.core.oauth.TokenCache; import com.databricks.sdk.core.utils.Environment; import java.io.IOException; import java.util.ArrayList; -import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; -import org.mockito.Mockito; public class DatabricksConfigTest { @Test @@ -199,43 +195,4 @@ public void testClone() { assert newWorkspaceConfig.getClientId().equals("my-client-id"); assert newWorkspaceConfig.getClientSecret().equals("my-client-secret"); } - - @Test - public void testGetTokenCache() { - // Test that getTokenCache returns a FileTokenCache by default - String testHost = "https://test-host.cloud.databricks.com"; - String testClientId = "test-client-id"; - List testScopes = Arrays.asList("offline_access", "clusters", "sql"); - - DatabricksConfig config = - new DatabricksConfig().setHost(testHost).setClientId(testClientId).setScopes(testScopes); - - TokenCache cache = config.getTokenCache(); - - assertNotNull(cache, "TokenCache from config should not be null"); - assertInstanceOf( - FileTokenCache.class, cache, "Default cache should be a FileTokenCache instance"); - } - - @Test - public void testSetTokenCache() { - // Test that a custom token cache can be set and is returned by getTokenCache - String testHost = "https://test-host.cloud.databricks.com"; - String testClientId = "test-client-id"; - List testScopes = Arrays.asList("offline_access", "clusters", "sql"); - - // Create a mock token cache - TokenCache mockCache = Mockito.mock(TokenCache.class); - - DatabricksConfig config = - new DatabricksConfig() - .setHost(testHost) - .setClientId(testClientId) - .setScopes(testScopes) - .setTokenCache(mockCache); - - // Verify the custom cache is returned - TokenCache returnedCache = config.getTokenCache(); - assertSame(mockCache, returnedCache, "Should return the custom token cache"); - } } 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 ff71a8661..73eb92146 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 @@ -238,13 +238,12 @@ void cacheWithValidTokenTest() throws IOException { new OpenIDConnectEndpoints( "https://test.databricks.com/token", "https://test.databricks.com/authorize"); - // Create our provider and mock the browser auth method + // Create our provider with the mock token cache and mock the browser auth method ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - // Spy on the config to inject the mock token cache and endpoints + // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -261,7 +260,8 @@ void cacheWithValidTokenTest() throws IOException { Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); // Verify performBrowserAuth was NOT called since refresh succeeded - Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); + Mockito.verify(provider, Mockito.never()) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -315,13 +315,12 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { new OpenIDConnectEndpoints( "https://test.databricks.com/token", "https://test.databricks.com/authorize"); - // Create our provider and mock the browser auth method + // Create our provider with the mock token cache ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - // Spy on the config to inject the mock token cache and endpoints + // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -338,7 +337,8 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); // Verify performBrowserAuth was NOT called since refresh succeeded - Mockito.verify(provider, Mockito.never()).performBrowserAuth(any(DatabricksConfig.class)); + Mockito.verify(provider, Mockito.never()) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -406,14 +406,13 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(browserAuthCreds) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); - // Spy on the config to inject the mock token cache and endpoints + // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -427,7 +426,8 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { Mockito.verify(mockTokenCache, Mockito.times(1)).load(); // Verify performBrowserAuth was called since refresh failed - Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + Mockito.verify(provider, Mockito.times(1)) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -467,17 +467,13 @@ void cacheWithInvalidTokensTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider()); + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(browserAuthCreds) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class)); - - // Spy on the config to inject the mock token cache - DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Configure provider - HeaderFactory headerFactory = provider.configure(spyConfig); + HeaderFactory headerFactory = provider.configure(config); // Verify headers contain the browser auth token (fallback) Map headers = headerFactory.headers(); assertEquals("Bearer browser_access_token", headers.get("Authorization")); @@ -486,7 +482,8 @@ void cacheWithInvalidTokensTest() throws IOException { Mockito.verify(mockTokenCache, Mockito.times(1)).load(); // Verify performBrowserAuth was called since we had an invalid token - Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + Mockito.verify(provider, Mockito.times(1)) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -525,20 +522,19 @@ void disabledTokenCacheTest() throws IOException { new OpenIDConnectEndpoints( "https://test.databricks.com/token", "https://test.databricks.com/authorize"); - // Create our provider and mock the browser auth method - ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider()); - Mockito.doReturn(browserAuthCreds) - .when(provider) - .performBrowserAuth(any(DatabricksConfig.class)); - // Create a mock token cache that simulates being disabled by returning null on load TokenCache mockTokenCache = Mockito.mock(TokenCache.class); Mockito.doReturn(null).when(mockTokenCache).load(); - // Spy on the config to inject the mock token cache and endpoints + // Create our provider with the disabled token cache and mock the browser auth method + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); + Mockito.doReturn(browserAuthCreds) + .when(provider) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + + // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(mockTokenCache).when(spyConfig).getTokenCache(); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); // Configure provider @@ -549,7 +545,8 @@ void disabledTokenCacheTest() throws IOException { assertEquals("Bearer browser_access_token", headers.get("Authorization")); // Verify performBrowserAuth was called immediately (no attempt to use cache) - Mockito.verify(provider, Mockito.times(1)).performBrowserAuth(any(DatabricksConfig.class)); + Mockito.verify(provider, Mockito.times(1)) + .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); // Verify token was saved to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); From 9ce45effeed7a5a14fb705ee1d4174448512603f Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 16:27:31 +0530 Subject: [PATCH 09/13] remove unnecessary test --- ...xternalBrowserCredentialsProviderTest.java | 63 ------------------- 1 file changed, 63 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 73eb92146..5b430fd55 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 @@ -488,67 +488,4 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); } - - @Test - void disabledTokenCacheTest() throws IOException { - // Create mock HTTP client for token operations - HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - - // Setup browser auth result - Token browserAuthToken = - new Token( - "browser_access_token", - "Bearer", - "browser_refresh_token", - LocalDateTime.now().plusHours(1)); - - SessionCredentials browserAuthCreds = - new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); - - // Create config with browser auth type - DatabricksConfig config = - new DatabricksConfig() - .setAuthType("external-browser") - .setHost("https://test.databricks.com") - .setClientId("test-client-id") - .setHttpClient(mockHttpClient); - - // We need to provide OIDC endpoints - OpenIDConnectEndpoints endpoints = - new OpenIDConnectEndpoints( - "https://test.databricks.com/token", "https://test.databricks.com/authorize"); - - // Create a mock token cache that simulates being disabled by returning null on load - TokenCache mockTokenCache = Mockito.mock(TokenCache.class); - Mockito.doReturn(null).when(mockTokenCache).load(); - - // Create our provider with the disabled token cache and mock the browser auth method - ExternalBrowserCredentialsProvider provider = - Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) - .when(provider) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); - - // Spy on the config to inject the endpoints - DatabricksConfig spyConfig = Mockito.spy(config); - Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); - - // Configure provider - HeaderFactory headerFactory = provider.configure(spyConfig); - - // Verify headers contain the browser auth token - Map headers = headerFactory.headers(); - assertEquals("Bearer browser_access_token", headers.get("Authorization")); - - // Verify performBrowserAuth was called immediately (no attempt to use cache) - Mockito.verify(provider, Mockito.times(1)) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); - - // Verify token was saved to cache - Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); - } } From 174c3c22681a1c0b43f3a4fd120ebbcabe889774 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 16:37:37 +0530 Subject: [PATCH 10/13] use databricks-cli --- .../sdk/core/oauth/ExternalBrowserCredentialsProvider.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 fb7795b03..d6e42bbf1 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 com.databricks.sdk.core.HeaderFactory; import java.io.IOException; import java.nio.file.Path; +import java.util.Objects; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,9 +45,12 @@ public String authType() { @Override public HeaderFactory configure(DatabricksConfig config) { - if (config.getHost() == null || config.getAuthType() != "external-browser") { + if (config.getHost() == null || !Objects.equals(config.getAuthType(), "external-browser")) { return null; } + if (config.getClientId() == null && config.getAzureClientId() == null) { + config.setClientId("databricks-cli"); + } try { if (tokenCache == null) { // Create a default FileTokenCache based on config From ab64d0fde1b2bb50999c1beec822b2fb5364e9cb Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 17:45:14 +0530 Subject: [PATCH 11/13] update logic used to get client id --- .../ExternalBrowserCredentialsProvider.java | 34 ++++--- .../sdk/core/oauth/OAuthClient.java | 11 --- .../sdk/core/oauth/OAuthClientUtils.java | 58 +++++++++++ ...xternalBrowserCredentialsProviderTest.java | 30 ++++-- .../sdk/core/oauth/OAuthClientUtilsTest.java | 97 +++++++++++++++++++ 5 files changed, 200 insertions(+), 30 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java 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 d6e42bbf1..36e731755 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 @@ -48,22 +48,24 @@ public HeaderFactory configure(DatabricksConfig config) { if (config.getHost() == null || !Objects.equals(config.getAuthType(), "external-browser")) { return null; } - if (config.getClientId() == null && config.getAzureClientId() == null) { - config.setClientId("databricks-cli"); - } + + // Use the utility class to resolve client ID and client secret + String[] clientCreds = OAuthClientUtils.resolveClientCredentials(config); + String clientId = clientCreds[0]; + String clientSecret = clientCreds[1]; + try { if (tokenCache == null) { // Create a default FileTokenCache based on config Path cachePath = - TokenCacheUtils.getCacheFilePath( - config.getHost(), config.getClientId(), config.getScopes()); + TokenCacheUtils.getCacheFilePath(config.getHost(), clientId, config.getScopes()); tokenCache = new FileTokenCache(cachePath); } // First try to use the cached token if available (will return null if disabled) Token cachedToken = tokenCache.load(); if (cachedToken != null && cachedToken.getRefreshToken() != null) { - LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId()); + LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId); try { // Create SessionCredentials with the cached token and try to refresh if needed @@ -71,8 +73,8 @@ public HeaderFactory configure(DatabricksConfig config) { new SessionCredentials.Builder() .withToken(cachedToken) .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) + .withClientId(clientId) + .withClientSecret(clientSecret) .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) .withTokenCache(tokenCache) @@ -88,7 +90,8 @@ public HeaderFactory configure(DatabricksConfig config) { } // If no cached token or refresh failed, perform browser auth - SessionCredentials credentials = performBrowserAuth(config, tokenCache); + SessionCredentials credentials = + performBrowserAuth(config, clientId, clientSecret, tokenCache); tokenCache.save(credentials.getToken()); return credentials.configure(config); } catch (IOException | DatabricksException e) { @@ -97,10 +100,19 @@ public HeaderFactory configure(DatabricksConfig config) { } } - SessionCredentials performBrowserAuth(DatabricksConfig config, TokenCache tokenCache) + SessionCredentials performBrowserAuth( + DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache) throws IOException { LOGGER.debug("Performing browser authentication"); - OAuthClient client = new OAuthClient(config); + OAuthClient client = + new OAuthClient.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(clientId) + .withClientSecret(clientSecret) + .withHost(config.getHost()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .build(); Consent consent = client.initiateConsent(); // Use the existing browser flow to get credentials diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index 632fc831c..cf65ba71a 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java @@ -85,17 +85,6 @@ public OAuthClient build() throws IOException { private final boolean isAws; private final boolean isAzure; - public OAuthClient(DatabricksConfig config) throws IOException { - this( - new Builder() - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withHost(config.getHost()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withScopes(config.getScopes())); - } - private OAuthClient(Builder b) throws IOException { this.clientId = Objects.requireNonNull(b.clientId); this.clientSecret = b.clientSecret; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java new file mode 100644 index 000000000..adc464646 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java @@ -0,0 +1,58 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksConfig; + +/** + * Utility methods for OAuth client credentials resolution. + */ +public class OAuthClientUtils { + + /** + * Default client ID to use when no client ID is specified. + */ + private static final String DEFAULT_CLIENT_ID = "databricks-cli"; + + /** + * Resolves the OAuth client ID from the configuration. + * Prioritizes regular OAuth client ID, then Azure client ID, and falls back to default client ID. + * + * @param config The Databricks configuration + * @return The resolved client ID + */ + public static String resolveClientId(DatabricksConfig config) { + if (config.getClientId() != null) { + return config.getClientId(); + } else if (config.getAzureClientId() != null) { + return config.getAzureClientId(); + } + return DEFAULT_CLIENT_ID; + } + + /** + * Resolves the OAuth client secret from the configuration. + * Prioritizes regular OAuth client secret, then Azure client secret. + * + * @param config The Databricks configuration + * @return The resolved client secret, or null if not present + */ + public static String resolveClientSecret(DatabricksConfig config) { + if (config.getClientSecret() != null) { + return config.getClientSecret(); + } else if (config.getAzureClientSecret() != null) { + return config.getAzureClientSecret(); + } + return null; + } + + /** + * Resolves both client ID and client secret from the configuration. + * + * @param config The Databricks configuration + * @return An array containing the client ID and client secret (may be null) + */ + public static String[] resolveClientCredentials(DatabricksConfig config) { + String clientId = resolveClientId(config); + String clientSecret = resolveClientSecret(config); + return new String[]{clientId, clientSecret}; + } +} \ No newline at end of file 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 5b430fd55..2be4264ed 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 @@ -42,7 +42,14 @@ void clientAndConsentTest() throws IOException { assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint()); - OAuthClient testClient = new OAuthClient(config); + OAuthClient testClient = new OAuthClient.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withHost(config.getHost()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .build(); assertEquals("test-client-id", testClient.getClientId()); Consent testConsent = testClient.initiateConsent(); @@ -79,7 +86,14 @@ void clientAndConsentTestWithCustomRedirectUrl() throws IOException { assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint()); - OAuthClient testClient = new OAuthClient(config); + OAuthClient testClient = new OAuthClient.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withHost(config.getHost()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .build(); assertEquals("test-client-id", testClient.getClientId()); Consent testConsent = testClient.initiateConsent(); @@ -261,7 +275,7 @@ void cacheWithValidTokenTest() throws IOException { // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -338,7 +352,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -409,7 +423,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(browserAuthCreds) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); // Spy on the config to inject the endpoints DatabricksConfig spyConfig = Mockito.spy(config); @@ -427,7 +441,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Verify performBrowserAuth was called since refresh failed Mockito.verify(provider, Mockito.times(1)) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -470,7 +484,7 @@ void cacheWithInvalidTokensTest() throws IOException { Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); Mockito.doReturn(browserAuthCreds) .when(provider) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); // Configure provider HeaderFactory headerFactory = provider.configure(config); @@ -483,7 +497,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Verify performBrowserAuth was called since we had an invalid token Mockito.verify(provider, Mockito.times(1)) - .performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class)); + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); // Verify token was saved after browser auth (for the new token) Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java new file mode 100644 index 000000000..14ec42cac --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java @@ -0,0 +1,97 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import com.databricks.sdk.core.DatabricksConfig; +import org.junit.jupiter.api.Test; + +public class OAuthClientUtilsTest { + + @Test + void resolveClientIdTest() { + // Test with regular client ID + DatabricksConfig config = new DatabricksConfig() + .setClientId("test-client-id") + .setAzureClientId("azure-client-id"); + assertEquals("test-client-id", OAuthClientUtils.resolveClientId(config)); + + // Test with only Azure client ID + config = new DatabricksConfig() + .setClientId(null) + .setAzureClientId("azure-client-id"); + assertEquals("azure-client-id", OAuthClientUtils.resolveClientId(config)); + + // Test with no client IDs + config = new DatabricksConfig() + .setClientId(null) + .setAzureClientId(null); + assertEquals("databricks-cli", OAuthClientUtils.resolveClientId(config)); + } + + @Test + void resolveClientSecretTest() { + // Test with regular client secret + DatabricksConfig config = new DatabricksConfig() + .setClientSecret("test-client-secret") + .setAzureClientSecret("azure-client-secret"); + assertEquals("test-client-secret", OAuthClientUtils.resolveClientSecret(config)); + + // Test with only Azure client secret + config = new DatabricksConfig() + .setClientSecret(null) + .setAzureClientSecret("azure-client-secret"); + assertEquals("azure-client-secret", OAuthClientUtils.resolveClientSecret(config)); + + // Test with no client secrets + config = new DatabricksConfig() + .setClientSecret(null) + .setAzureClientSecret(null); + assertNull(OAuthClientUtils.resolveClientSecret(config)); + } + + @Test + void resolveClientCredentialsTest() { + // Test with both client ID and secret + DatabricksConfig config = new DatabricksConfig() + .setClientId("test-client-id") + .setClientSecret("test-client-secret"); + String[] credentials = OAuthClientUtils.resolveClientCredentials(config); + assertEquals("test-client-id", credentials[0]); + assertEquals("test-client-secret", credentials[1]); + + // Test with only client ID + config = new DatabricksConfig() + .setClientId("test-client-id") + .setClientSecret(null); + credentials = OAuthClientUtils.resolveClientCredentials(config); + assertEquals("test-client-id", credentials[0]); + assertNull(credentials[1]); + + // Test with Azure credentials + config = new DatabricksConfig() + .setClientId(null) + .setClientSecret(null) + .setAzureClientId("azure-client-id") + .setAzureClientSecret("azure-client-secret"); + credentials = OAuthClientUtils.resolveClientCredentials(config); + assertEquals("azure-client-id", credentials[0]); + assertEquals("azure-client-secret", credentials[1]); + + // Test with no credentials + config = new DatabricksConfig(); + credentials = OAuthClientUtils.resolveClientCredentials(config); + assertEquals("databricks-cli", credentials[0]); + assertNull(credentials[1]); + + // Test mixed credentials preference + config = new DatabricksConfig() + .setClientId("test-client-id") + .setClientSecret(null) + .setAzureClientId("azure-client-id") + .setAzureClientSecret("azure-client-secret"); + credentials = OAuthClientUtils.resolveClientCredentials(config); + assertEquals("test-client-id", credentials[0]); + assertEquals("azure-client-secret", credentials[1]); + } +} \ No newline at end of file From 23b5bd152b83957a724f6ca150b6713b892e30c2 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 17:53:04 +0530 Subject: [PATCH 12/13] fmt --- .../ExternalBrowserCredentialsProvider.java | 4 +- .../sdk/core/oauth/OAuthClientUtils.java | 22 +++---- ...xternalBrowserCredentialsProviderTest.java | 46 ++++++++------ .../sdk/core/oauth/OAuthClientUtilsTest.java | 62 ++++++++----------- 4 files changed, 66 insertions(+), 68 deletions(-) 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 36e731755..d50f9265c 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 @@ -48,12 +48,12 @@ public HeaderFactory configure(DatabricksConfig config) { if (config.getHost() == null || !Objects.equals(config.getAuthType(), "external-browser")) { return null; } - + // Use the utility class to resolve client ID and client secret String[] clientCreds = OAuthClientUtils.resolveClientCredentials(config); String clientId = clientCreds[0]; String clientSecret = clientCreds[1]; - + try { if (tokenCache == null) { // Create a default FileTokenCache based on config diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java index adc464646..2001bde21 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java @@ -2,19 +2,15 @@ import com.databricks.sdk.core.DatabricksConfig; -/** - * Utility methods for OAuth client credentials resolution. - */ +/** Utility methods for OAuth client credentials resolution. */ public class OAuthClientUtils { - /** - * Default client ID to use when no client ID is specified. - */ + /** Default client ID to use when no client ID is specified. */ private static final String DEFAULT_CLIENT_ID = "databricks-cli"; /** - * Resolves the OAuth client ID from the configuration. - * Prioritizes regular OAuth client ID, then Azure client ID, and falls back to default client ID. + * Resolves the OAuth client ID from the configuration. Prioritizes regular OAuth client ID, then + * Azure client ID, and falls back to default client ID. * * @param config The Databricks configuration * @return The resolved client ID @@ -29,8 +25,8 @@ public static String resolveClientId(DatabricksConfig config) { } /** - * Resolves the OAuth client secret from the configuration. - * Prioritizes regular OAuth client secret, then Azure client secret. + * Resolves the OAuth client secret from the configuration. Prioritizes regular OAuth client + * secret, then Azure client secret. * * @param config The Databricks configuration * @return The resolved client secret, or null if not present @@ -46,13 +42,13 @@ public static String resolveClientSecret(DatabricksConfig config) { /** * Resolves both client ID and client secret from the configuration. - * + * * @param config The Databricks configuration * @return An array containing the client ID and client secret (may be null) */ public static String[] resolveClientCredentials(DatabricksConfig config) { String clientId = resolveClientId(config); String clientSecret = resolveClientSecret(config); - return new String[]{clientId, clientSecret}; + return new String[] {clientId, clientSecret}; } -} \ No newline at end of file +} 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 2be4264ed..1714b731c 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 @@ -42,14 +42,15 @@ void clientAndConsentTest() throws IOException { assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint()); - OAuthClient testClient = new OAuthClient.Builder() - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withHost(config.getHost()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withScopes(config.getScopes()) - .build(); + OAuthClient testClient = + new OAuthClient.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withHost(config.getHost()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .build(); assertEquals("test-client-id", testClient.getClientId()); Consent testConsent = testClient.initiateConsent(); @@ -86,14 +87,15 @@ void clientAndConsentTestWithCustomRedirectUrl() throws IOException { assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint()); - OAuthClient testClient = new OAuthClient.Builder() - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withHost(config.getHost()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withScopes(config.getScopes()) - .build(); + OAuthClient testClient = + new OAuthClient.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withHost(config.getHost()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withScopes(config.getScopes()) + .build(); assertEquals("test-client-id", testClient.getClientId()); Consent testConsent = testClient.initiateConsent(); @@ -275,7 +277,11 @@ void cacheWithValidTokenTest() throws IOException { // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()) - .performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class)); + .performBrowserAuth( + any(DatabricksConfig.class), + any(String.class), + any(String.class), + any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); @@ -352,7 +358,11 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // Verify performBrowserAuth was NOT called since refresh succeeded Mockito.verify(provider, Mockito.never()) - .performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class)); + .performBrowserAuth( + any(DatabricksConfig.class), + any(String.class), + any(String.class), + any(TokenCache.class)); // Verify token was saved back to cache Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java index 14ec42cac..a4b384b72 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java @@ -11,69 +11,60 @@ public class OAuthClientUtilsTest { @Test void resolveClientIdTest() { // Test with regular client ID - DatabricksConfig config = new DatabricksConfig() - .setClientId("test-client-id") - .setAzureClientId("azure-client-id"); + DatabricksConfig config = + new DatabricksConfig().setClientId("test-client-id").setAzureClientId("azure-client-id"); assertEquals("test-client-id", OAuthClientUtils.resolveClientId(config)); // Test with only Azure client ID - config = new DatabricksConfig() - .setClientId(null) - .setAzureClientId("azure-client-id"); + config = new DatabricksConfig().setClientId(null).setAzureClientId("azure-client-id"); assertEquals("azure-client-id", OAuthClientUtils.resolveClientId(config)); // Test with no client IDs - config = new DatabricksConfig() - .setClientId(null) - .setAzureClientId(null); + config = new DatabricksConfig().setClientId(null).setAzureClientId(null); assertEquals("databricks-cli", OAuthClientUtils.resolveClientId(config)); } @Test void resolveClientSecretTest() { // Test with regular client secret - DatabricksConfig config = new DatabricksConfig() - .setClientSecret("test-client-secret") - .setAzureClientSecret("azure-client-secret"); + DatabricksConfig config = + new DatabricksConfig() + .setClientSecret("test-client-secret") + .setAzureClientSecret("azure-client-secret"); assertEquals("test-client-secret", OAuthClientUtils.resolveClientSecret(config)); // Test with only Azure client secret - config = new DatabricksConfig() - .setClientSecret(null) - .setAzureClientSecret("azure-client-secret"); + config = + new DatabricksConfig().setClientSecret(null).setAzureClientSecret("azure-client-secret"); assertEquals("azure-client-secret", OAuthClientUtils.resolveClientSecret(config)); // Test with no client secrets - config = new DatabricksConfig() - .setClientSecret(null) - .setAzureClientSecret(null); + config = new DatabricksConfig().setClientSecret(null).setAzureClientSecret(null); assertNull(OAuthClientUtils.resolveClientSecret(config)); } @Test void resolveClientCredentialsTest() { // Test with both client ID and secret - DatabricksConfig config = new DatabricksConfig() - .setClientId("test-client-id") - .setClientSecret("test-client-secret"); + DatabricksConfig config = + new DatabricksConfig().setClientId("test-client-id").setClientSecret("test-client-secret"); String[] credentials = OAuthClientUtils.resolveClientCredentials(config); assertEquals("test-client-id", credentials[0]); assertEquals("test-client-secret", credentials[1]); // Test with only client ID - config = new DatabricksConfig() - .setClientId("test-client-id") - .setClientSecret(null); + config = new DatabricksConfig().setClientId("test-client-id").setClientSecret(null); credentials = OAuthClientUtils.resolveClientCredentials(config); assertEquals("test-client-id", credentials[0]); assertNull(credentials[1]); // Test with Azure credentials - config = new DatabricksConfig() - .setClientId(null) - .setClientSecret(null) - .setAzureClientId("azure-client-id") - .setAzureClientSecret("azure-client-secret"); + config = + new DatabricksConfig() + .setClientId(null) + .setClientSecret(null) + .setAzureClientId("azure-client-id") + .setAzureClientSecret("azure-client-secret"); credentials = OAuthClientUtils.resolveClientCredentials(config); assertEquals("azure-client-id", credentials[0]); assertEquals("azure-client-secret", credentials[1]); @@ -85,13 +76,14 @@ void resolveClientCredentialsTest() { assertNull(credentials[1]); // Test mixed credentials preference - config = new DatabricksConfig() - .setClientId("test-client-id") - .setClientSecret(null) - .setAzureClientId("azure-client-id") - .setAzureClientSecret("azure-client-secret"); + config = + new DatabricksConfig() + .setClientId("test-client-id") + .setClientSecret(null) + .setAzureClientId("azure-client-id") + .setAzureClientSecret("azure-client-secret"); credentials = OAuthClientUtils.resolveClientCredentials(config); assertEquals("test-client-id", credentials[0]); assertEquals("azure-client-secret", credentials[1]); } -} \ No newline at end of file +} From ce75413f2d7e81ad84a47ab1f47be90aa0eb6676 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Thu, 17 Apr 2025 20:24:49 +0530 Subject: [PATCH 13/13] remove extra function --- .../ExternalBrowserCredentialsProvider.java | 5 +-- .../sdk/core/oauth/OAuthClientUtils.java | 12 ----- .../sdk/core/oauth/OAuthClientUtilsTest.java | 44 ------------------- 3 files changed, 2 insertions(+), 59 deletions(-) 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 d50f9265c..b8aa4c66f 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 @@ -50,9 +50,8 @@ public HeaderFactory configure(DatabricksConfig config) { } // Use the utility class to resolve client ID and client secret - String[] clientCreds = OAuthClientUtils.resolveClientCredentials(config); - String clientId = clientCreds[0]; - String clientSecret = clientCreds[1]; + String clientId = OAuthClientUtils.resolveClientId(config); + String clientSecret = OAuthClientUtils.resolveClientSecret(config); try { if (tokenCache == null) { diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java index 2001bde21..5908eff79 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java @@ -39,16 +39,4 @@ public static String resolveClientSecret(DatabricksConfig config) { } return null; } - - /** - * Resolves both client ID and client secret from the configuration. - * - * @param config The Databricks configuration - * @return An array containing the client ID and client secret (may be null) - */ - public static String[] resolveClientCredentials(DatabricksConfig config) { - String clientId = resolveClientId(config); - String clientSecret = resolveClientSecret(config); - return new String[] {clientId, clientSecret}; - } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java index a4b384b72..b9af3928f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java @@ -42,48 +42,4 @@ void resolveClientSecretTest() { config = new DatabricksConfig().setClientSecret(null).setAzureClientSecret(null); assertNull(OAuthClientUtils.resolveClientSecret(config)); } - - @Test - void resolveClientCredentialsTest() { - // Test with both client ID and secret - DatabricksConfig config = - new DatabricksConfig().setClientId("test-client-id").setClientSecret("test-client-secret"); - String[] credentials = OAuthClientUtils.resolveClientCredentials(config); - assertEquals("test-client-id", credentials[0]); - assertEquals("test-client-secret", credentials[1]); - - // Test with only client ID - config = new DatabricksConfig().setClientId("test-client-id").setClientSecret(null); - credentials = OAuthClientUtils.resolveClientCredentials(config); - assertEquals("test-client-id", credentials[0]); - assertNull(credentials[1]); - - // Test with Azure credentials - config = - new DatabricksConfig() - .setClientId(null) - .setClientSecret(null) - .setAzureClientId("azure-client-id") - .setAzureClientSecret("azure-client-secret"); - credentials = OAuthClientUtils.resolveClientCredentials(config); - assertEquals("azure-client-id", credentials[0]); - assertEquals("azure-client-secret", credentials[1]); - - // Test with no credentials - config = new DatabricksConfig(); - credentials = OAuthClientUtils.resolveClientCredentials(config); - assertEquals("databricks-cli", credentials[0]); - assertNull(credentials[1]); - - // Test mixed credentials preference - config = - new DatabricksConfig() - .setClientId("test-client-id") - .setClientSecret(null) - .setAzureClientId("azure-client-id") - .setAzureClientSecret("azure-client-secret"); - credentials = OAuthClientUtils.resolveClientCredentials(config); - assertEquals("test-client-id", credentials[0]); - assertEquals("azure-client-secret", credentials[1]); - } }