From cb324042201ea16a7cbb30b65441b1c3c9921e7d Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Tue, 1 Apr 2025 12:16:06 +0530 Subject: [PATCH 1/9] Add token cache for U2M OAuth --- pom.xml | 5 + .../api/IDatabricksConnectionContext.java | 3 + .../api/impl/DatabricksConnectionContext.java | 5 + ...ingExternalBrowserCredentialsProvider.java | 122 ++++++++++++++ .../com/databricks/jdbc/auth/TokenCache.java | 119 ++++++++++++++ .../jdbc/common/DatabricksJdbcUrlParams.java | 3 +- .../impl/common/ClientConfigurator.java | 7 +- ...xternalBrowserCredentialsProviderTest.java | 154 ++++++++++++++++++ .../databricks/jdbc/auth/TokenCacheTest.java | 74 +++++++++ .../impl/common/ClientConfiguratorTest.java | 37 +++++ 10 files changed, 527 insertions(+), 2 deletions(-) create mode 100644 src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java create mode 100644 src/main/java/com/databricks/jdbc/auth/TokenCache.java create mode 100644 src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java create mode 100644 src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java diff --git a/pom.xml b/pom.xml index c11c2a5295..c761030921 100644 --- a/pom.xml +++ b/pom.xml @@ -188,6 +188,11 @@ jackson-core ${jackson.version} + + com.fasterxml.jackson.datatype + jackson-datatype-jsr310 + ${jackson.version} + org.immutables value diff --git a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java index 9d81b40d39..2f723971d5 100644 --- a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java @@ -231,4 +231,7 @@ public interface IDatabricksConnectionContext { /** Returns the size for HTTP connection pool */ int getHttpConnectionPoolSize(); + + /** Returns the passphrase used for encrypting/decrypting token cache */ + String getTokenCachePassPhrase(); } diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index 1601600fbe..cc20db6cd5 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -734,6 +734,11 @@ public int getHttpConnectionPoolSize() { return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HTTP_CONNECTION_POOL_SIZE)); } + @Override + public String getTokenCachePassPhrase() { + return getParameter(DatabricksJdbcUrlParams.TOKEN_CACHE_PASS_PHRASE); + } + private static boolean nullOrEmptyString(String s) { return s == null || s.isEmpty(); } diff --git a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java new file mode 100644 index 0000000000..42e3d8f7c6 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java @@ -0,0 +1,122 @@ +package com.databricks.jdbc.auth; + +import com.databricks.jdbc.log.JdbcLogger; +import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.databricks.sdk.core.CredentialsProvider; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.HeaderFactory; +import com.databricks.sdk.core.oauth.Consent; +import com.databricks.sdk.core.oauth.OAuthClient; +import com.databricks.sdk.core.oauth.SessionCredentials; +import com.databricks.sdk.core.oauth.Token; +import com.google.common.annotations.VisibleForTesting; +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import org.apache.http.HttpHeaders; + +/** + * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow with token + * caching and automatic token refresh. This provider encrypts and caches tokens in the user's home + * directory and reuses them when possible. + */ +public class CachingExternalBrowserCredentialsProvider implements CredentialsProvider { + + private static final String AUTH_TYPE = "external-browser-with-cache"; + private static final JdbcLogger LOGGER = + JdbcLoggerFactory.getLogger(CachingExternalBrowserCredentialsProvider.class); + private final TokenCache tokenCache; + private final DatabricksConfig config; + private Token currentToken; + + public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, String passphrase) { + this(config, new TokenCache(config.getUsername(), passphrase)); + } + + @VisibleForTesting + public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, TokenCache tokenCache) { + this.config = config; + this.tokenCache = tokenCache; + } + + @Override + public String authType() { + return AUTH_TYPE; + } + + @Override + public HeaderFactory configure(DatabricksConfig config) { + if (config.getHost() == null || !config.getAuthType().equals(AUTH_TYPE)) { + return null; + } + + try { + ensureValidToken(); + return createHeaderFactory(); + } catch (Exception e) { + LOGGER.error("Failed to configure external browser credentials", e); + throw new DatabricksException("Failed to configure external browser credentials", e); + } + } + + private void ensureValidToken() throws IOException, DatabricksException { + // Try to load from cache first + currentToken = tokenCache.load(); + + if (currentToken != null) { + LOGGER.debug("Cached token found"); + if (!currentToken.isExpired()) { + // Use cached access token if it's still valid + LOGGER.debug("Use cached token since it is still valid"); + return; + } + + // Try to refresh the token if we have a refresh token + if (currentToken.getRefreshToken() != null) { + try { + LOGGER.debug("Using cached refresh token to get new access token"); + currentToken = refreshAccessToken(); + tokenCache.save(currentToken); + return; + } catch (Exception e) { + LOGGER.info("Failed to refresh access token, will restart browser auth", e); + // If refresh fails, fall through to browser auth + } + } + } + + // If we get here, we need to do browser auth + LOGGER.debug( + "Cached token not found or unable to refresh access token, will restart browser auth"); + currentToken = performBrowserAuth(); + tokenCache.save(currentToken); + } + + @VisibleForTesting + Token refreshAccessToken() throws IOException, DatabricksException { + OAuthClient client = new OAuthClient(config); + Consent consent = client.initiateConsent(); + SessionCredentials creds = + consent.exchangeCallbackParameters(Map.of("refresh_token", currentToken.getRefreshToken())); + return creds.getToken(); + } + + @VisibleForTesting + Token performBrowserAuth() throws IOException, DatabricksException { + OAuthClient client = new OAuthClient(config); + Consent consent = client.initiateConsent(); + SessionCredentials creds = consent.launchExternalBrowser(); + return creds.getToken(); + } + + private HeaderFactory createHeaderFactory() { + return () -> { + Map headers = new HashMap<>(); + headers.put( + HttpHeaders.AUTHORIZATION, + currentToken.getTokenType() + " " + currentToken.getAccessToken()); + return headers; + }; + } +} diff --git a/src/main/java/com/databricks/jdbc/auth/TokenCache.java b/src/main/java/com/databricks/jdbc/auth/TokenCache.java new file mode 100644 index 0000000000..3b15677f98 --- /dev/null +++ b/src/main/java/com/databricks/jdbc/auth/TokenCache.java @@ -0,0 +1,119 @@ +package com.databricks.jdbc.auth; + +import com.databricks.sdk.core.oauth.Token; +import com.databricks.sdk.core.utils.ClockSupplier; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.security.spec.KeySpec; +import java.time.LocalDateTime; +import java.util.Base64; +import javax.crypto.Cipher; +import javax.crypto.SecretKey; +import javax.crypto.SecretKeyFactory; +import javax.crypto.spec.PBEKeySpec; +import javax.crypto.spec.SecretKeySpec; + +public class TokenCache { + private static final String CACHE_DIR = ".databricks"; + private static final String CACHE_FILE_SUFFIX = ".databricks_jdbc_token_cache"; + 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 Path cacheFile; + private final String passphrase; + private final ObjectMapper mapper; + + @JsonIgnoreProperties(ignoreUnknown = true) + static class SerializableToken extends Token { + public SerializableToken(String accessToken, String tokenType, LocalDateTime expiry) { + super(accessToken, tokenType, expiry); + } + + public SerializableToken( + String accessToken, String tokenType, LocalDateTime expiry, ClockSupplier clockSupplier) { + super(accessToken, tokenType, expiry, clockSupplier); + } + + @JsonCreator + public SerializableToken( + @JsonProperty("accessToken") String accessToken, + @JsonProperty("tokenType") String tokenType, + @JsonProperty("refreshToken") String refreshToken, + @JsonProperty("expiry") LocalDateTime expiry) { + super(accessToken, tokenType, refreshToken, expiry); + } + + public SerializableToken( + String accessToken, + String tokenType, + String refreshToken, + LocalDateTime expiry, + ClockSupplier clockSupplier) { + super(accessToken, tokenType, refreshToken, expiry, clockSupplier); + } + } + + public TokenCache(String host, String passphrase) { + if (passphrase == null || passphrase.isEmpty()) { + throw new IllegalArgumentException( + "Required setting TokenCachePassPhrase has not been provided in connection settings"); + } + this.passphrase = passphrase; + this.cacheFile = + Paths.get(System.getProperty("java.io.tmpdir"), CACHE_DIR, host + CACHE_FILE_SUFFIX); + this.mapper = new ObjectMapper(); + this.mapper.registerModule(new JavaTimeModule()); + } + + public void save(Token token) throws IOException { + try { + Files.createDirectories(cacheFile.getParent()); + String json = mapper.writeValueAsString(token); + byte[] encrypted = encrypt(json.getBytes()); + Files.write(cacheFile, encrypted); + } catch (Exception e) { + throw new IOException("Failed to save token cache: " + e.getMessage(), e); + } + } + + public Token load() throws IOException { + try { + if (!Files.exists(cacheFile)) { + return null; + } + byte[] encrypted = Files.readAllBytes(cacheFile); + byte[] decrypted = decrypt(encrypted); + return mapper.readValue(decrypted, SerializableToken.class); + } catch (Exception e) { + throw new IOException("Failed to load token cache: " + e.getMessage(), e); + } + } + + 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); + } + + 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)); + } + + 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)); + } +} diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index 2bdece4ab4..40c7207ae8 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -90,7 +90,8 @@ public enum DatabricksJdbcUrlParams { "EnableComplexDatatypeSupport", "flag to enable native support of complex data types as java objects", "0"), - AZURE_TENANT_ID("AzureTenantId", "Azure tenant ID"); + AZURE_TENANT_ID("AzureTenantId", "Azure tenant ID"), + TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"); private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java index c9772f3960..1d1b0fbd1f 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java @@ -4,6 +4,7 @@ import static com.databricks.jdbc.common.util.DatabricksAuthUtil.initializeConfigWithToken; import com.databricks.jdbc.api.IDatabricksConnectionContext; +import com.databricks.jdbc.auth.CachingExternalBrowserCredentialsProvider; import com.databricks.jdbc.auth.OAuthRefreshCredentialsProvider; import com.databricks.jdbc.auth.PrivateKeyClientCredentialProvider; import com.databricks.jdbc.common.AuthMech; @@ -125,8 +126,12 @@ public void setupOAuthConfig() throws DatabricksParsingException { /** Setup the OAuth U2M authentication settings in the databricks config. */ public void setupU2MConfig() throws DatabricksParsingException { + CredentialsProvider provider = + new CachingExternalBrowserCredentialsProvider( + databricksConfig, connectionContext.getTokenCachePassPhrase()); databricksConfig - .setAuthType(DatabricksJdbcConstants.U2M_AUTH_TYPE) + .setAuthType(provider.authType()) + .setCredentialsProvider(provider) .setHost(connectionContext.getHostForOAuth()) .setClientId(connectionContext.getClientId()) .setClientSecret(connectionContext.getClientSecret()) diff --git a/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java b/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java new file mode 100644 index 0000000000..2cd513f618 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java @@ -0,0 +1,154 @@ +package com.databricks.jdbc.auth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.HeaderFactory; +import com.databricks.sdk.core.oauth.Token; +import java.io.IOException; +import java.time.LocalDateTime; +import java.util.Arrays; +import java.util.Map; +import org.apache.http.HttpHeaders; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +public class CachingExternalBrowserCredentialsProviderTest { + private static final String TEST_HOST = "test-host"; + + @Mock private TokenCache tokenCache; + + private CachingExternalBrowserCredentialsProvider provider; + private DatabricksConfig config; + + @BeforeEach + void setUp() { + config = + new DatabricksConfig() + .setAuthType("external-browser-with-cache") + .setHost(TEST_HOST) + .setClientId("test-client-id") + .setUsername("test-user") + .setScopes(Arrays.asList("offline_access", "clusters", "sql")); + + provider = spy(new CachingExternalBrowserCredentialsProvider(config, tokenCache)); + } + + @Test + void testAuthType() { + assertEquals("external-browser-with-cache", provider.authType()); + } + + @Test + void testConfigureWithInvalidConfig() { + DatabricksConfig invalidConfig = new DatabricksConfig().setAuthType("invalid-type"); + assertNull(provider.configure(invalidConfig)); + } + + @Test + void testUseValidTokenFromCache() throws IOException { + // Setup valid token in cache + Token validToken = + new Token( + "cached-token", "Bearer", "cached-refresh-token", LocalDateTime.now().plusHours(1)); + when(tokenCache.load()).thenReturn(validToken); + + // Should use cached token without any refresh or browser auth + HeaderFactory headerFactory = provider.configure(config); + assertNotNull(headerFactory); + + Map headers = headerFactory.headers(); + assertEquals("Bearer cached-token", headers.get(HttpHeaders.AUTHORIZATION)); + + // Verify no refresh or browser auth was attempted + verify(provider, never()).refreshAccessToken(); + verify(provider, never()).performBrowserAuth(); + } + + @Test + void testRefreshExpiredTokenSuccess() throws IOException { + // Setup expired token in cache + Token expiredToken = + new Token( + "expired-token", "Bearer", "cached-refresh-token", LocalDateTime.now().minusMinutes(5)); + when(tokenCache.load()).thenReturn(expiredToken); + + // Setup successful token refresh + Token refreshedToken = + new Token( + "refreshed-token", "Bearer", "new-refresh-token", LocalDateTime.now().plusHours(1)); + doReturn(refreshedToken).when(provider).refreshAccessToken(); + + // Should refresh token using cached refresh token + HeaderFactory headerFactory = provider.configure(config); + assertNotNull(headerFactory); + + Map headers = headerFactory.headers(); + assertEquals("Bearer refreshed-token", headers.get(HttpHeaders.AUTHORIZATION)); + + // Verify refresh was attempted but not browser auth + verify(provider).refreshAccessToken(); + verify(provider, never()).performBrowserAuth(); + verify(tokenCache).save(refreshedToken); + } + + @Test + void testRefreshTokenFailureFallbackToBrowserAuth() throws IOException { + // Setup expired token in cache + Token expiredToken = + new Token( + "expired-token", + "Bearer", + "invalid-refresh-token", + LocalDateTime.now().minusMinutes(5)); + when(tokenCache.load()).thenReturn(expiredToken); + + // Setup failed token refresh + doThrow(new IOException("Invalid refresh token")).when(provider).refreshAccessToken(); + + // Setup successful browser auth + Token newToken = + new Token("new-token", "Bearer", "new-refresh-token", LocalDateTime.now().plusHours(1)); + doReturn(newToken).when(provider).performBrowserAuth(); + + // Should fall back to browser auth after refresh fails + HeaderFactory headerFactory = provider.configure(config); + assertNotNull(headerFactory); + + Map headers = headerFactory.headers(); + assertEquals("Bearer new-token", headers.get(HttpHeaders.AUTHORIZATION)); + + // Verify both refresh and browser auth were attempted + verify(provider).refreshAccessToken(); + verify(provider).performBrowserAuth(); + verify(tokenCache).save(newToken); + } + + @Test + void testEmptyCacheFallbackToBrowserAuth() throws IOException { + // Setup empty cache + when(tokenCache.load()).thenReturn(null); + + // Setup successful browser auth + Token newToken = + new Token("new-token", "Bearer", "new-refresh-token", LocalDateTime.now().plusHours(1)); + doReturn(newToken).when(provider).performBrowserAuth(); + + // Should directly use browser auth + HeaderFactory headerFactory = provider.configure(config); + assertNotNull(headerFactory); + + Map headers = headerFactory.headers(); + assertEquals("Bearer new-token", headers.get(HttpHeaders.AUTHORIZATION)); + + // Verify only browser auth was attempted + verify(provider, never()).refreshAccessToken(); + verify(provider).performBrowserAuth(); + verify(tokenCache).save(newToken); + } +} diff --git a/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java new file mode 100644 index 0000000000..6b55bce591 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java @@ -0,0 +1,74 @@ +package com.databricks.jdbc.auth; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.oauth.Token; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.time.LocalDateTime; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TokenCacheTest { + private static final String TEST_HOST = "test-host"; + private static final String TEST_PASSPHRASE = "test-passphrase"; + private Path cacheFile; + private TokenCache tokenCache; + + @BeforeEach + void setUp() { + tokenCache = new TokenCache(TEST_HOST, TEST_PASSPHRASE); + cacheFile = + Paths.get( + System.getProperty("java.io.tmpdir"), + ".databricks", + TEST_HOST + ".databricks_jdbc_token_cache_test"); + } + + @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); + assertEquals("access-token", loadedToken.getAccessToken()); + assertEquals("Bearer", loadedToken.getTokenType()); + assertEquals("refresh-token", loadedToken.getRefreshToken()); + assertFalse(loadedToken.isExpired()); + } + + @Test + void testInvalidPassphrase() { + assertThrows(IllegalArgumentException.class, () -> new TokenCache(TEST_HOST, null)); + assertThrows(IllegalArgumentException.class, () -> new TokenCache(TEST_HOST, "")); + } + + @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(); + assertEquals("token2", loadedToken.getAccessToken()); + assertEquals("refresh2", loadedToken.getRefreshToken()); + } +} diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 1f2343c3a1..031037e93f 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -8,6 +8,7 @@ import com.databricks.jdbc.api.IDatabricksConnectionContext; import com.databricks.jdbc.api.impl.DatabricksConnectionContextFactory; +import com.databricks.jdbc.auth.CachingExternalBrowserCredentialsProvider; import com.databricks.jdbc.auth.PrivateKeyClientCredentialProvider; import com.databricks.jdbc.common.AuthFlow; import com.databricks.jdbc.common.AuthMech; @@ -209,6 +210,42 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals(DatabricksJdbcConstants.U2M_AUTH_TYPE, config.getAuthType()); } + @Test + void getWorkspaceClient_OAuthWithCachingExternalBrowser_AuthenticatesCorrectly() + throws DatabricksParsingException { + when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); + when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); + when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); + when(mockContext.getClientId()).thenReturn("client-id"); + when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); + configurator = new ClientConfigurator(mockContext); + + WorkspaceClient client = configurator.getWorkspaceClient(); + assertNotNull(client); + DatabricksConfig config = client.config(); + + assertEquals("https://oauth-browser.databricks.com", config.getHost()); + assertEquals("client-id", config.getClientId()); + assertInstanceOf( + CachingExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); + } + + @Test + void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsException() + throws DatabricksParsingException { + when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); + when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); + when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); + when(mockContext.getClientId()).thenReturn("client-id"); + when(mockContext.getTokenCachePassPhrase()).thenReturn(null); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); + + assertThrows(IllegalArgumentException.class, () -> new ClientConfigurator(mockContext)); + } + @Test void testNonOauth() { when(mockContext.getAuthMech()).thenReturn(AuthMech.OTHER); From 276786f24b84dc96c2b599ebf166796e31f0bdd1 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Tue, 1 Apr 2025 12:20:20 +0530 Subject: [PATCH 2/9] fix --- .../jdbc/auth/CachingExternalBrowserCredentialsProvider.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java index 42e3d8f7c6..68bc1ab6cb 100644 --- a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java +++ b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java @@ -31,7 +31,7 @@ public class CachingExternalBrowserCredentialsProvider implements CredentialsPro private Token currentToken; public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, String passphrase) { - this(config, new TokenCache(config.getUsername(), passphrase)); + this(config, new TokenCache(config.getHost(), passphrase)); } @VisibleForTesting From e1ab9c82d342f890da9c4ce42a120ce5efa4d94a Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 13:29:07 +0530 Subject: [PATCH 3/9] fix tests --- .../dbclient/impl/common/ClientConfiguratorTest.java | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 031037e93f..294da35b12 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -170,6 +170,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.getClientSecret()).thenReturn("browser-client-secret"); when(mockContext.getOAuthScopesForU2M()).thenReturn(List.of(new String[] {"scope1", "scope2"})); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.getTokenCachePassPhrase()).thenReturn("token-cache-passphrase"); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -180,7 +181,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals("browser-client-secret", config.getClientSecret()); assertEquals(List.of(new String[] {"scope1", "scope2"}), config.getScopes()); assertEquals(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL, config.getOAuthRedirectUrl()); - assertEquals(DatabricksJdbcConstants.U2M_AUTH_TYPE, config.getAuthType()); + assertEquals("external-browser-with-cache", config.getAuthType()); } @Test @@ -196,6 +197,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); when(mockContext.getOAuthDiscoveryURL()).thenReturn(TEST_DISCOVERY_URL); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.getTokenCachePassPhrase()).thenReturn("token-cache-passphrase"); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -207,7 +209,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals("browser-client-secret", config.getClientSecret()); assertEquals(List.of(new String[] {"scope1", "scope2"}), config.getScopes()); assertEquals(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL, config.getOAuthRedirectUrl()); - assertEquals(DatabricksJdbcConstants.U2M_AUTH_TYPE, config.getAuthType()); + assertEquals("external-browser-with-cache", config.getAuthType()); } @Test @@ -233,12 +235,9 @@ void getWorkspaceClient_OAuthWithCachingExternalBrowser_AuthenticatesCorrectly() } @Test - void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsException() - throws DatabricksParsingException { + void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsException() { when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); - when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); - when(mockContext.getClientId()).thenReturn("client-id"); when(mockContext.getTokenCachePassPhrase()).thenReturn(null); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); From 9db852fe51af6589bdc2d45e3ddf0b267806b280 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 13:38:10 +0530 Subject: [PATCH 4/9] add EnableTokenCache --- .../api/IDatabricksConnectionContext.java | 3 ++ .../api/impl/DatabricksConnectionContext.java | 5 +++ .../jdbc/common/DatabricksJdbcUrlParams.java | 3 +- .../impl/common/ClientConfigurator.java | 16 +++++-- .../impl/DatabricksConnectionContextTest.java | 42 ++++++++++++++---- .../impl/common/ClientConfiguratorTest.java | 43 +++++++++++++++++++ 6 files changed, 100 insertions(+), 12 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java index e16c4d770f..4b7f10dcf9 100644 --- a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java @@ -250,4 +250,7 @@ public interface IDatabricksConnectionContext { /** Returns the passphrase used for encrypting/decrypting token cache */ String getTokenCachePassPhrase(); + + /** Returns whether token caching is enabled for OAuth authentication */ + boolean isTokenCacheEnabled(); } diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index d695b3516a..9655aea757 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -797,6 +797,11 @@ public String getTokenCachePassPhrase() { return getParameter(DatabricksJdbcUrlParams.TOKEN_CACHE_PASS_PHRASE); } + @Override + public boolean isTokenCacheEnabled() { + return getParameter(DatabricksJdbcUrlParams.ENABLE_TOKEN_CACHE).equals("1"); + } + private static boolean nullOrEmptyString(String s) { return s == null || s.isEmpty(); } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index e13d447b24..0174669a28 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -107,7 +107,8 @@ public enum DatabricksJdbcUrlParams { "DefaultStringColumnLength", "Maximum number of characters that can be contained in STRING columns", "255"), - TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"); + TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"), + ENABLE_TOKEN_CACHE("EnableTokenCache", "Enable caching OAuth tokens", "1"); private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java index e1aca91fd7..4f8017735c 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java @@ -21,6 +21,7 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.ProxyConfig; import com.databricks.sdk.core.commons.CommonsHttpClient; +import com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider; import com.databricks.sdk.core.utils.Cloud; import java.security.cert.*; import java.util.Arrays; @@ -130,9 +131,18 @@ public void setupOAuthConfig() throws DatabricksParsingException { /** Setup the OAuth U2M authentication settings in the databricks config. */ public void setupU2MConfig() throws DatabricksParsingException { - CredentialsProvider provider = - new CachingExternalBrowserCredentialsProvider( - databricksConfig, connectionContext.getTokenCachePassPhrase()); + CredentialsProvider provider; + + if (connectionContext.isTokenCacheEnabled()) { + LOGGER.debug("Using CachingExternalBrowserCredentialsProvider as token caching is enabled"); + provider = + new CachingExternalBrowserCredentialsProvider( + databricksConfig, connectionContext.getTokenCachePassPhrase()); + } else { + LOGGER.debug("Using ExternalBrowserCredentialsProvider as token caching is disabled"); + provider = new ExternalBrowserCredentialsProvider(); + } + databricksConfig .setAuthType(provider.authType()) .setCredentialsProvider(provider) diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java index fb46338f37..d51601f023 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java @@ -441,13 +441,39 @@ public void testParsingOfUrlWithSpecifiedCatalogAndSchema() throws DatabricksSQL @Test void testLogLevels() { - assertEquals(getLogLevel(123), LogLevel.OFF); - assertEquals(getLogLevel(0), LogLevel.OFF); - assertEquals(getLogLevel(1), LogLevel.FATAL); - assertEquals(getLogLevel(2), LogLevel.ERROR); - assertEquals(getLogLevel(3), LogLevel.WARN); - assertEquals(getLogLevel(4), LogLevel.INFO); - assertEquals(getLogLevel(5), LogLevel.DEBUG); - assertEquals(getLogLevel(6), LogLevel.TRACE); + assertEquals(LogLevel.OFF, getLogLevel(0)); + assertEquals(LogLevel.FATAL, getLogLevel(1)); + assertEquals(LogLevel.ERROR, getLogLevel(2)); + assertEquals(LogLevel.WARN, getLogLevel(3)); + assertEquals(LogLevel.INFO, getLogLevel(4)); + assertEquals(LogLevel.DEBUG, getLogLevel(5)); + assertEquals(LogLevel.TRACE, getLogLevel(6)); + assertEquals(LogLevel.INFO, getLogLevel(7)); + } + + @Test + public void testIsTokenCacheEnabled() throws DatabricksSQLException { + // Test with EnableTokenCache=1 (default) + Properties properties1 = new Properties(); + DatabricksConnectionContext connectionContext1 = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties1); + assertTrue(connectionContext1.isTokenCacheEnabled()); + + // Test with EnableTokenCache=0 + Properties properties2 = new Properties(); + properties2.setProperty("EnableTokenCache", "0"); + DatabricksConnectionContext connectionContext2 = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties2); + assertFalse(connectionContext2.isTokenCacheEnabled()); + + // Test with EnableTokenCache=1 explicitly set + Properties properties3 = new Properties(); + properties3.setProperty("EnableTokenCache", "1"); + DatabricksConnectionContext connectionContext3 = + (DatabricksConnectionContext) + DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties3); + assertTrue(connectionContext3.isTokenCacheEnabled()); } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 294da35b12..073f2becc0 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -222,6 +222,7 @@ void getWorkspaceClient_OAuthWithCachingExternalBrowser_AuthenticatesCorrectly() when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); + when(mockContext.isTokenCacheEnabled()).thenReturn(true); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); @@ -245,6 +246,48 @@ void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsExcep assertThrows(IllegalArgumentException.class, () -> new ClientConfigurator(mockContext)); } + @Test + void getWorkspaceClient_OAuthBrowserAuth_WithTokenCacheEnabled_UsesCachedProvider() + throws DatabricksParsingException { + when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); + when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); + when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); + when(mockContext.getClientId()).thenReturn("client-id"); + when(mockContext.getClientSecret()).thenReturn("client-secret"); + when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.isTokenCacheEnabled()).thenReturn(true); + + configurator = new ClientConfigurator(mockContext); + WorkspaceClient client = configurator.getWorkspaceClient(); + DatabricksConfig config = client.config(); + + assertEquals("external-browser-with-cache", config.getAuthType()); + assertInstanceOf( + CachingExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); + } + + @Test + void getWorkspaceClient_OAuthBrowserAuth_WithTokenCacheDisabled_UsesStandardProvider() + throws DatabricksParsingException { + when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); + when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); + when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); + when(mockContext.getClientId()).thenReturn("client-id"); + when(mockContext.getClientSecret()).thenReturn("client-secret"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.isTokenCacheEnabled()).thenReturn(false); + + configurator = new ClientConfigurator(mockContext); + WorkspaceClient client = configurator.getWorkspaceClient(); + DatabricksConfig config = client.config(); + + assertEquals("external-browser", config.getAuthType()); + assertInstanceOf( + com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider.class, + config.getCredentialsProvider()); + } + @Test void testNonOauth() { when(mockContext.getAuthMech()).thenReturn(AuthMech.OTHER); From 63be61129e23e6620701787af12237ed2dbc4c3d Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 13:42:51 +0530 Subject: [PATCH 5/9] add javadoc --- ...ingExternalBrowserCredentialsProvider.java | 80 ++++++++++++++++++- .../com/databricks/jdbc/auth/TokenCache.java | 53 ++++++++++++ 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java index 68bc1ab6cb..94de38822b 100644 --- a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java +++ b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java @@ -18,8 +18,21 @@ /** * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow with token - * caching and automatic token refresh. This provider encrypts and caches tokens in the user's home - * directory and reuses them when possible. + * caching and automatic token refresh. This provider encrypts and caches tokens in the user's + * temporary directory and reuses them when possible. + * + *

This provider extends the standard OAuth browser-based authentication flow by adding + * persistence of tokens between sessions. When a token is obtained after successful authentication, + * it is encrypted and stored locally. On subsequent connection attempts, the provider will: + * + *

    + *
  1. Try to load and use a cached token if available + *
  2. If the cached token is expired but has a refresh token, attempt to refresh it + *
  3. If no cached token exists or refresh fails, initiate the browser-based OAuth flow + *
+ * + *

This approach minimizes the need for users to repeatedly authenticate through the browser, + * improving the user experience while maintaining security through encryption of the cached tokens. */ public class CachingExternalBrowserCredentialsProvider implements CredentialsProvider { @@ -30,21 +43,49 @@ public class CachingExternalBrowserCredentialsProvider implements CredentialsPro private final DatabricksConfig config; private Token currentToken; + /** + * Creates a new CachingExternalBrowserCredentialsProvider with the specified configuration and + * passphrase for token encryption. + * + * @param config The Databricks configuration to use for authentication + * @param passphrase The passphrase to use for encrypting and decrypting cached tokens + * @throws IllegalArgumentException if the passphrase is null or empty + */ public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, String passphrase) { this(config, new TokenCache(config.getHost(), passphrase)); } + /** + * Creates a new CachingExternalBrowserCredentialsProvider with the specified configuration and + * token cache. This constructor is primarily used for testing. + * + * @param config The Databricks configuration to use for authentication + * @param tokenCache The token cache to use for storing and retrieving tokens + */ @VisibleForTesting public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, TokenCache tokenCache) { this.config = config; this.tokenCache = tokenCache; } + /** + * Returns the authentication type identifier for this provider. + * + * @return The string "external-browser-with-cache" + */ @Override public String authType() { return AUTH_TYPE; } + /** + * Configures the authentication by setting up the necessary headers for authenticated requests. + * This method implements the core OAuth flow with caching logic. + * + * @param config The Databricks configuration to use + * @return A HeaderFactory that adds the OAuth authentication header to requests, or null if the + * configuration is not valid for this provider + */ @Override public HeaderFactory configure(DatabricksConfig config) { if (config.getHost() == null || !config.getAuthType().equals(AUTH_TYPE)) { @@ -60,6 +101,22 @@ public HeaderFactory configure(DatabricksConfig config) { } } + /** + * Ensures that a valid token is available for authentication. This method implements the token + * loading, validation, refreshing, and browser authentication flow. + * + *

The method follows this sequence: + * + *

    + *
  1. Try to load a token from the cache + *
  2. If a token is found and not expired, use it + *
  3. If a token is found and expired but has a refresh token, try to refresh it + *
  4. If no token is found or refresh fails, perform browser-based authentication + *
+ * + * @throws IOException If there is an error reading from or writing to the token cache + * @throws DatabricksException If there is an error during the authentication process + */ private void ensureValidToken() throws IOException, DatabricksException { // Try to load from cache first currentToken = tokenCache.load(); @@ -93,6 +150,13 @@ private void ensureValidToken() throws IOException, DatabricksException { tokenCache.save(currentToken); } + /** + * Refreshes an access token using the refresh token from the current token. + * + * @return A new token with a refreshed access token + * @throws IOException If there is an error during the refresh process + * @throws DatabricksException If the Databricks API returns an error + */ @VisibleForTesting Token refreshAccessToken() throws IOException, DatabricksException { OAuthClient client = new OAuthClient(config); @@ -102,6 +166,13 @@ Token refreshAccessToken() throws IOException, DatabricksException { return creds.getToken(); } + /** + * Performs browser-based authentication to obtain a new token. + * + * @return A new token obtained through browser authentication + * @throws IOException If there is an error during the authentication process + * @throws DatabricksException If the Databricks API returns an error + */ @VisibleForTesting Token performBrowserAuth() throws IOException, DatabricksException { OAuthClient client = new OAuthClient(config); @@ -110,6 +181,11 @@ Token performBrowserAuth() throws IOException, DatabricksException { return creds.getToken(); } + /** + * Creates a HeaderFactory that adds the OAuth Bearer token to requests. + * + * @return A HeaderFactory that adds the Authorization header with the OAuth Bearer token + */ private HeaderFactory createHeaderFactory() { return () -> { Map headers = new HashMap<>(); diff --git a/src/main/java/com/databricks/jdbc/auth/TokenCache.java b/src/main/java/com/databricks/jdbc/auth/TokenCache.java index 3b15677f98..da0416a277 100644 --- a/src/main/java/com/databricks/jdbc/auth/TokenCache.java +++ b/src/main/java/com/databricks/jdbc/auth/TokenCache.java @@ -20,6 +20,16 @@ import javax.crypto.spec.PBEKeySpec; import javax.crypto.spec.SecretKeySpec; +/** + * A secure cache for OAuth tokens that encrypts and persists tokens to the local filesystem. + * + *

This class provides functionality to securely store and retrieve OAuth tokens between + * application sessions. Tokens are encrypted using AES encryption with a key derived from the + * provided passphrase using PBKDF2 key derivation. + * + *

The tokens are stored in a file within the system's temporary directory under a '.databricks' + * subdirectory. Each host connection has its own token cache file. + */ public class TokenCache { private static final String CACHE_DIR = ".databricks"; private static final String CACHE_FILE_SUFFIX = ".databricks_jdbc_token_cache"; @@ -33,6 +43,10 @@ public class TokenCache { private final String passphrase; private final ObjectMapper mapper; + /** + * A serializable version of the Token class that can be serialized/deserialized by Jackson. This + * class extends the Token class from the SDK and adds JSON annotations for proper serialization. + */ @JsonIgnoreProperties(ignoreUnknown = true) static class SerializableToken extends Token { public SerializableToken(String accessToken, String tokenType, LocalDateTime expiry) { @@ -63,6 +77,13 @@ public SerializableToken( } } + /** + * Constructs a new TokenCache instance for the specified host with encryption. + * + * @param host The Databricks host address for which to cache tokens + * @param passphrase The passphrase used to encrypt/decrypt the token cache + * @throws IllegalArgumentException if the passphrase is null or empty + */ public TokenCache(String host, String passphrase) { if (passphrase == null || passphrase.isEmpty()) { throw new IllegalArgumentException( @@ -75,6 +96,12 @@ public TokenCache(String host, String passphrase) { this.mapper.registerModule(new JavaTimeModule()); } + /** + * Saves a token to the cache file, encrypting it with the configured passphrase. + * + * @param token The token to save to the cache + * @throws IOException If an error occurs writing the token to the file or during encryption + */ public void save(Token token) throws IOException { try { Files.createDirectories(cacheFile.getParent()); @@ -86,6 +113,12 @@ public void save(Token token) throws IOException { } } + /** + * Loads a token from the cache file, decrypting it with the configured passphrase. + * + * @return The decrypted token from the cache or null if the cache file doesn't exist + * @throws IOException If an error occurs reading the token from the file or during decryption + */ public Token load() throws IOException { try { if (!Files.exists(cacheFile)) { @@ -99,18 +132,38 @@ public Token load() throws IOException { } } + /** + * 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, Base64 encoded + * @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()); From 6c6eabb840be0b7c16e4b3b64139a841c178b74b Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 13:46:22 +0530 Subject: [PATCH 6/9] fix tests --- .../jdbc/api/impl/DatabricksConnectionContextTest.java | 2 +- .../jdbc/dbclient/impl/common/ClientConfiguratorTest.java | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java index d51601f023..9b99db53ba 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java @@ -448,7 +448,7 @@ void testLogLevels() { assertEquals(LogLevel.INFO, getLogLevel(4)); assertEquals(LogLevel.DEBUG, getLogLevel(5)); assertEquals(LogLevel.TRACE, getLogLevel(6)); - assertEquals(LogLevel.INFO, getLogLevel(7)); + assertEquals(LogLevel.OFF, getLogLevel(123)); } @Test diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 073f2becc0..9674e1dfed 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -170,7 +170,6 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.getClientSecret()).thenReturn("browser-client-secret"); when(mockContext.getOAuthScopesForU2M()).thenReturn(List.of(new String[] {"scope1", "scope2"})); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); - when(mockContext.getTokenCachePassPhrase()).thenReturn("token-cache-passphrase"); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -181,7 +180,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals("browser-client-secret", config.getClientSecret()); assertEquals(List.of(new String[] {"scope1", "scope2"}), config.getScopes()); assertEquals(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL, config.getOAuthRedirectUrl()); - assertEquals("external-browser-with-cache", config.getAuthType()); + assertEquals("external-browser", config.getAuthType()); } @Test @@ -198,6 +197,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.getOAuthDiscoveryURL()).thenReturn(TEST_DISCOVERY_URL); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); when(mockContext.getTokenCachePassPhrase()).thenReturn("token-cache-passphrase"); + when(mockContext.isTokenCacheEnabled()).thenReturn(true); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -242,6 +242,7 @@ void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsExcep when(mockContext.getTokenCachePassPhrase()).thenReturn(null); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); + when(mockContext.isTokenCacheEnabled()).thenReturn(true); assertThrows(IllegalArgumentException.class, () -> new ClientConfigurator(mockContext)); } From 0f60c106b40c84064e6a0e679f727589b201172a Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 15:47:40 +0530 Subject: [PATCH 7/9] cache file must use user name instead of host --- ...ingExternalBrowserCredentialsProvider.java | 2 +- .../com/databricks/jdbc/auth/TokenCache.java | 22 ++++++++++++++----- .../databricks/jdbc/auth/TokenCacheTest.java | 10 ++++----- 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java index 94de38822b..c7de20196e 100644 --- a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java +++ b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java @@ -52,7 +52,7 @@ public class CachingExternalBrowserCredentialsProvider implements CredentialsPro * @throws IllegalArgumentException if the passphrase is null or empty */ public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, String passphrase) { - this(config, new TokenCache(config.getHost(), passphrase)); + this(config, new TokenCache(passphrase)); } /** diff --git a/src/main/java/com/databricks/jdbc/auth/TokenCache.java b/src/main/java/com/databricks/jdbc/auth/TokenCache.java index da0416a277..c6c668bab0 100644 --- a/src/main/java/com/databricks/jdbc/auth/TokenCache.java +++ b/src/main/java/com/databricks/jdbc/auth/TokenCache.java @@ -28,7 +28,7 @@ * provided passphrase using PBKDF2 key derivation. * *

The tokens are stored in a file within the system's temporary directory under a '.databricks' - * subdirectory. Each host connection has its own token cache file. + * subdirectory. Each user has their own token cache file identified by the sanitized username. */ public class TokenCache { private static final String CACHE_DIR = ".databricks"; @@ -78,24 +78,36 @@ public SerializableToken( } /** - * Constructs a new TokenCache instance for the specified host with encryption. + * Constructs a new TokenCache instance with encryption using the user's system username to + * identify the cache file. * - * @param host The Databricks host address for which to cache tokens * @param passphrase The passphrase used to encrypt/decrypt the token cache * @throws IllegalArgumentException if the passphrase is null or empty */ - public TokenCache(String host, String passphrase) { + public TokenCache(String passphrase) { if (passphrase == null || passphrase.isEmpty()) { throw new IllegalArgumentException( "Required setting TokenCachePassPhrase has not been provided in connection settings"); } this.passphrase = passphrase; this.cacheFile = - Paths.get(System.getProperty("java.io.tmpdir"), CACHE_DIR, host + CACHE_FILE_SUFFIX); + Paths.get( + System.getProperty("java.io.tmpdir"), + CACHE_DIR, + sanitizeUsernameForFile(System.getProperty("user.name")) + CACHE_FILE_SUFFIX); this.mapper = new ObjectMapper(); this.mapper.registerModule(new JavaTimeModule()); } + /** + * Sanitizes the given username so it can be safely used in a file name. Replaces all characters + * that are not a-z, A-Z, 0-9, or underscore (_) with an underscore. + */ + private static String sanitizeUsernameForFile(String username) { + if (username == null) return "unknown_user"; + return username.replaceAll("[^a-zA-Z0-9_]", "_"); + } + /** * Saves a token to the cache file, encrypting it with the configured passphrase. * diff --git a/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java index 6b55bce591..2b2faaa73c 100644 --- a/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java +++ b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java @@ -13,19 +13,19 @@ import org.junit.jupiter.api.Test; public class TokenCacheTest { - private static final String TEST_HOST = "test-host"; private static final String TEST_PASSPHRASE = "test-passphrase"; private Path cacheFile; private TokenCache tokenCache; @BeforeEach void setUp() { - tokenCache = new TokenCache(TEST_HOST, TEST_PASSPHRASE); + tokenCache = new TokenCache(TEST_PASSPHRASE); + String sanitizedUsername = System.getProperty("user.name").replaceAll("[^a-zA-Z0-9_]", "_"); cacheFile = Paths.get( System.getProperty("java.io.tmpdir"), ".databricks", - TEST_HOST + ".databricks_jdbc_token_cache_test"); + sanitizedUsername + ".databricks_jdbc_token_cache"); } @AfterEach @@ -55,8 +55,8 @@ void testSaveAndLoadToken() throws IOException { @Test void testInvalidPassphrase() { - assertThrows(IllegalArgumentException.class, () -> new TokenCache(TEST_HOST, null)); - assertThrows(IllegalArgumentException.class, () -> new TokenCache(TEST_HOST, "")); + assertThrows(IllegalArgumentException.class, () -> new TokenCache(null)); + assertThrows(IllegalArgumentException.class, () -> new TokenCache("")); } @Test From e109469634d12fd0933dd8c2557fccdbc04af78a Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 21:15:00 +0530 Subject: [PATCH 8/9] fix implementation --- ...ingExternalBrowserCredentialsProvider.java | 197 +++++++++--------- .../com/databricks/jdbc/auth/TokenCache.java | 34 +-- .../jdbc/common/util/StringUtil.java | 9 + .../impl/common/ClientConfigurator.java | 36 ++-- ...xternalBrowserCredentialsProviderTest.java | 99 ++++++--- .../databricks/jdbc/auth/TokenCacheTest.java | 10 +- .../jdbc/common/util/StringUtilTest.java | 40 ++++ .../impl/common/ClientConfiguratorTest.java | 73 +++---- 8 files changed, 295 insertions(+), 203 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java index c7de20196e..3d3b955cea 100644 --- a/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java +++ b/src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java @@ -1,17 +1,22 @@ package com.databricks.jdbc.auth; +import static com.databricks.jdbc.auth.AuthConstants.GRANT_TYPE_KEY; +import static com.databricks.jdbc.auth.AuthConstants.GRANT_TYPE_REFRESH_TOKEN_KEY; + +import com.databricks.jdbc.api.IDatabricksConnectionContext; +import com.databricks.jdbc.common.DatabricksJdbcConstants; +import com.databricks.jdbc.common.util.DatabricksAuthUtil; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; import com.databricks.sdk.core.CredentialsProvider; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.HeaderFactory; -import com.databricks.sdk.core.oauth.Consent; -import com.databricks.sdk.core.oauth.OAuthClient; -import com.databricks.sdk.core.oauth.SessionCredentials; -import com.databricks.sdk.core.oauth.Token; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.oauth.*; import com.google.common.annotations.VisibleForTesting; import java.io.IOException; +import java.time.LocalDateTime; import java.util.HashMap; import java.util.Map; import org.apache.http.HttpHeaders; @@ -21,51 +26,69 @@ * caching and automatic token refresh. This provider encrypts and caches tokens in the user's * temporary directory and reuses them when possible. * - *

This provider extends the standard OAuth browser-based authentication flow by adding - * persistence of tokens between sessions. When a token is obtained after successful authentication, - * it is encrypted and stored locally. On subsequent connection attempts, the provider will: + *

This provider extends {@code RefreshableTokenSource} to handle token refreshing in a + * standardized way. When a token is obtained after successful authentication, it is encrypted and + * stored locally. On subsequent connection attempts, the provider will: * *

    *
  1. Try to load and use a cached token if available - *
  2. If the cached token is expired but has a refresh token, attempt to refresh it + *
  3. If the cached token is expired but has a refresh token, attempt to refresh it using the + * OAuth2 token endpoint *
  4. If no cached token exists or refresh fails, initiate the browser-based OAuth flow *
* *

This approach minimizes the need for users to repeatedly authenticate through the browser, * improving the user experience while maintaining security through encryption of the cached tokens. */ -public class CachingExternalBrowserCredentialsProvider implements CredentialsProvider { +public class CachingExternalBrowserCredentialsProvider extends RefreshableTokenSource + implements CredentialsProvider { private static final String AUTH_TYPE = "external-browser-with-cache"; private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(CachingExternalBrowserCredentialsProvider.class); private final TokenCache tokenCache; private final DatabricksConfig config; - private Token currentToken; - - /** - * Creates a new CachingExternalBrowserCredentialsProvider with the specified configuration and - * passphrase for token encryption. - * - * @param config The Databricks configuration to use for authentication - * @param passphrase The passphrase to use for encrypting and decrypting cached tokens - * @throws IllegalArgumentException if the passphrase is null or empty - */ - public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, String passphrase) { - this(config, new TokenCache(passphrase)); - } + private final String tokenEndpoint; + private HttpClient hc; /** - * Creates a new CachingExternalBrowserCredentialsProvider with the specified configuration and - * token cache. This constructor is primarily used for testing. + * Creates a new CachingExternalBrowserCredentialsProvider with the specified configuration, + * connection context, and token cache. * * @param config The Databricks configuration to use for authentication + * @param context The connection context containing OAuth configuration parameters * @param tokenCache The token cache to use for storing and retrieving tokens */ - @VisibleForTesting - public CachingExternalBrowserCredentialsProvider(DatabricksConfig config, TokenCache tokenCache) { + public CachingExternalBrowserCredentialsProvider( + DatabricksConfig config, IDatabricksConnectionContext context, TokenCache tokenCache) { this.config = config; this.tokenCache = tokenCache; + this.tokenEndpoint = DatabricksAuthUtil.getTokenEndpoint(config, context); + try { + // Initialize token from cache + this.token = tokenCache.load(); + if (this.token == null) { + LOGGER.debug("No cached token found"); + // Initialize with an expired token to force authentication on first use + this.token = + new Token( + DatabricksJdbcConstants.EMPTY_STRING, + DatabricksJdbcConstants.EMPTY_STRING, + null, + LocalDateTime.now().minusMinutes(1)); + } else { + LOGGER.debug("Cached token found"); + } + } catch (IOException e) { + LOGGER.debug("Failed to load token from cache", e); + // Initialize with an expired token to force authentication on first use + this.token = + new Token( + DatabricksJdbcConstants.EMPTY_STRING, + DatabricksJdbcConstants.EMPTY_STRING, + null, + LocalDateTime.now().minusMinutes(1)); + } } /** @@ -92,82 +115,85 @@ public HeaderFactory configure(DatabricksConfig config) { return null; } - try { - ensureValidToken(); - return createHeaderFactory(); - } catch (Exception e) { - LOGGER.error("Failed to configure external browser credentials", e); - throw new DatabricksException("Failed to configure external browser credentials", e); + if (this.hc == null) { + this.hc = config.getHttpClient(); } + + return () -> { + Map headers = new HashMap<>(); + headers.put( + HttpHeaders.AUTHORIZATION, getToken().getTokenType() + " " + getToken().getAccessToken()); + return headers; + }; } /** - * Ensures that a valid token is available for authentication. This method implements the token - * loading, validation, refreshing, and browser authentication flow. - * - *

The method follows this sequence: - * - *

    - *
  1. Try to load a token from the cache - *
  2. If a token is found and not expired, use it - *
  3. If a token is found and expired but has a refresh token, try to refresh it - *
  4. If no token is found or refresh fails, perform browser-based authentication - *
+ * Implements the token refresh logic as required by RefreshableTokenSource. This method handles + * token refreshing, falling back to browser authentication if needed. * - * @throws IOException If there is an error reading from or writing to the token cache + * @return A new or refreshed token * @throws DatabricksException If there is an error during the authentication process */ - private void ensureValidToken() throws IOException, DatabricksException { - // Try to load from cache first - currentToken = tokenCache.load(); - - if (currentToken != null) { - LOGGER.debug("Cached token found"); - if (!currentToken.isExpired()) { - // Use cached access token if it's still valid - LOGGER.debug("Use cached token since it is still valid"); - return; - } - - // Try to refresh the token if we have a refresh token - if (currentToken.getRefreshToken() != null) { + @Override + protected Token refresh() { + try { + // Try to refresh if we have a refresh token + if (this.token != null && this.token.getRefreshToken() != null) { try { - LOGGER.debug("Using cached refresh token to get new access token"); - currentToken = refreshAccessToken(); - tokenCache.save(currentToken); - return; + LOGGER.debug("Using refresh token to get new access token"); + Token refreshedToken = refreshAccessToken(); + tokenCache.save(refreshedToken); + return refreshedToken; } catch (Exception e) { LOGGER.info("Failed to refresh access token, will restart browser auth", e); // If refresh fails, fall through to browser auth } } - } - // If we get here, we need to do browser auth - LOGGER.debug( - "Cached token not found or unable to refresh access token, will restart browser auth"); - currentToken = performBrowserAuth(); - tokenCache.save(currentToken); + // If we get here, we need to do browser auth + LOGGER.debug("Performing browser authentication to get new access token"); + Token newToken = performBrowserAuth(); + tokenCache.save(newToken); + return newToken; + } catch (Exception e) { + String errorMessage = "Failed to refresh or obtain new token"; + LOGGER.error(errorMessage, e); + throw new DatabricksException(errorMessage, e); + } } /** - * Refreshes an access token using the refresh token from the current token. + * Refreshes an access token using the refresh token from the current token. This method follows + * the OAuth 2.0 refresh token flow by sending a request to the token endpoint with the refresh + * token grant type. * * @return A new token with a refreshed access token - * @throws IOException If there is an error during the refresh process - * @throws DatabricksException If the Databricks API returns an error + * @throws DatabricksException If there is an error during the refresh process or if the token or + * refresh token is not available */ @VisibleForTesting - Token refreshAccessToken() throws IOException, DatabricksException { - OAuthClient client = new OAuthClient(config); - Consent consent = client.initiateConsent(); - SessionCredentials creds = - consent.exchangeCallbackParameters(Map.of("refresh_token", currentToken.getRefreshToken())); - return creds.getToken(); + Token refreshAccessToken() throws DatabricksException { + if (this.token == null || this.token.getRefreshToken() == null) { + throw new DatabricksException("oauth2: token is not set or refresh token is not available"); + } + + Map params = new HashMap<>(); + params.put(GRANT_TYPE_KEY, GRANT_TYPE_REFRESH_TOKEN_KEY); + params.put(GRANT_TYPE_REFRESH_TOKEN_KEY, this.token.getRefreshToken()); + Map headers = new HashMap<>(); + return retrieveToken( + hc, + config.getClientId(), + config.getClientSecret(), + tokenEndpoint, + params, + headers, + AuthParameterPosition.BODY); } /** - * Performs browser-based authentication to obtain a new token. + * Performs browser-based authentication to obtain a new token. This method launches a browser + * window to allow the user to authenticate and authorize the application. * * @return A new token obtained through browser authentication * @throws IOException If there is an error during the authentication process @@ -180,19 +206,4 @@ Token performBrowserAuth() throws IOException, DatabricksException { SessionCredentials creds = consent.launchExternalBrowser(); return creds.getToken(); } - - /** - * Creates a HeaderFactory that adds the OAuth Bearer token to requests. - * - * @return A HeaderFactory that adds the Authorization header with the OAuth Bearer token - */ - private HeaderFactory createHeaderFactory() { - return () -> { - Map headers = new HashMap<>(); - headers.put( - HttpHeaders.AUTHORIZATION, - currentToken.getTokenType() + " " + currentToken.getAccessToken()); - return headers; - }; - } } diff --git a/src/main/java/com/databricks/jdbc/auth/TokenCache.java b/src/main/java/com/databricks/jdbc/auth/TokenCache.java index c6c668bab0..96a1ebaa7b 100644 --- a/src/main/java/com/databricks/jdbc/auth/TokenCache.java +++ b/src/main/java/com/databricks/jdbc/auth/TokenCache.java @@ -1,5 +1,6 @@ package com.databricks.jdbc.auth; +import com.databricks.jdbc.common.util.StringUtil; import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.utils.ClockSupplier; import com.fasterxml.jackson.annotation.JsonCreator; @@ -7,6 +8,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import com.google.common.annotations.VisibleForTesting; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; @@ -85,27 +87,31 @@ public SerializableToken( * @throws IllegalArgumentException if the passphrase is null or empty */ public TokenCache(String passphrase) { - if (passphrase == null || passphrase.isEmpty()) { - throw new IllegalArgumentException( - "Required setting TokenCachePassPhrase has not been provided in connection settings"); - } - this.passphrase = passphrase; - this.cacheFile = + this( Paths.get( System.getProperty("java.io.tmpdir"), CACHE_DIR, - sanitizeUsernameForFile(System.getProperty("user.name")) + CACHE_FILE_SUFFIX); - this.mapper = new ObjectMapper(); - this.mapper.registerModule(new JavaTimeModule()); + StringUtil.sanitizeUsernameForFile(System.getProperty("user.name")) + + CACHE_FILE_SUFFIX), + passphrase); } /** - * Sanitizes the given username so it can be safely used in a file name. Replaces all characters - * that are not a-z, A-Z, 0-9, or underscore (_) with an underscore. + * Constructs a new TokenCache instance with encryption using the cache file path provided. + * + * @param cacheFile The cache file path + * @param passphrase The passphrase used to encrypt/decrypt the token cache */ - private static String sanitizeUsernameForFile(String username) { - if (username == null) return "unknown_user"; - return username.replaceAll("[^a-zA-Z0-9_]", "_"); + @VisibleForTesting + public TokenCache(Path cacheFile, String passphrase) { + if (passphrase == null || passphrase.isEmpty()) { + throw new IllegalArgumentException( + "Required setting TokenCachePassPhrase has not been provided in connection settings"); + } + this.passphrase = passphrase; + this.cacheFile = cacheFile; + this.mapper = new ObjectMapper(); + this.mapper.registerModule(new JavaTimeModule()); } /** diff --git a/src/main/java/com/databricks/jdbc/common/util/StringUtil.java b/src/main/java/com/databricks/jdbc/common/util/StringUtil.java index e9b7f655d2..3fb1a0bf50 100644 --- a/src/main/java/com/databricks/jdbc/common/util/StringUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/StringUtil.java @@ -59,4 +59,13 @@ public static String getVolumePath(String catalog, String schema, String volume) // We need to escape '' to prevent SQL injection return escapeStringLiteral(String.format("/Volumes/%s/%s/%s/", catalog, schema, volume)); } + + /** + * Sanitizes the given username so it can be safely used in a file name. Replaces all characters + * that are not a-z, A-Z, 0-9, or underscore (_) with an underscore. + */ + public static String sanitizeUsernameForFile(String username) { + if (username == null) return "unknown_user"; + return username.replaceAll("[^a-zA-Z0-9_]", "_"); + } } diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java index 4f8017735c..a6fbb04c54 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java @@ -4,10 +4,7 @@ import static com.databricks.jdbc.common.util.DatabricksAuthUtil.initializeConfigWithToken; import com.databricks.jdbc.api.IDatabricksConnectionContext; -import com.databricks.jdbc.auth.AzureMSICredentialProvider; -import com.databricks.jdbc.auth.CachingExternalBrowserCredentialsProvider; -import com.databricks.jdbc.auth.OAuthRefreshCredentialsProvider; -import com.databricks.jdbc.auth.PrivateKeyClientCredentialProvider; +import com.databricks.jdbc.auth.*; import com.databricks.jdbc.common.AuthMech; import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.jdbc.common.util.DriverUtil; @@ -131,28 +128,28 @@ public void setupOAuthConfig() throws DatabricksParsingException { /** Setup the OAuth U2M authentication settings in the databricks config. */ public void setupU2MConfig() throws DatabricksParsingException { - CredentialsProvider provider; + databricksConfig + .setHost(connectionContext.getHostForOAuth()) + .setClientId(connectionContext.getClientId()) + .setClientSecret(connectionContext.getClientSecret()) + .setOAuthRedirectUrl(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL); + if (!databricksConfig.isAzure()) { + databricksConfig.setScopes(connectionContext.getOAuthScopesForU2M()); + } + CredentialsProvider provider; if (connectionContext.isTokenCacheEnabled()) { LOGGER.debug("Using CachingExternalBrowserCredentialsProvider as token caching is enabled"); + TokenCache tokenCache = new TokenCache(connectionContext.getTokenCachePassPhrase()); provider = new CachingExternalBrowserCredentialsProvider( - databricksConfig, connectionContext.getTokenCachePassPhrase()); + databricksConfig, connectionContext, tokenCache); } else { LOGGER.debug("Using ExternalBrowserCredentialsProvider as token caching is disabled"); provider = new ExternalBrowserCredentialsProvider(); } - databricksConfig - .setAuthType(provider.authType()) - .setCredentialsProvider(provider) - .setHost(connectionContext.getHostForOAuth()) - .setClientId(connectionContext.getClientId()) - .setClientSecret(connectionContext.getClientSecret()) - .setOAuthRedirectUrl(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL); - if (!databricksConfig.isAzure()) { - databricksConfig.setScopes(connectionContext.getOAuthScopesForU2M()); - } + databricksConfig.setCredentialsProvider(provider).setAuthType(provider.authType()); } /** Setup the PAT authentication settings in the databricks config. */ @@ -177,14 +174,13 @@ public void resetAccessTokenInConfig(String newAccessToken) { /** Setup the OAuth U2M refresh token authentication settings in the databricks config. */ public void setupU2MRefreshConfig() throws DatabricksParsingException { - CredentialsProvider provider = - new OAuthRefreshCredentialsProvider(connectionContext, databricksConfig); databricksConfig .setHost(connectionContext.getHostForOAuth()) - .setAuthType(provider.authType()) // oauth-refresh - .setCredentialsProvider(provider) .setClientId(connectionContext.getClientId()) .setClientSecret(connectionContext.getClientSecret()); + CredentialsProvider provider = + new OAuthRefreshCredentialsProvider(connectionContext, databricksConfig); + databricksConfig.setAuthType(provider.authType()).setCredentialsProvider(provider); } /** Setup the OAuth M2M authentication settings in the databricks config. */ diff --git a/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java b/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java index 2cd513f618..ead75942e5 100644 --- a/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java +++ b/src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java @@ -1,14 +1,19 @@ package com.databricks.jdbc.auth; +import static com.databricks.jdbc.TestConstants.TEST_AUTH_URL; +import static com.databricks.jdbc.TestConstants.TEST_TOKEN_URL; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; +import com.databricks.jdbc.api.IDatabricksConnectionContext; +import com.databricks.jdbc.common.DatabricksJdbcConstants; import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.HeaderFactory; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; import com.databricks.sdk.core.oauth.Token; import java.io.IOException; import java.time.LocalDateTime; -import java.util.Arrays; import java.util.Map; import org.apache.http.HttpHeaders; import org.junit.jupiter.api.BeforeEach; @@ -20,44 +25,50 @@ @ExtendWith(MockitoExtension.class) public class CachingExternalBrowserCredentialsProviderTest { private static final String TEST_HOST = "test-host"; + private static final String AUTH_TYPE = "external-browser-with-cache"; @Mock private TokenCache tokenCache; + @Mock private IDatabricksConnectionContext connectionContext; + @Mock private DatabricksConfig config; private CachingExternalBrowserCredentialsProvider provider; - private DatabricksConfig config; @BeforeEach - void setUp() { - config = - new DatabricksConfig() - .setAuthType("external-browser-with-cache") - .setHost(TEST_HOST) - .setClientId("test-client-id") - .setUsername("test-user") - .setScopes(Arrays.asList("offline_access", "clusters", "sql")); - - provider = spy(new CachingExternalBrowserCredentialsProvider(config, tokenCache)); + void setUp() throws IOException { + // Set up necessary mocks + doReturn(new OpenIDConnectEndpoints(TEST_TOKEN_URL, TEST_AUTH_URL)) + .when(config) + .getOidcEndpoints(); } @Test void testAuthType() { - assertEquals("external-browser-with-cache", provider.authType()); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); + assertEquals(AUTH_TYPE, provider.authType()); } @Test void testConfigureWithInvalidConfig() { DatabricksConfig invalidConfig = new DatabricksConfig().setAuthType("invalid-type"); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); assertNull(provider.configure(invalidConfig)); } @Test void testUseValidTokenFromCache() throws IOException { + when(config.getHost()).thenReturn(TEST_HOST); + when(config.getAuthType()).thenReturn(AUTH_TYPE); // Setup valid token in cache Token validToken = new Token( "cached-token", "Bearer", "cached-refresh-token", LocalDateTime.now().plusHours(1)); when(tokenCache.load()).thenReturn(validToken); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); + // Should use cached token without any refresh or browser auth HeaderFactory headerFactory = provider.configure(config); assertNotNull(headerFactory); @@ -66,18 +77,23 @@ void testUseValidTokenFromCache() throws IOException { assertEquals("Bearer cached-token", headers.get(HttpHeaders.AUTHORIZATION)); // Verify no refresh or browser auth was attempted - verify(provider, never()).refreshAccessToken(); - verify(provider, never()).performBrowserAuth(); + verify(provider, never()).refresh(); } @Test - void testRefreshExpiredTokenSuccess() throws IOException { + void testRefreshExpiredTokenSuccess() throws IOException, DatabricksException { + when(config.getHost()).thenReturn(TEST_HOST); + when(config.getAuthType()).thenReturn(AUTH_TYPE); + // Setup expired token in cache Token expiredToken = new Token( "expired-token", "Bearer", "cached-refresh-token", LocalDateTime.now().minusMinutes(5)); when(tokenCache.load()).thenReturn(expiredToken); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); + // Setup successful token refresh Token refreshedToken = new Token( @@ -91,14 +107,16 @@ void testRefreshExpiredTokenSuccess() throws IOException { Map headers = headerFactory.headers(); assertEquals("Bearer refreshed-token", headers.get(HttpHeaders.AUTHORIZATION)); - // Verify refresh was attempted but not browser auth + // Verify refresh was attempted verify(provider).refreshAccessToken(); - verify(provider, never()).performBrowserAuth(); verify(tokenCache).save(refreshedToken); } @Test - void testRefreshTokenFailureFallbackToBrowserAuth() throws IOException { + void testRefreshTokenFailureFallbackToBrowserAuth() throws IOException, DatabricksException { + when(config.getHost()).thenReturn(TEST_HOST); + when(config.getAuthType()).thenReturn(AUTH_TYPE); + // Setup expired token in cache Token expiredToken = new Token( @@ -108,15 +126,18 @@ void testRefreshTokenFailureFallbackToBrowserAuth() throws IOException { LocalDateTime.now().minusMinutes(5)); when(tokenCache.load()).thenReturn(expiredToken); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); + // Setup failed token refresh - doThrow(new IOException("Invalid refresh token")).when(provider).refreshAccessToken(); + doThrow(new DatabricksException("Invalid refresh token")).when(provider).refreshAccessToken(); // Setup successful browser auth Token newToken = new Token("new-token", "Bearer", "new-refresh-token", LocalDateTime.now().plusHours(1)); doReturn(newToken).when(provider).performBrowserAuth(); - // Should fall back to browser auth after refresh fails + // Configure and get the token to trigger refresh HeaderFactory headerFactory = provider.configure(config); assertNotNull(headerFactory); @@ -130,25 +151,51 @@ void testRefreshTokenFailureFallbackToBrowserAuth() throws IOException { } @Test - void testEmptyCacheFallbackToBrowserAuth() throws IOException { - // Setup empty cache + void testEmptyCacheFallbackToBrowserAuth() throws IOException, DatabricksException { + when(config.getHost()).thenReturn(TEST_HOST); + when(config.getAuthType()).thenReturn(AUTH_TYPE); + + // Setup empty cache with null token + // RefreshableTokenSource initialization will create empty token if null is returned when(tokenCache.load()).thenReturn(null); + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); // Setup successful browser auth Token newToken = new Token("new-token", "Bearer", "new-refresh-token", LocalDateTime.now().plusHours(1)); doReturn(newToken).when(provider).performBrowserAuth(); - // Should directly use browser auth + // Configure and get the token to trigger refresh HeaderFactory headerFactory = provider.configure(config); assertNotNull(headerFactory); Map headers = headerFactory.headers(); assertEquals("Bearer new-token", headers.get(HttpHeaders.AUTHORIZATION)); - // Verify only browser auth was attempted - verify(provider, never()).refreshAccessToken(); + // Verify only browser auth was attempted (refresh will not be attempted with null refresh + // token) verify(provider).performBrowserAuth(); verify(tokenCache).save(newToken); } + + @Test + void testRefreshAccessToken() throws DatabricksException, IOException { + // Create a token with a refresh token + Token token = + new Token( + DatabricksJdbcConstants.EMPTY_STRING, + DatabricksJdbcConstants.EMPTY_STRING, + "test-refresh-token", + LocalDateTime.now().minusMinutes(1)); + when(tokenCache.load()).thenReturn(token); + + provider = + spy(new CachingExternalBrowserCredentialsProvider(config, connectionContext, tokenCache)); + + // Test that refreshAccessToken is called by refresh() + doThrow(new DatabricksException("Test exception")).when(provider).refreshAccessToken(); + + assertThrows(DatabricksException.class, () -> provider.refresh()); + } } diff --git a/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java index 2b2faaa73c..61c5e8cd84 100644 --- a/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java +++ b/src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.*; +import com.databricks.jdbc.common.util.StringUtil; import com.databricks.sdk.core.oauth.Token; import java.io.IOException; import java.nio.file.Files; @@ -18,14 +19,15 @@ public class TokenCacheTest { private TokenCache tokenCache; @BeforeEach - void setUp() { - tokenCache = new TokenCache(TEST_PASSPHRASE); - String sanitizedUsername = System.getProperty("user.name").replaceAll("[^a-zA-Z0-9_]", "_"); + void setUp() throws IOException { + String sanitizedUsername = StringUtil.sanitizeUsernameForFile(System.getProperty("user.name")); cacheFile = Paths.get( System.getProperty("java.io.tmpdir"), ".databricks", - sanitizedUsername + ".databricks_jdbc_token_cache"); + sanitizedUsername + ".databricks_jdbc_token_cache_test"); + tokenCache = new TokenCache(cacheFile, TEST_PASSPHRASE); + Files.deleteIfExists(cacheFile); } @AfterEach diff --git a/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java index b8592198fe..435934dab4 100644 --- a/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java @@ -73,4 +73,44 @@ public void testEscapeStringLiteral() { String expected = "''1'';select * from other-table"; assertEquals(expected, StringUtil.escapeStringLiteral(sqlValue)); } + + @Test + public void testNullUsername() { + assertEquals("unknown_user", StringUtil.sanitizeUsernameForFile(null)); + } + + @Test + public void testEmptyUsername() { + assertEquals("", StringUtil.sanitizeUsernameForFile("")); + } + + @Test + public void testAlreadySanitizedUsername() { + assertEquals("john_doe123", StringUtil.sanitizeUsernameForFile("john_doe123")); + } + + @Test + public void testUsernameWithSpaces() { + assertEquals("john_doe", StringUtil.sanitizeUsernameForFile("john doe")); + } + + @Test + public void testUsernameWithSpecialCharacters() { + assertEquals("john_doe_123", StringUtil.sanitizeUsernameForFile("john.doe@123")); + } + + @Test + public void testUsernameWithMixedCharacters() { + assertEquals("John_Doe_123_user", StringUtil.sanitizeUsernameForFile("John-Doe#123!user")); + } + + @Test + public void testUsernameWithUnicodeCharacters() { + assertEquals("_ser_123", StringUtil.sanitizeUsernameForFile("üser★123")); + } + + @Test + public void testUsernameWithAllInvalidCharacters() { + assertEquals("______", StringUtil.sanitizeUsernameForFile("!@#$%^")); + } } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 9674e1dfed..dfe6d2033a 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -13,6 +13,7 @@ import com.databricks.jdbc.common.AuthFlow; import com.databricks.jdbc.common.AuthMech; import com.databricks.jdbc.common.DatabricksJdbcConstants; +import com.databricks.jdbc.common.util.DatabricksAuthUtil; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.sdk.WorkspaceClient; @@ -21,6 +22,7 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.ProxyConfig; import com.databricks.sdk.core.commons.CommonsHttpClient; +import com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider; import com.databricks.sdk.core.utils.Cloud; import java.io.IOException; import java.util.List; @@ -29,6 +31,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) @@ -180,7 +183,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals("browser-client-secret", config.getClientSecret()); assertEquals(List.of(new String[] {"scope1", "scope2"}), config.getScopes()); assertEquals(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL, config.getOAuthRedirectUrl()); - assertEquals("external-browser", config.getAuthType()); + assertEquals(DatabricksJdbcConstants.U2M_AUTH_TYPE, config.getAuthType()); } @Test @@ -196,8 +199,6 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); when(mockContext.getOAuthDiscoveryURL()).thenReturn(TEST_DISCOVERY_URL); when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); - when(mockContext.getTokenCachePassPhrase()).thenReturn("token-cache-passphrase"); - when(mockContext.isTokenCacheEnabled()).thenReturn(true); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -209,30 +210,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect assertEquals("browser-client-secret", config.getClientSecret()); assertEquals(List.of(new String[] {"scope1", "scope2"}), config.getScopes()); assertEquals(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL, config.getOAuthRedirectUrl()); - assertEquals("external-browser-with-cache", config.getAuthType()); - } - - @Test - void getWorkspaceClient_OAuthWithCachingExternalBrowser_AuthenticatesCorrectly() - throws DatabricksParsingException { - when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); - when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); - when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); - when(mockContext.getClientId()).thenReturn("client-id"); - when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); - when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); - when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); - when(mockContext.isTokenCacheEnabled()).thenReturn(true); - configurator = new ClientConfigurator(mockContext); - - WorkspaceClient client = configurator.getWorkspaceClient(); - assertNotNull(client); - DatabricksConfig config = client.config(); - - assertEquals("https://oauth-browser.databricks.com", config.getHost()); - assertEquals("client-id", config.getClientId()); - assertInstanceOf( - CachingExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); + assertEquals(DatabricksJdbcConstants.U2M_AUTH_TYPE, config.getAuthType()); } @Test @@ -250,22 +228,27 @@ void getWorkspaceClient_OAuthWithCachingExternalBrowser_NoPassphrase_ThrowsExcep @Test void getWorkspaceClient_OAuthBrowserAuth_WithTokenCacheEnabled_UsesCachedProvider() throws DatabricksParsingException { - when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); - when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); - when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); - when(mockContext.getClientId()).thenReturn("client-id"); - when(mockContext.getClientSecret()).thenReturn("client-secret"); - when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); - when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); - when(mockContext.isTokenCacheEnabled()).thenReturn(true); - - configurator = new ClientConfigurator(mockContext); - WorkspaceClient client = configurator.getWorkspaceClient(); - DatabricksConfig config = client.config(); - - assertEquals("external-browser-with-cache", config.getAuthType()); - assertInstanceOf( - CachingExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); + try (MockedStatic mockedStatic = mockStatic(DatabricksAuthUtil.class)) { + when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH); + when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION); + when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com"); + when(mockContext.getClientId()).thenReturn("client-id"); + when(mockContext.getClientSecret()).thenReturn("client-secret"); + when(mockContext.getTokenCachePassPhrase()).thenReturn("test-passphrase"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); + when(mockContext.isTokenCacheEnabled()).thenReturn(true); + mockedStatic + .when(() -> DatabricksAuthUtil.getTokenEndpoint(any(), any())) + .thenReturn("https://oauth-browser.databricks.com"); + + configurator = new ClientConfigurator(mockContext); + WorkspaceClient client = configurator.getWorkspaceClient(); + DatabricksConfig config = client.config(); + + assertEquals("external-browser-with-cache", config.getAuthType()); + assertInstanceOf( + CachingExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); + } } @Test @@ -284,9 +267,7 @@ void getWorkspaceClient_OAuthBrowserAuth_WithTokenCacheDisabled_UsesStandardProv DatabricksConfig config = client.config(); assertEquals("external-browser", config.getAuthType()); - assertInstanceOf( - com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider.class, - config.getCredentialsProvider()); + assertInstanceOf(ExternalBrowserCredentialsProvider.class, config.getCredentialsProvider()); } @Test From 6f188b0fd3a75773bc65eb5ebe1e29b6172b7d1c Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 2 Apr 2025 21:20:38 +0530 Subject: [PATCH 9/9] remove unnecessary test --- .../java/com/databricks/jdbc/common/util/StringUtilTest.java | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java index 435934dab4..fb9b124835 100644 --- a/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/StringUtilTest.java @@ -104,11 +104,6 @@ public void testUsernameWithMixedCharacters() { assertEquals("John_Doe_123_user", StringUtil.sanitizeUsernameForFile("John-Doe#123!user")); } - @Test - public void testUsernameWithUnicodeCharacters() { - assertEquals("_ser_123", StringUtil.sanitizeUsernameForFile("üser★123")); - } - @Test public void testUsernameWithAllInvalidCharacters() { assertEquals("______", StringUtil.sanitizeUsernameForFile("!@#$%^"));