Refactor Token Caching with CachedTokenSource Wrapper

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

@@ -1,5 +1,6 @@
 package com.databricks.sdk.core;
 
+import com.databricks.sdk.core.oauth.CachedTokenSource;
 import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.oauth.Token;
 import com.databricks.sdk.core.utils.AzureUtils;
@@ -19,7 +20,7 @@ public String authType() {
     return AZURE_CLI;
   }
 
-  public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
+  public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
     String azPath =
         Optional.ofNullable(config.getEnv()).map(env -> env.get("AZ_PATH")).orElse("az");
 
@@ -35,7 +36,7 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
       List<String> extendedCmd = new ArrayList<>(cmd);
       extendedCmd.addAll(Arrays.asList("--subscription", subscription.get()));
       try {
-        return getToken(config, extendedCmd);
+        return getTokenSource(config, extendedCmd);
       } catch (DatabricksException ex) {
         LOG.warn("Failed to get token for subscription. Using resource only token.");
       }
@@ -45,14 +46,15 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
               + "It is recommended to specify this field in the Databricks configuration to avoid authentication errors.");
     }
 
-    return getToken(config, cmd);
+    return getTokenSource(config, cmd);
   }
 
-  protected CliTokenSource getToken(DatabricksConfig config, List<String> cmd) {
-    CliTokenSource token =
+  protected CachedTokenSource getTokenSource(DatabricksConfig config, List<String> cmd) {
+    CliTokenSource tokenSource =
         new CliTokenSource(cmd, "tokenType", "accessToken", "expiresOn", config.getEnv());
-    token.getToken(); // We need this to check if the CLI is installed and to validate the config.
-    return token;
+    CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+    cachedTokenSource.getToken(); // Check if the CLI is installed and to validate the config.
+    return cachedTokenSource;
   }
 
   private Optional<String> getSubscription(DatabricksConfig config) {
@@ -77,16 +79,16 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
     try {
       AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor);
       String resource = config.getEffectiveAzureLoginAppId();
-      CliTokenSource tokenSource = tokenSourceFor(config, resource);
-      CliTokenSource mgmtTokenSource;
+      CachedTokenSource tokenSource = tokenSourceFor(config, resource);
+      CachedTokenSource mgmtTokenSource;
       try {
         mgmtTokenSource =
             tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint());
       } catch (Exception e) {
         LOG.debug("Not including service management token in headers", e);
         mgmtTokenSource = null;
       }
-      CliTokenSource finalMgmtTokenSource = mgmtTokenSource;
+      CachedTokenSource finalMgmtTokenSource = mgmtTokenSource;
       return OAuthHeaderFactory.fromSuppliers(
           tokenSource::getToken,
           () -> {

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

@@ -1,7 +1,7 @@
 package com.databricks.sdk.core;
 
-import com.databricks.sdk.core.oauth.RefreshableTokenSource;
 import com.databricks.sdk.core.oauth.Token;
+import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.Environment;
 import com.databricks.sdk.core.utils.OSUtils;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -18,7 +18,7 @@
 import java.util.List;
 import org.apache.commons.io.IOUtils;
 
-public class CliTokenSource extends RefreshableTokenSource {
+public class CliTokenSource implements TokenSource {
   private List<String> cmd;
   private String tokenTypeField;
   private String accessTokenField;
@@ -86,7 +86,7 @@ private String getProcessStream(InputStream stream) throws IOException {
   }
 
   @Override
-  protected Token refresh() {
+  public Token getToken() {
     try {
       ProcessBuilder processBuilder = new ProcessBuilder(cmd);
       processBuilder.environment().putAll(env.getEnv());

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

@@ -1,5 +1,6 @@
 package com.databricks.sdk.core;
 
+import com.databricks.sdk.core.oauth.CachedTokenSource;
 import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.utils.OSUtils;
 import java.util.*;
@@ -47,8 +48,11 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
       if (tokenSource == null) {
         return null;
       }
-      tokenSource.getToken(); // We need this for checking if databricks CLI is installed.
-      return OAuthHeaderFactory.fromTokenSource(tokenSource);
+
+      CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+      cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
+
+      return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
     } catch (DatabricksException e) {
       String stderr = e.getMessage();
       if (stderr.contains("not found")) {

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

@@ -46,8 +46,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
             config.getEffectiveAzureLoginAppId(),
             idToken.get(),
             "urn:ietf:params:oauth:client-assertion-type:jwt-bearer");
-
-    return OAuthHeaderFactory.fromTokenSource(tokenSource);
+    CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+    return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
   }
 
   /**

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

@@ -28,8 +28,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
     }
     AzureUtils.ensureHostPresent(
         config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor);
-    RefreshableTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId());
-    RefreshableTokenSource cloud =
+    CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId());
+    CachedTokenSource cloud =
         tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint());
 
     return OAuthHeaderFactory.fromSuppliers(
@@ -44,29 +44,32 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
   }
 
   /**
-   * Creates a RefreshableTokenSource for the specified Azure resource.
+   * Creates a CachedTokenSource for the specified Azure resource.
    *
-   * <p>This function constructs a RefreshableTokenSource instance that fetches OAuth tokens for the
+   * <p>This function constructs a CachedTokenSource instance that fetches OAuth tokens for the
    * given Azure resource. It uses the authentication parameters provided by the DatabricksConfig
    * instance to generate the tokens.
    *
    * @param config The DatabricksConfig instance containing the required authentication parameters.
    * @param resource The Azure resource for which OAuth tokens need to be fetched.
-   * @return A RefreshableTokenSource instance capable of fetching OAuth tokens for the specified
-   *     Azure resource.
+   * @return A CachedTokenSource instance capable of fetching OAuth tokens for the specified Azure
+   *     resource.
    */
-  private static RefreshableTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
+  private static CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
     String aadEndpoint = config.getAzureEnvironment().getActiveDirectoryEndpoint();
     String tokenUrl = aadEndpoint + config.getAzureTenantId() + "/oauth2/token";
     Map<String, String> endpointParams = new HashMap<>();
     endpointParams.put("resource", resource);
-    return new ClientCredentials.Builder()
-        .withHttpClient(config.getHttpClient())
-        .withClientId(config.getAzureClientId())
-        .withClientSecret(config.getAzureClientSecret())
-        .withTokenUrl(tokenUrl)
-        .withEndpointParametersSupplier(() -> endpointParams)
-        .withAuthParameterPosition(AuthParameterPosition.BODY)
-        .build();
+
+    ClientCredentials clientCredentials =
+        new ClientCredentials.Builder()
+            .withHttpClient(config.getHttpClient())
+            .withClientId(config.getAzureClientId())
+            .withClientSecret(config.getAzureClientSecret())
+            .withTokenUrl(tokenUrl)
+            .withEndpointParametersSupplier(() -> endpointParams)
+            .withAuthParameterPosition(AuthParameterPosition.BODY)
+            .build();
+    return new CachedTokenSource.Builder(clientCredentials).build();
   }
 }