Add Configuration Option to Disable Async Token Refresh (#478)

## What changes are proposed in this pull request?
This PR adds a new configuration option to control asynchronous token
refresh behavior in the Databricks SDK for Java.

## Changes
- Added `disableAsyncTokenRefresh` configuration to `DatabricksConfig`
with environment variable `DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH`
- **Async token refresh is enabled by default** for optimal performance

## To disable async token refresh
**Environment Variable:**
```bash
export DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH=true
```

**Programmatic:**
```java
DatabricksConfig config = new DatabricksConfig()
    .setDisableAsyncTokenRefresh(true);
```

## How is this tested?
Existing unit and integration tests.

---------

Co-authored-by: Parth Bansal <[email protected]>

Author: emmyzhou-db

NEXT_CHANGELOG.md

@@ -3,6 +3,9 @@
 ## Release v0.55.0
 
 ### New Features and Improvements
+* Enabled asynchronous token refreshes by default. A new `disable_async_token_refresh` configuration option has been added to allow disabling this feature if necessary.
+  To disable asynchronous token refresh, set the environment variable `DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH=true` or configure it within your configuration object.
+  The previous `DATABRICKS_ENABLE_EXPERIMENTAL_ASYNC_TOKEN_REFRESH` option has been removed as asynchronous refresh is now the default behavior.
 
 ### Bug Fixes
 

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

@@ -52,7 +52,10 @@ public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource
   protected CachedTokenSource getTokenSource(DatabricksConfig config, List<String> cmd) {
     CliTokenSource tokenSource =
         new CliTokenSource(cmd, "tokenType", "accessToken", "expiresOn", config.getEnv());
-    CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+    CachedTokenSource cachedTokenSource =
+        new CachedTokenSource.Builder(tokenSource)
+            .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+            .build();
     cachedTokenSource.getToken(); // Check if the CLI is installed and to validate the config.
     return cachedTokenSource;
   }

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

@@ -49,7 +49,10 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
         return null;
       }
 
-      CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+      CachedTokenSource cachedTokenSource =
+          new CachedTokenSource.Builder(tokenSource)
+              .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+              .build();
       cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed.
 
       return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);

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

@@ -159,6 +159,10 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_ENV", auth = "env-oidc")
   private String oidcTokenEnv;
 
+  /** Disable asynchronous token refresh when set to true. */
+  @ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
+  private Boolean disableAsyncTokenRefresh;
+
   public Environment getEnv() {
     return env;
   }
@@ -575,6 +579,15 @@ public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) {
     return this;
   }
 
+  public boolean getDisableAsyncTokenRefresh() {
+    return disableAsyncTokenRefresh != null && disableAsyncTokenRefresh;
+  }
+
+  public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRefresh) {
+    this.disableAsyncTokenRefresh = disableAsyncTokenRefresh;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

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

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

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

@@ -70,6 +70,8 @@ private static CachedTokenSource tokenSourceFor(DatabricksConfig config, String
             .withEndpointParametersSupplier(() -> endpointParams)
             .withAuthParameterPosition(AuthParameterPosition.BODY)
             .build();
-    return new CachedTokenSource.Builder(clientCredentials).build();
+    return new CachedTokenSource.Builder(clientCredentials)
+        .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+        .build();
   }
 }

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

@@ -40,8 +40,7 @@ private enum TokenState {
   // The token source to use for refreshing the token.
   private final TokenSource tokenSource;
   // Whether asynchronous refresh is enabled.
-  private boolean asyncEnabled =
-      Boolean.parseBoolean(System.getenv("DATABRICKS_ENABLE_EXPERIMENTAL_ASYNC_TOKEN_REFRESH"));
+  private boolean asyncDisabled = false;
   // Duration before expiry to consider a token as 'stale'.
   private final Duration staleDuration;
   // Additional buffer before expiry to consider a token as expired.
@@ -50,15 +49,15 @@ private enum TokenState {
   private final ClockSupplier clockSupplier;
 
   // The current OAuth token. May be null if not yet fetched.
-  private volatile Token token;
+  protected volatile Token token;
   // Whether a refresh is currently in progress (for async refresh).
   private boolean refreshInProgress = false;
   // Whether the last refresh attempt succeeded.
   private boolean lastRefreshSucceeded = true;
 
   private CachedTokenSource(Builder builder) {
     this.tokenSource = builder.tokenSource;
-    this.asyncEnabled = builder.asyncEnabled;
+    this.asyncDisabled = builder.asyncDisabled;
     this.staleDuration = builder.staleDuration;
     this.expiryBuffer = builder.expiryBuffer;
     this.clockSupplier = builder.clockSupplier;
@@ -73,7 +72,7 @@ private CachedTokenSource(Builder builder) {
    */
   public static class Builder {
     private final TokenSource tokenSource;
-    private boolean asyncEnabled = false;
+    private boolean asyncDisabled = false;
     private Duration staleDuration = DEFAULT_STALE_DURATION;
     private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER;
     private ClockSupplier clockSupplier = new UtcClockSupplier();
@@ -110,11 +109,11 @@ public Builder setToken(Token token) {
      * current token. When disabled, all refreshes are performed synchronously and will block the
      * calling thread.
      *
-     * @param asyncEnabled True to enable asynchronous refresh, false to disable.
+     * @param asyncDisabled True to disable asynchronous refresh, false to enable.
      * @return This builder instance for method chaining.
      */
-    public Builder setAsyncEnabled(boolean asyncEnabled) {
-      this.asyncEnabled = asyncEnabled;
+    public Builder setAsyncDisabled(boolean asyncDisabled) {
+      this.asyncDisabled = asyncDisabled;
       return this;
     }
 
@@ -182,10 +181,10 @@ public CachedTokenSource build() {
    * @return The current valid token
    */
   public Token getToken() {
-    if (asyncEnabled) {
-      return getTokenAsync();
+    if (asyncDisabled) {
+      return getTokenBlocking();
     }
-    return getTokenBlocking();
+    return getTokenAsync();
   }
 
   /**

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

@@ -17,6 +17,7 @@ public class DataPlaneTokenSource {
   private final HttpClient httpClient;
   private final TokenSource cpTokenSource;
   private final String host;
+  private final boolean asyncDisabled;
   private final ConcurrentHashMap<TokenSourceKey, CachedTokenSource> sourcesCache;
   /**
    * Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This
@@ -42,11 +43,13 @@ static TokenSourceKey create(String endpoint, String authDetails) {
    * @throws NullPointerException if any parameter is null.
    * @throws IllegalArgumentException if the host is empty.
    */
-  public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) {
+  public DataPlaneTokenSource(
+      HttpClient httpClient, TokenSource cpTokenSource, String host, boolean asyncDisabled) {
     this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
     this.cpTokenSource =
         Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
     this.host = Objects.requireNonNull(host, "Host cannot be null");
+    this.asyncDisabled = asyncDisabled;
 
     if (host.isEmpty()) {
       throw new IllegalArgumentException("Host cannot be empty");
@@ -85,6 +88,7 @@ public Token getToken(String endpoint, String authDetails) {
                 new CachedTokenSource.Builder(
                         new EndpointTokenSource(
                             this.cpTokenSource, k.authDetails(), this.httpClient, this.host))
+                    .setAsyncDisabled(asyncDisabled)
                     .build());
 
     return specificSource.getToken();

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

@@ -78,7 +78,10 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
                   Optional.of(config.getEffectiveOAuthRedirectUrl()),
                   Optional.of(tokenCache));
 
-          CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build();
+          CachedTokenSource cachedTokenSource =
+              new CachedTokenSource.Builder(tokenSource)
+                  .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+                  .build();
           LOGGER.debug("Using cached token, will immediately refresh");
           cachedTokenSource.getToken();
           return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
@@ -128,6 +131,9 @@ CachedTokenSource performBrowserAuth(
             Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()),
             Optional.ofNullable(tokenCache));
 
-    return new CachedTokenSource.Builder(tokenSource).setToken(token).build();
+    return new CachedTokenSource.Builder(tokenSource)
+        .setToken(token)
+        .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+        .build();
   }
 }

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

@@ -58,7 +58,10 @@ public HeaderFactory configure(DatabricksConfig config) throws DatabricksExcepti
                         .build())
             .build();
 
-    CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(clientCredentials).build();
+    CachedTokenSource cachedTokenSource =
+        new CachedTokenSource.Builder(clientCredentials)
+            .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+            .build();
 
     return () -> {
       Map<String, String> headers = new HashMap<>();