[FEATURE] Adding token cache to u2m oauth

NEXT_CHANGELOG.md

@@ -3,7 +3,8 @@
 ## Release v0.46.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

@@ -669,4 +669,14 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
                 "headerFactory"));
     return clone(fieldsToSkip).setHost(host);
   }
+
+  /**
+   * 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";
+  }
 }

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

@@ -5,12 +5,38 @@
 import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.HeaderFactory;
 import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Objects;
+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 a specified TokenCache or creates a
+ * default one if none is provided.
  */
 public class ExternalBrowserCredentialsProvider implements CredentialsProvider {
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class);
+
+  private TokenCache tokenCache;
+
+  /**
+   * Creates a new ExternalBrowserCredentialsProvider with the specified TokenCache.
+   *
+   * @param tokenCache the TokenCache to use for caching tokens
+   */
+  public ExternalBrowserCredentialsProvider(TokenCache tokenCache) {
+    this.tokenCache = tokenCache;
+  }
+
+  /**
+   * Creates a new ExternalBrowserCredentialsProvider with a default TokenCache. A FileTokenCache
+   * will be created when credentials are configured.
+   */
+  public ExternalBrowserCredentialsProvider() {
+    this(null);
+  }
 
   @Override
   public String authType() {
@@ -19,16 +45,87 @@ public String authType() {
 
   @Override
   public HeaderFactory configure(DatabricksConfig config) {
-    if (config.getHost() == null || config.getAuthType() != "external-browser") {
+    if (config.getHost() == null || !Objects.equals(config.getAuthType(), "external-browser")) {
       return null;
     }
+
+    // Use the utility class to resolve client ID and client secret
+    String clientId = OAuthClientUtils.resolveClientId(config);
+    String clientSecret = OAuthClientUtils.resolveClientSecret(config);
+
     try {
-      OAuthClient client = new OAuthClient(config);
-      Consent consent = client.initiateConsent();
-      SessionCredentials creds = consent.launchExternalBrowser();
-      return creds.configure(config);
+      if (tokenCache == null) {
+        // Create a default FileTokenCache based on config
+        Path cachePath =
+            TokenCacheUtils.getCacheFilePath(config.getHost(), clientId, config.getScopes());
+        tokenCache = new FileTokenCache(cachePath);
+      }
+
+      // 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(), clientId);
+
+        try {
+          // Create SessionCredentials with the cached token and try to refresh if needed
+          SessionCredentials cachedCreds =
+              new SessionCredentials.Builder()
+                  .withToken(cachedToken)
+                  .withHttpClient(config.getHttpClient())
+                  .withClientId(clientId)
+                  .withClientSecret(clientSecret)
+                  .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
+                  .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
+                  .withTokenCache(tokenCache)
+                  .build();
+
+          LOGGER.debug("Using cached token, will immediately refresh");
+          cachedCreds.token = cachedCreds.refresh();
+          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, clientId, clientSecret, tokenCache);
+      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, String clientId, String clientSecret, TokenCache tokenCache)
+      throws IOException {
+    LOGGER.debug("Performing browser authentication");
+    OAuthClient client =
+        new OAuthClient.Builder()
+            .withHttpClient(config.getHttpClient())
+            .withClientId(clientId)
+            .withClientSecret(clientSecret)
+            .withHost(config.getHost())
+            .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
+            .withScopes(config.getScopes())
+            .build();
+    Consent consent = client.initiateConsent();
+
+    // Use the existing browser flow to get credentials
+    SessionCredentials credentials = consent.launchExternalBrowser();
+
+    // Create a new SessionCredentials with the same token but with our token cache
+    return new SessionCredentials.Builder()
+        .withToken(credentials.getToken())
+        .withHttpClient(config.getHttpClient())
+        .withClientId(config.getClientId())
+        .withClientSecret(config.getClientSecret())
+        .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
+        .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
+        .withTokenCache(tokenCache)
+        .build();
+  }
 }

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

@@ -85,20 +85,6 @@ public OAuthClient build() throws IOException {
   private final boolean isAws;
   private final boolean isAzure;
 
-  public OAuthClient(DatabricksConfig config) throws IOException {
-    this(
-        new Builder()
-            .withHttpClient(config.getHttpClient())
-            .withClientId(config.getClientId())
-            .withClientSecret(config.getClientSecret())
-            .withHost(config.getHost())
-            .withRedirectUrl(
-                config.getOAuthRedirectUrl() != null
-                    ? config.getOAuthRedirectUrl()
-                    : "http://localhost:8080/callback")
-            .withScopes(config.getScopes()));
-  }
-
   private OAuthClient(Builder b) throws IOException {
     this.clientId = Objects.requireNonNull(b.clientId);
     this.clientSecret = b.clientSecret;

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

@@ -0,0 +1,42 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.DatabricksConfig;
+
+/** Utility methods for OAuth client credentials resolution. */
+public class OAuthClientUtils {
+
+  /** Default client ID to use when no client ID is specified. */
+  private static final String DEFAULT_CLIENT_ID = "databricks-cli";
+
+  /**
+   * Resolves the OAuth client ID from the configuration. Prioritizes regular OAuth client ID, then
+   * Azure client ID, and falls back to default client ID.
+   *
+   * @param config The Databricks configuration
+   * @return The resolved client ID
+   */
+  public static String resolveClientId(DatabricksConfig config) {
+    if (config.getClientId() != null) {
+      return config.getClientId();
+    } else if (config.getAzureClientId() != null) {
+      return config.getAzureClientId();
+    }
+    return DEFAULT_CLIENT_ID;
+  }
+
+  /**
+   * Resolves the OAuth client secret from the configuration. Prioritizes regular OAuth client
+   * secret, then Azure client secret.
+   *
+   * @param config The Databricks configuration
+   * @return The resolved client secret, or null if not present
+   */
+  public static String resolveClientSecret(DatabricksConfig config) {
+    if (config.getClientSecret() != null) {
+      return config.getClientSecret();
+    } else if (config.getAzureClientSecret() != null) {
+      return config.getAzureClientSecret();
+    }
+    return null;
+  }
+}

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() {
@@ -43,6 +46,7 @@ static class Builder {
     private String redirectUrl;
     private String clientId;
     private String clientSecret;
+    private TokenCache tokenCache;
 
     public Builder withHttpClient(HttpClient hc) {
       this.hc = hc;
@@ -74,6 +78,11 @@ public Builder withClientSecret(String clientSecret) {
       return this;
     }
 
+    public Builder withTokenCache(TokenCache tokenCache) {
+      this.tokenCache = tokenCache;
+      return this;
+    }
+
     public SessionCredentials build() {
       return new SessionCredentials(this);
     }
@@ -84,6 +93,7 @@ public SessionCredentials build() {
   private final String redirectUrl;
   private final String clientId;
   private final String clientSecret;
+  private final TokenCache tokenCache;
 
   private SessionCredentials(Builder b) {
     super(b.token);
@@ -92,6 +102,7 @@ private SessionCredentials(Builder b) {
     this.redirectUrl = b.redirectUrl;
     this.clientId = b.clientId;
     this.clientSecret = b.clientSecret;
+    this.tokenCache = b.tokenCache;
   }
 
   @Override
@@ -113,7 +124,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;
   }
 }