diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index a58fc01d6..23b1bfc2a 100644 --- a/NEXT_CHANGELOG.md +++ b/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 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 index 4ec58480e..7087147cf 100644 --- 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 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 index 6d5a2eb9f..8c319cd63 100644 --- 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 index de6548982..5f5fdc471 100644 --- 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 @@ -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; 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 index 281d2d9a4..e661903b5 100644 --- 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 index a7a041f41..6e6df86eb 100644 --- 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 index 0dbabbb1a..564bda97f 100644 --- 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. @@ -50,7 +49,7 @@ 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. @@ -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 index 37e7d707e..75ffed81a 100644 --- 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 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 index 0766d8a1b..3708887ef 100644 --- 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 index c517b8a3e..3b1294797 100644 --- 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 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 index ec94c014b..fff6e351c 100644 --- 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 index d3edae491..ca3221253 100644 --- 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 index 588b0e971..032f3f3bb 100755 --- 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 @@ -27,7 +27,10 @@ public ServingEndpointsDataPlaneImpl( this.servingEndpointsAPI = servingEndpointsAPI; this.dataPlaneTokenSource = new DataPlaneTokenSource( - apiClient.getHttpClient(), config.getTokenSource(), config.getHost()); + apiClient.getHttpClient(), + config.getTokenSource(), + config.getHost(), + config.getDisableAsyncTokenRefresh()); this.infos = new ConcurrentHashMap<>(); } 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 index 0a29f5a7c..7afbfb2b3 100644 --- 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,12 +23,12 @@ public class CachedTokenSourceTest { private static Stream 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}") @@ -36,7 +36,7 @@ private static Stream 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 index bac591212..5e35bf63d 100644 --- 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 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);