Add Configuration Option to Disable Async Token Refresh

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<>();

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

@@ -38,7 +38,9 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
               .build();
 
       CachedTokenSource cachedTokenSource =
-          new CachedTokenSource.Builder(clientCredentials).build();
+          new CachedTokenSource.Builder(clientCredentials)
+              .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+              .build();
 
       return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
     } catch (IOException e) {

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

@@ -41,7 +41,9 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
     TokenSource cachedTokenSource =
         (tokenSource instanceof CachedTokenSource)
             ? tokenSource
-            : new CachedTokenSource.Builder(tokenSource).build();
+            : new CachedTokenSource.Builder(tokenSource)
+                .setAsyncDisabled(config.getDisableAsyncTokenRefresh())
+                .build();
 
     try {
       // Validate that we can get a token before returning a HeaderFactory

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java

@@ -23,20 +23,20 @@ public class CachedTokenSourceTest {
 
   private static Stream<Arguments> provideAsyncRefreshScenarios() {
     return Stream.of(
-        Arguments.of("Fresh token, async enabled", FRESH_MINUTES, true, false, INITIAL_TOKEN),
-        Arguments.of("Stale token, async enabled", STALE_MINUTES, true, true, INITIAL_TOKEN),
-        Arguments.of("Expired token, async enabled", EXPIRED_MINUTES, true, true, REFRESH_TOKEN),
-        Arguments.of("Fresh token, async disabled", FRESH_MINUTES, false, false, INITIAL_TOKEN),
-        Arguments.of("Stale token, async disabled", STALE_MINUTES, false, false, INITIAL_TOKEN),
-        Arguments.of("Expired token, async disabled", EXPIRED_MINUTES, false, true, REFRESH_TOKEN));
+        Arguments.of("Fresh token, async enabled", FRESH_MINUTES, false, false, INITIAL_TOKEN),
+        Arguments.of("Stale token, async enabled", STALE_MINUTES, false, true, INITIAL_TOKEN),
+        Arguments.of("Expired token, async enabled", EXPIRED_MINUTES, false, true, REFRESH_TOKEN),
+        Arguments.of("Fresh token, async disabled", FRESH_MINUTES, true, false, INITIAL_TOKEN),
+        Arguments.of("Stale token, async disabled", STALE_MINUTES, true, false, INITIAL_TOKEN),
+        Arguments.of("Expired token, async disabled", EXPIRED_MINUTES, true, true, REFRESH_TOKEN));
   }
 
   @ParameterizedTest(name = "{0}")
   @MethodSource("provideAsyncRefreshScenarios")
   void testAsyncRefreshParametrized(
       String testName,
       long minutesUntilExpiry,
-      boolean asyncEnabled,
+      boolean asyncDisabled,
       boolean expectRefresh,
       String expectedToken)
       throws Exception {
@@ -67,7 +67,7 @@ public Token getToken() {
 
     CachedTokenSource source =
         new CachedTokenSource.Builder(tokenSource)
-            .setAsyncEnabled(asyncEnabled)
+            .setAsyncDisabled(asyncDisabled)
             .setToken(initialToken)
             .build();
 
@@ -127,7 +127,7 @@ public Token getToken() {
     TestSource testSource = new TestSource();
     CachedTokenSource source =
         new CachedTokenSource.Builder(testSource)
-            .setAsyncEnabled(true)
+            .setAsyncDisabled(false)
             .setToken(staleToken)
             .setClockSupplier(clockSupplier)
             .build();

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java

@@ -184,11 +184,13 @@ void testDataPlaneTokenSource(
       assertThrows(
           expectedException,
           () -> {
-            DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host);
+            DataPlaneTokenSource source =
+                new DataPlaneTokenSource(httpClient, cpTokenSource, host, false);
             source.getToken(endpoint, authDetails);
           });
     } else {
-      DataPlaneTokenSource source = new DataPlaneTokenSource(httpClient, cpTokenSource, host);
+      DataPlaneTokenSource source =
+          new DataPlaneTokenSource(httpClient, cpTokenSource, host, false);
       Token token = source.getToken(endpoint, authDetails);
       assertNotNull(token);
       assertEquals(expectedToken.getAccessToken(), token.getAccessToken());
@@ -214,7 +216,7 @@ void testEndpointTokenSourceCaching() throws Exception {
     try (MockedConstruction<EndpointTokenSource> mockedConstruction =
         mockConstruction(EndpointTokenSource.class)) {
       DataPlaneTokenSource source =
-          new DataPlaneTokenSource(mockHttpClient, mockCpTokenSource, TEST_HOST);
+          new DataPlaneTokenSource(mockHttpClient, mockCpTokenSource, TEST_HOST, false);
 
       // First call - should create new EndpointTokenSource
       source.getToken(TEST_ENDPOINT_1, TEST_AUTH_DETAILS_1);