Enabled async by default via config

emmyzhou-db · emmyzhou-db · commit 28bcb0353865 · 2025-06-26T14:08:27.000Z
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java b/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;
   }
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/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);
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
@@ -161,7 +161,7 @@ public class DatabricksConfig {
 
   /** Disable asynchronous token refresh when set to true. */
   @ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
-  private boolean disableAsyncTokenRefresh = false;
+  private Boolean disableAsyncTokenRefresh;
 
   public Environment getEnv() {
     return env;
@@ -584,7 +584,7 @@ public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) {
   }
 
   public boolean getDisableAsyncTokenRefresh() {
-    return disableAsyncTokenRefresh;
+    return disableAsyncTokenRefresh != null && disableAsyncTokenRefresh;
   }
 
   public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRefresh) {
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureGithubOidcCredentialsProvider.java b/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);
   }
 
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProvider.java b/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();
   }
 }
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/CachedTokenSource.java b/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.
@@ -58,7 +57,7 @@ private enum TokenState {
 
   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();
   }
 
   /**
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DataPlaneTokenSource.java b/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();
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/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();
   }
 }
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubOidcCredentialsProvider.java b/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<>();
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java b/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) {
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java b/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
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/service/serving/ServingEndpointsDataPlaneImpl.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/service/serving/ServingEndpointsDataPlaneImpl.java
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java b/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();
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DataPlaneTokenSourceTest.java b/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);

Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,7 @@ public class DatabricksConfig {`
`161`	`161`
`162`	`162`	`/** Disable asynchronous token refresh when set to true. */`
`163`	`163`	`@ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")`
`164`		`- private boolean disableAsyncTokenRefresh = false;`
	`164`	`+ private Boolean disableAsyncTokenRefresh;`
`165`	`165`
`166`	`166`	`public Environment getEnv() {`
`167`	`167`	`return env;`
`@@ -584,7 +584,7 @@ public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) {`
`584`	`584`	`}`
`585`	`585`
`586`	`586`	`public boolean getDisableAsyncTokenRefresh() {`
`587`		`- return disableAsyncTokenRefresh;`
	`587`	`+ return disableAsyncTokenRefresh != null && disableAsyncTokenRefresh;`
`588`	`588`	`}`
`589`	`589`
`590`	`590`	`public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRefresh) {`
Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,8 @@ private static CachedTokenSource tokenSourceFor(DatabricksConfig config, String`
`70`	`70`	`.withEndpointParametersSupplier(() -> endpointParams)`
`71`	`71`	`.withAuthParameterPosition(AuthParameterPosition.BODY)`
`72`	`72`	`.build();`
`73`		`- return new CachedTokenSource.Builder(clientCredentials).build();`
	`73`	`+ return new CachedTokenSource.Builder(clientCredentials)`
	`74`	`+ .setAsyncDisabled(config.getDisableAsyncTokenRefresh())`
	`75`	`+ .build();`
`74`	`76`	`}`
`75`	`77`	`}`