[FEATURE] Adding token cache to u2m oauth #429

vikrantpuppala · 2025-04-10T13:03:35Z

What changes are proposed in this pull request?

What: Introduce a token cache to be used in browser based u2m oauth
Why: Makes user experience better by using valid access/refresh tokens from the cache instead of requiring a new browser auth

How is this tested?

In addition to the unit tests,
Manual testing:

ran example program without the cache, performed browser action and saw that a cache was created
reconnected the program and program worked without the need for browser action
changed the expiry in cache to past time and observed that token was refreshed
changed the refresh token in cache to be garbage, saw refresh failing and performed browser action

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

gopalldb · 2025-04-10T14:27:36Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

+   * @return The OAuth redirect URL to use
+   */
+  public String getEffectiveOAuthRedirectUrl() {
+    return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback";


this is hardcoded at 2 places, making this error prone for future

i made this change to remove the hardcoded value from 2 places. where else do you mean is this hardcoded?

Can we separate unrelated refactors such as this into separate PRs?

I actually added code in ExternalBrowserCredentialsProvider where we would have to define this string "http://localhost:8080/callback" again without the refactor. Thought it would be in scope for this change?

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

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

mgyucht

Overall this seems like a nice contribution, left a few bits of feedback. Do we need to support disabling the cache & encrypting it? We don't have either of those features in the Python SDK and haven't heard any feedback/issues about this.

NEXT_CHANGELOG.md

databricks-sdk-java/pom.xml

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

mgyucht · 2025-04-11T08:28:53Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

+   * @return The OAuth redirect URL to use
+   */
+  public String getEffectiveOAuthRedirectUrl() {
+    return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback";


Can we separate unrelated refactors such as this into separate PRs?

mgyucht · 2025-04-11T08:30:49Z

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

  public HeaderFactory configure(DatabricksConfig config) {
-    if (config.getHost() == null || config.getAuthType() != "external-browser") {
+    if (config.getHost() == null
+        || config.getClientId() == null


Why is this needed? Are customers explicitly requesting external-browser but also configuring client IDs and that's causing issues?

Based on the testing I did, we do need the client ID for external browser auth?

Take a look at how it is done in the Python SDK: https://github.com/databricks/databricks-sdk-py/blob/main/databricks/sdk/credentials_provider.py#L213-L220. Let's preserve that behavior, defaulting to databricks-cli if not set.

mgyucht · 2025-04-11T08:47:44Z

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java

+  }
+
+  @Test
+  void testDatabricksConfigDisabledTokenCache() throws IOException {


This test belongs in DatabricksConfigTest, and we only need to test that the resulting token cache has caching disabled. (i.e. we can remove from line 198 onwards.)

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenCacheTest.java

...java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java

mgyucht · 2025-04-17T10:16:27Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

+   *
+   * @return A TokenCache instance
+   */
+  public synchronized TokenCache getTokenCache() {


Just talking with @renaudhartert-db. Is it possible to inline this token cache entirely into ExternalBrowserCredentialsProvider? We're trying to break apart DatabricksConfig and stop using it as a dependency injection mechanism. I think you could do this by

accepting the TokenCache as a constructor parameter in ExternalBrowserCredentialsProvider, falling back to this implementation if not specified

Adding the token cache to SessionCredentials via its builder.

mgyucht

One suggestion, otherwise this looks good to me!

mgyucht · 2025-04-17T13:34:58Z

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

+  /**
+   * Resolves both client ID and client secret from the configuration.
+   *
+   * @param config The Databricks configuration
+   * @return An array containing the client ID and client secret (may be null)
+   */
+  public static String[] resolveClientCredentials(DatabricksConfig config) {
+    String clientId = resolveClientId(config);
+    String clientSecret = resolveClientSecret(config);
+    return new String[] {clientId, clientSecret};
+  }


Is this really necessary? I think callers can just call resolveClientId() and resolveClientSecret() as needed.

…n-cache

github-actions · 2025-04-17T14:56:22Z

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-java

Inputs:

PR number: 429
Commit SHA: ef3c0001be029b3ce58c253cda8c6c0c1c3d3055

Checks will be approved automatically on success.

vikrantpuppala added 2 commits April 10, 2025 14:44

Adding token cache to u2m oauth

eb0fad0

add a way to enable/disable token cache

e8080b3

vikrantpuppala temporarily deployed to test-trigger-is April 10, 2025 13:03 — with GitHub Actions Inactive

gopalldb reviewed Apr 10, 2025

View reviewed changes

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java Outdated Show resolved Hide resolved

gopalldb reviewed Apr 10, 2025

View reviewed changes

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java Outdated Show resolved Hide resolved

gopalldb reviewed Apr 10, 2025

View reviewed changes

...sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java Outdated Show resolved Hide resolved

fmt

b57f6e2

vikrantpuppala temporarily deployed to test-trigger-is April 10, 2025 16:01 — with GitHub Actions Inactive

add changelog

6f1b69d

vikrantpuppala temporarily deployed to test-trigger-is April 10, 2025 16:04 — with GitHub Actions Inactive

vikrantpuppala mentioned this pull request Apr 11, 2025

[PECOBLR-131][PECOBLR-180] Add token cache for U2M OAuth by using sdk databricks/databricks-jdbc#783

Merged

mgyucht reviewed Apr 11, 2025

View reviewed changes

update implementation based on discussion

d8c366e

vikrantpuppala temporarily deployed to test-trigger-is April 11, 2025 13:01 — with GitHub Actions Inactive

fmt

3dad770

vikrantpuppala temporarily deployed to test-trigger-is April 11, 2025 13:14 — with GitHub Actions Inactive

fix

c9fc008

vikrantpuppala temporarily deployed to test-trigger-is April 11, 2025 13:41 — with GitHub Actions Inactive

mgyucht reviewed Apr 17, 2025

View reviewed changes

address comments

c4433de

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 10:51 — with GitHub Actions Inactive

vikrantpuppala added 2 commits April 17, 2025 16:27

remove unnecessary test

9ce45ef

use databricks-cli

174c3c2

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 11:07 — with GitHub Actions Inactive

update logic used to get client id

ab64d0f

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 12:15 — with GitHub Actions Inactive

fmt

23b5bd1

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 12:23 — with GitHub Actions Inactive

mgyucht approved these changes Apr 17, 2025

View reviewed changes

remove extra function

ce75413

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 14:55 — with GitHub Actions Inactive

Merge remote-tracking branch 'refs/remotes/databricks/main' into toke…

ef3c000

…n-cache

vikrantpuppala temporarily deployed to test-trigger-is April 17, 2025 14:56 — with GitHub Actions Inactive

mgyucht added this pull request to the merge queue Apr 17, 2025

Merged via the queue into databricks:main with commit 55cc9ed Apr 17, 2025
15 checks passed

[FEATURE] Adding token cache to u2m oauth #429

[FEATURE] Adding token cache to u2m oauth #429

Uh oh!

Conversation

vikrantpuppala commented Apr 10, 2025 • edited by mgyucht Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes are proposed in this pull request?

How is this tested?

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vikrantpuppala Apr 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Uh oh!

Uh oh!

mgyucht left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mgyucht left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Apr 17, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vikrantpuppala commented Apr 10, 2025 •

edited by mgyucht

Loading

vikrantpuppala Apr 10, 2025 •

edited

Loading