[FEATURE] Adding token cache to u2m oauth

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.44.0
 
 ### New Features and Improvements
+Added `TokenCache` to `ExternalBrowserCredentialsProvider` to reduce number of authentications needed for U2M OAuth 
 
 ### Bug Fixes
 

databricks-sdk-java/pom.xml

@@ -97,5 +97,10 @@
       <artifactId>google-auth-library-oauth2-http</artifactId>
       <version>1.20.0</version>
     </dependency>
+    <dependency>
+      <groupId>com.fasterxml.jackson.datatype</groupId>
+      <artifactId>jackson-datatype-jsr310</artifactId>
+      <version>${jackson.version}</version>
+    </dependency>
   </dependencies>
 </project>

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

@@ -5,6 +5,7 @@
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
 import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.core.oauth.TokenCache;
 import com.databricks.sdk.core.utils.Cloud;
 import com.databricks.sdk.core.utils.Environment;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -38,6 +39,23 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_REDIRECT_URL", auth = "oauth")
   private String redirectUrl;
 
+  /**
+   * The passphrase used to encrypt the OAuth token cache. This is optional and only used with token
+   * caching.
+   */
+  @ConfigAttribute(
+      env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE",
+      auth = "oauth",
+      sensitive = true)
+  private String oAuthTokenCachePassphrase;
+
+  /**
+   * Controls whether OAuth token caching is enabled. When set to false, tokens will not be cached
+   * or loaded from cache.
+   */
+  @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_ENABLED", auth = "oauth")
+  private Boolean isTokenCacheEnabled;
+
   /**
    * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints.
    *
@@ -141,6 +159,9 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  // Lazily initialized OAuth token cache
+  private transient TokenCache tokenCache;
+
   public Environment getEnv() {
     return env;
   }
@@ -669,4 +690,62 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
                 "headerFactory"));
     return clone(fieldsToSkip).setHost(host);
   }
+
+  public String getOAuthTokenCachePassphrase() {
+    return oAuthTokenCachePassphrase;
+  }
+
+  public DatabricksConfig setOAuthTokenCachePassphrase(String oAuthPassphrase) {
+    this.oAuthTokenCachePassphrase = oAuthPassphrase;
+    return this;
+  }
+
+  /**
+   * Gets whether OAuth token caching is enabled. Default is true.
+   *
+   * @return true if token caching is enabled, false otherwise
+   */
+  public boolean isTokenCacheEnabled() {
+    return isTokenCacheEnabled == null || isTokenCacheEnabled;
+  }
+
+  /**
+   * Sets whether OAuth token caching is enabled.
+   *
+   * @param enabled true to enable token caching, false to disable
+   * @return this config instance
+   */
+  public DatabricksConfig setTokenCacheEnabled(boolean enabled) {
+    this.isTokenCacheEnabled = enabled;
+    return this;
+  }
+
+  /**
+   * Gets the default OAuth redirect URL. If one is not provided explicitly, uses
+   * http://localhost:8080/callback
+   *
+   * @return The OAuth redirect URL to use
+   */
+  public String getEffectiveOAuthRedirectUrl() {
+    return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback";
+  }
+
+  /**
+   * Gets the TokenCache instance for the current configuration. Creates it if it doesn't exist yet.
+   * When token caching is disabled, the TokenCache will be created but operations will be no-ops.
+   *
+   * @return A TokenCache instance for the current host, client ID, and scopes
+   */
+  public synchronized TokenCache getTokenCache() {
+    if (tokenCache == null) {
+      tokenCache =
+          new TokenCache(
+              getHost(),
+              getClientId(),
+              getScopes(),
+              getOAuthTokenCachePassphrase(),
+              isTokenCacheEnabled());
+    }
+    return tokenCache;
+  }
 }

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

@@ -5,12 +5,19 @@
 import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.HeaderFactory;
 import java.io.IOException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a
- * browser for the user to authorize the application.
+ * browser for the user to authorize the application. When cache support is enabled with {@link
+ * DatabricksConfig#setOAuthTokenCachePassphrase} and {@link
+ * DatabricksConfig#setTokenCacheEnabled(boolean)}, tokens will be cached to avoid repeated
+ * authentication.
  */
 public class ExternalBrowserCredentialsProvider implements CredentialsProvider {
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class);
 
   @Override
   public String authType() {
@@ -19,16 +26,57 @@ public String authType() {
 
   @Override
   public HeaderFactory configure(DatabricksConfig config) {
-    if (config.getHost() == null || config.getAuthType() != "external-browser") {
+    if (config.getHost() == null
+        || config.getClientId() == null
+        || !config.getAuthType().equals("external-browser")) {
       return null;
     }
+
     try {
-      OAuthClient client = new OAuthClient(config);
-      Consent consent = client.initiateConsent();
-      SessionCredentials creds = consent.launchExternalBrowser();
-      return creds.configure(config);
+      // Get the token cache from config
+      TokenCache tokenCache = config.getTokenCache();
+
+      // First try to use the cached token if available (will return null if disabled)
+      Token cachedToken = tokenCache.load();
+      if (cachedToken != null && cachedToken.getRefreshToken() != null) {
+        LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId());
+
+        try {
+          // Create SessionCredentials with the cached token and try to refresh if needed
+          SessionCredentials cachedCreds =
+              new SessionCredentials.Builder()
+                  .withToken(cachedToken)
+                  .withHttpClient(config.getHttpClient())
+                  .withClientId(config.getClientId())
+                  .withClientSecret(config.getClientSecret())
+                  .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
+                  .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
+                  .build();
+
+          LOGGER.debug("Using cached token, will immediately refresh");
+          cachedCreds.token = cachedCreds.refresh();
+          tokenCache.save(cachedToken);
+          return cachedCreds.configure(config);
+        } catch (Exception e) {
+          // If token refresh fails, log and continue to browser auth
+          LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage());
+        }
+      }
+
+      // If no cached token or refresh failed, perform browser auth
+      SessionCredentials credentials = performBrowserAuth(config);
+      tokenCache.save(credentials.getToken());
+      return credentials.configure(config);
     } catch (IOException | DatabricksException e) {
+      LOGGER.error("Failed to authenticate: {}", e.getMessage());
       return null;
     }
   }
+
+  SessionCredentials performBrowserAuth(DatabricksConfig config) throws IOException {
+    LOGGER.debug("Performing browser authentication");
+    OAuthClient client = new OAuthClient(config);
+    Consent consent = client.initiateConsent();
+    return consent.launchExternalBrowser();
+  }
 }

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

@@ -92,10 +92,7 @@ public OAuthClient(DatabricksConfig config) throws IOException {
             .withClientId(config.getClientId())
             .withClientSecret(config.getClientSecret())
             .withHost(config.getHost())
-            .withRedirectUrl(
-                config.getOAuthRedirectUrl() != null
-                    ? config.getOAuthRedirectUrl()
-                    : "http://localhost:8080/callback")
+            .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
             .withScopes(config.getScopes()));
   }
 

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

@@ -9,6 +9,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.http.HttpHeaders;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * An implementation of RefreshableTokenSource implementing the refresh_token OAuth grant type.
@@ -20,6 +22,7 @@
 public class SessionCredentials extends RefreshableTokenSource
     implements CredentialsProvider, Serializable {
   private static final long serialVersionUID = 3083941540130596650L;
+  private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentials.class);
 
   @Override
   public String authType() {
@@ -28,6 +31,8 @@ public String authType() {
 
   @Override
   public HeaderFactory configure(DatabricksConfig config) {
+    this.tokenCache = config.getTokenCache();
+
     return () -> {
       Map<String, String> headers = new HashMap<>();
       headers.put(
@@ -84,6 +89,7 @@ public SessionCredentials build() {
   private final String redirectUrl;
   private final String clientId;
   private final String clientSecret;
+  private transient TokenCache tokenCache;
 
   private SessionCredentials(Builder b) {
     super(b.token);
@@ -113,7 +119,20 @@ protected Token refresh() {
       // cross-origin requests
       headers.put("Origin", redirectUrl);
     }
-    return retrieveToken(
-        hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY);
+    Token newToken =
+        retrieveToken(
+            hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY);
+
+    // Save the refreshed token directly to cache
+    if (tokenCache != null) {
+      try {
+        tokenCache.save(newToken);
+        LOGGER.debug("Saved refreshed token to cache");
+      } catch (Exception e) {
+        LOGGER.warn("Failed to save token to cache: {}", e.getMessage());
+      }
+    }
+
+    return newToken;
   }
 }

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

@@ -2,6 +2,7 @@
 
 import com.databricks.sdk.core.utils.ClockSupplier;
 import com.databricks.sdk.core.utils.SystemClockSupplier;
+import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import java.time.LocalDateTime;
 import java.time.temporal.ChronoUnit;
@@ -37,7 +38,12 @@ public Token(
   }
 
   /** Constructor for refreshable tokens. */
-  public Token(String accessToken, String tokenType, String refreshToken, LocalDateTime expiry) {
+  @JsonCreator
+  public Token(
+      @JsonProperty("accessToken") String accessToken,
+      @JsonProperty("tokenType") String tokenType,
+      @JsonProperty("refreshToken") String refreshToken,
+      @JsonProperty("expiry") LocalDateTime expiry) {
     this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier());
   }