[PECOBLR-131][PECOBLR-180] Add token cache for U2M OAuth #768

vikrantpuppala · 2025-04-01T06:48:05Z

Description

Adds support for token cache to the JDBC driver, specifically adding two parameters: EnableTokenCache and TokenCachePassPhrase

Testing

Manual and unit tests

Additional Notes to the Reviewer

Copilot

Pull Request Overview

This PR adds token cache support for U2M OAuth in the JDBC driver. The changes include:

Introducing a new TokenCache implementation with encryption to persist tokens.
Adding comprehensive tests for token caching behavior and the caching credentials provider.
Updating configuration classes and URL parameters to support the new TokenCachePassPhrase parameter.

Reviewed Changes

Copilot reviewed 9 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java	Added tests verifying behavior with a valid passphrase and error handling when missing.
src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java	New tests for saving, loading, and error conditions in token caching.
src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java	New tests to verify both the token refresh and browser authentication fallback logic.
src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java	Modified to use the caching credentials provider for U2M OAuth.
src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java	Added a new URL parameter for TokenCachePassPhrase.
src/main/java/com/databricks/jdbc/auth/TokenCache.java	New implementation for encrypted token caching.
src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java	New credentials provider implementation with token caching and refresh support.
src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java and src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java	Added a getter for the token cache passphrase.

Files not reviewed (1)

pom.xml: Language not supported

src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java

# Conflicts: # src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java # src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java # src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java # src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java

…n-cache # Conflicts: # src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java # src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java # src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java

Copilot

Pull Request Overview

This PR introduces support for token caching in the JDBC driver’s OAuth flow, enabling caching of tokens through the addition of two new URL parameters and corresponding tests. Key changes include:

Adding new URL parameters (EnableTokenCache and TokenCachePassPhrase) in the connection context.
Implementing a TokenCache class that encrypts and persists OAuth tokens.
Extending authentication flows in the ClientConfigurator and CachingExternalBrowserCredentialsProvider to leverage cached tokens.

Reviewed Changes

Copilot reviewed 10 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java	Updated tests to expect "external-browser-with-cache" auth type and a token passphrase setting
src/test/java/com/databricks/jdbc/auth/TokenCacheTest.java	Added tests for saving, loading, and validating token cache behavior
src/test/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProviderTest.java	Added tests confirming correct provider behavior with valid, expired, and missing tokens
src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java	Added tests to verify token cache enablement based on URL parameters
src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java	Updated to use the caching provider based on token cache enablement
src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java	Introduced TOKEN_CACHE_PASS_PHRASE and ENABLE_TOKEN_CACHE parameters
src/main/java/com/databricks/jdbc/auth/TokenCache.java	New implementation to encrypt and cache tokens locally
src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java	New provider implementation which leverages cached tokens
src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java	Implemented accessors for token cache passphrase and enablement flag
src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java	Updated interface to include token cache methods

Files not reviewed (1)

pom.xml: Language not supported

Comments suppressed due to low confidence (2)

src/main/java/com/databricks/jdbc/auth/TokenCache.java:109

Consider specifying an explicit cipher transformation (e.g., 'AES/CBC/PKCS5Padding') instead of relying on default behavior to ensure consistency across environments.

Cipher cipher = Cipher.getInstance(ALGORITHM);

src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java:110

[nitpick] Consider renaming 'TokenCachePassPhrase' to 'TokenCachePassphrase' for consistency with common terminology and with the rest of the codebase.

TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"),

gopalldb · 2025-04-09T08:20:28Z

src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java

+   * @param context The connection context containing OAuth configuration parameters
+   * @param tokenCache The token cache to use for storing and retrieving tokens
+   */
+  public CachingExternalBrowserCredentialsProvider(


We are reusing lots of code from SDK. Can we do a better structure of class to reuse ExternalBrowserCredentialsProvider from SDK?

ExternalBrowserCredentialsProvider does the hand-off to SessionCredentials internally. I was initially trying to overload those classes but it wouldn't just work

gopalldb · 2025-04-09T08:21:02Z

src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java

+ * <p>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 extends RefreshableTokenSource


Did we check with SDK team if we should add this in SDK? Looks useful for other clients as well.

gopalldb · 2025-04-09T08:22:02Z

src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java

+        try {
+          LOGGER.debug("Using refresh token to get new access token");
+          Token refreshedToken = refreshAccessToken();
+          tokenCache.save(refreshedToken);


So, is this TokenCache the only diff with SDK implementation?

Essentially, yes, + the logic on how/when to use token cache

gopalldb · 2025-04-09T08:24:01Z

src/main/java/com/databricks/jdbc/auth/TokenCache.java

+  @VisibleForTesting
+  public TokenCache(Path cacheFile, String passphrase) {
+    if (passphrase == null || passphrase.isEmpty()) {
+      throw new IllegalArgumentException(


log this error

gopalldb · 2025-04-09T08:24:51Z

src/main/java/com/databricks/jdbc/auth/TokenCache.java

+      byte[] encrypted = encrypt(json.getBytes());
+      Files.write(cacheFile, encrypted);
+    } catch (Exception e) {
+      throw new IOException("Failed to save token cache: " + e.getMessage(), e);


Add logging

gopalldb · 2025-04-09T08:24:58Z

src/main/java/com/databricks/jdbc/auth/TokenCache.java

+      byte[] decrypted = decrypt(encrypted);
+      return mapper.readValue(decrypted, SerializableToken.class);
+    } catch (Exception e) {
+      throw new IOException("Failed to load token cache: " + e.getMessage(), e);


gopalldb · 2025-04-09T08:25:32Z

src/main/java/com/databricks/jdbc/auth/TokenCache.java

+      byte[] decrypted = decrypt(encrypted);
+      return mapper.readValue(decrypted, SerializableToken.class);
+    } catch (Exception e) {
+      throw new IOException("Failed to load token cache: " + e.getMessage(), e);


What happens to IOException? Do we catch it any upstream class and handle it?

vikrantpuppala · 2025-04-23T10:49:22Z

Closing in favour of #783

Add token cache for U2M OAuth

cb32404

vikrantpuppala requested a review from Copilot April 1, 2025 06:48

Copilot AI reviewed Apr 1, 2025

View reviewed changes

src/main/java/com/databricks/jdbc/auth/CachingExternalBrowserCredentialsProvider.java Outdated Show resolved Hide resolved

vikrantpuppala added 2 commits April 1, 2025 12:19

fix

276786f

vikrantpuppala temporarily deployed to azure-prod April 1, 2025 06:51 — with GitHub Actions Inactive

vikrantpuppala added 2 commits April 2, 2025 13:28

fix tests

e1ab9c8

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 08:00 — with GitHub Actions Inactive

add EnableTokenCache

9db852f

vikrantpuppala requested a review from Copilot April 2, 2025 08:08

Copilot AI reviewed Apr 2, 2025

View reviewed changes

vikrantpuppala requested review from gopalldb, jayantsing-db, jprakash-db, madhav-db, samikshya-db and shivam2680 April 2, 2025 08:09

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 08:09 — with GitHub Actions Inactive

add javadoc

63be611

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 08:14 — with GitHub Actions Inactive

fix tests

6c6eabb

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 08:17 — with GitHub Actions Inactive

cache file must use user name instead of host

0f60c10

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 10:18 — with GitHub Actions Inactive

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 10:19 — with GitHub Actions Inactive

fix implementation

e109469

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 15:46 — with GitHub Actions Inactive

remove unnecessary test

6f188b0

vikrantpuppala temporarily deployed to azure-prod April 2, 2025 15:51 — with GitHub Actions Inactive

vikrantpuppala changed the title ~~[PECOBLR-131] Add token cache for U2M OAuth~~ [PECOBLR-131][PECOBLR-180] Add token cache for U2M OAuth Apr 7, 2025

gopalldb reviewed Apr 9, 2025

View reviewed changes

vikrantpuppala closed this Apr 23, 2025

[PECOBLR-131][PECOBLR-180] Add token cache for U2M OAuth #768

[PECOBLR-131][PECOBLR-180] Add token cache for U2M OAuth #768

Uh oh!

Conversation

vikrantpuppala commented Apr 1, 2025

Description

Testing

Additional Notes to the Reviewer

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vikrantpuppala commented Apr 23, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants