[FEATURE] Adding token cache to u2m oauth

NEXT_CHANGELOG.md

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

databricks-sdk-java/pom.xml

@@ -97,5 +97,11 @@
       <artifactId>google-auth-library-oauth2-http</artifactId>
       <version>1.20.0</version>
     </dependency>
+    <!-- Jackson JSR310 module needed to serialize/deserialize java.time classes in TokenCache -->
+    <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

@@ -4,13 +4,17 @@
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
+import com.databricks.sdk.core.oauth.FileTokenCache;
 import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.core.oauth.TokenCache;
+import com.databricks.sdk.core.oauth.TokenCacheUtils;
 import com.databricks.sdk.core.utils.Cloud;
 import com.databricks.sdk.core.utils.Environment;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.nio.file.Path;
 import java.util.*;
 import org.apache.http.HttpMessage;
 
@@ -141,6 +145,9 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  // Lazily initialized OAuth token cache
+  private transient TokenCache tokenCache;
+
   public Environment getEnv() {
     return env;
   }
@@ -669,4 +676,41 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
                 "headerFactory"));
     return clone(fieldsToSkip).setHost(host);
   }
+
+  /**
+   * Sets a custom TokenCache implementation.
+   *
+   * @param tokenCache the TokenCache implementation to use
+   * @return this config instance
+   */
+  public DatabricksConfig setTokenCache(TokenCache tokenCache) {
+    this.tokenCache = tokenCache;
+    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.
+   *
+   * <p>If a custom TokenCache has been set, it will be returned. Otherwise, a SimpleFileTokenCache
+   * will be created based on the configuration properties.
+   *
+   * @return A TokenCache instance
+   */
+  public synchronized TokenCache getTokenCache() {
+    if (tokenCache == null) {
+      Path cachePath = TokenCacheUtils.getCacheFilePath(getHost(), getClientId(), getScopes());
+      tokenCache = new FileTokenCache(cachePath);
+    }
+    return tokenCache;
+  }
 }

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

@@ -5,12 +5,17 @@
 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. Uses the cache from {@link
+ * DatabricksConfig#getTokenCache()}, 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 +24,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(cachedCreds.getToken());
+          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/FileTokenCache.java

@@ -0,0 +1,77 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.utils.SerDeUtils;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.File;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Objects;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** A TokenCache implementation that stores tokens as plain files. */
+public class FileTokenCache implements TokenCache {
+  private static final Logger LOGGER = LoggerFactory.getLogger(FileTokenCache.class);
+
+  private final Path cacheFile;
+  private final ObjectMapper mapper;
+
+  /**
+   * Constructs a new SimpleFileTokenCache instance.
+   *
+   * @param cacheFilePath The path where the token cache will be stored
+   */
+  public FileTokenCache(Path cacheFilePath) {
+    Objects.requireNonNull(cacheFilePath, "cacheFilePath must be defined");
+
+    this.cacheFile = cacheFilePath;
+    this.mapper = SerDeUtils.createMapper();
+  }
+
+  @Override
+  public void save(Token token) {
+    try {
+      Files.createDirectories(cacheFile.getParent());
+
+      // Serialize token to JSON
+      String json = mapper.writeValueAsString(token);
+      byte[] dataToWrite = json.getBytes(StandardCharsets.UTF_8);
+
+      Files.write(cacheFile, dataToWrite);
+      // Set file permissions to be readable only by the owner (equivalent to 0600)
+      File file = cacheFile.toFile();
+      file.setReadable(false, false);
+      file.setReadable(true, true);
+      file.setWritable(false, false);
+      file.setWritable(true, true);
+
+      LOGGER.debug("Successfully saved token to cache: {}", cacheFile);
+    } catch (Exception e) {
+      LOGGER.warn("Failed to save token to cache: {}", cacheFile, e);
+    }
+  }
+
+  @Override
+  public Token load() {
+    try {
+      if (!Files.exists(cacheFile)) {
+        LOGGER.debug("No token cache file found at: {}", cacheFile);
+        return null;
+      }
+
+      byte[] fileContent = Files.readAllBytes(cacheFile);
+
+      // Deserialize token from JSON
+      String json = new String(fileContent, StandardCharsets.UTF_8);
+      Token token = mapper.readValue(json, Token.class);
+      LOGGER.debug("Successfully loaded token from cache: {}", cacheFile);
+      return token;
+    } catch (Exception e) {
+      // If there's any issue loading the token, return null
+      // to allow a fresh token to be obtained
+      LOGGER.warn("Failed to load token from cache: {}", e.getMessage());
+      return null;
+    }
+  }
+}

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,15 @@ 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) {
+      tokenCache.save(newToken);
+      LOGGER.debug("Saved refreshed token to cache");
+    }
+    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());
   }
 

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

@@ -0,0 +1,21 @@
+package com.databricks.sdk.core.oauth;
+
+/**
+ * TokenCache interface for storing and retrieving OAuth tokens. Implementations can use different
+ * storage mechanisms and security approaches.
+ */
+public interface TokenCache {
+  /**
+   * Saves a Token to the cache.
+   *
+   * @param token The Token to save
+   */
+  void save(Token token);
+
+  /**
+   * Loads a Token from the cache.
+   *
+   * @return The Token from the cache or null if the cache doesn't exist or is invalid
+   */
+  Token load();
+}