Adding token cache to u2m oauth

Author: vikrantpuppala

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,8 @@
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
 import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.core.oauth.Token;
+import com.databricks.sdk.core.oauth.TokenCache;
 import com.databricks.sdk.core.utils.Cloud;
 import com.databricks.sdk.core.utils.Environment;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -38,6 +40,13 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_REDIRECT_URL", auth = "oauth")
   private String redirectUrl;
 
+  /**
+   * The passphrase used to encrypt the OAuth token cache.
+   * This is optional and only used with token caching.
+   */
+  @ConfigAttribute(env = "DATABRICKS_OAUTH_TOKEN_CACHE_PASSPHRASE", auth = "oauth", sensitive = true)
+  private String oAuthTokenCachePassphrase;
+
   /**
    * The OpenID Connect discovery URL used to retrieve OIDC configuration and endpoints.
    *
@@ -141,6 +150,9 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  // Lazily initialized OAuth token cache
+  private transient TokenCache tokenCache;
+
   public Environment getEnv() {
     return env;
   }
@@ -669,4 +681,63 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
                 "headerFactory"));
     return clone(fieldsToSkip).setHost(host);
   }
+
+  public String getOAuthPassphrase() {
+    return oAuthTokenCachePassphrase;
+  }
+
+  public DatabricksConfig setOAuthPassphrase(String oAuthPassphrase) {
+    this.oAuthTokenCachePassphrase = oAuthPassphrase;
+    return this;
+  }
+
+  /**
+   * Gets the default OAuth redirect URL.
+   * If one is not provided explicitly, uses http://localhost:8080/callback
+   * 
+   * @return The OAuth redirect URL to use
+   */
+  public String getEffectiveOAuthRedirectUrl() {
+    return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback";
+  }
+  
+  /**
+   * Gets the TokenCache instance for the current configuration.
+   * Creates it if it doesn't exist yet.
+   * 
+   * @return A TokenCache instance for the current host, client ID, and scopes
+   */
+  public synchronized TokenCache getTokenCache() {
+    if (tokenCache == null && getHost() != null && getClientId() != null) {
+      tokenCache = new TokenCache(getHost(), getClientId(), getScopes(), getOAuthPassphrase());
+    }
+    return tokenCache;
+  }
+  
+  /**
+   * Loads a cached token if one exists for the current configuration.
+   * 
+   * @return The cached token, or null if no valid token is cached
+   */
+  public Token loadCachedToken() {
+    TokenCache cache = getTokenCache();
+    if (cache == null) {
+      return null;
+    }
+    return cache.load();
+  }
+  
+  /**
+   * Saves a token to the cache for the current configuration.
+   * 
+   * @param token The token to cache
+   * @throws IOException If there is an error saving the token
+   */
+  public void saveTokenToCache(Token token) throws IOException {
+    TokenCache cache = getTokenCache();
+    if (cache == null) {
+      return;
+    }
+    cache.save(token);
+  }
 }

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

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

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,19 @@ protected Token refresh() {
       // cross-origin requests
       headers.put("Origin", redirectUrl);
     }
-    return retrieveToken(
+    Token newToken = retrieveToken(
         hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY);
+
+    // Save the refreshed token directly to cache
+    if (tokenCache != null) {
+      try {
+        tokenCache.save(newToken);
+        LOGGER.debug("Saved refreshed token to cache");
+      } catch (Exception e) {
+        LOGGER.warn("Failed to save token to cache: {}", e.getMessage());
+      }
+    }
+    
+    return newToken;
   }
 }

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());
   }