Add Support for Asynchronous Token Refresh via Config

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

@@ -9,9 +9,9 @@
 import org.slf4j.LoggerFactory;
 
 public class AzureCliCredentialsProvider implements CredentialsProvider {
+
   private final ObjectMapper mapper = new ObjectMapper();
   private static final Logger LOG = LoggerFactory.getLogger(AzureCliCredentialsProvider.class);
-
   public static final String AZURE_CLI = "azure-cli";
 
   @Override
@@ -27,11 +27,11 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
         new ArrayList<>(
             Arrays.asList(
                 azPath, "account", "get-access-token", "--resource", resource, "--output", "json"));
+
     Optional<String> subscription = getSubscription(config);
     if (subscription.isPresent()) {
       // This will fail if the user has access to the workspace, but not to the subscription
-      // itself.
-      // In such case, we fall back to not using the subscription.
+      // itself. In such case, we fall back to not using the subscription.
       List<String> extendedCmd = new ArrayList<>(cmd);
       extendedCmd.addAll(Arrays.asList("--subscription", subscription.get()));
       try {
@@ -60,11 +60,13 @@ private Optional<String> getSubscription(DatabricksConfig config) {
     if (resourceId == null || resourceId.equals("")) {
       return Optional.empty();
     }
+
     String[] components = resourceId.split("/");
     if (components.length < 3) {
       LOG.warn("Invalid azure workspace resource ID");
       return Optional.empty();
     }
+
     return Optional.of(components[2]);
   }
 
@@ -78,14 +80,18 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
       AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor);
       String resource = config.getEffectiveAzureLoginAppId();
       CliTokenSource tokenSource = tokenSourceFor(config, resource);
+      tokenSource.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
+
       CliTokenSource mgmtTokenSource;
       try {
         mgmtTokenSource =
             tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint());
+        mgmtTokenSource.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
       } catch (Exception e) {
         LOG.debug("Not including service management token in headers", e);
         mgmtTokenSource = null;
       }
+
       CliTokenSource finalMgmtTokenSource = mgmtTokenSource;
       return OAuthHeaderFactory.fromSuppliers(
           tokenSource::getToken,

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

@@ -9,7 +9,6 @@
 public class DatabricksCliCredentialsProvider implements CredentialsProvider {
 
   private static final Logger LOG = LoggerFactory.getLogger(DatabricksCliCredentialsProvider.class);
-
   public static final String DATABRICKS_CLI = "databricks-cli";
 
   @Override
@@ -26,12 +25,15 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
       LOG.debug("Databricks CLI could not be found");
       return null;
     }
+
     List<String> cmd =
         new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
+
     if (config.isAccountClient()) {
       cmd.add("--account-id");
       cmd.add(config.getAccountId());
     }
+
     return new CliTokenSource(cmd, "token_type", "access_token", "expiry", config.getEnv());
   }
 
@@ -47,7 +49,9 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
       if (tokenSource == null) {
         return null;
       }
+
       tokenSource.getToken(); // We need this for checking if databricks CLI is installed.
+      tokenSource.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
       return OAuthHeaderFactory.fromTokenSource(tokenSource);
     } catch (DatabricksException e) {
       String stderr = e.getMessage();

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

@@ -159,6 +159,15 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_ENV", auth = "env-oidc")
   private String oidcTokenEnv;
 
+  /**
+   * Whether to enable experimental async token refresh.
+   *
+   * <p><b>Experimental:</b> This feature and its setting are experimental and may be removed in
+   * future releases.
+   */
+  @ConfigAttribute(env = "DATABRICKS_ENABLE_EXPERIMENTAL_ASYNC_TOKEN_REFRESH")
+  private Boolean enableExperimentalAsyncTokenRefresh;
+
   public Environment getEnv() {
     return env;
   }
@@ -575,6 +584,16 @@ public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) {
     return this;
   }
 
+  public Boolean getEnableExperimentalAsyncTokenRefresh() {
+    return enableExperimentalAsyncTokenRefresh;
+  }
+
+  public DatabricksConfig setEnableExperimentalAsyncTokenRefresh(
+      Boolean enableExperimentalAsyncTokenRefresh) {
+    this.enableExperimentalAsyncTokenRefresh = enableExperimentalAsyncTokenRefresh;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

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

@@ -36,7 +36,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
       return null;
     }
 
-    TokenSource tokenSource =
+    OidcTokenSource tokenSource =
         new OidcTokenSource(
             config.getHttpClient(),
             config.getDatabricksEnvironment().getAzureEnvironment().getActiveDirectoryEndpoint()
@@ -47,6 +47,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
             idToken.get(),
             "urn:ietf:params:oauth:client-assertion-type:jwt-bearer");
 
+    tokenSource.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
     return OAuthHeaderFactory.fromTokenSource(tokenSource);
   }
 

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

@@ -26,11 +26,15 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
         || config.getAzureTenantId() == null) {
       return null;
     }
+
     AzureUtils.ensureHostPresent(
         config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor);
+
     RefreshableTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId());
+    inner.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
     RefreshableTokenSource cloud =
         tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint());
+    cloud.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
 
     return OAuthHeaderFactory.fromSuppliers(
         inner::getToken,

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

@@ -37,6 +37,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
               .withAuthParameterPosition(AuthParameterPosition.HEADER)
               .build();
 
+      tokenSource.withAsyncRefresh(config.getEnableExperimentalAsyncTokenRefresh());
       return OAuthHeaderFactory.fromTokenSource(tokenSource);
     } catch (IOException e) {
       // TODO: Log exception