From 23746118951e37b4ddad7bf971c252acc4051319 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 28 May 2025 16:58:06 +0000 Subject: [PATCH 01/47] First draft --- .../core/oauth/RefreshableTokenSource.java | 96 ++++++++++++++- .../com/databricks/sdk/core/oauth/Token.java | 9 ++ .../oauth/RefreshableTokenSourceTest.java | 114 ++++++++++++++++++ 3 files changed, 216 insertions(+), 3 deletions(-) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index e93f91ae5..efdc29469 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -5,10 +5,12 @@ import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; +import java.time.Duration; import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; import java.util.Base64; import java.util.Map; +import java.util.concurrent.CompletableFuture; import org.apache.http.HttpHeaders; /** @@ -18,7 +20,20 @@ * at least 10 seconds until expiry). If not, refresh() is called first to refresh the token. */ public abstract class RefreshableTokenSource implements TokenSource { + + private enum TokenState { + FRESH, // The token is valid. + STALE, // The token is valid but will expire soon. + EXPIRED // The token has expired and cannot be used without refreshing. + } + + private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); + protected Token token; + private boolean asyncEnabled = false; + private Duration staleDuration = DEFAULT_STALE_DURATION; + private boolean refreshInProgress = false; + private boolean lastRefreshSucceeded = true; public RefreshableTokenSource() {} @@ -26,6 +41,11 @@ public RefreshableTokenSource(Token token) { this.token = token; } + public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { + this.asyncEnabled = enabled; + return this; + } + /** * Helper method implementing OAuth token refresh. * @@ -80,9 +100,79 @@ protected static Token retrieveToken( protected abstract Token refresh(); public synchronized Token getToken() { - if (token == null || !token.isValid()) { - token = refresh(); + if (!asyncEnabled) { + return getTokenBlocking(); + } + return getTokenAsync(); + } + + protected TokenState getTokenState() { + if (token == null) { + return TokenState.EXPIRED; + } + Duration lifeTime = token.getLifetime(); + if (lifeTime.compareTo(Duration.ZERO) <= 0) { + return TokenState.EXPIRED; + } + if (lifeTime.compareTo(staleDuration) <= 0) { + return TokenState.STALE; + } + return TokenState.FRESH; + } + + protected synchronized Token getTokenBlocking() { + refreshInProgress = false; + TokenState state = getTokenState(); + if (state != TokenState.EXPIRED) { + return token; + } + try { + Token newToken = refresh(); + token = newToken; + return newToken; + } catch (Exception e) { + lastRefreshSucceeded = false; + throw e; + } + } + + protected Token getTokenAsync() { + TokenState state; + Token currentToken; + synchronized (this) { + state = getTokenState(); + currentToken = token; + } + if (state == TokenState.FRESH) { + return currentToken; + } + if (state == TokenState.STALE) { + triggerAsyncRefresh(); + return token; + } else { + return getTokenBlocking(); + } + } + + protected synchronized void triggerAsyncRefresh() { + if (!refreshInProgress && lastRefreshSucceeded) { + refreshInProgress = true; + CompletableFuture.runAsync( + () -> { + try { + Token newToken = refresh(); + synchronized (this) { + token = newToken; + lastRefreshSucceeded = true; + refreshInProgress = false; + } + } catch (Exception e) { + synchronized (this) { + lastRefreshSucceeded = false; + refreshInProgress = false; + } + } + }); } - return token; } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index f0fd72f68..5aaa012da 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -4,6 +4,7 @@ import com.databricks.sdk.core.utils.SystemClockSupplier; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import java.time.Duration; import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; import java.util.Objects; @@ -91,4 +92,12 @@ public String getRefreshToken() { public String getAccessToken() { return accessToken; } + + public LocalDateTime getExpiry() { + return this.expiry; + } + + public Duration getLifetime() { + return Duration.between(LocalDateTime.now(clockSupplier.getClock()), this.expiry); + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java new file mode 100644 index 000000000..e543117c9 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -0,0 +1,114 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.utils.FakeClockSupplier; +import java.time.Instant; +import java.time.LocalDateTime; +import java.time.ZoneId; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; + +public class RefreshableTokenSourceTest { + @Test + void testAsyncRefresh() throws Exception { + // Set up a fake clock and initial token that is about to become stale + Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); + ZoneId zoneId = ZoneId.of("UTC"); + FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + + // Token expires in 2 minutes (less than the default stale duration of 3 minutes) + Token initialToken = + new Token("initial-token", "Bearer", null, currentTime.plusMinutes(2), fakeClock); + + CountDownLatch refreshCalled = new CountDownLatch(1); + Token refreshedToken = + new Token("refreshed-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + + // Subclass with a refresh() that signals when called + RefreshableTokenSource source = + new RefreshableTokenSource(initialToken) { + @Override + protected Token refresh() { + refreshCalled.countDown(); + return refreshedToken; + } + }.enableAsyncRefresh(true); + + // First call should return the stale token and trigger async refresh + Token token1 = source.getToken(); + assertEquals( + "initial-token", token1.getAccessToken(), "Should return the stale token immediately"); + + // Wait for async refresh to complete (with timeout) + boolean refreshed = refreshCalled.await(2, TimeUnit.SECONDS); + assertTrue(refreshed, "Async refresh should have been triggered"); + + // After refresh, getToken should return the refreshed token + // (may need to wait a bit for the token to be set) + for (int i = 0; i < 10; i++) { + Token token2 = source.getToken(); + if ("refreshed-token".equals(token2.getAccessToken())) { + return; // Success + } + Thread.sleep(100); + } + fail("Token was not refreshed asynchronously"); + } + + @Test + void testSyncRefreshWhenExpired() { + // Set up a fake clock and an expired token + Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); + ZoneId zoneId = ZoneId.of("UTC"); + FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + + Token expiredToken = + new Token("expired-token", "Bearer", null, currentTime.minusMinutes(1), fakeClock); + Token refreshedToken = + new Token("refreshed-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + + final boolean[] refreshCalled = {false}; + RefreshableTokenSource source = + new RefreshableTokenSource(expiredToken) { + @Override + protected Token refresh() { + refreshCalled[0] = true; + return refreshedToken; + } + }.enableAsyncRefresh(true); + + // Should call refresh synchronously and return the refreshed token + Token token = source.getToken(); + assertTrue(refreshCalled[0], "refresh() should be called synchronously for expired token"); + assertEquals("refreshed-token", token.getAccessToken(), "Should return the refreshed token"); + } + + @Test + void testNoRefreshWhenTokenIsFresh() { + // Set up a fake clock and a fresh token + Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); + ZoneId zoneId = ZoneId.of("UTC"); + FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + + Token freshToken = + new Token("fresh-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + + RefreshableTokenSource source = + new RefreshableTokenSource(freshToken) { + @Override + protected Token refresh() { + fail("refresh() should not be called when token is fresh"); + return null; + } + }; + + // Should return the fresh token and never call refresh + Token token = source.getToken(); + assertEquals("fresh-token", token.getAccessToken(), "Should return the fresh token"); + } +} From 5148136c960d522333dcf1456343af993d3670f7 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 28 May 2025 17:24:32 +0000 Subject: [PATCH 02/47] Update test --- .../oauth/RefreshableTokenSourceTest.java | 121 ++++++++++++++---- 1 file changed, 97 insertions(+), 24 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index e543117c9..15c7768dd 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -11,21 +11,29 @@ import org.junit.jupiter.api.Test; public class RefreshableTokenSourceTest { + private static class MutableBoolean { + boolean value = false; + } + private static final String TOKEN_TYPE = "Bearer"; + private static final String INITIAL_TOKEN = "initial-token"; + private static final String REFRESHED_TOKEN = "refreshed-token"; + private static final String EXPIRED_TOKEN = "expired-token"; + private static final String FRESH_TOKEN = "fresh-token"; + private static final Instant FIXED_INSTANT = Instant.parse("2023-10-18T12:00:00.00Z"); + private static final ZoneId ZONE_ID = ZoneId.of("UTC"); + @Test void testAsyncRefresh() throws Exception { - // Set up a fake clock and initial token that is about to become stale - Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); - ZoneId zoneId = ZoneId.of("UTC"); - FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); // Token expires in 2 minutes (less than the default stale duration of 3 minutes) Token initialToken = - new Token("initial-token", "Bearer", null, currentTime.plusMinutes(2), fakeClock); + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); CountDownLatch refreshCalled = new CountDownLatch(1); Token refreshedToken = - new Token("refreshed-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + new Token(REFRESHED_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); // Subclass with a refresh() that signals when called RefreshableTokenSource source = @@ -40,7 +48,7 @@ protected Token refresh() { // First call should return the stale token and trigger async refresh Token token1 = source.getToken(); assertEquals( - "initial-token", token1.getAccessToken(), "Should return the stale token immediately"); + INITIAL_TOKEN, token1.getAccessToken(), "Should return the stale token immediately"); // Wait for async refresh to complete (with timeout) boolean refreshed = refreshCalled.await(2, TimeUnit.SECONDS); @@ -50,7 +58,7 @@ protected Token refresh() { // (may need to wait a bit for the token to be set) for (int i = 0; i < 10; i++) { Token token2 = source.getToken(); - if ("refreshed-token".equals(token2.getAccessToken())) { + if (REFRESHED_TOKEN.equals(token2.getAccessToken())) { return; // Success } Thread.sleep(100); @@ -60,43 +68,37 @@ protected Token refresh() { @Test void testSyncRefreshWhenExpired() { - // Set up a fake clock and an expired token - Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); - ZoneId zoneId = ZoneId.of("UTC"); - FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); Token expiredToken = - new Token("expired-token", "Bearer", null, currentTime.minusMinutes(1), fakeClock); + new Token(EXPIRED_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); Token refreshedToken = - new Token("refreshed-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + new Token(REFRESHED_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); - final boolean[] refreshCalled = {false}; + MutableBoolean refreshCalled = new MutableBoolean(); RefreshableTokenSource source = new RefreshableTokenSource(expiredToken) { @Override protected Token refresh() { - refreshCalled[0] = true; + refreshCalled.value = true; return refreshedToken; } }.enableAsyncRefresh(true); // Should call refresh synchronously and return the refreshed token Token token = source.getToken(); - assertTrue(refreshCalled[0], "refresh() should be called synchronously for expired token"); - assertEquals("refreshed-token", token.getAccessToken(), "Should return the refreshed token"); + assertTrue(refreshCalled.value, "refresh() should be called synchronously for expired token"); + assertEquals(REFRESHED_TOKEN, token.getAccessToken(), "Should return the refreshed token"); } @Test void testNoRefreshWhenTokenIsFresh() { - // Set up a fake clock and a fresh token - Instant now = Instant.parse("2023-10-18T12:00:00.00Z"); - ZoneId zoneId = ZoneId.of("UTC"); - FakeClockSupplier fakeClock = new FakeClockSupplier(now, zoneId); + FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); Token freshToken = - new Token("fresh-token", "Bearer", null, currentTime.plusMinutes(10), fakeClock); + new Token(FRESH_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); RefreshableTokenSource source = new RefreshableTokenSource(freshToken) { @@ -109,6 +111,77 @@ protected Token refresh() { // Should return the fresh token and never call refresh Token token = source.getToken(); - assertEquals("fresh-token", token.getAccessToken(), "Should return the fresh token"); + assertEquals(FRESH_TOKEN, token.getAccessToken(), "Should return the fresh token"); + } + + @Test + void testRefreshThrowsException() { + FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); + LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + + Token expiredToken = + new Token(EXPIRED_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); + + RefreshableTokenSource source = + new RefreshableTokenSource(expiredToken) { + @Override + protected Token refresh() { + throw new RuntimeException("Simulated refresh failure"); + } + }; + + RuntimeException thrown = assertThrows( + RuntimeException.class, + source::getToken, + "getToken() should propagate exception from refresh() when token is expired"); + assertEquals("Simulated refresh failure", thrown.getMessage()); + } + + @Test + void testFailedAsyncRefreshForcesNextRefreshToBeSynchronous() throws Exception { + FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); + LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + // Token is stale (expires in 2 minutes) + Token staleToken = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); + + class TestSource extends RefreshableTokenSource { + int refreshCallCount = 0; + boolean failFirst = true; + TestSource(Token token) { super(token); } + @Override + protected Token refresh() { + refreshCallCount++; + if (failFirst) { + failFirst = false; + throw new RuntimeException("Simulated async failure"); + } + throw new RuntimeException("Simulated sync failure"); + } + } + + TestSource source = new TestSource(staleToken); + source.enableAsyncRefresh(true); + + // First call triggers async refresh, which fails + source.getToken(); + Thread.sleep(300); // Give time for async refresh to run + assertEquals(1, source.refreshCallCount, "refresh() should have been called once (async, failed)"); + + // Token is still stale, so next call should NOT trigger another refresh + source.getToken(); + Thread.sleep(200); + assertEquals(1, source.refreshCallCount, "refresh() should NOT be called again while stale after async failure"); + + // Advance the clock so the token is now expired + source.token = new Token( + INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); + + // Now getToken() should call refresh synchronously and throw + RuntimeException thrown = assertThrows( + RuntimeException.class, + source::getToken, + "getToken() should call refresh synchronously and propagate exception when expired"); + assertEquals("Simulated sync failure", thrown.getMessage()); + assertEquals(2, source.refreshCallCount, "refresh() should have been called synchronously after expiry"); } } From 78c03a8a9fea6a9ea816ea91fb73988de90a5e07 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 30 May 2025 09:32:17 +0000 Subject: [PATCH 03/47] Clean up unit tests --- .../core/oauth/RefreshableTokenSource.java | 66 +++++- .../oauth/RefreshableTokenSourceTest.java | 220 ++++++++---------- 2 files changed, 159 insertions(+), 127 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index efdc29469..74097906c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -16,31 +16,54 @@ /** * An OAuth TokenSource which can be refreshed. * - *

Calls to getToken() will first check if the token is still valid (currently defined by having - * at least 10 seconds until expiry). If not, refresh() is called first to refresh the token. + *

This class supports both synchronous and asynchronous token refresh. When async is enabled, + * stale tokens will trigger a background refresh, while expired tokens will block until a new token + * is fetched. */ public abstract class RefreshableTokenSource implements TokenSource { + /** + * Enum representing the state of the token. FRESH: Token is valid and not close to expiry. STALE: + * Token is valid but will expire soon. EXPIRED: Token has expired and must be refreshed. + */ private enum TokenState { - FRESH, // The token is valid. - STALE, // The token is valid but will expire soon. - EXPIRED // The token has expired and cannot be used without refreshing. + FRESH, + STALE, + EXPIRED } + // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); + /** The current OAuth token. May be null if not yet fetched. */ protected Token token; + /** Whether asynchronous refresh is enabled. */ private boolean asyncEnabled = false; + /** Duration before expiry to consider a token as 'stale'. */ private Duration staleDuration = DEFAULT_STALE_DURATION; + /** Whether a refresh is currently in progress (for async refresh). */ private boolean refreshInProgress = false; + /** Whether the last refresh attempt succeeded. */ private boolean lastRefreshSucceeded = true; + /** Default constructor. */ public RefreshableTokenSource() {} + /** + * Constructor with initial token. + * + * @param token The initial token to use. + */ public RefreshableTokenSource(Token token) { this.token = token; } + /** + * Enable or disable asynchronous token refresh. + * + * @param enabled true to enable async refresh, false to disable + * @return this instance for chaining + */ public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { this.asyncEnabled = enabled; return this; @@ -49,6 +72,7 @@ public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { /** * Helper method implementing OAuth token refresh. * + * @param hc The HTTP client to use for the request. * @param clientId The client ID to authenticate with. * @param clientSecret The client secret to authenticate with. * @param tokenUrl The authorization URL for fetching tokens. @@ -56,6 +80,7 @@ public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { * @param headers Additional headers. * @param position The position of the authentication parameters in the request. * @return The newly fetched Token. + * @throws DatabricksException if the refresh fails */ protected static Token retrieveToken( HttpClient hc, @@ -99,6 +124,12 @@ protected static Token retrieveToken( protected abstract Token refresh(); + /** + * Get the current token, refreshing if necessary. If async refresh is enabled, may return a stale + * token while a refresh is in progress. + * + * @return The current valid token + */ public synchronized Token getToken() { if (!asyncEnabled) { return getTokenBlocking(); @@ -106,6 +137,11 @@ public synchronized Token getToken() { return getTokenAsync(); } + /** + * Determine the state of the current token (fresh, stale, or expired). + * + * @return The token state + */ protected TokenState getTokenState() { if (token == null) { return TokenState.EXPIRED; @@ -120,6 +156,11 @@ protected TokenState getTokenState() { return TokenState.FRESH; } + /** + * Get the current token, blocking to refresh if expired. + * + * @return The current valid token + */ protected synchronized Token getTokenBlocking() { refreshInProgress = false; TokenState state = getTokenState(); @@ -129,6 +170,7 @@ protected synchronized Token getTokenBlocking() { try { Token newToken = refresh(); token = newToken; + lastRefreshSucceeded = true; return newToken; } catch (Exception e) { lastRefreshSucceeded = false; @@ -136,6 +178,12 @@ protected synchronized Token getTokenBlocking() { } } + /** + * Get the current token, possibly triggering an async refresh if stale. If the token is expired, + * blocks to refresh. + * + * @return The current valid or stale token + */ protected Token getTokenAsync() { TokenState state; Token currentToken; @@ -147,23 +195,29 @@ protected Token getTokenAsync() { return currentToken; } if (state == TokenState.STALE) { + // Trigger background refresh, return current token triggerAsyncRefresh(); return token; } else { + // Token is expired, block to refresh return getTokenBlocking(); } } + /** + * Trigger an asynchronous refresh of the token if not already in progress and last refresh + * succeeded. + */ protected synchronized void triggerAsyncRefresh() { if (!refreshInProgress && lastRefreshSucceeded) { refreshInProgress = true; CompletableFuture.runAsync( () -> { try { + // Attempt to refresh the token in the background Token newToken = refresh(); synchronized (this) { token = newToken; - lastRefreshSucceeded = true; refreshInProgress = false; } } catch (Exception e) { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index 15c7768dd..be177cc03 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -8,34 +8,41 @@ import java.time.ZoneId; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; public class RefreshableTokenSourceTest { - private static class MutableBoolean { - boolean value = false; - } private static final String TOKEN_TYPE = "Bearer"; private static final String INITIAL_TOKEN = "initial-token"; private static final String REFRESHED_TOKEN = "refreshed-token"; - private static final String EXPIRED_TOKEN = "expired-token"; - private static final String FRESH_TOKEN = "fresh-token"; private static final Instant FIXED_INSTANT = Instant.parse("2023-10-18T12:00:00.00Z"); private static final ZoneId ZONE_ID = ZoneId.of("UTC"); - @Test - void testAsyncRefresh() throws Exception { + @ParameterizedTest(name = "{0}") + @MethodSource("provideAsyncRefreshScenarios") + void testAsyncRefreshParametrized( + String testName, + long minutesUntilExpiry, + boolean asyncEnabled, + boolean expectRefresh, + boolean expectRefreshedToken) + throws Exception { FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); - - // Token expires in 2 minutes (less than the default stale duration of 3 minutes) Token initialToken = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); - - CountDownLatch refreshCalled = new CountDownLatch(1); + new Token( + INITIAL_TOKEN, + TOKEN_TYPE, + null, + currentTime.plusMinutes(minutesUntilExpiry), + fakeClock); Token refreshedToken = new Token(REFRESHED_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); + CountDownLatch refreshCalled = new CountDownLatch(1); - // Subclass with a refresh() that signals when called RefreshableTokenSource source = new RefreshableTokenSource(initialToken) { @Override @@ -43,119 +50,74 @@ protected Token refresh() { refreshCalled.countDown(); return refreshedToken; } - }.enableAsyncRefresh(true); + }.enableAsyncRefresh(asyncEnabled); - // First call should return the stale token and trigger async refresh Token token1 = source.getToken(); - assertEquals( - INITIAL_TOKEN, token1.getAccessToken(), "Should return the stale token immediately"); - - // Wait for async refresh to complete (with timeout) - boolean refreshed = refreshCalled.await(2, TimeUnit.SECONDS); - assertTrue(refreshed, "Async refresh should have been triggered"); - - // After refresh, getToken should return the refreshed token - // (may need to wait a bit for the token to be set) - for (int i = 0; i < 10; i++) { - Token token2 = source.getToken(); - if (REFRESHED_TOKEN.equals(token2.getAccessToken())) { - return; // Success + if (expectRefresh) { + // Wait for async refresh if enabled, otherwise refresh is sync + boolean refreshed = refreshCalled.await(2, TimeUnit.SECONDS); + assertTrue(refreshed, "Refresh should have been triggered"); + } else { + assertEquals(1, refreshCalled.getCount(), "Refresh should NOT have been triggered"); + } + if (expectRefreshedToken) { + // Wait for async to complete if needed + for (int i = 0; i < 10; i++) { + Token token2 = source.getToken(); + if (REFRESHED_TOKEN.equals(token2.getAccessToken())) { + return; // Success + } + Thread.sleep(100); } - Thread.sleep(100); + fail("Token was not refreshed as expected"); + } else { + assertEquals(INITIAL_TOKEN, token1.getAccessToken(), "Should return the initial token"); } - fail("Token was not refreshed asynchronously"); } - @Test - void testSyncRefreshWhenExpired() { - FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); - LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); - - Token expiredToken = - new Token(EXPIRED_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); - Token refreshedToken = - new Token(REFRESHED_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); - - MutableBoolean refreshCalled = new MutableBoolean(); - RefreshableTokenSource source = - new RefreshableTokenSource(expiredToken) { - @Override - protected Token refresh() { - refreshCalled.value = true; - return refreshedToken; - } - }.enableAsyncRefresh(true); - - // Should call refresh synchronously and return the refreshed token - Token token = source.getToken(); - assertTrue(refreshCalled.value, "refresh() should be called synchronously for expired token"); - assertEquals(REFRESHED_TOKEN, token.getAccessToken(), "Should return the refreshed token"); + private static Stream provideAsyncRefreshScenarios() { + return Stream.of( + Arguments.of("Fresh token, async enabled", 10, true, false, false), + Arguments.of("Stale token, async enabled", 1, true, true, true), + Arguments.of("Expired token, async enabled", -1, true, true, true), + Arguments.of("Fresh token, async disabled", 10, false, false, false), + Arguments.of("Stale token, async disabled", 1, false, false, false), + Arguments.of("Expired token, async disabled", -1, false, true, true)); } + /** + * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is forced to be synchronous. + * It ensures that after an async failure, the system does not repeatedly attempt async refreshes while the token is stale, + * and only performs a synchronous refresh when the token is expired. After a successful sync refresh, async refreshes resume as normal. + */ @Test - void testNoRefreshWhenTokenIsFresh() { + void testAsyncRefreshFailureFallback() throws Exception { FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); - - Token freshToken = - new Token(FRESH_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); - - RefreshableTokenSource source = - new RefreshableTokenSource(freshToken) { - @Override - protected Token refresh() { - fail("refresh() should not be called when token is fresh"); - return null; - } - }; - - // Should return the fresh token and never call refresh - Token token = source.getToken(); - assertEquals(FRESH_TOKEN, token.getAccessToken(), "Should return the fresh token"); - } - - @Test - void testRefreshThrowsException() { - FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); - LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); - - Token expiredToken = - new Token(EXPIRED_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); - - RefreshableTokenSource source = - new RefreshableTokenSource(expiredToken) { - @Override - protected Token refresh() { - throw new RuntimeException("Simulated refresh failure"); - } - }; - - RuntimeException thrown = assertThrows( - RuntimeException.class, - source::getToken, - "getToken() should propagate exception from refresh() when token is expired"); - assertEquals("Simulated refresh failure", thrown.getMessage()); - } - - @Test - void testFailedAsyncRefreshForcesNextRefreshToBeSynchronous() throws Exception { - FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); - LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); - // Token is stale (expires in 2 minutes) - Token staleToken = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); + Token staleToken = + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); class TestSource extends RefreshableTokenSource { int refreshCallCount = 0; - boolean failFirst = true; - TestSource(Token token) { super(token); } + boolean isFirstRefresh = true; + + TestSource(Token token) { + super(token); + } + @Override protected Token refresh() { refreshCallCount++; - if (failFirst) { - failFirst = false; + if (isFirstRefresh) { + isFirstRefresh = false; throw new RuntimeException("Simulated async failure"); } - throw new RuntimeException("Simulated sync failure"); + return new Token( + REFRESHED_TOKEN, + TOKEN_TYPE, + null, + LocalDateTime.now(fakeClock.getClock()).plusMinutes(10), + fakeClock); } } @@ -164,24 +126,40 @@ protected Token refresh() { // First call triggers async refresh, which fails source.getToken(); - Thread.sleep(300); // Give time for async refresh to run - assertEquals(1, source.refreshCallCount, "refresh() should have been called once (async, failed)"); + Thread.sleep(300); + assertEquals( + 1, source.refreshCallCount, "refresh() should have been called once (async, failed)"); - // Token is still stale, so next call should NOT trigger another refresh + // Token is still stale, so next call should NOT trigger another refresh since the last refresh + // failed source.getToken(); Thread.sleep(200); - assertEquals(1, source.refreshCallCount, "refresh() should NOT be called again while stale after async failure"); + assertEquals( + 1, + source.refreshCallCount, + "refresh() should NOT be called again while stale after async failure"); // Advance the clock so the token is now expired - source.token = new Token( - INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); - - // Now getToken() should call refresh synchronously and throw - RuntimeException thrown = assertThrows( - RuntimeException.class, - source::getToken, - "getToken() should call refresh synchronously and propagate exception when expired"); - assertEquals("Simulated sync failure", thrown.getMessage()); - assertEquals(2, source.refreshCallCount, "refresh() should have been called synchronously after expiry"); + source.token = + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); + + // Now getToken() should call refresh synchronously and return the refreshed token + Token token = source.getToken(); + assertEquals( + REFRESHED_TOKEN, + token.getAccessToken(), + "Should return the refreshed token after sync refresh"); + assertEquals( + 2, source.refreshCallCount, "refresh() should have been called synchronously after expiry"); + + // Make the token stale again and trigger async refresh since the last sync refresh succeeded + source.token = + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); + source.getToken(); + Thread.sleep(300); + assertEquals( + 3, + source.refreshCallCount, + "refresh() should have been called again asynchronously after making token stale"); } } From 5a37c59833ed2261251b03b26605ea1fdada55e4 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 30 May 2025 12:43:04 +0000 Subject: [PATCH 04/47] Clean up comments --- .../core/oauth/RefreshableTokenSource.java | 125 ++++++++++-------- .../oauth/RefreshableTokenSourceTest.java | 8 +- 2 files changed, 77 insertions(+), 56 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 74097906c..e1168222f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -24,7 +24,8 @@ public abstract class RefreshableTokenSource implements TokenSource { /** * Enum representing the state of the token. FRESH: Token is valid and not close to expiry. STALE: - * Token is valid but will expire soon. EXPIRED: Token has expired and must be refreshed. + * Token is valid but will expire soon - an async refresh will be triggered if enabled. EXPIRED: + * Token has expired and must be refreshed using a blocking call. */ private enum TokenState { FRESH, @@ -70,63 +71,21 @@ public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { } /** - * Helper method implementing OAuth token refresh. + * Refresh the OAuth token. Subclasses must implement this to define how the token is refreshed. * - * @param hc The HTTP client to use for the request. - * @param clientId The client ID to authenticate with. - * @param clientSecret The client secret to authenticate with. - * @param tokenUrl The authorization URL for fetching tokens. - * @param params Additional request parameters. - * @param headers Additional headers. - * @param position The position of the authentication parameters in the request. - * @return The newly fetched Token. - * @throws DatabricksException if the refresh fails + *

This method may throw an exception if the token cannot be refreshed. The specific exception + * type depends on the implementation. + * + * @return The newly refreshed Token. */ - protected static Token retrieveToken( - HttpClient hc, - String clientId, - String clientSecret, - String tokenUrl, - Map params, - Map headers, - AuthParameterPosition position) { - switch (position) { - case BODY: - if (clientId != null) { - params.put("client_id", clientId); - } - if (clientSecret != null) { - params.put("client_secret", clientSecret); - } - break; - case HEADER: - String authHeaderValue = - "Basic " - + Base64.getEncoder().encodeToString((clientId + ":" + clientSecret).getBytes()); - headers.put(HttpHeaders.AUTHORIZATION, authHeaderValue); - break; - } - headers.put("Content-Type", "application/x-www-form-urlencoded"); - Request req = new Request("POST", tokenUrl, FormRequest.wrapValuesInList(params)); - req.withHeaders(headers); - try { - ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build(); - OAuthResponse resp = apiClient.execute(req, OAuthResponse.class); - if (resp.getErrorCode() != null) { - throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); - } - LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); - return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); - } catch (Exception e) { - throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); - } - } - protected abstract Token refresh(); /** - * Get the current token, refreshing if necessary. If async refresh is enabled, may return a stale - * token while a refresh is in progress. + * Gets the current token, refreshing if necessary. If async refresh is enabled, may return a + * stale token while a refresh is in progress. + * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. * * @return The current valid token */ @@ -159,6 +118,9 @@ protected TokenState getTokenState() { /** * Get the current token, blocking to refresh if expired. * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. + * * @return The current valid token */ protected synchronized Token getTokenBlocking() { @@ -182,6 +144,9 @@ protected synchronized Token getTokenBlocking() { * Get the current token, possibly triggering an async refresh if stale. If the token is expired, * blocks to refresh. * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. + * * @return The current valid or stale token */ protected Token getTokenAsync() { @@ -229,4 +194,58 @@ protected synchronized void triggerAsyncRefresh() { }); } } + + /** + * Helper method implementing OAuth token refresh. + * + * @param hc The HTTP client to use for the request. + * @param clientId The client ID to authenticate with. + * @param clientSecret The client secret to authenticate with. + * @param tokenUrl The authorization URL for fetching tokens. + * @param params Additional request parameters. + * @param headers Additional headers. + * @param position The position of the authentication parameters in the request. + * @return The newly fetched Token. + * @throws DatabricksException if the refresh fails + * @throws IllegalArgumentException if the OAuth response contains an error + */ + protected static Token retrieveToken( + HttpClient hc, + String clientId, + String clientSecret, + String tokenUrl, + Map params, + Map headers, + AuthParameterPosition position) { + switch (position) { + case BODY: + if (clientId != null) { + params.put("client_id", clientId); + } + if (clientSecret != null) { + params.put("client_secret", clientSecret); + } + break; + case HEADER: + String authHeaderValue = + "Basic " + + Base64.getEncoder().encodeToString((clientId + ":" + clientSecret).getBytes()); + headers.put(HttpHeaders.AUTHORIZATION, authHeaderValue); + break; + } + headers.put("Content-Type", "application/x-www-form-urlencoded"); + Request req = new Request("POST", tokenUrl, FormRequest.wrapValuesInList(params)); + req.withHeaders(headers); + try { + ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build(); + OAuthResponse resp = apiClient.execute(req, OAuthResponse.class); + if (resp.getErrorCode() != null) { + throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); + } + LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); + return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); + } catch (Exception e) { + throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); + } + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index be177cc03..e4da02959 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -86,9 +86,11 @@ private static Stream provideAsyncRefreshScenarios() { } /** - * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is forced to be synchronous. - * It ensures that after an async failure, the system does not repeatedly attempt async refreshes while the token is stale, - * and only performs a synchronous refresh when the token is expired. After a successful sync refresh, async refreshes resume as normal. + * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is + * forced to be synchronous. It ensures that after an async failure, the system does not + * repeatedly attempt async refreshes while the token is stale, and only performs a synchronous + * refresh when the token is expired. After a successful sync refresh, async refreshes resume as + * normal. */ @Test void testAsyncRefreshFailureFallback() throws Exception { From f5e4f8abb57df8916e7c2ae317f2447c8f2a0a56 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 30 May 2025 12:50:04 +0000 Subject: [PATCH 05/47] Add Javadoc to Token.java --- .../com/databricks/sdk/core/oauth/Token.java | 37 +++++++++++++++++++ .../databricks/sdk/core/oauth/TokenTest.java | 20 ++++++++++ 2 files changed, 57 insertions(+) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index 5aaa012da..ee5df2f5f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -66,6 +66,12 @@ public Token( this.clockSupplier = clockSupplier; } + /** + * Checks if the token is expired. Tokens are considered expired 40 seconds before their actual + * expiry time to account for Azure Databricks rejecting tokens that expire in 30 seconds or less. + * + * @return true if the token is expired or about to expire, false otherwise + */ public boolean isExpired() { if (expiry == null) { return false; @@ -77,26 +83,57 @@ public boolean isExpired() { return potentiallyExpired.isBefore(now); } + /** + * Checks if the token is valid. A token is valid if it has a non-null access token and is not + * expired. + * + * @return true if the token is valid, false otherwise + */ public boolean isValid() { return accessToken != null && !isExpired(); } + /** + * Returns the type of the token (e.g., "Bearer"). + * + * @return the token type + */ public String getTokenType() { return tokenType; } + /** + * Returns the refresh token, if available. May be null for non-refreshable tokens. + * + * @return the refresh token or null + */ public String getRefreshToken() { return refreshToken; } + /** + * Returns the access token string. + * + * @return the access token + */ public String getAccessToken() { return accessToken; } + /** + * Returns the expiry time of the token as a LocalDateTime. + * + * @return the expiry time + */ public LocalDateTime getExpiry() { return this.expiry; } + /** + * Returns the remaining lifetime of the token as a Duration. + * + * @return the duration between now and the token's expiry + */ public Duration getLifetime() { return Duration.between(LocalDateTime.now(clockSupplier.getClock()), this.expiry); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index 2d87a32c2..91d687b2f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -71,4 +71,24 @@ void expiredToken() { assertTrue(token.isExpired()); assertFalse(token.isValid()); } + + @Test + void tokenLifetimeInFuture() { + Token token = + new Token(accessToken, tokenType, currentLocalDateTime.plusMinutes(10), fakeClockSupplier); + assertEquals(java.time.Duration.ofMinutes(10), token.getLifetime()); + } + + @Test + void tokenLifetimeExpired() { + Token token = + new Token(accessToken, tokenType, currentLocalDateTime.minusMinutes(2), fakeClockSupplier); + assertEquals(java.time.Duration.ofMinutes(-2), token.getLifetime()); + } + + @Test + void tokenLifetimeZero() { + Token token = new Token(accessToken, tokenType, currentLocalDateTime, fakeClockSupplier); + assertEquals(java.time.Duration.ZERO, token.getLifetime()); + } } From 3e651092f882b5f6ca41a2b354248fa42f21dfc2 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 30 May 2025 14:00:19 +0000 Subject: [PATCH 06/47] Add a token expiry buffer field --- .../sdk/core/oauth/RefreshableTokenSource.java | 18 ++++++++++++++++-- .../core/oauth/RefreshableTokenSourceTest.java | 4 ++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index e1168222f..291da9e84 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -42,6 +42,8 @@ private enum TokenState { private boolean asyncEnabled = false; /** Duration before expiry to consider a token as 'stale'. */ private Duration staleDuration = DEFAULT_STALE_DURATION; + /** Additional buffer before expiry to consider a token as expired. */ + private Duration expiryBuffer = Duration.ZERO; /** Whether a refresh is currently in progress (for async refresh). */ private boolean refreshInProgress = false; /** Whether the last refresh attempt succeeded. */ @@ -65,11 +67,23 @@ public RefreshableTokenSource(Token token) { * @param enabled true to enable async refresh, false to disable * @return this instance for chaining */ - public RefreshableTokenSource enableAsyncRefresh(boolean enabled) { + public RefreshableTokenSource withAsyncRefresh(boolean enabled) { this.asyncEnabled = enabled; return this; } + /** + * Set the expiry buffer. If the token's lifetime is less than this buffer, it is considered + * expired. + * + * @param buffer the expiry buffer duration + * @return this instance for chaining + */ + public RefreshableTokenSource withExpiryBuffer(Duration buffer) { + this.expiryBuffer = buffer; + return this; + } + /** * Refresh the OAuth token. Subclasses must implement this to define how the token is refreshed. * @@ -106,7 +120,7 @@ protected TokenState getTokenState() { return TokenState.EXPIRED; } Duration lifeTime = token.getLifetime(); - if (lifeTime.compareTo(Duration.ZERO) <= 0) { + if (lifeTime.compareTo(expiryBuffer) <= 0) { return TokenState.EXPIRED; } if (lifeTime.compareTo(staleDuration) <= 0) { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index e4da02959..8d991e003 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -50,7 +50,7 @@ protected Token refresh() { refreshCalled.countDown(); return refreshedToken; } - }.enableAsyncRefresh(asyncEnabled); + }.withAsyncRefresh(asyncEnabled); Token token1 = source.getToken(); if (expectRefresh) { @@ -124,7 +124,7 @@ protected Token refresh() { } TestSource source = new TestSource(staleToken); - source.enableAsyncRefresh(true); + source.withAsyncRefresh(true); // First call triggers async refresh, which fails source.getToken(); From dfce414f347dbd7bd995012eb94eda1107907099 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 2 Jun 2025 15:39:23 +0000 Subject: [PATCH 07/47] Fix for comments --- .../core/oauth/RefreshableTokenSource.java | 50 ++++++++----------- .../com/databricks/sdk/core/oauth/Token.java | 28 ----------- .../core/oauth/DataPlaneTokenSourceTest.java | 1 - .../oauth/DatabricksOAuthTokenSourceTest.java | 1 - .../sdk/core/oauth/FileTokenCacheTest.java | 18 ------- .../databricks/sdk/core/oauth/TokenTest.java | 26 ---------- 6 files changed, 22 insertions(+), 102 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 291da9e84..5270bf6cd 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -36,20 +36,20 @@ private enum TokenState { // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); - /** The current OAuth token. May be null if not yet fetched. */ + // The current OAuth token. May be null if not yet fetched. protected Token token; - /** Whether asynchronous refresh is enabled. */ + // Whether asynchronous refresh is enabled. private boolean asyncEnabled = false; - /** Duration before expiry to consider a token as 'stale'. */ + // Duration before expiry to consider a token as 'stale'. private Duration staleDuration = DEFAULT_STALE_DURATION; - /** Additional buffer before expiry to consider a token as expired. */ - private Duration expiryBuffer = Duration.ZERO; - /** Whether a refresh is currently in progress (for async refresh). */ + // Additional buffer before expiry to consider a token as expired. + private Duration expiryBuffer = Duration.ofSeconds(40); + // Whether a refresh is currently in progress (for async refresh). private boolean refreshInProgress = false; - /** Whether the last refresh attempt succeeded. */ + // Whether the last refresh attempt succeeded. private boolean lastRefreshSucceeded = true; - /** Default constructor. */ + /** Constructs a new {@code RefreshableTokenSource} with no initial token. */ public RefreshableTokenSource() {} /** @@ -103,7 +103,7 @@ public RefreshableTokenSource withExpiryBuffer(Duration buffer) { * * @return The current valid token */ - public synchronized Token getToken() { + public Token getToken() { if (!asyncEnabled) { return getTokenBlocking(); } @@ -143,15 +143,10 @@ protected synchronized Token getTokenBlocking() { if (state != TokenState.EXPIRED) { return token; } - try { - Token newToken = refresh(); - token = newToken; - lastRefreshSucceeded = true; - return newToken; - } catch (Exception e) { - lastRefreshSucceeded = false; - throw e; - } + lastRefreshSucceeded = false; + token = refresh(); // May throw an exception + lastRefreshSucceeded = true; + return token; } /** @@ -170,16 +165,15 @@ protected Token getTokenAsync() { state = getTokenState(); currentToken = token; } - if (state == TokenState.FRESH) { - return currentToken; - } - if (state == TokenState.STALE) { - // Trigger background refresh, return current token - triggerAsyncRefresh(); - return token; - } else { - // Token is expired, block to refresh - return getTokenBlocking(); + switch (state) { + case FRESH: + return currentToken; + case STALE: + triggerAsyncRefresh(); + return currentToken; + case EXPIRED: + default: + return getTokenBlocking(); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index ee5df2f5f..59afddd84 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -6,7 +6,6 @@ import com.fasterxml.jackson.annotation.JsonProperty; import java.time.Duration; import java.time.LocalDateTime; -import java.time.temporal.ChronoUnit; import java.util.Objects; public class Token { @@ -66,33 +65,6 @@ public Token( this.clockSupplier = clockSupplier; } - /** - * Checks if the token is expired. Tokens are considered expired 40 seconds before their actual - * expiry time to account for Azure Databricks rejecting tokens that expire in 30 seconds or less. - * - * @return true if the token is expired or about to expire, false otherwise - */ - public boolean isExpired() { - if (expiry == null) { - return false; - } - // Azure Databricks rejects tokens that expire in 30 seconds or less, - // so we refresh the token 40 seconds before it expires. - LocalDateTime potentiallyExpired = expiry.minus(40, ChronoUnit.SECONDS); - LocalDateTime now = LocalDateTime.now(clockSupplier.getClock()); - return potentiallyExpired.isBefore(now); - } - - /** - * Checks if the token is valid. A token is valid if it has a non-null access token and is not - * expired. - * - * @return true if the token is valid, false otherwise - */ - public boolean isValid() { - return accessToken != null && !isExpired(); - } - /** * Returns the type of the token (e.g., "Bearer"). * 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 5887c4ee1..d50452b2f 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 @@ -196,7 +196,6 @@ void testDataPlaneTokenSource( assertEquals(expectedToken.getAccessToken(), token.getAccessToken()); assertEquals(expectedToken.getTokenType(), token.getTokenType()); assertEquals(expectedToken.getRefreshToken(), token.getRefreshToken()); - assertTrue(token.isValid()); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java index 8217179f2..ee226cd42 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java @@ -330,7 +330,6 @@ void testTokenSource(TestCase testCase) { assertEquals(TOKEN, token.getAccessToken()); assertEquals(TOKEN_TYPE, token.getTokenType()); assertEquals(REFRESH_TOKEN, token.getRefreshToken()); - assertFalse(token.isExpired()); // Verify correct audience was used verify(testCase.idTokenSource, atLeastOnce()).getIDToken(testCase.expectedAudience); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java index ede6cfd11..fb89e5b81 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -54,24 +54,6 @@ void testSaveAndLoadToken() { assertEquals("access-token", loadedToken.getAccessToken()); assertEquals("Bearer", loadedToken.getTokenType()); assertEquals("refresh-token", loadedToken.getRefreshToken()); - assertFalse(loadedToken.isExpired(), "Token should not be expired"); - } - - @Test - void testTokenExpiry() { - // Create an expired token - LocalDateTime pastTime = LocalDateTime.now().minusHours(1); - Token expiredToken = new Token("access-token", "Bearer", "refresh-token", pastTime); - - // Verify it's marked as expired - assertTrue(expiredToken.isExpired(), "Token should be expired"); - - // Create a valid token - LocalDateTime futureTime = LocalDateTime.now().plusMinutes(30); - Token validToken = new Token("access-token", "Bearer", "refresh-token", futureTime); - - // Verify it's not marked as expired - assertFalse(validToken.isExpired(), "Token should not be expired"); } @Test diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index 91d687b2f..ed0e4b3bd 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -30,7 +30,6 @@ void createNonRefreshableToken() { assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertNull(token.getRefreshToken()); - assertTrue(token.isValid()); } @Test @@ -45,31 +44,6 @@ void createRefreshableToken() { assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertEquals(refreshToken, token.getRefreshToken()); - assertTrue(token.isValid()); - } - - @Test - void tokenExpiryMoreThan40Seconds() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusSeconds(50), fakeClockSupplier); - assertFalse(token.isExpired()); - assertTrue(token.isValid()); - } - - @Test - void tokenExpiryLessThan40Seconds() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusSeconds(30), fakeClockSupplier); - assertTrue(token.isExpired()); - assertFalse(token.isValid()); - } - - @Test - void expiredToken() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.minusSeconds(10), fakeClockSupplier); - assertTrue(token.isExpired()); - assertFalse(token.isValid()); } @Test From 5bd4215a686793a2eb438a6049c829882a6f4490 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 12:01:48 +0000 Subject: [PATCH 08/47] Update tests --- .../core/oauth/RefreshableTokenSource.java | 24 ++++- .../com/databricks/sdk/core/oauth/Token.java | 34 +------ .../oauth/RefreshableTokenSourceTest.java | 90 +++++++------------ .../databricks/sdk/core/oauth/TokenTest.java | 43 +-------- 4 files changed, 58 insertions(+), 133 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 5270bf6cd..50e596348 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -5,6 +5,8 @@ import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.utils.ClockSupplier; +import com.databricks.sdk.core.utils.SystemClockSupplier; import java.time.Duration; import java.time.LocalDateTime; import java.time.temporal.ChronoUnit; @@ -48,6 +50,8 @@ private enum TokenState { private boolean refreshInProgress = false; // Whether the last refresh attempt succeeded. private boolean lastRefreshSucceeded = true; + // Clock supplier for current time, for testing purposes. + private ClockSupplier clockSupplier = new SystemClockSupplier(); /** Constructs a new {@code RefreshableTokenSource} with no initial token. */ public RefreshableTokenSource() {} @@ -61,6 +65,17 @@ public RefreshableTokenSource(Token token) { this.token = token; } + /** + * Set the clock supplier for current time. + * + * @param clockSupplier The clock supplier to use. + * @return this instance for chaining + */ + public RefreshableTokenSource withClockSupplier(ClockSupplier clockSupplier) { + this.clockSupplier = clockSupplier; + return this; + } + /** * Enable or disable asynchronous token refresh. * @@ -119,7 +134,8 @@ protected TokenState getTokenState() { if (token == null) { return TokenState.EXPIRED; } - Duration lifeTime = token.getLifetime(); + Duration lifeTime = + Duration.between(LocalDateTime.now(clockSupplier.getClock()), token.getExpiry()); if (lifeTime.compareTo(expiryBuffer) <= 0) { return TokenState.EXPIRED; } @@ -138,7 +154,6 @@ protected TokenState getTokenState() { * @return The current valid token */ protected synchronized Token getTokenBlocking() { - refreshInProgress = false; TokenState state = getTokenState(); if (state != TokenState.EXPIRED) { return token; @@ -160,7 +175,7 @@ protected synchronized Token getTokenBlocking() { */ protected Token getTokenAsync() { TokenState state; - Token currentToken; + Token currentToken = token; synchronized (this) { state = getTokenState(); currentToken = token; @@ -172,8 +187,9 @@ protected Token getTokenAsync() { triggerAsyncRefresh(); return currentToken; case EXPIRED: - default: return getTokenBlocking(); + default: + throw new IllegalStateException("Invalid token state: " + state); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index 59afddd84..296809f61 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -1,10 +1,7 @@ package com.databricks.sdk.core.oauth; -import com.databricks.sdk.core.utils.ClockSupplier; -import com.databricks.sdk.core.utils.SystemClockSupplier; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import java.time.Duration; import java.time.LocalDateTime; import java.util.Objects; @@ -24,17 +21,9 @@ public class Token { */ @JsonProperty private LocalDateTime expiry; - private final ClockSupplier clockSupplier; - /** Constructor for non-refreshable tokens (e.g. M2M). */ public Token(String accessToken, String tokenType, LocalDateTime expiry) { - this(accessToken, tokenType, null, expiry, new SystemClockSupplier()); - } - - /** Constructor for non-refreshable tokens (e.g. M2M) with ClockSupplier */ - public Token( - String accessToken, String tokenType, LocalDateTime expiry, ClockSupplier clockSupplier) { - this(accessToken, tokenType, null, expiry, clockSupplier); + this(accessToken, tokenType, null, expiry); } /** Constructor for refreshable tokens. */ @@ -44,25 +33,13 @@ public Token( @JsonProperty("tokenType") String tokenType, @JsonProperty("refreshToken") String refreshToken, @JsonProperty("expiry") LocalDateTime expiry) { - this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier()); - } - - /** Constructor for refreshable tokens with ClockSupplier. */ - public Token( - String accessToken, - String tokenType, - String refreshToken, - LocalDateTime expiry, - ClockSupplier clockSupplier) { Objects.requireNonNull(accessToken, "accessToken must be defined"); Objects.requireNonNull(tokenType, "tokenType must be defined"); Objects.requireNonNull(expiry, "expiry must be defined"); - Objects.requireNonNull(clockSupplier, "clockSupplier must be defined"); this.accessToken = accessToken; this.tokenType = tokenType; this.refreshToken = refreshToken; this.expiry = expiry; - this.clockSupplier = clockSupplier; } /** @@ -100,13 +77,4 @@ public String getAccessToken() { public LocalDateTime getExpiry() { return this.expiry; } - - /** - * Returns the remaining lifetime of the token as a Duration. - * - * @return the duration between now and the token's expiry - */ - public Duration getLifetime() { - return Duration.between(LocalDateTime.now(clockSupplier.getClock()), this.expiry); - } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index 8d991e003..aece009f8 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -2,10 +2,7 @@ import static org.junit.jupiter.api.Assertions.*; -import com.databricks.sdk.core.utils.FakeClockSupplier; -import java.time.Instant; import java.time.LocalDateTime; -import java.time.ZoneId; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.stream.Stream; @@ -17,9 +14,20 @@ public class RefreshableTokenSourceTest { private static final String TOKEN_TYPE = "Bearer"; private static final String INITIAL_TOKEN = "initial-token"; - private static final String REFRESHED_TOKEN = "refreshed-token"; - private static final Instant FIXED_INSTANT = Instant.parse("2023-10-18T12:00:00.00Z"); - private static final ZoneId ZONE_ID = ZoneId.of("UTC"); + private static final String REFRESH_TOKEN = "refreshed-token"; + private static final long FRESH_TIME_MINUTES = 10; + private static final long STALE_TIME_MINUTES = 1; + private static final long EXPIRED_TIME_MINUTES = -1; + + private static Stream provideAsyncRefreshScenarios() { + return Stream.of( + Arguments.of("Fresh token, async enabled", FRESH_TIME_MINUTES, true, false, false), + Arguments.of("Stale token, async enabled", STALE_TIME_MINUTES, true, true, false), + Arguments.of("Expired token, async enabled", EXPIRED_TIME_MINUTES, true, true, true), + Arguments.of("Fresh token, async disabled", FRESH_TIME_MINUTES, false, false, false), + Arguments.of("Stale token, async disabled", STALE_TIME_MINUTES, false, false, false), + Arguments.of("Expired token, async disabled", EXPIRED_TIME_MINUTES, false, true, true)); + } @ParameterizedTest(name = "{0}") @MethodSource("provideAsyncRefreshScenarios") @@ -30,17 +38,12 @@ void testAsyncRefreshParametrized( boolean expectRefresh, boolean expectRefreshedToken) throws Exception { - FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); - LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); + Token initialToken = new Token( - INITIAL_TOKEN, - TOKEN_TYPE, - null, - currentTime.plusMinutes(minutesUntilExpiry), - fakeClock); + INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(minutesUntilExpiry)); Token refreshedToken = - new Token(REFRESHED_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(10), fakeClock); + new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); CountDownLatch refreshCalled = new CountDownLatch(1); RefreshableTokenSource source = @@ -48,43 +51,27 @@ void testAsyncRefreshParametrized( @Override protected Token refresh() { refreshCalled.countDown(); + try { + Thread.sleep(500); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } return refreshedToken; } }.withAsyncRefresh(asyncEnabled); - Token token1 = source.getToken(); - if (expectRefresh) { - // Wait for async refresh if enabled, otherwise refresh is sync - boolean refreshed = refreshCalled.await(2, TimeUnit.SECONDS); - assertTrue(refreshed, "Refresh should have been triggered"); - } else { - assertEquals(1, refreshCalled.getCount(), "Refresh should NOT have been triggered"); - } + Token token = source.getToken(); + + boolean refreshed = refreshCalled.await(1, TimeUnit.SECONDS); + assertEquals(expectRefresh, refreshed, "Refresh should have been triggered"); + if (expectRefreshedToken) { - // Wait for async to complete if needed - for (int i = 0; i < 10; i++) { - Token token2 = source.getToken(); - if (REFRESHED_TOKEN.equals(token2.getAccessToken())) { - return; // Success - } - Thread.sleep(100); - } - fail("Token was not refreshed as expected"); + assertEquals(REFRESH_TOKEN, token.getAccessToken(), "Token was not refreshed as expected"); } else { - assertEquals(INITIAL_TOKEN, token1.getAccessToken(), "Should return the initial token"); + assertEquals(INITIAL_TOKEN, token.getAccessToken(), "Should return the initial token"); } } - private static Stream provideAsyncRefreshScenarios() { - return Stream.of( - Arguments.of("Fresh token, async enabled", 10, true, false, false), - Arguments.of("Stale token, async enabled", 1, true, true, true), - Arguments.of("Expired token, async enabled", -1, true, true, true), - Arguments.of("Fresh token, async disabled", 10, false, false, false), - Arguments.of("Stale token, async disabled", 1, false, false, false), - Arguments.of("Expired token, async disabled", -1, false, true, true)); - } - /** * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is * forced to be synchronous. It ensures that after an async failure, the system does not @@ -94,10 +81,8 @@ private static Stream provideAsyncRefreshScenarios() { */ @Test void testAsyncRefreshFailureFallback() throws Exception { - FakeClockSupplier fakeClock = new FakeClockSupplier(FIXED_INSTANT, ZONE_ID); - LocalDateTime currentTime = LocalDateTime.now(fakeClock.getClock()); Token staleToken = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); class TestSource extends RefreshableTokenSource { int refreshCallCount = 0; @@ -114,12 +99,7 @@ protected Token refresh() { isFirstRefresh = false; throw new RuntimeException("Simulated async failure"); } - return new Token( - REFRESHED_TOKEN, - TOKEN_TYPE, - null, - LocalDateTime.now(fakeClock.getClock()).plusMinutes(10), - fakeClock); + return new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); } } @@ -142,21 +122,19 @@ protected Token refresh() { "refresh() should NOT be called again while stale after async failure"); // Advance the clock so the token is now expired - source.token = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.minusMinutes(1), fakeClock); + source.token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().minusMinutes(1)); // Now getToken() should call refresh synchronously and return the refreshed token Token token = source.getToken(); assertEquals( - REFRESHED_TOKEN, + REFRESH_TOKEN, token.getAccessToken(), "Should return the refreshed token after sync refresh"); assertEquals( 2, source.refreshCallCount, "refresh() should have been called synchronously after expiry"); // Make the token stale again and trigger async refresh since the last sync refresh succeeded - source.token = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, currentTime.plusMinutes(2), fakeClock); + source.token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); source.getToken(); Thread.sleep(300); assertEquals( diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index ed0e4b3bd..d82dfab82 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -2,10 +2,7 @@ import static org.junit.jupiter.api.Assertions.*; -import com.databricks.sdk.core.utils.FakeClockSupplier; -import java.time.Instant; import java.time.LocalDateTime; -import java.time.ZoneId; import org.junit.jupiter.api.Test; class TokenTest { @@ -13,20 +10,11 @@ class TokenTest { private static final String accessToken = "testAccessToken"; private static final String refreshToken = "testRefreshToken"; private static final String tokenType = "testTokenType"; - private final LocalDateTime currentLocalDateTime; - private final FakeClockSupplier fakeClockSupplier; - - TokenTest() { - Instant instant = Instant.parse("2023-10-18T12:00:00.00Z"); - ZoneId zoneId = ZoneId.of("UTC"); - fakeClockSupplier = new FakeClockSupplier(instant, zoneId); - currentLocalDateTime = LocalDateTime.now(fakeClockSupplier.getClock()); - } + private static final LocalDateTime currentLocalDateTime = LocalDateTime.now(); @Test void createNonRefreshableToken() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusMinutes(5), fakeClockSupplier); + Token token = new Token(accessToken, tokenType, currentLocalDateTime.plusMinutes(5)); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertNull(token.getRefreshToken()); @@ -35,34 +23,9 @@ void createNonRefreshableToken() { @Test void createRefreshableToken() { Token token = - new Token( - accessToken, - tokenType, - refreshToken, - currentLocalDateTime.plusMinutes(5), - fakeClockSupplier); + new Token(accessToken, tokenType, refreshToken, currentLocalDateTime.plusMinutes(5)); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertEquals(refreshToken, token.getRefreshToken()); } - - @Test - void tokenLifetimeInFuture() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusMinutes(10), fakeClockSupplier); - assertEquals(java.time.Duration.ofMinutes(10), token.getLifetime()); - } - - @Test - void tokenLifetimeExpired() { - Token token = - new Token(accessToken, tokenType, currentLocalDateTime.minusMinutes(2), fakeClockSupplier); - assertEquals(java.time.Duration.ofMinutes(-2), token.getLifetime()); - } - - @Test - void tokenLifetimeZero() { - Token token = new Token(accessToken, tokenType, currentLocalDateTime, fakeClockSupplier); - assertEquals(java.time.Duration.ZERO, token.getLifetime()); - } } From 93e0bafcb95c16892cf65bf5a7db2cc69287ce76 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 12:14:19 +0000 Subject: [PATCH 09/47] Clean up tests --- .../oauth/RefreshableTokenSourceTest.java | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index aece009f8..11a791bc9 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -21,12 +21,16 @@ public class RefreshableTokenSourceTest { private static Stream provideAsyncRefreshScenarios() { return Stream.of( - Arguments.of("Fresh token, async enabled", FRESH_TIME_MINUTES, true, false, false), - Arguments.of("Stale token, async enabled", STALE_TIME_MINUTES, true, true, false), - Arguments.of("Expired token, async enabled", EXPIRED_TIME_MINUTES, true, true, true), - Arguments.of("Fresh token, async disabled", FRESH_TIME_MINUTES, false, false, false), - Arguments.of("Stale token, async disabled", STALE_TIME_MINUTES, false, false, false), - Arguments.of("Expired token, async disabled", EXPIRED_TIME_MINUTES, false, true, true)); + Arguments.of("Fresh token, async enabled", FRESH_TIME_MINUTES, true, false, INITIAL_TOKEN), + Arguments.of("Stale token, async enabled", STALE_TIME_MINUTES, true, true, INITIAL_TOKEN), + Arguments.of( + "Expired token, async enabled", EXPIRED_TIME_MINUTES, true, true, REFRESH_TOKEN), + Arguments.of( + "Fresh token, async disabled", FRESH_TIME_MINUTES, false, false, INITIAL_TOKEN), + Arguments.of( + "Stale token, async disabled", STALE_TIME_MINUTES, false, false, INITIAL_TOKEN), + Arguments.of( + "Expired token, async disabled", EXPIRED_TIME_MINUTES, false, true, REFRESH_TOKEN)); } @ParameterizedTest(name = "{0}") @@ -36,7 +40,7 @@ void testAsyncRefreshParametrized( long minutesUntilExpiry, boolean asyncEnabled, boolean expectRefresh, - boolean expectRefreshedToken) + String expectedToken) throws Exception { Token initialToken = @@ -64,12 +68,7 @@ protected Token refresh() { boolean refreshed = refreshCalled.await(1, TimeUnit.SECONDS); assertEquals(expectRefresh, refreshed, "Refresh should have been triggered"); - - if (expectRefreshedToken) { - assertEquals(REFRESH_TOKEN, token.getAccessToken(), "Token was not refreshed as expected"); - } else { - assertEquals(INITIAL_TOKEN, token.getAccessToken(), "Should return the initial token"); - } + assertEquals(expectedToken, token.getAccessToken(), "Token value did not match expected"); } /** From 15e221de95db5352512ff63afdd3f1e1056ccfe8 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 12:24:03 +0000 Subject: [PATCH 10/47] Add logging --- .../core/oauth/RefreshableTokenSource.java | 11 +++++++++- .../oauth/RefreshableTokenSourceTest.java | 22 ++++++++----------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 50e596348..24f88bbd0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -14,6 +14,8 @@ import java.util.Map; import java.util.concurrent.CompletableFuture; import org.apache.http.HttpHeaders; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An OAuth TokenSource which can be refreshed. @@ -35,6 +37,7 @@ private enum TokenState { EXPIRED } + private static final Logger logger = LoggerFactory.getLogger(RefreshableTokenSource.class); // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); @@ -159,7 +162,12 @@ protected synchronized Token getTokenBlocking() { return token; } lastRefreshSucceeded = false; - token = refresh(); // May throw an exception + try { + token = refresh(); // May throw an exception + } catch (Exception e) { + logger.error("Failed to refresh token synchronously", e); + throw e; + } lastRefreshSucceeded = true; return token; } @@ -213,6 +221,7 @@ protected synchronized void triggerAsyncRefresh() { synchronized (this) { lastRefreshSucceeded = false; refreshInProgress = false; + logger.error("Async token refresh failed", e); } } }); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java index 11a791bc9..c33e7fda5 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java @@ -15,22 +15,18 @@ public class RefreshableTokenSourceTest { private static final String TOKEN_TYPE = "Bearer"; private static final String INITIAL_TOKEN = "initial-token"; private static final String REFRESH_TOKEN = "refreshed-token"; - private static final long FRESH_TIME_MINUTES = 10; - private static final long STALE_TIME_MINUTES = 1; - private static final long EXPIRED_TIME_MINUTES = -1; + private static final long FRESH_MINUTES = 10; + private static final long STALE_MINUTES = 1; + private static final long EXPIRED_MINUTES = -1; private static Stream provideAsyncRefreshScenarios() { return Stream.of( - Arguments.of("Fresh token, async enabled", FRESH_TIME_MINUTES, true, false, INITIAL_TOKEN), - Arguments.of("Stale token, async enabled", STALE_TIME_MINUTES, true, true, INITIAL_TOKEN), - Arguments.of( - "Expired token, async enabled", EXPIRED_TIME_MINUTES, true, true, REFRESH_TOKEN), - Arguments.of( - "Fresh token, async disabled", FRESH_TIME_MINUTES, false, false, INITIAL_TOKEN), - Arguments.of( - "Stale token, async disabled", STALE_TIME_MINUTES, false, false, INITIAL_TOKEN), - Arguments.of( - "Expired token, async disabled", EXPIRED_TIME_MINUTES, false, true, REFRESH_TOKEN)); + 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)); } @ParameterizedTest(name = "{0}") From 7589dab5f40bd1cf86f46984b0b2268d47e02638 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 12:39:02 +0000 Subject: [PATCH 11/47] Performance optimization --- .../core/oauth/RefreshableTokenSource.java | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 24f88bbd0..769b88731 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -42,7 +42,7 @@ private enum TokenState { private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); // The current OAuth token. May be null if not yet fetched. - protected Token token; + protected volatile Token token; // Whether asynchronous refresh is enabled. private boolean asyncEnabled = false; // Duration before expiry to consider a token as 'stale'. @@ -133,12 +133,12 @@ public Token getToken() { * * @return The token state */ - protected TokenState getTokenState() { - if (token == null) { + protected TokenState getTokenState(Token t) { + if (t == null) { return TokenState.EXPIRED; } Duration lifeTime = - Duration.between(LocalDateTime.now(clockSupplier.getClock()), token.getExpiry()); + Duration.between(LocalDateTime.now(clockSupplier.getClock()), t.getExpiry()); if (lifeTime.compareTo(expiryBuffer) <= 0) { return TokenState.EXPIRED; } @@ -157,13 +157,12 @@ protected TokenState getTokenState() { * @return The current valid token */ protected synchronized Token getTokenBlocking() { - TokenState state = getTokenState(); - if (state != TokenState.EXPIRED) { + if (getTokenState(token) != TokenState.EXPIRED) { return token; } lastRefreshSucceeded = false; try { - token = refresh(); // May throw an exception + token = refresh(); } catch (Exception e) { logger.error("Failed to refresh token synchronously", e); throw e; @@ -182,13 +181,9 @@ protected synchronized Token getTokenBlocking() { * @return The current valid or stale token */ protected Token getTokenAsync() { - TokenState state; Token currentToken = token; - synchronized (this) { - state = getTokenState(); - currentToken = token; - } - switch (state) { + + switch (getTokenState(currentToken)) { case FRESH: return currentToken; case STALE: @@ -197,7 +192,7 @@ protected Token getTokenAsync() { case EXPIRED: return getTokenBlocking(); default: - throw new IllegalStateException("Invalid token state: " + state); + throw new IllegalStateException("Invalid token state."); } } From d97c7348bcc05c4da7d7c785e7a671b71cc8f5f2 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 13:24:41 +0000 Subject: [PATCH 12/47] Furter optimizations --- .../core/oauth/RefreshableTokenSource.java | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index 769b88731..ce0f05657 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -71,6 +71,8 @@ public RefreshableTokenSource(Token token) { /** * Set the clock supplier for current time. * + *

Experimental: This method may change or be removed in future releases. + * * @param clockSupplier The clock supplier to use. * @return this instance for chaining */ @@ -82,6 +84,8 @@ public RefreshableTokenSource withClockSupplier(ClockSupplier clockSupplier) { /** * Enable or disable asynchronous token refresh. * + *

Experimental: This method may change or be removed in future releases. + * * @param enabled true to enable async refresh, false to disable * @return this instance for chaining */ @@ -94,6 +98,8 @@ public RefreshableTokenSource withAsyncRefresh(boolean enabled) { * Set the expiry buffer. If the token's lifetime is less than this buffer, it is considered * expired. * + *

Experimental: This method may change or be removed in future releases. + * * @param buffer the expiry buffer duration * @return this instance for chaining */ @@ -156,19 +162,28 @@ protected TokenState getTokenState(Token t) { * * @return The current valid token */ - protected synchronized Token getTokenBlocking() { + protected Token getTokenBlocking() { + // Use double-checked locking to minimize synchronization overhead on reads: + // 1. Check if the token is expired without locking. + // 2. If expired, synchronize and check again (another thread may have refreshed it). + // 3. If still expired, perform the refresh. if (getTokenState(token) != TokenState.EXPIRED) { return token; } - lastRefreshSucceeded = false; - try { - token = refresh(); - } catch (Exception e) { - logger.error("Failed to refresh token synchronously", e); - throw e; + synchronized (this) { + if (getTokenState(token) != TokenState.EXPIRED) { + return token; + } + lastRefreshSucceeded = false; + try { + token = refresh(); + } catch (Exception e) { + logger.error("Failed to refresh token synchronously", e); + throw e; + } + lastRefreshSucceeded = true; + return token; } - lastRefreshSucceeded = true; - return token; } /** From 105bc99b03eb346918c7047d3b3a550215ca9f06 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 14:02:53 +0000 Subject: [PATCH 13/47] Add extra token state check in async refresh --- .../com/databricks/sdk/core/oauth/RefreshableTokenSource.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index ce0f05657..c62dccad7 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -216,7 +216,8 @@ protected Token getTokenAsync() { * succeeded. */ protected synchronized void triggerAsyncRefresh() { - if (!refreshInProgress && lastRefreshSucceeded) { + // Check token state to avoid triggering a refresh if another thread has already refreshed it + if (!refreshInProgress && lastRefreshSucceeded && getTokenState(token) != TokenState.FRESH) { refreshInProgress = true; CompletableFuture.runAsync( () -> { From b24a0fce8da59fee5bfe24f35b3e615ab4b91485 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 3 Jun 2025 17:37:45 +0000 Subject: [PATCH 14/47] Change LocalDateTime to Instant --- .../databricks/sdk/core/CliTokenSource.java | 32 +++++++++++++---- .../oauth/DatabricksOAuthTokenSource.java | 4 +-- .../sdk/core/oauth/EndpointTokenSource.java | 4 +-- .../sdk/core/oauth/OidcTokenSource.java | 4 +-- .../core/oauth/RefreshableTokenSource.java | 5 ++- .../com/databricks/sdk/core/oauth/Token.java | 22 ++++++------ .../core/AzureCliCredentialsProviderTest.java | 3 +- .../sdk/core/CliTokenSourceTest.java | 35 +++++++++++++++---- .../sdk/core/DatabricksConfigTest.java | 3 +- .../core/oauth/DataPlaneTokenSourceTest.java | 14 ++++---- .../core/oauth/EndpointTokenSourceTest.java | 4 +-- ...xternalBrowserCredentialsProviderTest.java | 16 ++++----- .../sdk/core/oauth/FileTokenCacheTest.java | 14 ++++---- .../core/oauth/OAuthHeaderFactoryTest.java | 6 ++-- .../TokenSourceCredentialsProviderTest.java | 4 +-- .../databricks/sdk/core/oauth/TokenTest.java | 15 ++++---- 16 files changed, 109 insertions(+), 76 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 8a7328904..3b3845e51 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -8,7 +8,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.io.InputStream; +import java.time.Instant; import java.time.LocalDateTime; +import java.time.ZoneOffset; import java.time.format.DateTimeFormatter; import java.time.format.DateTimeParseException; import java.util.Arrays; @@ -36,21 +38,39 @@ public CliTokenSource( this.env = env; } - static LocalDateTime parseExpiry(String expiry) { + /** + * Parses an expiry time string and returns the corresponding {@link Instant}. + * + *

The expiry time string is verified to always be in UTC. Any time zone or offset information + * present in the input is ignored, and the value is parsed as a UTC time. + * + *

The method attempts to parse the input using several common date-time formats, including + * ISO-8601 and patterns with varying sub-second precision. + * + * @param expiry the expiry time string to parse, which must represent a UTC time + * @return the parsed {@link Instant} representing the expiry time in UTC + * @throws DateTimeParseException if the input string cannot be parsed as a valid date-time + */ + static Instant parseExpiry(String expiry) { + DateTimeParseException lastException = null; + try { + return Instant.parse(expiry); + } catch (DateTimeParseException e) { + lastException = e; + } + String multiplePrecisionPattern = "[SSSSSSSSS][SSSSSSSS][SSSSSSS][SSSSSS][SSSSS][SSSS][SSS][SS][S]"; List datePatterns = Arrays.asList( "yyyy-MM-dd HH:mm:ss", "yyyy-MM-dd HH:mm:ss." + multiplePrecisionPattern, - "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "XXX", - "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "'Z'"); - DateTimeParseException lastException = null; + "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "XXX"); for (String pattern : datePatterns) { try { DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern); LocalDateTime dateTime = LocalDateTime.parse(expiry, formatter); - return dateTime; + return dateTime.atZone(ZoneOffset.UTC).toInstant(); } catch (DateTimeParseException e) { lastException = e; } @@ -83,7 +103,7 @@ protected Token refresh() { String tokenType = jsonNode.get(tokenTypeField).asText(); String accessToken = jsonNode.get(accessTokenField).asText(); String expiry = jsonNode.get(expiryField).asText(); - LocalDateTime expiresOn = parseExpiry(expiry); + Instant expiresOn = parseExpiry(expiry); return new Token(accessToken, tokenType, expiresOn); } catch (DatabricksException e) { throw e; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java index f16ae2aed..484e0712e 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java @@ -3,7 +3,7 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.HttpClient; import com.google.common.base.Strings; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -166,7 +166,7 @@ public Token refresh() { throw e; } - LocalDateTime expiry = LocalDateTime.now().plusSeconds(response.getExpiresIn()); + Instant expiry = Instant.now().plusSeconds(response.getExpiresIn()); return new Token( response.getAccessToken(), response.getTokenType(), response.getRefreshToken(), expiry); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java index 3ca75c441..ed08f57d6 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java @@ -2,7 +2,7 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.HttpClient; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -87,7 +87,7 @@ protected Token refresh() { throw e; } - LocalDateTime expiry = LocalDateTime.now().plusSeconds(oauthResponse.getExpiresIn()); + Instant expiry = Instant.now().plusSeconds(oauthResponse.getExpiresIn()); return new Token( oauthResponse.getAccessToken(), oauthResponse.getTokenType(), diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java index 719544ebf..b15f55ded 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java @@ -8,7 +8,7 @@ import com.google.common.base.Strings; import com.google.common.collect.ImmutableMap; import java.io.IOException; -import java.time.LocalDateTime; +import java.time.Instant; /** * {@code OidcTokenSource} is responsible for obtaining OAuth tokens using the OpenID Connect (OIDC) @@ -77,7 +77,7 @@ protected Token refresh() { if (resp.getErrorCode() != null) { throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); } - LocalDateTime expiry = LocalDateTime.now().plusSeconds(resp.getExpiresIn()); + Instant expiry = Instant.now().plusSeconds(resp.getExpiresIn()); return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index e93f91ae5..750aae967 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -5,8 +5,7 @@ import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; import com.databricks.sdk.core.http.Request; -import java.time.LocalDateTime; -import java.time.temporal.ChronoUnit; +import java.time.Instant; import java.util.Base64; import java.util.Map; import org.apache.http.HttpHeaders; @@ -70,7 +69,7 @@ protected static Token retrieveToken( if (resp.getErrorCode() != null) { throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); } - LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); + Instant expiry = Instant.now().plusSeconds(resp.getExpiresIn()); return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); } catch (Exception e) { throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index f0fd72f68..ac6fbc3ac 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -4,8 +4,7 @@ import com.databricks.sdk.core.utils.SystemClockSupplier; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import java.time.LocalDateTime; -import java.time.temporal.ChronoUnit; +import java.time.Instant; import java.util.Objects; public class Token { @@ -19,21 +18,20 @@ public class Token { * The expiry time of the token. * *

OAuth token responses include the duration of the lifetime of the access token. When the - * token is retrieved, this is converted to a LocalDateTime tracking the expiry time of the token - * with respect to the current clock. + * token is retrieved, this is converted to an Instant tracking the expiry time of the token with + * respect to the current clock. */ - @JsonProperty private LocalDateTime expiry; + @JsonProperty private Instant expiry; private final ClockSupplier clockSupplier; /** Constructor for non-refreshable tokens (e.g. M2M). */ - public Token(String accessToken, String tokenType, LocalDateTime expiry) { + public Token(String accessToken, String tokenType, Instant expiry) { this(accessToken, tokenType, null, expiry, new SystemClockSupplier()); } /** Constructor for non-refreshable tokens (e.g. M2M) with ClockSupplier */ - public Token( - String accessToken, String tokenType, LocalDateTime expiry, ClockSupplier clockSupplier) { + public Token(String accessToken, String tokenType, Instant expiry, ClockSupplier clockSupplier) { this(accessToken, tokenType, null, expiry, clockSupplier); } @@ -43,7 +41,7 @@ public Token( @JsonProperty("accessToken") String accessToken, @JsonProperty("tokenType") String tokenType, @JsonProperty("refreshToken") String refreshToken, - @JsonProperty("expiry") LocalDateTime expiry) { + @JsonProperty("expiry") Instant expiry) { this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier()); } @@ -52,7 +50,7 @@ public Token( String accessToken, String tokenType, String refreshToken, - LocalDateTime expiry, + Instant expiry, ClockSupplier clockSupplier) { Objects.requireNonNull(accessToken, "accessToken must be defined"); Objects.requireNonNull(tokenType, "tokenType must be defined"); @@ -71,8 +69,8 @@ public boolean isExpired() { } // Azure Databricks rejects tokens that expire in 30 seconds or less, // so we refresh the token 40 seconds before it expires. - LocalDateTime potentiallyExpired = expiry.minus(40, ChronoUnit.SECONDS); - LocalDateTime now = LocalDateTime.now(clockSupplier.getClock()); + Instant potentiallyExpired = expiry.minusSeconds(40); + Instant now = Instant.now(clockSupplier.getClock()); return potentiallyExpired.isBefore(now); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java index 0212b7652..4e8a57b06 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java @@ -7,7 +7,6 @@ import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.oauth.TokenSource; -import java.time.LocalDateTime; import java.util.Arrays; import java.util.List; import org.junit.jupiter.api.Test; @@ -25,7 +24,7 @@ class AzureCliCredentialsProviderTest { private static CliTokenSource mockTokenSource() { CliTokenSource tokenSource = Mockito.mock(CliTokenSource.class); Mockito.when(tokenSource.getToken()) - .thenReturn(new Token(TOKEN, TOKEN_TYPE, LocalDateTime.now())); + .thenReturn(new Token(TOKEN, TOKEN_TYPE, java.time.Instant.now())); return tokenSource; } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java index 20e2f6095..abe609b01 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java @@ -2,7 +2,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; -import java.time.LocalDateTime; +import java.time.Instant; import java.time.format.DateTimeParseException; import org.junit.jupiter.api.Test; @@ -10,20 +10,20 @@ public class CliTokenSourceTest { @Test public void testParseExpiryWithoutTruncate() { - LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612218Z"); - assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612218), parsedDateTime); + Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612218Z"); + assertEquals(Instant.parse("2023-07-17T09:02:22.330612218Z"), parsedInstant); } @Test public void testParseExpiryWithTruncate() { - LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221Z"); - assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612210), parsedDateTime); + Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221Z"); + assertEquals(Instant.parse("2023-07-17T09:02:22.330612210Z"), parsedInstant); } @Test public void testParseExpiryWithTruncateAndLessNanoSecondDigits() { - LocalDateTime parsedDateTime = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612Z"); - assertEquals(LocalDateTime.of(2023, 7, 17, 9, 2, 22, 330612000), parsedDateTime); + Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612Z"); + assertEquals(Instant.parse("2023-07-17T09:02:22.330612000Z"), parsedInstant); } @Test @@ -34,4 +34,25 @@ public void testParseExpiryWithMoreThanNineNanoSecondDigits() { assert (e.getMessage().contains("could not be parsed")); } } + + @Test + public void testParseExpiryWithSpaceFormat() { + Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22"); + assertEquals(Instant.parse("2023-07-17T09:02:22Z"), parsedInstant); + } + + @Test + public void testParseExpiryWithSpaceFormatAndMillis() { + Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22.123"); + assertEquals(Instant.parse("2023-07-17T09:02:22.123Z"), parsedInstant); + } + + @Test + public void testParseExpiryWithInvalidFormat() { + try { + CliTokenSource.parseExpiry("17-07-2023 09:02:22"); + } catch (DateTimeParseException e) { + assert (e.getMessage().contains("could not be parsed")); + } + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 38b6fcd9c..b3ac333a3 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -12,7 +12,6 @@ import com.databricks.sdk.core.oauth.TokenSource; import com.databricks.sdk.core.utils.Environment; import java.io.IOException; -import java.time.LocalDateTime; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -232,7 +231,7 @@ public void testGetTokenSourceWithOAuth() { HttpClient httpClient = mock(HttpClient.class); TokenSource mockTokenSource = mock(TokenSource.class); when(mockTokenSource.getToken()) - .thenReturn(new Token("test-token", "Bearer", LocalDateTime.now().plusHours(1))); + .thenReturn(new Token("test-token", "Bearer", java.time.Instant.now().plusSeconds(3600))); OAuthHeaderFactory mockHeaderFactory = OAuthHeaderFactory.fromTokenSource(mockTokenSource); CredentialsProvider mockProvider = mock(CredentialsProvider.class); when(mockProvider.authType()).thenReturn("test"); 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 5887c4ee1..1fb96a559 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 @@ -8,7 +8,7 @@ import com.databricks.sdk.core.http.Response; import java.io.IOException; import java.net.URL; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.stream.Stream; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -29,8 +29,7 @@ public class DataPlaneTokenSourceTest { private static Stream provideDataPlaneTokenScenarios() throws Exception { // Mock DatabricksOAuthTokenSource for control plane token - Token cpToken = - new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(600)); + Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, Instant.now().plusSeconds(600)); DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); @@ -81,9 +80,8 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti "dp-access-token1", TEST_TOKEN_TYPE, TEST_REFRESH_TOKEN, - LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)), - null // No exception - ), + Instant.now().plusSeconds(TEST_EXPIRES_IN)), + null), Arguments.of( "Success: endpoint2/auth2 (different cache key)", TEST_ENDPOINT_2, @@ -95,7 +93,7 @@ private static Stream provideDataPlaneTokenScenarios() throws Excepti "dp-access-token2", TEST_TOKEN_TYPE, TEST_REFRESH_TOKEN, - LocalDateTime.now().plusSeconds(TEST_EXPIRES_IN)), + Instant.now().plusSeconds(TEST_EXPIRES_IN)), null), Arguments.of( "Error response from endpoint", @@ -203,7 +201,7 @@ void testDataPlaneTokenSource( @Test void testEndpointTokenSourceCaching() throws Exception { Token cpToken = - new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, LocalDateTime.now().plusSeconds(3600)); + new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, null, Instant.now().plusSeconds(3600)); DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java index a3af2254f..5fdac9f2d 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EndpointTokenSourceTest.java @@ -9,7 +9,7 @@ import com.databricks.sdk.core.http.Response; import java.io.IOException; import java.net.URL; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.stream.Stream; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -48,7 +48,7 @@ private static Stream provideEndpointTokenScenarios() throws Exceptio String malformedJson = "{not valid json}"; // Mock DatabricksOAuthTokenSource for control plane token - Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, LocalDateTime.now().plusMinutes(10)); + Token cpToken = new Token(TEST_CP_TOKEN, TEST_TOKEN_TYPE, Instant.now().plusSeconds(600)); DatabricksOAuthTokenSource mockCpTokenSource = mock(DatabricksOAuthTokenSource.class); when(mockCpTokenSource.getToken()).thenReturn(cpToken); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 1714b731c..932690bd7 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -13,7 +13,7 @@ import com.databricks.sdk.core.http.Response; import java.io.IOException; import java.net.URL; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.Arrays; import java.util.HashMap; import java.util.Map; @@ -210,7 +210,7 @@ void sessionCredentials() throws IOException { "originalAccessToken", "originalTokenType", "originalRefreshToken", - LocalDateTime.MAX)) + Instant.MAX)) .build(); Token token = sessionCredentials.refresh(); @@ -234,7 +234,7 @@ void cacheWithValidTokenTest() throws IOException { .execute(any(Request.class)); // Create an valid token with valid refresh token - LocalDateTime futureTime = LocalDateTime.now().plusHours(1); + Instant futureTime = Instant.now().plusSeconds(3600); Token validToken = new Token("valid_access_token", "Bearer", "valid_refresh_token", futureTime); // Create mock token cache that returns the valid token @@ -314,7 +314,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { .execute(any(Request.class)); // Create an expired token with valid refresh token - LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Instant pastTime = Instant.now().minusSeconds(3600); Token expiredToken = new Token("expired_access_token", "Bearer", "valid_refresh_token", pastTime); @@ -392,7 +392,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { .execute(any(Request.class)); // Create an expired token with invalid refresh token - LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Instant pastTime = Instant.now().minusSeconds(3600); Token expiredToken = new Token("expired_access_token", "Bearer", "invalid_refresh_token", pastTime); @@ -406,7 +406,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { "browser_access_token", "Bearer", "browser_refresh_token", - LocalDateTime.now().plusHours(1)); + Instant.now().plusSeconds(3600)); SessionCredentials browserAuthCreds = new SessionCredentials.Builder() @@ -460,7 +460,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { @Test void cacheWithInvalidTokensTest() throws IOException { // Create completely invalid token (no refresh token) - LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Instant pastTime = Instant.now().minusSeconds(3600); Token invalidToken = new Token("expired_access_token", "Bearer", null, pastTime); // Create mock token cache that returns the invalid token @@ -473,7 +473,7 @@ void cacheWithInvalidTokensTest() throws IOException { "browser_access_token", "Bearer", "browser_refresh_token", - LocalDateTime.now().plusHours(1)); + Instant.now().plusSeconds(3600)); SessionCredentials browserAuthCreds = new SessionCredentials.Builder() diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java index ede6cfd11..303f6de66 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -5,7 +5,7 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.Arrays; import java.util.List; import org.junit.jupiter.api.AfterEach; @@ -42,7 +42,7 @@ void testEmptyCache() { @Test void testSaveAndLoadToken() { // Given a token - LocalDateTime expiry = LocalDateTime.now().plusHours(1); + Instant expiry = Instant.now().plusSeconds(3600); Token token = new Token("access-token", "Bearer", "refresh-token", expiry); // When saving and loading the token @@ -60,14 +60,14 @@ void testSaveAndLoadToken() { @Test void testTokenExpiry() { // Create an expired token - LocalDateTime pastTime = LocalDateTime.now().minusHours(1); + Instant pastTime = Instant.now().minusSeconds(3600); Token expiredToken = new Token("access-token", "Bearer", "refresh-token", pastTime); // Verify it's marked as expired assertTrue(expiredToken.isExpired(), "Token should be expired"); // Create a valid token - LocalDateTime futureTime = LocalDateTime.now().plusMinutes(30); + Instant futureTime = Instant.now().plusSeconds(1800); Token validToken = new Token("access-token", "Bearer", "refresh-token", futureTime); // Verify it's not marked as expired @@ -86,8 +86,8 @@ void testNullPathRejection() { @Test void testOverwriteToken() { // Given two tokens saved in sequence - Token token1 = new Token("token1", "Bearer", "refresh1", LocalDateTime.now().plusHours(1)); - Token token2 = new Token("token2", "Bearer", "refresh2", LocalDateTime.now().plusHours(2)); + Token token1 = new Token("token1", "Bearer", "refresh1", Instant.now().plusSeconds(3600)); + Token token2 = new Token("token2", "Bearer", "refresh2", Instant.now().plusSeconds(7200)); tokenCache.save(token1); tokenCache.save(token2); @@ -110,7 +110,7 @@ void testWithCustomPath(@TempDir Path tempDir) { // And a token Token testToken = new Token( - "test-access-token", "Bearer", "test-refresh-token", LocalDateTime.now().plusHours(1)); + "test-access-token", "Bearer", "test-refresh-token", Instant.now().plusSeconds(3600)); // When saving and loading cache.save(testToken); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthHeaderFactoryTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthHeaderFactoryTest.java index d0530b2c1..f0b83153c 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthHeaderFactoryTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthHeaderFactoryTest.java @@ -3,7 +3,7 @@ import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -25,7 +25,7 @@ public class OAuthHeaderFactoryTest { @Mock private TokenSource tokenSource; private static Stream provideTokenSourceTestCases() { - LocalDateTime expiry = LocalDateTime.now().plusHours(1); + Instant expiry = Instant.now().plusSeconds(3600); Token token = new Token(TOKEN_VALUE, TOKEN_TYPE, expiry); return Stream.of( @@ -57,7 +57,7 @@ public void testFromTokenSourceFactoryMethod( } private static Stream provideSuppliersTestCases() { - LocalDateTime expiry = LocalDateTime.now().plusHours(1); + Instant expiry = Instant.now().plusSeconds(3600); Token token = new Token(TOKEN_VALUE, TOKEN_TYPE, expiry); Map standardHeaders = new HashMap<>(); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java index 14eb3fa40..8d2d68fd4 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java @@ -6,7 +6,7 @@ import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.HeaderFactory; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.Map; import java.util.stream.Stream; import org.junit.jupiter.params.ParameterizedTest; @@ -19,7 +19,7 @@ class TokenSourceCredentialsProviderTest { private static final String TEST_AUTH_TYPE = "test-auth-type"; private static final String TEST_ACCESS_TOKEN_VALUE = "test-access-token"; private static final Token TEST_TOKEN = - new Token(TEST_ACCESS_TOKEN_VALUE, "Bearer", LocalDateTime.now().plusHours(1)); + new Token(TEST_ACCESS_TOKEN_VALUE, "Bearer", Instant.now().plusSeconds(3600)); /** Tests token retrieval scenarios */ @ParameterizedTest(name = "{0}") diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index 2d87a32c2..16b726222 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -4,7 +4,6 @@ import com.databricks.sdk.core.utils.FakeClockSupplier; import java.time.Instant; -import java.time.LocalDateTime; import java.time.ZoneId; import org.junit.jupiter.api.Test; @@ -13,20 +12,20 @@ class TokenTest { private static final String accessToken = "testAccessToken"; private static final String refreshToken = "testRefreshToken"; private static final String tokenType = "testTokenType"; - private final LocalDateTime currentLocalDateTime; + private final Instant currentInstant; private final FakeClockSupplier fakeClockSupplier; TokenTest() { Instant instant = Instant.parse("2023-10-18T12:00:00.00Z"); ZoneId zoneId = ZoneId.of("UTC"); fakeClockSupplier = new FakeClockSupplier(instant, zoneId); - currentLocalDateTime = LocalDateTime.now(fakeClockSupplier.getClock()); + currentInstant = Instant.now(fakeClockSupplier.getClock()); } @Test void createNonRefreshableToken() { Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusMinutes(5), fakeClockSupplier); + new Token(accessToken, tokenType, currentInstant.plusSeconds(300), fakeClockSupplier); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertNull(token.getRefreshToken()); @@ -40,7 +39,7 @@ void createRefreshableToken() { accessToken, tokenType, refreshToken, - currentLocalDateTime.plusMinutes(5), + currentInstant.plusSeconds(300), fakeClockSupplier); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); @@ -51,7 +50,7 @@ void createRefreshableToken() { @Test void tokenExpiryMoreThan40Seconds() { Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusSeconds(50), fakeClockSupplier); + new Token(accessToken, tokenType, currentInstant.plusSeconds(50), fakeClockSupplier); assertFalse(token.isExpired()); assertTrue(token.isValid()); } @@ -59,7 +58,7 @@ void tokenExpiryMoreThan40Seconds() { @Test void tokenExpiryLessThan40Seconds() { Token token = - new Token(accessToken, tokenType, currentLocalDateTime.plusSeconds(30), fakeClockSupplier); + new Token(accessToken, tokenType, currentInstant.plusSeconds(30), fakeClockSupplier); assertTrue(token.isExpired()); assertFalse(token.isValid()); } @@ -67,7 +66,7 @@ void tokenExpiryLessThan40Seconds() { @Test void expiredToken() { Token token = - new Token(accessToken, tokenType, currentLocalDateTime.minusSeconds(10), fakeClockSupplier); + new Token(accessToken, tokenType, currentInstant.minusSeconds(10), fakeClockSupplier); assertTrue(token.isExpired()); assertFalse(token.isValid()); } From 6f81b4cff985b9ad06665189c39eac630d1a5806 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 4 Jun 2025 08:51:23 +0000 Subject: [PATCH 15/47] Update parseExpiry in CilTokenSource --- .../main/java/com/databricks/sdk/core/CliTokenSource.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 3b3845e51..9eec8930f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -41,8 +41,8 @@ public CliTokenSource( /** * Parses an expiry time string and returns the corresponding {@link Instant}. * - *

The expiry time string is verified to always be in UTC. Any time zone or offset information - * present in the input is ignored, and the value is parsed as a UTC time. + *

The expiry time string is always in UTC. Any time zone or offset information present in the + * input is ignored. * *

The method attempts to parse the input using several common date-time formats, including * ISO-8601 and patterns with varying sub-second precision. @@ -53,6 +53,7 @@ public CliTokenSource( */ static Instant parseExpiry(String expiry) { DateTimeParseException lastException = null; + // Try to parse the expiry as an ISO-8601 string in UTC first try { return Instant.parse(expiry); } catch (DateTimeParseException e) { From daac1b2f8ec1935c9ee4be4553fc615ac44915ff Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 4 Jun 2025 09:02:34 +0000 Subject: [PATCH 16/47] Update javadoc --- .../src/main/java/com/databricks/sdk/core/CliTokenSource.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 9eec8930f..d07bbf787 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -49,7 +49,7 @@ public CliTokenSource( * * @param expiry the expiry time string to parse, which must represent a UTC time * @return the parsed {@link Instant} representing the expiry time in UTC - * @throws DateTimeParseException if the input string cannot be parsed as a valid date-time + * @throws DateTimeParseException if the input string cannot be parsed */ static Instant parseExpiry(String expiry) { DateTimeParseException lastException = null; From f3d4b8a669d12f6bb3ffb56054e71c7f6cdb7eff Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 4 Jun 2025 09:03:51 +0000 Subject: [PATCH 17/47] Update javadoc --- .../src/main/java/com/databricks/sdk/core/CliTokenSource.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index d07bbf787..2e176c210 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -48,7 +48,7 @@ public CliTokenSource( * ISO-8601 and patterns with varying sub-second precision. * * @param expiry the expiry time string to parse, which must represent a UTC time - * @return the parsed {@link Instant} representing the expiry time in UTC + * @return the parsed {@link Instant} * @throws DateTimeParseException if the input string cannot be parsed */ static Instant parseExpiry(String expiry) { From 12123a957402ab4bc492de71cb240b7842211a84 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 5 Jun 2025 09:15:03 +0000 Subject: [PATCH 18/47] Retrigger tests From 8705ff53e7017d2212acb839edc2efdec6683410 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 6 Jun 2025 08:55:42 +0000 Subject: [PATCH 19/47] Save progress --- .../databricks/sdk/core/CliTokenSource.java | 6 +- .../databricks/sdk/core/DatabricksConfig.java | 8 +- ...reServicePrincipalCredentialsProvider.java | 6 +- .../sdk/core/oauth/CachedTokenSource.java | 220 ++++++++++++++ .../sdk/core/oauth/ClientCredentials.java | 7 +- .../databricks/sdk/core/oauth/Consent.java | 2 +- .../oauth/DatabricksOAuthTokenSource.java | 4 +- .../sdk/core/oauth/EndpointTokenSource.java | 4 +- .../ExternalBrowserCredentialsProvider.java | 2 +- .../sdk/core/oauth/OidcTokenSource.java | 5 +- .../core/oauth/RefreshableTokenSource.java | 287 +----------------- .../sdk/core/oauth/SessionCredentials.java | 10 +- .../sdk/core/oauth/TokenEndpointClient.java | 60 ++++ .../databricks/sdk/core/utils/AzureUtils.java | 2 +- ...ceTest.java => CachedTokenSourceTest.java} | 80 +++-- ...xternalBrowserCredentialsProviderTest.java | 4 +- 16 files changed, 348 insertions(+), 359 deletions(-) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/CachedTokenSource.java rename databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/{RefreshableTokenSourceTest.java => CachedTokenSourceTest.java} (64%) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 8a7328904..2eb5c13f7 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -1,7 +1,7 @@ package com.databricks.sdk.core; -import com.databricks.sdk.core.oauth.RefreshableTokenSource; import com.databricks.sdk.core.oauth.Token; +import com.databricks.sdk.core.oauth.TokenSource; import com.databricks.sdk.core.utils.Environment; import com.databricks.sdk.core.utils.OSUtils; import com.fasterxml.jackson.databind.JsonNode; @@ -15,7 +15,7 @@ import java.util.List; import org.apache.commons.io.IOUtils; -public class CliTokenSource extends RefreshableTokenSource { +public class CliTokenSource implements TokenSource { private List cmd; private String tokenTypeField; private String accessTokenField; @@ -64,7 +64,7 @@ private String getProcessStream(InputStream stream) throws IOException { } @Override - protected Token refresh() { + public Token getToken() { try { ProcessBuilder processBuilder = new ProcessBuilder(cmd); processBuilder.environment().putAll(env.getEnv()); 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..df16ebae3 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 @@ -410,13 +410,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** @deprecated Use {@link #getAzureUseMsi()} instead. */ + /** + * @deprecated Use {@link #getAzureUseMsi()} instead. + */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ + /** + * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. + */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; 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 c7c7bb672..d10fefd94 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 @@ -28,8 +28,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } AzureUtils.ensureHostPresent( config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor); - RefreshableTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); - RefreshableTokenSource cloud = + TokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); + TokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); return OAuthHeaderFactory.fromSuppliers( @@ -55,7 +55,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { * @return A RefreshableTokenSource instance capable of fetching OAuth tokens for the specified * Azure resource. */ - private static RefreshableTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + private static TokenSource tokenSourceFor(DatabricksConfig config, String resource) { String aadEndpoint = config.getAzureEnvironment().getActiveDirectoryEndpoint(); String tokenUrl = aadEndpoint + config.getAzureTenantId() + "/oauth2/token"; Map endpointParams = new HashMap<>(); 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 new file mode 100644 index 000000000..72aa38f74 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/CachedTokenSource.java @@ -0,0 +1,220 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.utils.ClockSupplier; +import com.databricks.sdk.core.utils.SystemClockSupplier; +import java.time.Duration; +import java.time.LocalDateTime; +import java.util.concurrent.CompletableFuture; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * An OAuth TokenSource which can be refreshed. + * + *

This class supports both synchronous and asynchronous token refresh. When async is enabled, + * stale tokens will trigger a background refresh, while expired tokens will block until a new token + * is fetched. + */ +public class CachedTokenSource implements TokenSource { + + /** + * Enum representing the state of the token. FRESH: Token is valid and not close to expiry. STALE: + * Token is valid but will expire soon - an async refresh will be triggered if enabled. EXPIRED: + * Token has expired and must be refreshed using a blocking call. + */ + private enum TokenState { + FRESH, + STALE, + EXPIRED + } + + private static final Logger logger = LoggerFactory.getLogger(RefreshableTokenSource.class); + // Default duration before expiry to consider a token as 'stale'. + private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); + private static final Duration DEFAULT_EXPIRY_BUFFER = Duration.ofSeconds(40); + + private final TokenSource tokenSource; + private final boolean asyncEnabled; + private final Duration staleDuration; + private final Duration expiryBuffer; + private final ClockSupplier clockSupplier; + + // The current OAuth token. May be null if not yet fetched. + private volatile Token token; + // Whether a refresh is currently in progress (for async refresh). + private volatile boolean refreshInProgress = false; + // Whether the last refresh attempt succeeded. + private volatile boolean lastRefreshSucceeded = true; + + private CachedTokenSource(Builder builder) { + this.tokenSource = builder.tokenSource; + this.asyncEnabled = builder.asyncEnabled; + this.staleDuration = builder.staleDuration; + this.expiryBuffer = builder.expiryBuffer; + this.clockSupplier = builder.clockSupplier; + this.token = builder.token; + } + + public static class Builder { + private final TokenSource tokenSource; + private Token token; + private boolean asyncEnabled = false; + private Duration staleDuration = DEFAULT_STALE_DURATION; + private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER; + private ClockSupplier clockSupplier = new SystemClockSupplier(); + + public Builder(TokenSource tokenSource) { + this.tokenSource = tokenSource; + } + + public Builder withToken(Token token) { + this.token = token; + return this; + } + + public Builder withAsyncEnabled(boolean asyncEnabled) { + this.asyncEnabled = asyncEnabled; + return this; + } + + public Builder withStaleDuration(Duration staleDuration) { + this.staleDuration = staleDuration; + return this; + } + + public Builder withExpiryBuffer(Duration expiryBuffer) { + this.expiryBuffer = expiryBuffer; + return this; + } + + public Builder withClockSupplier(ClockSupplier clockSupplier) { + this.clockSupplier = clockSupplier; + return this; + } + + public CachedTokenSource build() { + return new CachedTokenSource(this); + } + } + + /** + * Gets the current token, refreshing if necessary. If async refresh is enabled, may return a + * stale token while a refresh is in progress. + * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. + * + * @return The current valid token + */ + public Token getToken() { + if (!asyncEnabled) { + return getTokenBlocking(); + } + return getTokenAsync(); + } + + /** + * Determine the state of the current token (fresh, stale, or expired). + * + * @return The token state + */ + protected TokenState getTokenState(Token t) { + if (t == null) { + return TokenState.EXPIRED; + } + Duration lifeTime = + Duration.between(LocalDateTime.now(clockSupplier.getClock()), t.getExpiry()); + if (lifeTime.compareTo(expiryBuffer) <= 0) { + return TokenState.EXPIRED; + } + if (lifeTime.compareTo(staleDuration) <= 0) { + return TokenState.STALE; + } + return TokenState.FRESH; + } + + /** + * Get the current token, blocking to refresh if expired. + * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. + * + * @return The current valid token + */ + protected Token getTokenBlocking() { + // Use double-checked locking to minimize synchronization overhead on reads: + // 1. Check if the token is expired without locking. + // 2. If expired, synchronize and check again (another thread may have refreshed it). + // 3. If still expired, perform the refresh. + if (getTokenState(token) != TokenState.EXPIRED) { + return token; + } + synchronized (this) { + if (getTokenState(token) != TokenState.EXPIRED) { + return token; + } + lastRefreshSucceeded = false; + try { + token = tokenSource.getToken(); + } catch (Exception e) { + logger.error("Failed to refresh token synchronously", e); + throw e; + } + lastRefreshSucceeded = true; + return token; + } + } + + /** + * Get the current token, possibly triggering an async refresh if stale. If the token is expired, + * blocks to refresh. + * + *

This method may throw an exception if the token cannot be refreshed, depending on the + * implementation of {@link #refresh()}. + * + * @return The current valid or stale token + */ + protected Token getTokenAsync() { + Token currentToken = token; + + switch (getTokenState(currentToken)) { + case FRESH: + return currentToken; + case STALE: + triggerAsyncRefresh(); + return currentToken; + case EXPIRED: + return getTokenBlocking(); + default: + throw new IllegalStateException("Invalid token state."); + } + } + + /** + * Trigger an asynchronous refresh of the token if not already in progress and last refresh + * succeeded. + */ + protected synchronized void triggerAsyncRefresh() { + // Check token state to avoid triggering a refresh if another thread has already refreshed it + if (!refreshInProgress && lastRefreshSucceeded && getTokenState(token) != TokenState.FRESH) { + refreshInProgress = true; + CompletableFuture.runAsync( + () -> { + try { + // Attempt to refresh the token in the background + Token newToken = tokenSource.getToken(); + synchronized (this) { + token = newToken; + refreshInProgress = false; + } + } catch (Exception e) { + synchronized (this) { + lastRefreshSucceeded = false; + refreshInProgress = false; + logger.error("Async token refresh failed", e); + } + } + }); + } + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java index 1c4b7d6de..8cee3ef29 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ClientCredentials.java @@ -13,7 +13,7 @@ * support all OAuth endpoints, authentication parameters can be passed in the request body or in * the Authorization header. */ -public class ClientCredentials extends RefreshableTokenSource { +public class ClientCredentials implements TokenSource { public static class Builder { private String clientId; private String clientSecret; @@ -97,7 +97,7 @@ private ClientCredentials( } @Override - protected Token refresh() { + public Token getToken() { Map params = new HashMap<>(); params.put("grant_type", "client_credentials"); if (scopes != null) { @@ -106,6 +106,7 @@ protected Token refresh() { if (endpointParamsSupplier != null) { params.putAll(endpointParamsSupplier.get()); } - return retrieveToken(hc, clientId, clientSecret, tokenUrl, params, new HashMap<>(), position); + return TokenEndpointClient.retrieveToken( + hc, clientId, clientSecret, tokenUrl, params, new HashMap<>(), position); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 77045df97..68dc6f176 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -313,7 +313,7 @@ public SessionCredentials exchange(String code, String state) { headers.put("Origin", this.redirectUrl); } Token token = - RefreshableTokenSource.retrieveToken( + TokenEndpointClient.retrieveToken( this.hc, this.clientId, this.clientSecret, diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java index f16ae2aed..ac76234a0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java @@ -14,7 +14,7 @@ * Implementation of TokenSource that handles OAuth token exchange for Databricks authentication. * This class manages the OAuth token exchange flow using ID tokens to obtain access tokens. */ -public class DatabricksOAuthTokenSource extends RefreshableTokenSource { +public class DatabricksOAuthTokenSource implements TokenSource { private static final Logger LOG = LoggerFactory.getLogger(DatabricksOAuthTokenSource.class); /** OAuth client ID used for token exchange. */ @@ -128,7 +128,7 @@ public DatabricksOAuthTokenSource build() { * @throws NullPointerException when any of the required parameters are null. */ @Override - public Token refresh() { + public Token getToken() { Objects.requireNonNull(clientId, "ClientID cannot be null"); Objects.requireNonNull(host, "Host cannot be null"); Objects.requireNonNull(endpoints, "Endpoints cannot be null"); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java index 3ca75c441..7bcddacdd 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EndpointTokenSource.java @@ -13,7 +13,7 @@ * Represents a token source that exchanges a control plane token for an endpoint-specific dataplane * token. It utilizes an underlying {@link TokenSource} to obtain the initial control plane token. */ -public class EndpointTokenSource extends RefreshableTokenSource { +public class EndpointTokenSource implements TokenSource { private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class); private static final String JWT_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"; private static final String GRANT_TYPE_PARAM = "grant_type"; @@ -67,7 +67,7 @@ public EndpointTokenSource( * @throws NullPointerException if any of the parameters are null. */ @Override - protected Token refresh() { + public Token getToken() { Token cpToken = cpTokenSource.getToken(); Map params = new HashMap<>(); params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE); 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 7bae60022..494ad69cd 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 @@ -79,7 +79,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { .build(); LOGGER.debug("Using cached token, will immediately refresh"); - cachedCreds.token = cachedCreds.refresh(); + cachedCreds.token = cachedCreds.getToken(); return cachedCreds.configure(config); } catch (Exception e) { // If token refresh fails, log and continue to browser auth diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java index 719544ebf..1739c5c6c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OidcTokenSource.java @@ -15,7 +15,7 @@ * protocol. It communicates with an OAuth server to request access tokens using the client * credentials grant type instead of a client secret. */ -class OidcTokenSource extends RefreshableTokenSource { +class OidcTokenSource implements TokenSource { private final HttpClient httpClient; private final String tokenUrl; @@ -58,7 +58,8 @@ private static void putIfDefined( } } - protected Token refresh() { + @Override + public Token getToken() { Response rawResp; try { rawResp = httpClient.execute(new FormRequest(tokenUrl, params)); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java index c62dccad7..c0cf080ce 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java @@ -1,22 +1,5 @@ package com.databricks.sdk.core.oauth; -import com.databricks.sdk.core.ApiClient; -import com.databricks.sdk.core.DatabricksException; -import com.databricks.sdk.core.http.FormRequest; -import com.databricks.sdk.core.http.HttpClient; -import com.databricks.sdk.core.http.Request; -import com.databricks.sdk.core.utils.ClockSupplier; -import com.databricks.sdk.core.utils.SystemClockSupplier; -import java.time.Duration; -import java.time.LocalDateTime; -import java.time.temporal.ChronoUnit; -import java.util.Base64; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import org.apache.http.HttpHeaders; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - /** * An OAuth TokenSource which can be refreshed. * @@ -24,272 +7,4 @@ * stale tokens will trigger a background refresh, while expired tokens will block until a new token * is fetched. */ -public abstract class RefreshableTokenSource implements TokenSource { - - /** - * Enum representing the state of the token. FRESH: Token is valid and not close to expiry. STALE: - * Token is valid but will expire soon - an async refresh will be triggered if enabled. EXPIRED: - * Token has expired and must be refreshed using a blocking call. - */ - private enum TokenState { - FRESH, - STALE, - EXPIRED - } - - private static final Logger logger = LoggerFactory.getLogger(RefreshableTokenSource.class); - // Default duration before expiry to consider a token as 'stale'. - private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); - - // The current OAuth token. May be null if not yet fetched. - protected volatile Token token; - // Whether asynchronous refresh is enabled. - private boolean asyncEnabled = false; - // Duration before expiry to consider a token as 'stale'. - private Duration staleDuration = DEFAULT_STALE_DURATION; - // Additional buffer before expiry to consider a token as expired. - private Duration expiryBuffer = Duration.ofSeconds(40); - // Whether a refresh is currently in progress (for async refresh). - private boolean refreshInProgress = false; - // Whether the last refresh attempt succeeded. - private boolean lastRefreshSucceeded = true; - // Clock supplier for current time, for testing purposes. - private ClockSupplier clockSupplier = new SystemClockSupplier(); - - /** Constructs a new {@code RefreshableTokenSource} with no initial token. */ - public RefreshableTokenSource() {} - - /** - * Constructor with initial token. - * - * @param token The initial token to use. - */ - public RefreshableTokenSource(Token token) { - this.token = token; - } - - /** - * Set the clock supplier for current time. - * - *

Experimental: This method may change or be removed in future releases. - * - * @param clockSupplier The clock supplier to use. - * @return this instance for chaining - */ - public RefreshableTokenSource withClockSupplier(ClockSupplier clockSupplier) { - this.clockSupplier = clockSupplier; - return this; - } - - /** - * Enable or disable asynchronous token refresh. - * - *

Experimental: This method may change or be removed in future releases. - * - * @param enabled true to enable async refresh, false to disable - * @return this instance for chaining - */ - public RefreshableTokenSource withAsyncRefresh(boolean enabled) { - this.asyncEnabled = enabled; - return this; - } - - /** - * Set the expiry buffer. If the token's lifetime is less than this buffer, it is considered - * expired. - * - *

Experimental: This method may change or be removed in future releases. - * - * @param buffer the expiry buffer duration - * @return this instance for chaining - */ - public RefreshableTokenSource withExpiryBuffer(Duration buffer) { - this.expiryBuffer = buffer; - return this; - } - - /** - * Refresh the OAuth token. Subclasses must implement this to define how the token is refreshed. - * - *

This method may throw an exception if the token cannot be refreshed. The specific exception - * type depends on the implementation. - * - * @return The newly refreshed Token. - */ - protected abstract Token refresh(); - - /** - * Gets the current token, refreshing if necessary. If async refresh is enabled, may return a - * stale token while a refresh is in progress. - * - *

This method may throw an exception if the token cannot be refreshed, depending on the - * implementation of {@link #refresh()}. - * - * @return The current valid token - */ - public Token getToken() { - if (!asyncEnabled) { - return getTokenBlocking(); - } - return getTokenAsync(); - } - - /** - * Determine the state of the current token (fresh, stale, or expired). - * - * @return The token state - */ - protected TokenState getTokenState(Token t) { - if (t == null) { - return TokenState.EXPIRED; - } - Duration lifeTime = - Duration.between(LocalDateTime.now(clockSupplier.getClock()), t.getExpiry()); - if (lifeTime.compareTo(expiryBuffer) <= 0) { - return TokenState.EXPIRED; - } - if (lifeTime.compareTo(staleDuration) <= 0) { - return TokenState.STALE; - } - return TokenState.FRESH; - } - - /** - * Get the current token, blocking to refresh if expired. - * - *

This method may throw an exception if the token cannot be refreshed, depending on the - * implementation of {@link #refresh()}. - * - * @return The current valid token - */ - protected Token getTokenBlocking() { - // Use double-checked locking to minimize synchronization overhead on reads: - // 1. Check if the token is expired without locking. - // 2. If expired, synchronize and check again (another thread may have refreshed it). - // 3. If still expired, perform the refresh. - if (getTokenState(token) != TokenState.EXPIRED) { - return token; - } - synchronized (this) { - if (getTokenState(token) != TokenState.EXPIRED) { - return token; - } - lastRefreshSucceeded = false; - try { - token = refresh(); - } catch (Exception e) { - logger.error("Failed to refresh token synchronously", e); - throw e; - } - lastRefreshSucceeded = true; - return token; - } - } - - /** - * Get the current token, possibly triggering an async refresh if stale. If the token is expired, - * blocks to refresh. - * - *

This method may throw an exception if the token cannot be refreshed, depending on the - * implementation of {@link #refresh()}. - * - * @return The current valid or stale token - */ - protected Token getTokenAsync() { - Token currentToken = token; - - switch (getTokenState(currentToken)) { - case FRESH: - return currentToken; - case STALE: - triggerAsyncRefresh(); - return currentToken; - case EXPIRED: - return getTokenBlocking(); - default: - throw new IllegalStateException("Invalid token state."); - } - } - - /** - * Trigger an asynchronous refresh of the token if not already in progress and last refresh - * succeeded. - */ - protected synchronized void triggerAsyncRefresh() { - // Check token state to avoid triggering a refresh if another thread has already refreshed it - if (!refreshInProgress && lastRefreshSucceeded && getTokenState(token) != TokenState.FRESH) { - refreshInProgress = true; - CompletableFuture.runAsync( - () -> { - try { - // Attempt to refresh the token in the background - Token newToken = refresh(); - synchronized (this) { - token = newToken; - refreshInProgress = false; - } - } catch (Exception e) { - synchronized (this) { - lastRefreshSucceeded = false; - refreshInProgress = false; - logger.error("Async token refresh failed", e); - } - } - }); - } - } - - /** - * Helper method implementing OAuth token refresh. - * - * @param hc The HTTP client to use for the request. - * @param clientId The client ID to authenticate with. - * @param clientSecret The client secret to authenticate with. - * @param tokenUrl The authorization URL for fetching tokens. - * @param params Additional request parameters. - * @param headers Additional headers. - * @param position The position of the authentication parameters in the request. - * @return The newly fetched Token. - * @throws DatabricksException if the refresh fails - * @throws IllegalArgumentException if the OAuth response contains an error - */ - protected static Token retrieveToken( - HttpClient hc, - String clientId, - String clientSecret, - String tokenUrl, - Map params, - Map headers, - AuthParameterPosition position) { - switch (position) { - case BODY: - if (clientId != null) { - params.put("client_id", clientId); - } - if (clientSecret != null) { - params.put("client_secret", clientSecret); - } - break; - case HEADER: - String authHeaderValue = - "Basic " - + Base64.getEncoder().encodeToString((clientId + ":" + clientSecret).getBytes()); - headers.put(HttpHeaders.AUTHORIZATION, authHeaderValue); - break; - } - headers.put("Content-Type", "application/x-www-form-urlencoded"); - Request req = new Request("POST", tokenUrl, FormRequest.wrapValuesInList(params)); - req.withHeaders(headers); - try { - ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build(); - OAuthResponse resp = apiClient.execute(req, OAuthResponse.class); - if (resp.getErrorCode() != null) { - throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); - } - LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); - return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); - } catch (Exception e) { - throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); - } - } -} +public abstract class RefreshableTokenSource implements TokenSource {} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index 4d2d512e3..3bde86994 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -17,8 +17,7 @@ * requests to an API, and a long-lived refresh token, which can be used to fetch new access tokens. * Calling refresh() uses the refresh token to retrieve a new access token to authenticate to APIs. */ -public class SessionCredentials extends RefreshableTokenSource - implements CredentialsProvider, Serializable { +public class SessionCredentials implements TokenSource, CredentialsProvider, Serializable { private static final long serialVersionUID = 3083941540130596650L; private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentials.class); @@ -87,9 +86,10 @@ public SessionCredentials build() { private final String clientId; private final String clientSecret; private final TokenCache tokenCache; + protected Token token; private SessionCredentials(Builder b) { - super(b.token); + this.token = b.token; this.hc = b.hc; this.tokenUrl = b.tokenUrl; this.redirectUrl = b.redirectUrl; @@ -99,7 +99,7 @@ private SessionCredentials(Builder b) { } @Override - protected Token refresh() { + public Token getToken() { if (this.token == null) { throw new DatabricksException("oauth2: token is not set"); } @@ -118,7 +118,7 @@ protected Token refresh() { headers.put("Origin", redirectUrl); } Token newToken = - retrieveToken( + TokenEndpointClient.retrieveToken( hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); // Save the refreshed token directly to cache diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java index 69883dd24..ace0c9314 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java @@ -1,5 +1,6 @@ package com.databricks.sdk.core.oauth; +import com.databricks.sdk.core.ApiClient; import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; @@ -88,4 +89,63 @@ public static OAuthResponse requestToken( LOG.debug("Successfully obtained token response from {}", tokenEndpointUrl); return response; } + + /** + * Helper method implementing OAuth token refresh. + * + * @param hc The HTTP client to use for the request. + * @param clientId The client ID to authenticate with. + * @param clientSecret The client secret to authenticate with. + * @param tokenUrl The authorization URL for fetching tokens. + * @param params Additional request parameters. + * @param headers Additional headers. + * @param position The position of the authentication parameters in the request. + * @return The newly fetched Token. + * @throws DatabricksException if the refresh fails + * @throws IllegalArgumentException if the OAuth response contains an error + */ + public static Token retrieveToken( + HttpClient hc, + String clientId, + String clientSecret, + String tokenUrl, + Map params, + Map headers, + AuthParameterPosition position) { + switch (position) { + case BODY: + if (clientId != null) { + params.put("client_id", clientId); + } + if (clientSecret != null) { + params.put("client_secret", clientSecret); + } + break; + case HEADER: + String authHeaderValue = + "Basic " + + java.util.Base64.getEncoder() + .encodeToString((clientId + ":" + clientSecret).getBytes()); + headers.put(org.apache.http.HttpHeaders.AUTHORIZATION, authHeaderValue); + break; + } + headers.put("Content-Type", "application/x-www-form-urlencoded"); + com.databricks.sdk.core.http.Request req = + new com.databricks.sdk.core.http.Request( + "POST", tokenUrl, com.databricks.sdk.core.http.FormRequest.wrapValuesInList(params)); + req.withHeaders(headers); + try { + ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build(); + OAuthResponse resp = apiClient.execute(req, OAuthResponse.class); + if (resp.getErrorCode() != null) { + throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); + } + java.time.LocalDateTime expiry = + java.time.LocalDateTime.now() + .plus(resp.getExpiresIn(), java.time.temporal.ChronoUnit.SECONDS); + return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); + } catch (Exception e) { + throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); + } + } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java index 96dc116c2..09cea6e86 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/AzureUtils.java @@ -77,7 +77,7 @@ public static Map addWorkspaceResourceId( } public static Map addSpManagementToken( - RefreshableTokenSource tokenSource, Map headers) { + TokenSource tokenSource, Map headers) { headers.put("X-Databricks-Azure-SP-Management-Token", tokenSource.getToken().getAccessToken()); return headers; } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java similarity index 64% rename from databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java rename to databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java index c33e7fda5..ba9a41753 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/RefreshableTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/CachedTokenSourceTest.java @@ -11,7 +11,7 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -public class RefreshableTokenSourceTest { +public class CachedTokenSourceTest { private static final String TOKEN_TYPE = "Bearer"; private static final String INITIAL_TOKEN = "initial-token"; private static final String REFRESH_TOKEN = "refreshed-token"; @@ -46,10 +46,10 @@ void testAsyncRefreshParametrized( new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); CountDownLatch refreshCalled = new CountDownLatch(1); - RefreshableTokenSource source = - new RefreshableTokenSource(initialToken) { + TokenSource tokenSource = + new TokenSource() { @Override - protected Token refresh() { + public Token getToken() { refreshCalled.countDown(); try { Thread.sleep(500); @@ -58,7 +58,13 @@ protected Token refresh() { } return refreshedToken; } - }.withAsyncRefresh(asyncEnabled); + }; + + CachedTokenSource source = + new CachedTokenSource.Builder(tokenSource) + .withAsyncEnabled(asyncEnabled) + .withToken(initialToken) + .build(); Token token = source.getToken(); @@ -67,74 +73,56 @@ protected Token refresh() { assertEquals(expectedToken, token.getAccessToken(), "Token value did not match expected"); } - /** - * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is - * forced to be synchronous. It ensures that after an async failure, the system does not - * repeatedly attempt async refreshes while the token is stale, and only performs a synchronous - * refresh when the token is expired. After a successful sync refresh, async refreshes resume as - * normal. - */ @Test void testAsyncRefreshFailureFallback() throws Exception { - Token staleToken = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); - - class TestSource extends RefreshableTokenSource { + class MutableTokenSource implements TokenSource { int refreshCallCount = 0; boolean isFirstRefresh = true; - - TestSource(Token token) { - super(token); - } + Token token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); @Override - protected Token refresh() { + public Token getToken() { refreshCallCount++; if (isFirstRefresh) { isFirstRefresh = false; throw new RuntimeException("Simulated async failure"); } - return new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); + token = new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); + return token; } } - TestSource source = new TestSource(staleToken); - source.withAsyncRefresh(true); + MutableTokenSource mutableTokenSource = new MutableTokenSource(); + + CachedTokenSource source = + new CachedTokenSource.Builder(mutableTokenSource) + .withAsyncEnabled(true) + .withToken(mutableTokenSource.token) + .build(); // First call triggers async refresh, which fails source.getToken(); - Thread.sleep(300); - assertEquals( - 1, source.refreshCallCount, "refresh() should have been called once (async, failed)"); + assertEquals(1, mutableTokenSource.refreshCallCount); // Token is still stale, so next call should NOT trigger another refresh since the last refresh // failed source.getToken(); - Thread.sleep(200); - assertEquals( - 1, - source.refreshCallCount, - "refresh() should NOT be called again while stale after async failure"); + assertEquals(1, mutableTokenSource.refreshCallCount); // Advance the clock so the token is now expired - source.token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().minusMinutes(1)); + mutableTokenSource.token = + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().minusMinutes(1)); // Now getToken() should call refresh synchronously and return the refreshed token - Token token = source.getToken(); - assertEquals( - REFRESH_TOKEN, - token.getAccessToken(), - "Should return the refreshed token after sync refresh"); - assertEquals( - 2, source.refreshCallCount, "refresh() should have been called synchronously after expiry"); + Token token; + token = source.getToken(); + assertEquals(REFRESH_TOKEN, token.getAccessToken()); + assertEquals(2, mutableTokenSource.refreshCallCount); // Make the token stale again and trigger async refresh since the last sync refresh succeeded - source.token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); + mutableTokenSource.token = + new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); source.getToken(); - Thread.sleep(300); - assertEquals( - 3, - source.refreshCallCount, - "refresh() should have been called again asynchronously after making token stale"); + assertEquals(3, mutableTokenSource.refreshCallCount); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 1714b731c..e8f6faa2f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -186,7 +186,7 @@ void clientCredentials() throws IOException { .withClientSecret("abc") .withTokenUrl("https://tokenUrl") .build(); - Token token = clientCredentials.refresh(); + Token token = clientCredentials.getToken(); assertEquals("accessTokenFromServer", token.getAccessToken()); assertEquals("refreshTokenFromServer", token.getRefreshToken()); } @@ -212,7 +212,7 @@ void sessionCredentials() throws IOException { "originalRefreshToken", LocalDateTime.MAX)) .build(); - Token token = sessionCredentials.refresh(); + Token token = sessionCredentials.getToken(); // We check that we are actually getting the token from server response (that is defined // above) rather than what was given while creating session credentials From fdc50effa6c3d0c22bffe247e55ee3bb2acb18c8 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 6 Jun 2025 11:00:55 +0000 Subject: [PATCH 20/47] Removed redundant date formattters --- .../databricks/sdk/core/CliTokenSource.java | 42 ++-------- .../sdk/core/CliTokenSourceTest.java | 79 ++++++++----------- .../src/test/resources/testdata/az | 8 +- 3 files changed, 45 insertions(+), 84 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 2e176c210..0e68dc7fc 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -9,11 +9,8 @@ import java.io.IOException; import java.io.InputStream; import java.time.Instant; -import java.time.LocalDateTime; -import java.time.ZoneOffset; -import java.time.format.DateTimeFormatter; +import java.time.OffsetDateTime; import java.time.format.DateTimeParseException; -import java.util.Arrays; import java.util.List; import org.apache.commons.io.IOUtils; @@ -39,44 +36,15 @@ public CliTokenSource( } /** - * Parses an expiry time string and returns the corresponding {@link Instant}. + * Parses an expiry string in RFC 3339/ISO 8601 format (with or without offset) and returns the + * corresponding {@link Instant}. Any specified time zone or offset is converted to UTC. * - *

The expiry time string is always in UTC. Any time zone or offset information present in the - * input is ignored. - * - *

The method attempts to parse the input using several common date-time formats, including - * ISO-8601 and patterns with varying sub-second precision. - * - * @param expiry the expiry time string to parse, which must represent a UTC time + * @param expiry expiry time string in RFC 3339/ISO 8601 format * @return the parsed {@link Instant} * @throws DateTimeParseException if the input string cannot be parsed */ static Instant parseExpiry(String expiry) { - DateTimeParseException lastException = null; - // Try to parse the expiry as an ISO-8601 string in UTC first - try { - return Instant.parse(expiry); - } catch (DateTimeParseException e) { - lastException = e; - } - - String multiplePrecisionPattern = - "[SSSSSSSSS][SSSSSSSS][SSSSSSS][SSSSSS][SSSSS][SSSS][SSS][SS][S]"; - List datePatterns = - Arrays.asList( - "yyyy-MM-dd HH:mm:ss", - "yyyy-MM-dd HH:mm:ss." + multiplePrecisionPattern, - "yyyy-MM-dd'T'HH:mm:ss." + multiplePrecisionPattern + "XXX"); - for (String pattern : datePatterns) { - try { - DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern); - LocalDateTime dateTime = LocalDateTime.parse(expiry, formatter); - return dateTime.atZone(ZoneOffset.UTC).toInstant(); - } catch (DateTimeParseException e) { - lastException = e; - } - } - throw lastException; + return OffsetDateTime.parse(expiry).toInstant(); } private String getProcessStream(InputStream stream) throws IOException { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java index abe609b01..cc177881c 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java @@ -1,58 +1,49 @@ package com.databricks.sdk.core; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.time.Instant; import java.time.format.DateTimeParseException; -import org.junit.jupiter.api.Test; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; public class CliTokenSourceTest { - @Test - public void testParseExpiryWithoutTruncate() { - Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612218Z"); - assertEquals(Instant.parse("2023-07-17T09:02:22.330612218Z"), parsedInstant); + private static Stream expiryProvider() { + return Stream.of( + Arguments.of( + "2023-07-17T09:02:22.330612218Z", "2023-07-17T09:02:22.330612218Z", "9-digit nanos"), + Arguments.of( + "2023-07-17T09:02:22.33061221Z", "2023-07-17T09:02:22.330612210Z", "8-digit nanos"), + Arguments.of( + "2023-07-17T09:02:22.330612Z", "2023-07-17T09:02:22.330612000Z", "6-digit nanos"), + Arguments.of( + "2023-07-17T10:02:22.330612218+01:00", + "2023-07-17T09:02:22.330612218Z", + "+01:00 offset, 9-digit nanos"), + Arguments.of( + "2023-07-17T04:02:22.330612218-05:00", + "2023-07-17T09:02:22.330612218Z", + "-05:00 offset, 9-digit nanos"), + Arguments.of( + "2023-07-17T10:02:22.330612+01:00", + "2023-07-17T09:02:22.330612000Z", + "+01:00 offset, 6-digit nanos"), + Arguments.of("2023-07-17T09:02:22.33061221987Z", null, "Invalid: >9 nanos"), + Arguments.of("17-07-2023 09:02:22", null, "Invalid: non-ISO8601 format")); } - @Test - public void testParseExpiryWithTruncate() { - Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221Z"); - assertEquals(Instant.parse("2023-07-17T09:02:22.330612210Z"), parsedInstant); - } - - @Test - public void testParseExpiryWithTruncateAndLessNanoSecondDigits() { - Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17T09:02:22.330612Z"); - assertEquals(Instant.parse("2023-07-17T09:02:22.330612000Z"), parsedInstant); - } - - @Test - public void testParseExpiryWithMoreThanNineNanoSecondDigits() { - try { - CliTokenSource.parseExpiry("2023-07-17T09:02:22.33061221987Z"); - } catch (DateTimeParseException e) { - assert (e.getMessage().contains("could not be parsed")); - } - } - - @Test - public void testParseExpiryWithSpaceFormat() { - Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22"); - assertEquals(Instant.parse("2023-07-17T09:02:22Z"), parsedInstant); - } - - @Test - public void testParseExpiryWithSpaceFormatAndMillis() { - Instant parsedInstant = CliTokenSource.parseExpiry("2023-07-17 09:02:22.123"); - assertEquals(Instant.parse("2023-07-17T09:02:22.123Z"), parsedInstant); - } - - @Test - public void testParseExpiryWithInvalidFormat() { - try { - CliTokenSource.parseExpiry("17-07-2023 09:02:22"); - } catch (DateTimeParseException e) { - assert (e.getMessage().contains("could not be parsed")); + @ParameterizedTest(name = "{2}") + @MethodSource("expiryProvider") + public void testParseExpiry(String input, String expectedInstant, String description) { + if (expectedInstant == null) { + assertThrows(DateTimeParseException.class, () -> CliTokenSource.parseExpiry(input)); + } else { + Instant parsedInstant = CliTokenSource.parseExpiry(input); + assertEquals(Instant.parse(expectedInstant), parsedInstant); } } } diff --git a/databricks-sdk-java/src/test/resources/testdata/az b/databricks-sdk-java/src/test/resources/testdata/az index 29b824ed7..b24bddea9 100755 --- a/databricks-sdk-java/src/test/resources/testdata/az +++ b/databricks-sdk-java/src/test/resources/testdata/az @@ -22,14 +22,16 @@ for arg in "$@"; do fi done -# Macos -EXP="$(/bin/date -v+${EXPIRE:=10S} +'%F %T' 2>/dev/null)" +# MacOS +EXP="$(/bin/date -v+${EXPIRE:=10S} +'%FT%T%z' 2>/dev/null)" if [ -z "${EXP}" ]; then # Linux EXPIRE=$(/bin/echo $EXPIRE | /bin/sed 's/S/seconds/') EXPIRE=$(/bin/echo $EXPIRE | /bin/sed 's/M/minutes/') - EXP=$(/bin/date --date=+${EXPIRE:=10seconds} +'%F %T') + EXP=$(/bin/date --date=+${EXPIRE:=10seconds} +'%FT%T%z') fi +# Insert colon in timezone offset for ISO 8601 compliance +EXP="${EXP:0:19}${EXP:19:3}:${EXP:22:2}" if [ -z "${TF_AAD_TOKEN}" ]; then TF_AAD_TOKEN="..." From 70934b27ca4769d241bd3d02ff02e1d1e58c13d5 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 6 Jun 2025 14:08:29 +0000 Subject: [PATCH 21/47] Change clock supplier to use UTC time --- .../java/com/databricks/sdk/core/utils/SystemClockSupplier.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java index edac0cd94..c79e6884d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java @@ -5,6 +5,6 @@ public class SystemClockSupplier implements ClockSupplier { @Override public Clock getClock() { - return Clock.systemDefaultZone(); + return Clock.systemUTC(); } } From 8f19819e4b77c1e94e65e998cab592ed86ead20a Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sat, 7 Jun 2025 10:41:47 +0000 Subject: [PATCH 22/47] Update async tests --- .../sdk/core/oauth/CachedTokenSourceTest.java | 102 +++++++++++++----- 1 file changed, 77 insertions(+), 25 deletions(-) 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 ba9a41753..4a2bae8a4 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 @@ -2,7 +2,12 @@ import static org.junit.jupiter.api.Assertions.*; +import com.databricks.sdk.core.utils.ClockSupplier; +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; import java.time.LocalDateTime; +import java.time.ZoneId; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.stream.Stream; @@ -73,12 +78,28 @@ public Token getToken() { assertEquals(expectedToken, token.getAccessToken(), "Token value did not match expected"); } + /** + * This test verifies that if an asynchronous token refresh fails, the next refresh attempt is + * forced to be synchronous. It ensures that after an async failure, the system does not + * repeatedly attempt async refreshes while the token is stale, and only performs a synchronous + * refresh when the token is expired. After a successful sync refresh, async refreshes resume as + * normal. + */ @Test void testAsyncRefreshFailureFallback() throws Exception { - class MutableTokenSource implements TokenSource { + // Create a mutable clock supplier that we can control + MutableClockSupplier clockSupplier = new MutableClockSupplier(Instant.now()); + + // Create a token that will be stale (2 minutes until expiry) + Token staleToken = new Token( + INITIAL_TOKEN, + TOKEN_TYPE, + null, + LocalDateTime.now(clockSupplier.getClock()).plusMinutes(2)); + + class TestSource implements TokenSource { int refreshCallCount = 0; boolean isFirstRefresh = true; - Token token = new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); @Override public Token getToken() { @@ -87,42 +108,73 @@ public Token getToken() { isFirstRefresh = false; throw new RuntimeException("Simulated async failure"); } - token = new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); - return token; + // Return a token that expires in 10 minutes from current time + return new Token( + REFRESH_TOKEN, + TOKEN_TYPE, + null, + LocalDateTime.now(clockSupplier.getClock()).plusMinutes(10)); } } - MutableTokenSource mutableTokenSource = new MutableTokenSource(); - - CachedTokenSource source = - new CachedTokenSource.Builder(mutableTokenSource) - .withAsyncEnabled(true) - .withToken(mutableTokenSource.token) - .build(); + TestSource testSource = new TestSource(); + CachedTokenSource source = new CachedTokenSource.Builder(testSource) + .withAsyncEnabled(true) + .withToken(staleToken) + .withClockSupplier(clockSupplier) + .build(); // First call triggers async refresh, which fails source.getToken(); - assertEquals(1, mutableTokenSource.refreshCallCount); + Thread.sleep(300); + assertEquals( + 1, testSource.refreshCallCount, "refresh() should have been called once (async, failed)"); // Token is still stale, so next call should NOT trigger another refresh since the last refresh // failed source.getToken(); - assertEquals(1, mutableTokenSource.refreshCallCount); + Thread.sleep(200); + assertEquals( + 1, + testSource.refreshCallCount, + "refresh() should NOT be called again while stale after async failure"); - // Advance the clock so the token is now expired - mutableTokenSource.token = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().minusMinutes(1)); + // Advance time by 3 minutes to make the token expired + clockSupplier.advanceTime(Duration.ofMinutes(3)); // Now getToken() should call refresh synchronously and return the refreshed token - Token token; - token = source.getToken(); - assertEquals(REFRESH_TOKEN, token.getAccessToken()); - assertEquals(2, mutableTokenSource.refreshCallCount); - - // Make the token stale again and trigger async refresh since the last sync refresh succeeded - mutableTokenSource.token = - new Token(INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(2)); + Token token = source.getToken(); + assertEquals( + REFRESH_TOKEN, + token.getAccessToken(), + "Should return the refreshed token after sync refresh"); + assertEquals( + 2, testSource.refreshCallCount, "refresh() should have been called synchronously after expiry"); + + // Advance time by 8 minutes to make the token stale again + clockSupplier.advanceTime(Duration.ofMinutes(8)); source.getToken(); - assertEquals(3, mutableTokenSource.refreshCallCount); + Thread.sleep(300); + assertEquals( + 3, + testSource.refreshCallCount, + "refresh() should have been called again asynchronously after making token stale"); + } + + private static class MutableClockSupplier implements ClockSupplier { + private Instant instant; + + MutableClockSupplier(Instant instant) { + this.instant = instant; + } + + void advanceTime(Duration duration) { + this.instant = instant.plus(duration); + } + + @Override + public Clock getClock() { + return Clock.fixed(instant, ZoneId.of("UTC")); + } } } From 333d22e078f3e29c2d6105fde8e5d09231d155ef Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sat, 7 Jun 2025 11:39:04 +0000 Subject: [PATCH 23/47] Update async test --- .../sdk/core/oauth/CachedTokenSourceTest.java | 67 ++++++++++++------- 1 file changed, 41 insertions(+), 26 deletions(-) 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 4a2bae8a4..f3e67079a 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 @@ -88,14 +88,15 @@ public Token getToken() { @Test void testAsyncRefreshFailureFallback() throws Exception { // Create a mutable clock supplier that we can control - MutableClockSupplier clockSupplier = new MutableClockSupplier(Instant.now()); - + TestClockSupplier clockSupplier = new TestClockSupplier(Instant.now()); + // Create a token that will be stale (2 minutes until expiry) - Token staleToken = new Token( - INITIAL_TOKEN, - TOKEN_TYPE, - null, - LocalDateTime.now(clockSupplier.getClock()).plusMinutes(2)); + Token staleToken = + new Token( + INITIAL_TOKEN, + TOKEN_TYPE, + null, + LocalDateTime.now(clockSupplier.getClock()).plusMinutes(2)); class TestSource implements TokenSource { int refreshCallCount = 0; @@ -104,36 +105,45 @@ class TestSource implements TokenSource { @Override public Token getToken() { refreshCallCount++; + try { + // Sleep to simulate token fetch delay + Thread.sleep(500); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } if (isFirstRefresh) { isFirstRefresh = false; throw new RuntimeException("Simulated async failure"); } // Return a token that expires in 10 minutes from current time return new Token( - REFRESH_TOKEN, - TOKEN_TYPE, - null, + REFRESH_TOKEN + "-" + refreshCallCount, + TOKEN_TYPE, + null, LocalDateTime.now(clockSupplier.getClock()).plusMinutes(10)); } } TestSource testSource = new TestSource(); - CachedTokenSource source = new CachedTokenSource.Builder(testSource) - .withAsyncEnabled(true) - .withToken(staleToken) - .withClockSupplier(clockSupplier) - .build(); + CachedTokenSource source = + new CachedTokenSource.Builder(testSource) + .withAsyncEnabled(true) + .withToken(staleToken) + .withClockSupplier(clockSupplier) + .build(); // First call triggers async refresh, which fails - source.getToken(); - Thread.sleep(300); + // Should return stale token immediately (async refresh) + Token token = source.getToken(); + assertEquals(INITIAL_TOKEN, token.getAccessToken(), "Should return stale token immediately"); + Thread.sleep(600); assertEquals( 1, testSource.refreshCallCount, "refresh() should have been called once (async, failed)"); // Token is still stale, so next call should NOT trigger another refresh since the last refresh // failed - source.getToken(); - Thread.sleep(200); + token = source.getToken(); + assertEquals(INITIAL_TOKEN, token.getAccessToken(), "Should still return stale token"); assertEquals( 1, testSource.refreshCallCount, @@ -143,28 +153,33 @@ public Token getToken() { clockSupplier.advanceTime(Duration.ofMinutes(3)); // Now getToken() should call refresh synchronously and return the refreshed token - Token token = source.getToken(); + token = source.getToken(); assertEquals( - REFRESH_TOKEN, + REFRESH_TOKEN + "-2", token.getAccessToken(), "Should return the refreshed token after sync refresh"); assertEquals( - 2, testSource.refreshCallCount, "refresh() should have been called synchronously after expiry"); + 2, + testSource.refreshCallCount, + "refresh() should have been called synchronously after expiry"); // Advance time by 8 minutes to make the token stale again clockSupplier.advanceTime(Duration.ofMinutes(8)); - source.getToken(); - Thread.sleep(300); + // Should return stale token immediately (async refresh) + token = source.getToken(); + assertEquals( + REFRESH_TOKEN + "-2", token.getAccessToken(), "Should return stale token immediately"); + Thread.sleep(600); assertEquals( 3, testSource.refreshCallCount, "refresh() should have been called again asynchronously after making token stale"); } - private static class MutableClockSupplier implements ClockSupplier { + private static class TestClockSupplier implements ClockSupplier { private Instant instant; - MutableClockSupplier(Instant instant) { + TestClockSupplier(Instant instant) { this.instant = instant; } From 1bd052f1a4630b6c343fa7b01e6a474c45a2f8e9 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sat, 7 Jun 2025 18:03:50 +0000 Subject: [PATCH 24/47] Add support for space separated expiry strings --- .../databricks/sdk/core/CliTokenSource.java | 26 +++++++++++- .../sdk/core/CliTokenSourceTest.java | 40 ++++++++++++++----- .../src/test/resources/testdata/az | 10 ++--- 3 files changed, 60 insertions(+), 16 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java index 0e68dc7fc..9313a4b3f 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java @@ -9,8 +9,12 @@ import java.io.IOException; import java.io.InputStream; import java.time.Instant; +import java.time.LocalDateTime; import java.time.OffsetDateTime; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.time.format.DateTimeParseException; +import java.util.Arrays; import java.util.List; import org.apache.commons.io.IOUtils; @@ -44,7 +48,27 @@ public CliTokenSource( * @throws DateTimeParseException if the input string cannot be parsed */ static Instant parseExpiry(String expiry) { - return OffsetDateTime.parse(expiry).toInstant(); + DateTimeParseException lastException = null; + try { + return OffsetDateTime.parse(expiry).toInstant(); + } catch (DateTimeParseException e) { + lastException = e; + } + + String multiplePrecisionPattern = + "[SSSSSSSSS][SSSSSSSS][SSSSSSS][SSSSSS][SSSSS][SSSS][SSS][SS][S]"; + List datePatterns = + Arrays.asList("yyyy-MM-dd HH:mm:ss", "yyyy-MM-dd HH:mm:ss." + multiplePrecisionPattern); + for (String pattern : datePatterns) { + try { + DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern); + LocalDateTime dateTime = LocalDateTime.parse(expiry, formatter); + return dateTime.atZone(ZoneId.systemDefault()).toInstant(); + } catch (DateTimeParseException e) { + lastException = e; + } + } + throw lastException; } private String getProcessStream(InputStream stream) throws IOException { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java index cc177881c..d6c8fc740 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java @@ -4,6 +4,8 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import java.time.Instant; +import java.time.LocalDateTime; +import java.time.ZoneId; import java.time.format.DateTimeParseException; import java.util.stream.Stream; import org.junit.jupiter.params.ParameterizedTest; @@ -15,35 +17,55 @@ public class CliTokenSourceTest { private static Stream expiryProvider() { return Stream.of( Arguments.of( - "2023-07-17T09:02:22.330612218Z", "2023-07-17T09:02:22.330612218Z", "9-digit nanos"), + "2023-07-17T09:02:22.330612218Z", + Instant.parse("2023-07-17T09:02:22.330612218Z"), + "9-digit nanos"), Arguments.of( - "2023-07-17T09:02:22.33061221Z", "2023-07-17T09:02:22.330612210Z", "8-digit nanos"), + "2023-07-17T09:02:22.33061221Z", + Instant.parse("2023-07-17T09:02:22.330612210Z"), + "8-digit nanos"), Arguments.of( - "2023-07-17T09:02:22.330612Z", "2023-07-17T09:02:22.330612000Z", "6-digit nanos"), + "2023-07-17T09:02:22.330612Z", + Instant.parse("2023-07-17T09:02:22.330612000Z"), + "6-digit nanos"), Arguments.of( "2023-07-17T10:02:22.330612218+01:00", - "2023-07-17T09:02:22.330612218Z", + Instant.parse("2023-07-17T09:02:22.330612218Z"), "+01:00 offset, 9-digit nanos"), Arguments.of( "2023-07-17T04:02:22.330612218-05:00", - "2023-07-17T09:02:22.330612218Z", + Instant.parse("2023-07-17T09:02:22.330612218Z"), "-05:00 offset, 9-digit nanos"), Arguments.of( "2023-07-17T10:02:22.330612+01:00", - "2023-07-17T09:02:22.330612000Z", + Instant.parse("2023-07-17T09:02:22.330612000Z"), "+01:00 offset, 6-digit nanos"), Arguments.of("2023-07-17T09:02:22.33061221987Z", null, "Invalid: >9 nanos"), - Arguments.of("17-07-2023 09:02:22", null, "Invalid: non-ISO8601 format")); + Arguments.of("17-07-2023 09:02:22", null, "Invalid date format"), + Arguments.of( + "2023-07-17 09:02:22.330612218", + LocalDateTime.parse("2023-07-17T09:02:22.330612218") + .atZone(ZoneId.systemDefault()) + .toInstant(), + "Space separator, 9-digit nanos"), + Arguments.of( + "2023-07-17 09:02:22.330612", + LocalDateTime.parse("2023-07-17T09:02:22.330612") + .atZone(ZoneId.systemDefault()) + .toInstant(), + "Space separator, 6-digit nanos"), + Arguments.of( + "2023-07-17 09:02:22.33061221987", null, "Space separator, Invalid: >9 nanos")); } @ParameterizedTest(name = "{2}") @MethodSource("expiryProvider") - public void testParseExpiry(String input, String expectedInstant, String description) { + public void testParseExpiry(String input, Instant expectedInstant, String description) { if (expectedInstant == null) { assertThrows(DateTimeParseException.class, () -> CliTokenSource.parseExpiry(input)); } else { Instant parsedInstant = CliTokenSource.parseExpiry(input); - assertEquals(Instant.parse(expectedInstant), parsedInstant); + assertEquals(expectedInstant, parsedInstant); } } } diff --git a/databricks-sdk-java/src/test/resources/testdata/az b/databricks-sdk-java/src/test/resources/testdata/az index b24bddea9..00997a4cf 100755 --- a/databricks-sdk-java/src/test/resources/testdata/az +++ b/databricks-sdk-java/src/test/resources/testdata/az @@ -22,16 +22,14 @@ for arg in "$@"; do fi done -# MacOS -EXP="$(/bin/date -v+${EXPIRE:=10S} +'%FT%T%z' 2>/dev/null)" +# Macos +EXP="$(/bin/date -v+${EXPIRE:=10S} +'%F %T' 2>/dev/null)" if [ -z "${EXP}" ]; then # Linux EXPIRE=$(/bin/echo $EXPIRE | /bin/sed 's/S/seconds/') EXPIRE=$(/bin/echo $EXPIRE | /bin/sed 's/M/minutes/') - EXP=$(/bin/date --date=+${EXPIRE:=10seconds} +'%FT%T%z') + EXP=$(/bin/date --date=+${EXPIRE:=10seconds} +'%F %T') fi -# Insert colon in timezone offset for ISO 8601 compliance -EXP="${EXP:0:19}${EXP:19:3}:${EXP:22:2}" if [ -z "${TF_AAD_TOKEN}" ]; then TF_AAD_TOKEN="..." @@ -43,4 +41,4 @@ fi \"subscription\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\", \"tenant\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\", \"tokenType\": \"Bearer\" -}" +}" \ No newline at end of file From 408f3b465307952d7bad140f92d8a71ab1b78e2b Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Sat, 7 Jun 2025 18:09:32 +0000 Subject: [PATCH 25/47] revert test data --- databricks-sdk-java/src/test/resources/testdata/az | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/test/resources/testdata/az b/databricks-sdk-java/src/test/resources/testdata/az index 00997a4cf..29b824ed7 100755 --- a/databricks-sdk-java/src/test/resources/testdata/az +++ b/databricks-sdk-java/src/test/resources/testdata/az @@ -41,4 +41,4 @@ fi \"subscription\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\", \"tenant\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\", \"tokenType\": \"Bearer\" -}" \ No newline at end of file +}" From 3f06444ec9d3ac688a92490d62f4977111ba6599 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 13:27:45 +0000 Subject: [PATCH 26/47] Passing all tests --- .../sdk/core/AzureCliCredentialsProvider.java | 22 +++++----- .../DatabricksCliCredentialsProvider.java | 8 +++- ...reServicePrincipalCredentialsProvider.java | 29 +++++++------ .../ExternalBrowserCredentialsProvider.java | 43 ++++++++++++------- .../oauth/GithubOidcCredentialsProvider.java | 4 +- ...2MServicePrincipalCredentialsProvider.java | 4 +- .../sdk/core/oauth/TokenEndpointClient.java | 24 +++++------ .../core/AzureCliCredentialsProviderTest.java | 23 +++++----- ...xternalBrowserCredentialsProviderTest.java | 16 +++++-- 9 files changed, 104 insertions(+), 69 deletions(-) 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 6ed5b83d3..4ec58480e 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 @@ -1,5 +1,6 @@ package com.databricks.sdk.core; +import com.databricks.sdk.core.oauth.CachedTokenSource; import com.databricks.sdk.core.oauth.OAuthHeaderFactory; import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.utils.AzureUtils; @@ -19,7 +20,7 @@ public String authType() { return AZURE_CLI; } - public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { String azPath = Optional.ofNullable(config.getEnv()).map(env -> env.get("AZ_PATH")).orElse("az"); @@ -35,7 +36,7 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { List extendedCmd = new ArrayList<>(cmd); extendedCmd.addAll(Arrays.asList("--subscription", subscription.get())); try { - return getToken(config, extendedCmd); + return getTokenSource(config, extendedCmd); } catch (DatabricksException ex) { LOG.warn("Failed to get token for subscription. Using resource only token."); } @@ -45,14 +46,15 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + "It is recommended to specify this field in the Databricks configuration to avoid authentication errors."); } - return getToken(config, cmd); + return getTokenSource(config, cmd); } - protected CliTokenSource getToken(DatabricksConfig config, List cmd) { - CliTokenSource token = + protected CachedTokenSource getTokenSource(DatabricksConfig config, List cmd) { + CliTokenSource tokenSource = new CliTokenSource(cmd, "tokenType", "accessToken", "expiresOn", config.getEnv()); - token.getToken(); // We need this to check if the CLI is installed and to validate the config. - return token; + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build(); + cachedTokenSource.getToken(); // Check if the CLI is installed and to validate the config. + return cachedTokenSource; } private Optional getSubscription(DatabricksConfig config) { @@ -77,8 +79,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { try { AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); String resource = config.getEffectiveAzureLoginAppId(); - CliTokenSource tokenSource = tokenSourceFor(config, resource); - CliTokenSource mgmtTokenSource; + CachedTokenSource tokenSource = tokenSourceFor(config, resource); + CachedTokenSource mgmtTokenSource; try { mgmtTokenSource = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); @@ -86,7 +88,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOG.debug("Not including service management token in headers", e); mgmtTokenSource = null; } - CliTokenSource finalMgmtTokenSource = mgmtTokenSource; + CachedTokenSource finalMgmtTokenSource = mgmtTokenSource; return OAuthHeaderFactory.fromSuppliers( tokenSource::getToken, () -> { 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 655d0b599..6d5a2eb9f 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 @@ -1,5 +1,6 @@ package com.databricks.sdk.core; +import com.databricks.sdk.core.oauth.CachedTokenSource; import com.databricks.sdk.core.oauth.OAuthHeaderFactory; import com.databricks.sdk.core.utils.OSUtils; import java.util.*; @@ -47,8 +48,11 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { if (tokenSource == null) { return null; } - tokenSource.getToken(); // We need this for checking if databricks CLI is installed. - return OAuthHeaderFactory.fromTokenSource(tokenSource); + + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build(); + cachedTokenSource.getToken(); // We need this for checking if databricks CLI is installed. + + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (DatabricksException e) { String stderr = e.getMessage(); if (stderr.contains("not found")) { 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 d10fefd94..7b77a4dfe 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 @@ -44,29 +44,32 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } /** - * Creates a RefreshableTokenSource for the specified Azure resource. + * Creates a CachedTokenSource for the specified Azure resource. * - *

This function constructs a RefreshableTokenSource instance that fetches OAuth tokens for the + *

This function constructs a CachedTokenSource instance that fetches OAuth tokens for the * given Azure resource. It uses the authentication parameters provided by the DatabricksConfig * instance to generate the tokens. * * @param config The DatabricksConfig instance containing the required authentication parameters. * @param resource The Azure resource for which OAuth tokens need to be fetched. - * @return A RefreshableTokenSource instance capable of fetching OAuth tokens for the specified - * Azure resource. + * @return A CachedTokenSource instance capable of fetching OAuth tokens for the specified Azure + * resource. */ - private static TokenSource tokenSourceFor(DatabricksConfig config, String resource) { + private static CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { String aadEndpoint = config.getAzureEnvironment().getActiveDirectoryEndpoint(); String tokenUrl = aadEndpoint + config.getAzureTenantId() + "/oauth2/token"; Map endpointParams = new HashMap<>(); endpointParams.put("resource", resource); - return new ClientCredentials.Builder() - .withHttpClient(config.getHttpClient()) - .withClientId(config.getAzureClientId()) - .withClientSecret(config.getAzureClientSecret()) - .withTokenUrl(tokenUrl) - .withEndpointParametersSupplier(() -> endpointParams) - .withAuthParameterPosition(AuthParameterPosition.BODY) - .build(); + + ClientCredentials clientCredentials = + new ClientCredentials.Builder() + .withHttpClient(config.getHttpClient()) + .withClientId(config.getAzureClientId()) + .withClientSecret(config.getAzureClientSecret()) + .withTokenUrl(tokenUrl) + .withEndpointParametersSupplier(() -> endpointParams) + .withAuthParameterPosition(AuthParameterPosition.BODY) + .build(); + return new CachedTokenSource.Builder(clientCredentials).build(); } } 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 494ad69cd..03bec3f8a 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 @@ -67,7 +67,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { try { // Create SessionCredentials with the cached token and try to refresh if needed - SessionCredentials cachedCreds = + SessionCredentials sessionCredentials = new SessionCredentials.Builder() .withToken(cachedToken) .withHttpClient(config.getHttpClient()) @@ -79,26 +79,32 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { .build(); LOGGER.debug("Using cached token, will immediately refresh"); - cachedCreds.token = cachedCreds.getToken(); - return cachedCreds.configure(config); + sessionCredentials.token = sessionCredentials.getToken(); + + CachedTokenSource cachedTokenSource = + new CachedTokenSource.Builder(sessionCredentials) + .withToken(sessionCredentials.token) + .build(); + + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); } } - // If no cached token or refresh failed, perform browser auth - SessionCredentials credentials = + CachedTokenSource tokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); - tokenCache.save(credentials.getToken()); - return credentials.configure(config); + tokenCache.save(tokenSource.getToken()); + + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; } } - SessionCredentials performBrowserAuth( + CachedTokenSource performBrowserAuth( DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache) throws IOException { LOGGER.debug("Performing browser authentication"); @@ -117,14 +123,19 @@ SessionCredentials performBrowserAuth( SessionCredentials credentials = consent.launchExternalBrowser(); // Create a new SessionCredentials with the same token but with our token cache - return new SessionCredentials.Builder() - .withToken(credentials.getToken()) - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withTokenCache(tokenCache) + SessionCredentials sessionCredentials = + new SessionCredentials.Builder() + .withToken(credentials.getToken()) + .withHttpClient(config.getHttpClient()) + .withClientId(config.getClientId()) + .withClientSecret(config.getClientSecret()) + .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) + .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) + .withTokenCache(tokenCache) + .build(); + + return new CachedTokenSource.Builder(sessionCredentials) + .withToken(sessionCredentials.getToken()) .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 eeb70797e..c517b8a3e 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,9 +58,11 @@ public HeaderFactory configure(DatabricksConfig config) throws DatabricksExcepti .build()) .build(); + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(clientCredentials).build(); + return () -> { Map headers = new HashMap<>(); - headers.put("Authorization", "Bearer " + clientCredentials.getToken().getAccessToken()); + headers.put("Authorization", "Bearer " + cachedTokenSource.getToken().getAccessToken()); return headers; }; } 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 058fc268c..743680d52 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 @@ -27,7 +27,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { // https://login.microsoftonline.com/{cfg.azure_tenant_id}/.well-known/oauth-authorization-server try { OpenIDConnectEndpoints jsonResponse = config.getOidcEndpoints(); - ClientCredentials tokenSource = + ClientCredentials clientCredentials = new ClientCredentials.Builder() .withHttpClient(config.getHttpClient()) .withClientId(config.getClientId()) @@ -37,6 +37,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { .withAuthParameterPosition(AuthParameterPosition.HEADER) .build(); + CachedTokenSource tokenSource = new CachedTokenSource.Builder(clientCredentials).build(); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (IOException e) { // TODO: Log exception diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java index ace0c9314..2b7ab582e 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java @@ -4,11 +4,16 @@ import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.FormRequest; import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; +import java.time.LocalDateTime; +import java.time.temporal.ChronoUnit; +import java.util.Base64; import java.util.Map; import java.util.Objects; +import org.apache.http.HttpHeaders; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -93,7 +98,7 @@ public static OAuthResponse requestToken( /** * Helper method implementing OAuth token refresh. * - * @param hc The HTTP client to use for the request. + * @param hc The {@link HttpClient} to use for making the request. * @param clientId The client ID to authenticate with. * @param clientSecret The client secret to authenticate with. * @param tokenUrl The authorization URL for fetching tokens. @@ -101,8 +106,8 @@ public static OAuthResponse requestToken( * @param headers Additional headers. * @param position The position of the authentication parameters in the request. * @return The newly fetched Token. - * @throws DatabricksException if the refresh fails - * @throws IllegalArgumentException if the OAuth response contains an error + * @throws DatabricksException if the refresh fails. + * @throws IllegalArgumentException if the OAuth response contains an error. */ public static Token retrieveToken( HttpClient hc, @@ -124,15 +129,12 @@ public static Token retrieveToken( case HEADER: String authHeaderValue = "Basic " - + java.util.Base64.getEncoder() - .encodeToString((clientId + ":" + clientSecret).getBytes()); - headers.put(org.apache.http.HttpHeaders.AUTHORIZATION, authHeaderValue); + + Base64.getEncoder().encodeToString((clientId + ":" + clientSecret).getBytes()); + headers.put(HttpHeaders.AUTHORIZATION, authHeaderValue); break; } headers.put("Content-Type", "application/x-www-form-urlencoded"); - com.databricks.sdk.core.http.Request req = - new com.databricks.sdk.core.http.Request( - "POST", tokenUrl, com.databricks.sdk.core.http.FormRequest.wrapValuesInList(params)); + Request req = new Request("POST", tokenUrl, FormRequest.wrapValuesInList(params)); req.withHeaders(headers); try { ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build(); @@ -140,9 +142,7 @@ public static Token retrieveToken( if (resp.getErrorCode() != null) { throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); } - java.time.LocalDateTime expiry = - java.time.LocalDateTime.now() - .plus(resp.getExpiresIn(), java.time.temporal.ChronoUnit.SECONDS); + LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); } catch (Exception e) { throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java index 0212b7652..a57ae71df 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java @@ -5,6 +5,7 @@ import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.times; +import com.databricks.sdk.core.oauth.CachedTokenSource; import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.oauth.TokenSource; import java.time.LocalDateTime; @@ -22,18 +23,18 @@ class AzureCliCredentialsProviderTest { private static final String TOKEN = "t-123"; private static final String TOKEN_TYPE = "token-type"; - private static CliTokenSource mockTokenSource() { - CliTokenSource tokenSource = Mockito.mock(CliTokenSource.class); - Mockito.when(tokenSource.getToken()) + private static CachedTokenSource mockTokenSource() { + CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); + Mockito.when(cliTokenSource.getToken()) .thenReturn(new Token(TOKEN, TOKEN_TYPE, LocalDateTime.now())); - return tokenSource; + return new CachedTokenSource.Builder(cliTokenSource).build(); } private static AzureCliCredentialsProvider getAzureCliCredentialsProvider( TokenSource tokenSource) { AzureCliCredentialsProvider provider = Mockito.spy(new AzureCliCredentialsProvider()); - Mockito.doReturn(tokenSource).when(provider).getToken(any(), anyList()); + Mockito.doReturn(tokenSource).when(provider).getTokenSource(any(), anyList()); return provider; } @@ -52,7 +53,7 @@ void testWorkSpaceIDUsage() { String token = header.headers().get("Authorization"); assertEquals(token, TOKEN_TYPE + " " + TOKEN); - Mockito.verify(provider, times(2)).getToken(any(), argument.capture()); + Mockito.verify(provider, times(2)).getTokenSource(any(), argument.capture()); List value = argument.getValue(); value = value.subList(value.size() - 2, value.size()); @@ -62,11 +63,13 @@ void testWorkSpaceIDUsage() { @Test void testFallbackWhenTailsToGetTokenForSubscription() { - CliTokenSource tokenSource = mockTokenSource(); + CachedTokenSource tokenSource = mockTokenSource(); AzureCliCredentialsProvider provider = Mockito.spy(new AzureCliCredentialsProvider()); - Mockito.doThrow(new DatabricksException("error")).when(provider).getToken(any(), anyList()); - Mockito.doReturn(tokenSource).when(provider).getToken(any(), anyList()); + Mockito.doThrow(new DatabricksException("error")) + .when(provider) + .getTokenSource(any(), anyList()); + Mockito.doReturn(tokenSource).when(provider).getTokenSource(any(), anyList()); DatabricksConfig config = new DatabricksConfig() @@ -94,7 +97,7 @@ void testGetTokenWithoutWorkspaceResourceID() { String token = header.headers().get("Authorization"); assertEquals(token, TOKEN_TYPE + " " + TOKEN); - Mockito.verify(provider, times(2)).getToken(any(), argument.capture()); + Mockito.verify(provider, times(2)).getTokenSource(any(), argument.capture()); List value = argument.getValue(); assertFalse(value.contains("--subscription")); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index e8f6faa2f..f1e871277 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -408,13 +408,16 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { "browser_refresh_token", LocalDateTime.now().plusHours(1)); - SessionCredentials browserAuthCreds = + SessionCredentials sessionCredentials = new SessionCredentials.Builder() .withToken(browserAuthToken) .withClientId("test-client-id") .withTokenUrl("https://test-token-url") .build(); + CachedTokenSource browserAuthTokenSource = + new CachedTokenSource.Builder(sessionCredentials).withToken(browserAuthToken).build(); + // Create config with failing HTTP client and mock token cache DatabricksConfig config = new DatabricksConfig() @@ -431,7 +434,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); @@ -441,6 +444,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Configure provider HeaderFactory headerFactory = provider.configure(spyConfig); + assertNotNull(headerFactory); // Verify headers contain the browser auth token (fallback) Map headers = headerFactory.headers(); @@ -475,13 +479,16 @@ void cacheWithInvalidTokensTest() throws IOException { "browser_refresh_token", LocalDateTime.now().plusHours(1)); - SessionCredentials browserAuthCreds = + SessionCredentials sessionCredentials = new SessionCredentials.Builder() .withToken(browserAuthToken) .withClientId("test-client-id") .withTokenUrl("https://test-token-url") .build(); + CachedTokenSource browserAuthTokenSource = + new CachedTokenSource.Builder(sessionCredentials).withToken(browserAuthToken).build(); + // Create simple config DatabricksConfig config = new DatabricksConfig() @@ -492,12 +499,13 @@ void cacheWithInvalidTokensTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); // Configure provider HeaderFactory headerFactory = provider.configure(config); + assertNotNull(headerFactory); // Verify headers contain the browser auth token (fallback) Map headers = headerFactory.headers(); assertEquals("Bearer browser_access_token", headers.get("Authorization")); From f12024feb9c8c37ef31083f1f6b174537487c199 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 14:35:39 +0000 Subject: [PATCH 27/47] Wrapped all token sources in cached token source --- .../AzureGithubOidcCredentialsProvider.java | 4 ++-- ...ureServicePrincipalCredentialsProvider.java | 4 ++-- .../sdk/core/oauth/CachedTokenSource.java | 2 +- .../com/databricks/sdk/core/oauth/Consent.java | 2 +- .../sdk/core/oauth/DataPlaneTokenSource.java | 18 ++++++++++-------- .../sdk/core/oauth/RefreshableTokenSource.java | 10 ---------- .../oauth/TokenSourceCredentialsProvider.java | 11 ++++++++--- 7 files changed, 24 insertions(+), 27 deletions(-) delete mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java 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 b29b5aa0e..281d2d9a4 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,8 +46,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { config.getEffectiveAzureLoginAppId(), idToken.get(), "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"); - - return OAuthHeaderFactory.fromTokenSource(tokenSource); + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).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 7b77a4dfe..a7a041f41 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 @@ -28,8 +28,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } AzureUtils.ensureHostPresent( config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor); - TokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); - TokenSource cloud = + CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId()); + CachedTokenSource cloud = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); return OAuthHeaderFactory.fromSuppliers( 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 72aa38f74..a3eb6fe4b 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 @@ -28,7 +28,7 @@ private enum TokenState { EXPIRED } - private static final Logger logger = LoggerFactory.getLogger(RefreshableTokenSource.class); + private static final Logger logger = LoggerFactory.getLogger(CachedTokenSource.class); // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); private static final Duration DEFAULT_EXPIRY_BUFFER = Duration.ofSeconds(40); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 68dc6f176..f8e65d0a2 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -298,7 +298,7 @@ public SessionCredentials exchangeCallbackParameters(Map query) return exchange(query.get("code"), query.get("state")); } - public SessionCredentials exchange(String code, String state) { + private SessionCredentials exchange(String code, String state) { if (!this.state.equals(state)) { throw new DatabricksException( "state mismatch: original state: " + this.state + "; retrieved state: " + state); 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 8d48c2dff..80f68aaeb 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 @@ -7,16 +7,16 @@ /** * Manages and provides Databricks data plane tokens. This class is responsible for acquiring and * caching OAuth tokens that are specific to a particular Databricks data plane service endpoint and - * a set of authorization details. It utilizes a {@link DatabricksOAuthTokenSource} for obtaining - * control plane tokens, which may then be exchanged or used to authorize requests for data plane - * tokens. Cached {@link EndpointTokenSource} instances are used to efficiently reuse tokens for - * repeated requests to the same endpoint with the same authorization context. + * a set of authorization details. It utilizes a {@link TokenSource} for obtaining control plane + * tokens, which may then be exchanged or used to authorize requests for data plane tokens. Cached + * {@link EndpointTokenSource} instances are used to efficiently reuse tokens for repeated requests + * to the same endpoint with the same authorization context. */ public class DataPlaneTokenSource { private final HttpClient httpClient; private final TokenSource cpTokenSource; private final String host; - private final ConcurrentHashMap sourcesCache; + private final ConcurrentHashMap sourcesCache; /** * Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This * is a value object that uniquely identifies a token source configuration. @@ -103,12 +103,14 @@ public Token getToken(String endpoint, String authDetails) { TokenSourceKey key = new TokenSourceKey(endpoint, authDetails); - EndpointTokenSource specificSource = + CachedTokenSource specificSource = sourcesCache.computeIfAbsent( key, k -> - new EndpointTokenSource( - this.cpTokenSource, k.authDetails, this.httpClient, this.host)); + new CachedTokenSource.Builder( + new EndpointTokenSource( + this.cpTokenSource, k.authDetails, this.httpClient, this.host)) + .build()); return specificSource.getToken(); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java deleted file mode 100644 index c0cf080ce..000000000 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java +++ /dev/null @@ -1,10 +0,0 @@ -package com.databricks.sdk.core.oauth; - -/** - * An OAuth TokenSource which can be refreshed. - * - *

This class supports both synchronous and asynchronous token refresh. When async is enabled, - * stale tokens will trigger a background refresh, while expired tokens will block until a new token - * is fetched. - */ -public abstract class RefreshableTokenSource implements TokenSource {} 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 9a341b901..46f67eb78 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 @@ -37,11 +37,16 @@ public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType) */ @Override public OAuthHeaderFactory configure(DatabricksConfig config) { + // Check if the token source is already cached to prevent double caching + TokenSource source = + (tokenSource instanceof CachedTokenSource) + ? tokenSource + : new CachedTokenSource.Builder(tokenSource).build(); + try { // Validate that we can get a token before returning a HeaderFactory - tokenSource.getToken().getAccessToken(); - - return OAuthHeaderFactory.fromTokenSource(tokenSource); + source.getToken().getAccessToken(); + return OAuthHeaderFactory.fromTokenSource(source); } catch (Exception e) { return null; } From c32b7524ffa694ea9470c8ab51e799db1a16e737 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 14:43:37 +0000 Subject: [PATCH 28/47] update names --- .../sdk/core/oauth/ExternalBrowserCredentialsProvider.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 03bec3f8a..63cb2e1b7 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 @@ -81,12 +81,12 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOGGER.debug("Using cached token, will immediately refresh"); sessionCredentials.token = sessionCredentials.getToken(); - CachedTokenSource cachedTokenSource = + CachedTokenSource tokenSource = new CachedTokenSource.Builder(sessionCredentials) .withToken(sessionCredentials.token) .build(); - return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); From 7e441d6075df2a3e96a0c4d88fb42b784deeb542 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 14:57:25 +0000 Subject: [PATCH 29/47] Small change --- .../sdk/core/oauth/ExternalBrowserCredentialsProvider.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 63cb2e1b7..2a9de00fe 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 @@ -121,11 +121,12 @@ CachedTokenSource performBrowserAuth( // Use the existing browser flow to get credentials SessionCredentials credentials = consent.launchExternalBrowser(); + Token token = credentials.getToken(); // Create a new SessionCredentials with the same token but with our token cache SessionCredentials sessionCredentials = new SessionCredentials.Builder() - .withToken(credentials.getToken()) + .withToken(token) .withHttpClient(config.getHttpClient()) .withClientId(config.getClientId()) .withClientSecret(config.getClientSecret()) @@ -134,8 +135,6 @@ CachedTokenSource performBrowserAuth( .withTokenCache(tokenCache) .build(); - return new CachedTokenSource.Builder(sessionCredentials) - .withToken(sessionCredentials.getToken()) - .build(); + return new CachedTokenSource.Builder(sessionCredentials).withToken(token).build(); } } From 35cf2a4a746cc38f7d9041daeb11665f9482288c Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 15:10:43 +0000 Subject: [PATCH 30/47] Revert DatabricksConfig --- .../java/com/databricks/sdk/core/DatabricksConfig.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 df16ebae3..de6548982 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 @@ -410,17 +410,13 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** - * @deprecated Use {@link #getAzureUseMsi()} instead. - */ + /** @deprecated Use {@link #getAzureUseMsi()} instead. */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** - * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. - */ + /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; From 457d02c466ce4f7b6881cfc7898302d2aa141274 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 15:26:28 +0000 Subject: [PATCH 31/47] Merge branch 'emmyzhou-db/localdatetime-to-instant' into emmyzhou-db/async_token_cache_wrapper --- .../core/oauth/RefreshableTokenSource.java | 0 .../com/databricks/sdk/core/oauth/Token.java | 44 +++------------- .../sdk/core/oauth/TokenEndpointClient.java | 5 +- .../sdk/core/oauth/CachedTokenSourceTest.java | 9 ++-- .../sdk/core/oauth/FileTokenCacheTest.java | 20 +------- .../databricks/sdk/core/oauth/TokenTest.java | 50 ++----------------- 6 files changed, 19 insertions(+), 109 deletions(-) delete mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java deleted file mode 100644 index e69de29bb..000000000 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index a6dc0089f..2e78311f9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -21,16 +21,9 @@ public class Token { */ @JsonProperty private Instant expiry; - private final ClockSupplier clockSupplier; - /** Constructor for non-refreshable tokens (e.g. M2M). */ public Token(String accessToken, String tokenType, Instant expiry) { - this(accessToken, tokenType, null, expiry, new SystemClockSupplier()); - } - - /** Constructor for non-refreshable tokens (e.g. M2M) with ClockSupplier */ - public Token(String accessToken, String tokenType, Instant expiry, ClockSupplier clockSupplier) { - this(accessToken, tokenType, null, expiry, clockSupplier); + this(accessToken, tokenType, null, expiry); } /** Constructor for refreshable tokens. */ @@ -40,16 +33,6 @@ public Token( @JsonProperty("tokenType") String tokenType, @JsonProperty("refreshToken") String refreshToken, @JsonProperty("expiry") Instant expiry) { - this(accessToken, tokenType, refreshToken, expiry, new SystemClockSupplier()); - } - - /** Constructor for refreshable tokens with ClockSupplier. */ - public Token( - String accessToken, - String tokenType, - String refreshToken, - Instant expiry, - ClockSupplier clockSupplier) { Objects.requireNonNull(accessToken, "accessToken must be defined"); Objects.requireNonNull(tokenType, "tokenType must be defined"); Objects.requireNonNull(expiry, "expiry must be defined"); @@ -57,24 +40,13 @@ public Token( this.tokenType = tokenType; this.refreshToken = refreshToken; this.expiry = expiry; - this.clockSupplier = clockSupplier; - } - - public boolean isExpired() { - if (expiry == null) { - return false; - } - // Azure Databricks rejects tokens that expire in 30 seconds or less, - // so we refresh the token 40 seconds before it expires. - Instant potentiallyExpired = expiry.minusSeconds(40); - Instant now = Instant.now(clockSupplier.getClock()); - return potentiallyExpired.isBefore(now); - } - - public boolean isValid() { - return accessToken != null && !isExpired(); } + /** + * Returns the type of the token (e.g., "Bearer"). + * + * @return the token type + */ public String getTokenType() { return tokenType; } @@ -102,7 +74,7 @@ public String getAccessToken() { * * @return the expiry time */ - public LocalDateTime getExpiry() { + public Instant getExpiry() { return this.expiry; } -} +} \ No newline at end of file diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java index 2b7ab582e..3e8b13a4d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenEndpointClient.java @@ -8,8 +8,7 @@ import com.databricks.sdk.core.http.Response; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; -import java.time.LocalDateTime; -import java.time.temporal.ChronoUnit; +import java.time.Instant; import java.util.Base64; import java.util.Map; import java.util.Objects; @@ -142,7 +141,7 @@ public static Token retrieveToken( if (resp.getErrorCode() != null) { throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary()); } - LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS); + Instant expiry = Instant.now().plusSeconds(resp.getExpiresIn()); return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry); } catch (Exception e) { throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e); 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 f3e67079a..298c8786b 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 @@ -6,7 +6,6 @@ import java.time.Clock; import java.time.Duration; import java.time.Instant; -import java.time.LocalDateTime; import java.time.ZoneId; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -46,9 +45,9 @@ void testAsyncRefreshParametrized( Token initialToken = new Token( - INITIAL_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(minutesUntilExpiry)); + INITIAL_TOKEN, TOKEN_TYPE, null, Instant.now().plus(Duration.ofMinutes(minutesUntilExpiry))); Token refreshedToken = - new Token(REFRESH_TOKEN, TOKEN_TYPE, null, LocalDateTime.now().plusMinutes(10)); + new Token(REFRESH_TOKEN, TOKEN_TYPE, null, Instant.now().plus(Duration.ofMinutes(10))); CountDownLatch refreshCalled = new CountDownLatch(1); TokenSource tokenSource = @@ -96,7 +95,7 @@ void testAsyncRefreshFailureFallback() throws Exception { INITIAL_TOKEN, TOKEN_TYPE, null, - LocalDateTime.now(clockSupplier.getClock()).plusMinutes(2)); + Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(2))); class TestSource implements TokenSource { int refreshCallCount = 0; @@ -120,7 +119,7 @@ public Token getToken() { REFRESH_TOKEN + "-" + refreshCallCount, TOKEN_TYPE, null, - LocalDateTime.now(clockSupplier.getClock()).plusMinutes(10)); + Instant.now(clockSupplier.getClock()).plus(Duration.ofMinutes(10))); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java index 303f6de66..203f61e0f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -54,24 +54,6 @@ void testSaveAndLoadToken() { assertEquals("access-token", loadedToken.getAccessToken()); assertEquals("Bearer", loadedToken.getTokenType()); assertEquals("refresh-token", loadedToken.getRefreshToken()); - assertFalse(loadedToken.isExpired(), "Token should not be expired"); - } - - @Test - void testTokenExpiry() { - // Create an expired token - Instant pastTime = Instant.now().minusSeconds(3600); - Token expiredToken = new Token("access-token", "Bearer", "refresh-token", pastTime); - - // Verify it's marked as expired - assertTrue(expiredToken.isExpired(), "Token should be expired"); - - // Create a valid token - Instant futureTime = Instant.now().plusSeconds(1800); - Token validToken = new Token("access-token", "Bearer", "refresh-token", futureTime); - - // Verify it's not marked as expired - assertFalse(validToken.isExpired(), "Token should not be expired"); } @Test @@ -125,4 +107,4 @@ void testWithCustomPath(@TempDir Path tempDir) { // Verify the file exists assertTrue(Files.exists(tempPath), "Cache file should exist at custom path"); } -} +} \ No newline at end of file diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index ebe8bb1db..31cf728c6 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -2,9 +2,7 @@ import static org.junit.jupiter.api.Assertions.*; -import com.databricks.sdk.core.utils.FakeClockSupplier; import java.time.Instant; -import java.time.ZoneId; import org.junit.jupiter.api.Test; class TokenTest { @@ -12,20 +10,11 @@ class TokenTest { private static final String accessToken = "testAccessToken"; private static final String refreshToken = "testRefreshToken"; private static final String tokenType = "testTokenType"; - private final Instant currentInstant; - private final FakeClockSupplier fakeClockSupplier; - - TokenTest() { - Instant instant = Instant.parse("2023-10-18T12:00:00.00Z"); - ZoneId zoneId = ZoneId.of("UTC"); - fakeClockSupplier = new FakeClockSupplier(instant, zoneId); - currentInstant = Instant.now(fakeClockSupplier.getClock()); - } + private static final Instant currentInstant = Instant.now(); @Test void createNonRefreshableToken() { - Token token = - new Token(accessToken, tokenType, currentInstant.plusSeconds(300), fakeClockSupplier); + Token token = new Token(accessToken, tokenType, currentInstant.plusSeconds(300)); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertNull(token.getRefreshToken()); @@ -33,40 +22,9 @@ void createNonRefreshableToken() { @Test void createRefreshableToken() { - Token token = - new Token( - accessToken, - tokenType, - refreshToken, - currentInstant.plusSeconds(300), - fakeClockSupplier); + Token token = new Token(accessToken, tokenType, refreshToken, currentInstant.plusSeconds(300)); assertEquals(accessToken, token.getAccessToken()); assertEquals(tokenType, token.getTokenType()); assertEquals(refreshToken, token.getRefreshToken()); - assertTrue(token.isValid()); - } - - @Test - void tokenExpiryMoreThan40Seconds() { - Token token = - new Token(accessToken, tokenType, currentInstant.plusSeconds(50), fakeClockSupplier); - assertFalse(token.isExpired()); - assertTrue(token.isValid()); - } - - @Test - void tokenExpiryLessThan40Seconds() { - Token token = - new Token(accessToken, tokenType, currentInstant.plusSeconds(30), fakeClockSupplier); - assertTrue(token.isExpired()); - assertFalse(token.isValid()); - } - - @Test - void expiredToken() { - Token token = - new Token(accessToken, tokenType, currentInstant.minusSeconds(10), fakeClockSupplier); - assertTrue(token.isExpired()); - assertFalse(token.isValid()); } -} +} \ No newline at end of file From 1c99d4b4146db2a1a78334b090887393af76df11 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 15:27:21 +0000 Subject: [PATCH 32/47] Merge branch 'emmyzhou-db/localdatetime-to-instant' into emmyzhou-db/async_token_cache_wrapper --- .../src/main/java/com/databricks/sdk/core/oauth/Token.java | 2 +- .../com/databricks/sdk/core/oauth/CachedTokenSourceTest.java | 5 ++++- .../com/databricks/sdk/core/oauth/FileTokenCacheTest.java | 2 +- .../test/java/com/databricks/sdk/core/oauth/TokenTest.java | 2 +- 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java index 2e78311f9..18ef99880 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Token.java @@ -77,4 +77,4 @@ public String getAccessToken() { public Instant getExpiry() { return this.expiry; } -} \ No newline at end of file +} 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 298c8786b..3d0c3b12e 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 @@ -45,7 +45,10 @@ void testAsyncRefreshParametrized( Token initialToken = new Token( - INITIAL_TOKEN, TOKEN_TYPE, null, Instant.now().plus(Duration.ofMinutes(minutesUntilExpiry))); + INITIAL_TOKEN, + TOKEN_TYPE, + null, + Instant.now().plus(Duration.ofMinutes(minutesUntilExpiry))); Token refreshedToken = new Token(REFRESH_TOKEN, TOKEN_TYPE, null, Instant.now().plus(Duration.ofMinutes(10))); CountDownLatch refreshCalled = new CountDownLatch(1); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java index 203f61e0f..d08ddf583 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileTokenCacheTest.java @@ -107,4 +107,4 @@ void testWithCustomPath(@TempDir Path tempDir) { // Verify the file exists assertTrue(Files.exists(tempPath), "Cache file should exist at custom path"); } -} \ No newline at end of file +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java index 31cf728c6..23c1a418f 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenTest.java @@ -27,4 +27,4 @@ void createRefreshableToken() { assertEquals(tokenType, token.getTokenType()); assertEquals(refreshToken, token.getRefreshToken()); } -} \ No newline at end of file +} From 63d246cca9d1cc30fe5f687c3e7275ff8e2b089f Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 10 Jun 2025 15:32:40 +0000 Subject: [PATCH 33/47] Use Instant with CachedTokenSource --- .../com/databricks/sdk/core/oauth/CachedTokenSource.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 a3eb6fe4b..77b389333 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 @@ -3,7 +3,7 @@ import com.databricks.sdk.core.utils.ClockSupplier; import com.databricks.sdk.core.utils.SystemClockSupplier; import java.time.Duration; -import java.time.LocalDateTime; +import java.time.Instant; import java.util.concurrent.CompletableFuture; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -122,8 +122,7 @@ protected TokenState getTokenState(Token t) { if (t == null) { return TokenState.EXPIRED; } - Duration lifeTime = - Duration.between(LocalDateTime.now(clockSupplier.getClock()), t.getExpiry()); + Duration lifeTime = Duration.between(Instant.now(clockSupplier.getClock()), t.getExpiry()); if (lifeTime.compareTo(expiryBuffer) <= 0) { return TokenState.EXPIRED; } From 114982e7bcb79b1c93ca7fb5f698be3d03b0af78 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 11 Jun 2025 09:38:04 +0000 Subject: [PATCH 34/47] Update AzureCliCredentials --- .../sdk/core/AzureCliCredentialsProvider.java | 18 ++++++++--------- .../core/AzureCliCredentialsProviderTest.java | 20 +++++++++---------- 2 files changed, 18 insertions(+), 20 deletions(-) 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..17a098157 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 @@ -20,7 +20,7 @@ public String authType() { return AZURE_CLI; } - public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { String azPath = Optional.ofNullable(config.getEnv()).map(env -> env.get("AZ_PATH")).orElse("az"); @@ -49,12 +49,11 @@ public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource return getTokenSource(config, cmd); } - protected CachedTokenSource getTokenSource(DatabricksConfig config, List cmd) { - CliTokenSource tokenSource = + protected CliTokenSource getTokenSource(DatabricksConfig config, List cmd) { + CliTokenSource token = new CliTokenSource(cmd, "tokenType", "accessToken", "expiresOn", config.getEnv()); - CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build(); - cachedTokenSource.getToken(); // Check if the CLI is installed and to validate the config. - return cachedTokenSource; + token.getToken(); // We need this to check if the CLI is installed and to validate the config. + return token; } private Optional getSubscription(DatabricksConfig config) { @@ -79,8 +78,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { try { AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); String resource = config.getEffectiveAzureLoginAppId(); - CachedTokenSource tokenSource = tokenSourceFor(config, resource); - CachedTokenSource mgmtTokenSource; + CliTokenSource tokenSource = tokenSourceFor(config, resource); + CliTokenSource mgmtTokenSource; try { mgmtTokenSource = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); @@ -88,7 +87,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOG.debug("Not including service management token in headers", e); mgmtTokenSource = null; } - CachedTokenSource finalMgmtTokenSource = mgmtTokenSource; + CachedTokenSource finalMgmtTokenSource = + new CachedTokenSource.Builder(mgmtTokenSource).build(); return OAuthHeaderFactory.fromSuppliers( tokenSource::getToken, () -> { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java index 9585d2de2..c38385bf4 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java @@ -5,9 +5,9 @@ import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.times; -import com.databricks.sdk.core.oauth.CachedTokenSource; import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.oauth.TokenSource; +import java.time.Instant; import java.util.Arrays; import java.util.List; import org.junit.jupiter.api.Test; @@ -22,11 +22,11 @@ class AzureCliCredentialsProviderTest { private static final String TOKEN = "t-123"; private static final String TOKEN_TYPE = "token-type"; - private static CachedTokenSource mockTokenSource() { - CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); - Mockito.when(cliTokenSource.getToken()) - .thenReturn(new Token(TOKEN, TOKEN_TYPE, java.time.Instant.now())); - return new CachedTokenSource.Builder(cliTokenSource).build(); + private static CliTokenSource mockTokenSource() { + CliTokenSource tokenSource = Mockito.mock(CliTokenSource.class); + Mockito.when(tokenSource.getToken()) + .thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now())); + return tokenSource; } private static AzureCliCredentialsProvider getAzureCliCredentialsProvider( @@ -62,12 +62,10 @@ void testWorkSpaceIDUsage() { @Test void testFallbackWhenTailsToGetTokenForSubscription() { - CachedTokenSource tokenSource = mockTokenSource(); + CliTokenSource tokenSource = mockTokenSource(); AzureCliCredentialsProvider provider = Mockito.spy(new AzureCliCredentialsProvider()); - Mockito.doThrow(new DatabricksException("error")) - .when(provider) - .getTokenSource(any(), anyList()); + Mockito.doThrow(new DatabricksException("error")).when(provider).getTokenSource(any(), anyList()); Mockito.doReturn(tokenSource).when(provider).getTokenSource(any(), anyList()); DatabricksConfig config = @@ -102,4 +100,4 @@ void testGetTokenWithoutWorkspaceResourceID() { assertFalse(value.contains("--subscription")); assertFalse(value.contains(SUBSCRIPTION)); } -} +} \ No newline at end of file From 8417e1b5e88b735168b9b1100bb0deae13319c5e Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 11 Jun 2025 13:04:42 +0000 Subject: [PATCH 35/47] Revert AzureCliCredentialsProvider --- .../sdk/core/AzureCliCredentialsProvider.java | 18 +++++++++--------- .../core/AzureCliCredentialsProviderTest.java | 17 ++++++++++------- 2 files changed, 19 insertions(+), 16 deletions(-) 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 17a098157..4ec58480e 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 @@ -20,7 +20,7 @@ public String authType() { return AZURE_CLI; } - public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { + public CachedTokenSource tokenSourceFor(DatabricksConfig config, String resource) { String azPath = Optional.ofNullable(config.getEnv()).map(env -> env.get("AZ_PATH")).orElse("az"); @@ -49,11 +49,12 @@ public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) { return getTokenSource(config, cmd); } - protected CliTokenSource getTokenSource(DatabricksConfig config, List cmd) { - CliTokenSource token = + protected CachedTokenSource getTokenSource(DatabricksConfig config, List cmd) { + CliTokenSource tokenSource = new CliTokenSource(cmd, "tokenType", "accessToken", "expiresOn", config.getEnv()); - token.getToken(); // We need this to check if the CLI is installed and to validate the config. - return token; + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource).build(); + cachedTokenSource.getToken(); // Check if the CLI is installed and to validate the config. + return cachedTokenSource; } private Optional getSubscription(DatabricksConfig config) { @@ -78,8 +79,8 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { try { AzureUtils.ensureHostPresent(config, mapper, this::tokenSourceFor); String resource = config.getEffectiveAzureLoginAppId(); - CliTokenSource tokenSource = tokenSourceFor(config, resource); - CliTokenSource mgmtTokenSource; + CachedTokenSource tokenSource = tokenSourceFor(config, resource); + CachedTokenSource mgmtTokenSource; try { mgmtTokenSource = tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint()); @@ -87,8 +88,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOG.debug("Not including service management token in headers", e); mgmtTokenSource = null; } - CachedTokenSource finalMgmtTokenSource = - new CachedTokenSource.Builder(mgmtTokenSource).build(); + CachedTokenSource finalMgmtTokenSource = mgmtTokenSource; return OAuthHeaderFactory.fromSuppliers( tokenSource::getToken, () -> { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java index c38385bf4..ba2e485f9 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java @@ -5,6 +5,7 @@ import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.times; +import com.databricks.sdk.core.oauth.CachedTokenSource; import com.databricks.sdk.core.oauth.Token; import com.databricks.sdk.core.oauth.TokenSource; import java.time.Instant; @@ -22,11 +23,11 @@ class AzureCliCredentialsProviderTest { private static final String TOKEN = "t-123"; private static final String TOKEN_TYPE = "token-type"; - private static CliTokenSource mockTokenSource() { - CliTokenSource tokenSource = Mockito.mock(CliTokenSource.class); - Mockito.when(tokenSource.getToken()) + private static CachedTokenSource mockTokenSource() { + CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); + Mockito.when(cliTokenSource.getToken()) .thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now())); - return tokenSource; + return new CachedTokenSource.Builder(cliTokenSource).build(); } private static AzureCliCredentialsProvider getAzureCliCredentialsProvider( @@ -62,10 +63,12 @@ void testWorkSpaceIDUsage() { @Test void testFallbackWhenTailsToGetTokenForSubscription() { - CliTokenSource tokenSource = mockTokenSource(); + CachedTokenSource tokenSource = mockTokenSource(); AzureCliCredentialsProvider provider = Mockito.spy(new AzureCliCredentialsProvider()); - Mockito.doThrow(new DatabricksException("error")).when(provider).getTokenSource(any(), anyList()); + Mockito.doThrow(new DatabricksException("error")) + .when(provider) + .getTokenSource(any(), anyList()); Mockito.doReturn(tokenSource).when(provider).getTokenSource(any(), anyList()); DatabricksConfig config = @@ -100,4 +103,4 @@ void testGetTokenWithoutWorkspaceResourceID() { assertFalse(value.contains("--subscription")); assertFalse(value.contains(SUBSCRIPTION)); } -} \ No newline at end of file +} From e843dc739e72682bf2cb50b3547dfab61a15082a Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 11 Jun 2025 14:01:34 +0000 Subject: [PATCH 36/47] Update variable names --- .../core/oauth/ExternalBrowserCredentialsProvider.java | 10 +++++----- .../OAuthM2MServicePrincipalCredentialsProvider.java | 5 +++-- .../sdk/core/oauth/TokenSourceCredentialsProvider.java | 6 +++--- .../sdk/core/AzureCliCredentialsProviderTest.java | 3 +-- 4 files changed, 12 insertions(+), 12 deletions(-) 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 2a9de00fe..605385f9e 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 @@ -81,23 +81,23 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOGGER.debug("Using cached token, will immediately refresh"); sessionCredentials.token = sessionCredentials.getToken(); - CachedTokenSource tokenSource = + CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(sessionCredentials) .withToken(sessionCredentials.token) .build(); - return OAuthHeaderFactory.fromTokenSource(tokenSource); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); } } // If no cached token or refresh failed, perform browser auth - CachedTokenSource tokenSource = + CachedTokenSource cachedTokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); - tokenCache.save(tokenSource.getToken()); + tokenCache.save(cachedTokenSource.getToken()); - return OAuthHeaderFactory.fromTokenSource(tokenSource); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; 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 743680d52..ec94c014b 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 @@ -37,9 +37,10 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { .withAuthParameterPosition(AuthParameterPosition.HEADER) .build(); - CachedTokenSource tokenSource = new CachedTokenSource.Builder(clientCredentials).build(); + CachedTokenSource cachedTokenSource = + new CachedTokenSource.Builder(clientCredentials).build(); - return OAuthHeaderFactory.fromTokenSource(tokenSource); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (IOException e) { // TODO: Log exception throw new DatabricksException("Unable to fetch OIDC endpoint: " + e.getMessage(), 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 46f67eb78..d3edae491 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 @@ -38,15 +38,15 @@ public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType) @Override public OAuthHeaderFactory configure(DatabricksConfig config) { // Check if the token source is already cached to prevent double caching - TokenSource source = + TokenSource cachedTokenSource = (tokenSource instanceof CachedTokenSource) ? tokenSource : new CachedTokenSource.Builder(tokenSource).build(); try { // Validate that we can get a token before returning a HeaderFactory - source.getToken().getAccessToken(); - return OAuthHeaderFactory.fromTokenSource(source); + cachedTokenSource.getToken().getAccessToken(); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (Exception e) { return null; } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java index ba2e485f9..051a87643 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java @@ -25,8 +25,7 @@ class AzureCliCredentialsProviderTest { private static CachedTokenSource mockTokenSource() { CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); - Mockito.when(cliTokenSource.getToken()) - .thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now())); + Mockito.when(cliTokenSource.getToken()).thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now())); return new CachedTokenSource.Builder(cliTokenSource).build(); } From b774e1adb3669cfed969e3285540c11d0b9d6f85 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 16 Jun 2025 15:25:21 +0000 Subject: [PATCH 37/47] clean up CachedTokenSource --- .../sdk/core/oauth/CachedTokenSource.java | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) 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 77b389333..d0aebf358 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 @@ -31,20 +31,26 @@ private enum TokenState { private static final Logger logger = LoggerFactory.getLogger(CachedTokenSource.class); // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); + // Default additional buffer before expiry to consider a token as expired. private static final Duration DEFAULT_EXPIRY_BUFFER = Duration.ofSeconds(40); + // The token source to use for refreshing the token. private final TokenSource tokenSource; + // Whether asynchronous refresh is enabled. private final boolean asyncEnabled; + // Duration before expiry to consider a token as 'stale'. private final Duration staleDuration; + // Additional buffer before expiry to consider a token as expired. private final Duration expiryBuffer; + // Clock supplier for current time, can be overridden for testing purposes. 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 volatile boolean refreshInProgress = false; + private boolean refreshInProgress = false; // Whether the last refresh attempt succeeded. - private volatile boolean lastRefreshSucceeded = true; + private boolean lastRefreshSucceeded = true; private CachedTokenSource(Builder builder) { this.tokenSource = builder.tokenSource; @@ -193,14 +199,15 @@ protected Token getTokenAsync() { * Trigger an asynchronous refresh of the token if not already in progress and last refresh * succeeded. */ - protected synchronized void triggerAsyncRefresh() { - // Check token state to avoid triggering a refresh if another thread has already refreshed it + private synchronized void triggerAsyncRefresh() { + // Check token state again inside the synchronized block to avoid triggering a refresh if + // another thread updated the token in the meantime. if (!refreshInProgress && lastRefreshSucceeded && getTokenState(token) != TokenState.FRESH) { refreshInProgress = true; CompletableFuture.runAsync( () -> { try { - // Attempt to refresh the token in the background + // Attempt to fetch the token in the background Token newToken = tokenSource.getToken(); synchronized (this) { token = newToken; @@ -210,7 +217,7 @@ protected synchronized void triggerAsyncRefresh() { synchronized (this) { lastRefreshSucceeded = false; refreshInProgress = false; - logger.error("Async token refresh failed", e); + logger.error("Asynchronous token refresh failed", e); } } }); From e846dd71daa9fae25b255f717b77375c9f082587 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Mon, 16 Jun 2025 15:41:08 +0000 Subject: [PATCH 38/47] Rename SystemClockSupplier to UtcClockSupplier --- .../com/databricks/sdk/core/oauth/CachedTokenSource.java | 6 +++--- .../{SystemClockSupplier.java => UtcClockSupplier.java} | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) rename databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/{SystemClockSupplier.java => UtcClockSupplier.java} (70%) 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 d0aebf358..64357e2e3 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 @@ -1,7 +1,7 @@ package com.databricks.sdk.core.oauth; import com.databricks.sdk.core.utils.ClockSupplier; -import com.databricks.sdk.core.utils.SystemClockSupplier; +import com.databricks.sdk.core.utils.UtcClockSupplier; import java.time.Duration; import java.time.Instant; import java.util.concurrent.CompletableFuture; @@ -42,7 +42,7 @@ private enum TokenState { private final Duration staleDuration; // Additional buffer before expiry to consider a token as expired. private final Duration expiryBuffer; - // Clock supplier for current time, can be overridden for testing purposes. + // Clock supplier for current time. private final ClockSupplier clockSupplier; // The current OAuth token. May be null if not yet fetched. @@ -67,7 +67,7 @@ public static class Builder { private boolean asyncEnabled = false; private Duration staleDuration = DEFAULT_STALE_DURATION; private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER; - private ClockSupplier clockSupplier = new SystemClockSupplier(); + private ClockSupplier clockSupplier = new UtcClockSupplier(); public Builder(TokenSource tokenSource) { this.tokenSource = tokenSource; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/UtcClockSupplier.java similarity index 70% rename from databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java rename to databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/UtcClockSupplier.java index c79e6884d..73e6d2bc4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/SystemClockSupplier.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/utils/UtcClockSupplier.java @@ -2,7 +2,7 @@ import java.time.Clock; -public class SystemClockSupplier implements ClockSupplier { +public class UtcClockSupplier implements ClockSupplier { @Override public Clock getClock() { return Clock.systemUTC(); From 8970e1b2cc4879ae6540d61c20f4ed9f213c3d3c Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Tue, 17 Jun 2025 11:52:57 +0000 Subject: [PATCH 39/47] Refactor --- .../databricks/sdk/core/oauth/Consent.java | 22 ++++++++++--------- .../ExternalBrowserCredentialsProvider.java | 9 +++----- .../sdk/core/oauth/SessionCredentials.java | 10 +++++---- ...xternalBrowserCredentialsProviderTest.java | 2 +- 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index f8e65d0a2..730ca4739 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -265,7 +265,7 @@ public Map getParams() { * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. * @throws IOException if the webserver cannot be started, or if the browser cannot be opened */ - public SessionCredentials launchExternalBrowser() throws IOException { + public CachedTokenSource launchExternalBrowser() throws IOException { URL redirect = new URL(getRedirectUrl()); if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) { throw new IllegalArgumentException( @@ -288,7 +288,7 @@ protected void desktopBrowser() throws IOException { Desktop.getDesktop().browse(URI.create(this.authUrl)); } - public SessionCredentials exchangeCallbackParameters(Map query) { + public CachedTokenSource exchangeCallbackParameters(Map query) { if (query.containsKey("error")) { throw new DatabricksException(query.get("error") + ": " + query.get("error_description")); } @@ -298,7 +298,7 @@ public SessionCredentials exchangeCallbackParameters(Map query) return exchange(query.get("code"), query.get("state")); } - private SessionCredentials exchange(String code, String state) { + public CachedTokenSource exchange(String code, String state) { if (!this.state.equals(state)) { throw new DatabricksException( "state mismatch: original state: " + this.state + "; retrieved state: " + state); @@ -321,12 +321,14 @@ private SessionCredentials exchange(String code, String state) { params, headers, AuthParameterPosition.BODY); - return new SessionCredentials.Builder() - .withHttpClient(this.hc) - .withClientId(this.clientId) - .withClientSecret(this.clientSecret) - .withTokenUrl(this.tokenUrl) - .withToken(token) - .build(); + SessionCredentials sessionCredentials = + new SessionCredentials.Builder() + .withHttpClient(this.hc) + .withClientId(this.clientId) + .withClientSecret(this.clientSecret) + .withTokenUrl(this.tokenUrl) + .withToken(token) + .build(); + return new CachedTokenSource.Builder(sessionCredentials).withToken(token).build(); } } 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 605385f9e..e66ee99d1 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 @@ -79,24 +79,22 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { .build(); LOGGER.debug("Using cached token, will immediately refresh"); - sessionCredentials.token = sessionCredentials.getToken(); - + sessionCredentials.getToken(); CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(sessionCredentials) .withToken(sessionCredentials.token) .build(); - return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); } } + // If no cached token or refresh failed, perform browser auth CachedTokenSource cachedTokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); tokenCache.save(cachedTokenSource.getToken()); - return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); @@ -120,9 +118,8 @@ CachedTokenSource performBrowserAuth( Consent consent = client.initiateConsent(); // Use the existing browser flow to get credentials - SessionCredentials credentials = consent.launchExternalBrowser(); + CachedTokenSource credentials = consent.launchExternalBrowser(); Token token = credentials.getToken(); - // Create a new SessionCredentials with the same token but with our token cache SessionCredentials sessionCredentials = new SessionCredentials.Builder() diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index 3bde86994..871c74021 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -28,7 +28,9 @@ public String authType() { @Override public OAuthHeaderFactory configure(DatabricksConfig config) { - return OAuthHeaderFactory.fromTokenSource(this); + CachedTokenSource cachedTokenSource = + new CachedTokenSource.Builder(this).withToken(this.token).build(); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } static class Builder { @@ -117,15 +119,15 @@ public Token getToken() { // cross-origin requests headers.put("Origin", redirectUrl); } - Token newToken = + this.token = TokenEndpointClient.retrieveToken( hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); // Save the refreshed token directly to cache if (tokenCache != null) { - tokenCache.save(newToken); + tokenCache.save(this.token); LOGGER.debug("Saved refreshed token to cache"); } - return newToken; + return this.token; } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index c4e7fe709..51a27111b 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -166,7 +166,7 @@ void exchangeTest() throws IOException { Map queryCreds = new HashMap<>(); queryCreds.put("code", "testCode"); queryCreds.put("state", "testState"); - SessionCredentials creds = testConsent.exchangeCallbackParameters(queryCreds); + CachedTokenSource creds = testConsent.exchangeCallbackParameters(queryCreds); assertEquals("accessTokenFromServer", creds.token.getAccessToken()); assertEquals("refreshTokenFromServer", creds.token.getRefreshToken()); } From c880323c0f47d4aef72f93d503f7cc75186c1eb6 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 18 Jun 2025 09:51:01 +0000 Subject: [PATCH 40/47] Clean up --- .../sdk/core/oauth/CachedTokenSource.java | 10 ++++---- .../sdk/core/oauth/CachedTokenSourceTest.java | 23 +++---------------- 2 files changed, 8 insertions(+), 25 deletions(-) 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 64357e2e3..076188113 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 @@ -63,11 +63,11 @@ private CachedTokenSource(Builder builder) { public static class Builder { private final TokenSource tokenSource; - private Token token; private boolean asyncEnabled = false; private Duration staleDuration = DEFAULT_STALE_DURATION; private Duration expiryBuffer = DEFAULT_EXPIRY_BUFFER; private ClockSupplier clockSupplier = new UtcClockSupplier(); + private Token token; public Builder(TokenSource tokenSource) { this.tokenSource = tokenSource; @@ -113,10 +113,10 @@ public CachedTokenSource build() { * @return The current valid token */ public Token getToken() { - if (!asyncEnabled) { - return getTokenBlocking(); + if (asyncEnabled) { + return getTokenAsync(); } - return getTokenAsync(); + return getTokenBlocking(); } /** @@ -207,7 +207,7 @@ private synchronized void triggerAsyncRefresh() { CompletableFuture.runAsync( () -> { try { - // Attempt to fetch the token in the background + // Attempt to refresh the token in the background Token newToken = tokenSource.getToken(); synchronized (this) { token = newToken; 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 3d0c3b12e..2cad04c9a 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 @@ -2,11 +2,9 @@ import static org.junit.jupiter.api.Assertions.*; -import com.databricks.sdk.core.utils.ClockSupplier; -import java.time.Clock; +import com.databricks.sdk.core.utils.TestClockSupplier; import java.time.Duration; import java.time.Instant; -import java.time.ZoneId; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.stream.Stream; @@ -146,6 +144,7 @@ public Token getToken() { // failed token = source.getToken(); assertEquals(INITIAL_TOKEN, token.getAccessToken(), "Should still return stale token"); + Thread.sleep(600); assertEquals( 1, testSource.refreshCallCount, @@ -160,6 +159,7 @@ public Token getToken() { REFRESH_TOKEN + "-2", token.getAccessToken(), "Should return the refreshed token after sync refresh"); + Thread.sleep(600); assertEquals( 2, testSource.refreshCallCount, @@ -177,21 +177,4 @@ public Token getToken() { testSource.refreshCallCount, "refresh() should have been called again asynchronously after making token stale"); } - - private static class TestClockSupplier implements ClockSupplier { - private Instant instant; - - TestClockSupplier(Instant instant) { - this.instant = instant; - } - - void advanceTime(Duration duration) { - this.instant = instant.plus(duration); - } - - @Override - public Clock getClock() { - return Clock.fixed(instant, ZoneId.of("UTC")); - } - } } From 75367bfc7b977d4eda43daa779d3f1c988bec2e0 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Wed, 18 Jun 2025 15:15:14 +0000 Subject: [PATCH 41/47] Add javadoc to builder class --- .../sdk/core/oauth/CachedTokenSource.java | 72 ++++++++++++++++++- 1 file changed, 70 insertions(+), 2 deletions(-) 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 076188113..ce1569738 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 @@ -4,6 +4,7 @@ import com.databricks.sdk.core.utils.UtcClockSupplier; import java.time.Duration; import java.time.Instant; +import java.util.Objects; import java.util.concurrent.CompletableFuture; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -32,6 +33,8 @@ private enum TokenState { // Default duration before expiry to consider a token as 'stale'. private static final Duration DEFAULT_STALE_DURATION = Duration.ofMinutes(3); // Default additional buffer before expiry to consider a token as expired. + // This is 40 seconds by default since Azure Databricks rejects tokens that are within 30 seconds + // of expiry. private static final Duration DEFAULT_EXPIRY_BUFFER = Duration.ofSeconds(40); // The token source to use for refreshing the token. @@ -61,6 +64,12 @@ private CachedTokenSource(Builder builder) { this.token = builder.token; } + /** + * Builder for creating instances of {@link CachedTokenSource}. + * + *

This builder allows configuration of various aspects of token caching behavior, including + * asynchronous refresh, timing parameters, and initial token state. + */ public static class Builder { private final TokenSource tokenSource; private boolean asyncEnabled = false; @@ -69,35 +78,94 @@ public static class Builder { private ClockSupplier clockSupplier = new UtcClockSupplier(); private Token token; + /** + * Creates a new builder with the specified token source. + * + * @param tokenSource The underlying token source to use for refreshing tokens. + * @throws NullPointerException If the token source is null. + */ public Builder(TokenSource tokenSource) { - this.tokenSource = tokenSource; + this.tokenSource = Objects.requireNonNull(tokenSource); } + /** + * Sets an initial token to use in the cache. + * + *

If provided, this token will be used immediately without requiring an initial refresh from + * the underlying token source. + * + * @param token The initial token to cache. + * @return This builder instance for method chaining. + */ public Builder withToken(Token token) { this.token = token; return this; } + /** + * Enables or disables asynchronous token refresh. + * + *

When enabled, stale tokens will trigger a background refresh while continuing to serve the + * 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. + * @return This builder instance for method chaining. + */ public Builder withAsyncEnabled(boolean asyncEnabled) { this.asyncEnabled = asyncEnabled; return this; } + /** + * Sets the duration before token expiry at which the token is considered stale. + * + *

When asynchronous refresh is enabled, tokens that are stale but not yet expired will + * trigger a background refresh while continuing to serve the current token. + * + * @param staleDuration The duration before expiry to consider a token stale. Must be greater + * than the expiry buffer duration. + * @return This builder instance for method chaining. + */ public Builder withStaleDuration(Duration staleDuration) { this.staleDuration = staleDuration; return this; } + /** + * Sets the buffer duration before token expiry at which the token is considered expired. + * + *

Tokens within this buffer of their expiry time will be considered expired and require + * synchronous refresh. + * + * @param expiryBuffer The buffer duration before expiry to consider a token expired. Must be + * less than the stale duration. + * @return This builder instance for method chaining. + */ public Builder withExpiryBuffer(Duration expiryBuffer) { this.expiryBuffer = expiryBuffer; return this; } + /** + * Sets the clock supplier to use for time-based operations. + * + *

This is primarily useful for testing scenarios where you need to control the current time. + * In production, the default UTC clock supplier should be used. + * + * @param clockSupplier The clock supplier to use for determining current time. + * @return This builder instance for method chaining. + */ public Builder withClockSupplier(ClockSupplier clockSupplier) { this.clockSupplier = clockSupplier; return this; } + /** + * Builds and returns a new {@link CachedTokenSource} instance with the configured parameters. + * + * @return A new CachedTokenSource instance. + */ public CachedTokenSource build() { return new CachedTokenSource(this); } @@ -207,7 +275,7 @@ private synchronized void triggerAsyncRefresh() { CompletableFuture.runAsync( () -> { try { - // Attempt to refresh the token in the background + // Attempt to refresh the token in the background. Token newToken = tokenSource.getToken(); synchronized (this) { token = newToken; From 68edacc26b52dba38919992c12d825592c705ee8 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 20 Jun 2025 13:39:44 +0000 Subject: [PATCH 42/47] Add new attribute to config --- .../databricks/sdk/core/DatabricksConfig.java | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) 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..316f40afd 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 = false; + public Environment getEnv() { return env; } @@ -410,13 +414,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** @deprecated Use {@link #getAzureUseMsi()} instead. */ + /** + * @deprecated Use {@link #getAzureUseMsi()} instead. + */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ + /** + * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. + */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; @@ -575,6 +583,15 @@ public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) { return this; } + public boolean getDisableAsyncTokenRefresh() { + return disableAsyncTokenRefresh; + } + + public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRefresh) { + this.disableAsyncTokenRefresh = disableAsyncTokenRefresh; + return this; + } + public boolean isAzure() { if (azureWorkspaceResourceId != null) { return true; From bc66554bfc84fa0296552620f5aeb034933d84a9 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 26 Jun 2025 12:42:23 +0000 Subject: [PATCH 43/47] Use set instead of with for cached token source builder --- .../databricks/sdk/core/oauth/CachedTokenSource.java | 10 +++++----- .../core/oauth/ExternalBrowserCredentialsProvider.java | 2 +- .../databricks/sdk/core/oauth/SessionCredentials.java | 2 +- .../sdk/core/oauth/CachedTokenSourceTest.java | 10 +++++----- .../oauth/ExternalBrowserCredentialsProviderTest.java | 4 ++-- 5 files changed, 14 insertions(+), 14 deletions(-) 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 a0681ce9e..866d2092d 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 @@ -98,7 +98,7 @@ public Builder(TokenSource tokenSource) { * @param token The initial token to cache. * @return This builder instance for method chaining. */ - public Builder withToken(Token token) { + public Builder setToken(Token token) { this.token = token; return this; } @@ -113,7 +113,7 @@ public Builder withToken(Token token) { * @param asyncEnabled True to enable asynchronous refresh, false to disable. * @return This builder instance for method chaining. */ - public Builder withAsyncEnabled(boolean asyncEnabled) { + public Builder setAsyncEnabled(boolean asyncEnabled) { this.asyncEnabled = asyncEnabled; return this; } @@ -128,7 +128,7 @@ public Builder withAsyncEnabled(boolean asyncEnabled) { * than the expiry buffer duration. * @return This builder instance for method chaining. */ - public Builder withStaleDuration(Duration staleDuration) { + public Builder setStaleDuration(Duration staleDuration) { this.staleDuration = staleDuration; return this; } @@ -143,7 +143,7 @@ public Builder withStaleDuration(Duration staleDuration) { * less than the stale duration. * @return This builder instance for method chaining. */ - public Builder withExpiryBuffer(Duration expiryBuffer) { + public Builder setExpiryBuffer(Duration expiryBuffer) { this.expiryBuffer = expiryBuffer; return this; } @@ -157,7 +157,7 @@ public Builder withExpiryBuffer(Duration expiryBuffer) { * @param clockSupplier The clock supplier to use for determining current time. * @return This builder instance for method chaining. */ - public Builder withClockSupplier(ClockSupplier clockSupplier) { + public Builder setClockSupplier(ClockSupplier clockSupplier) { this.clockSupplier = clockSupplier; return this; } 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 3de7e0a18..e65c9a312 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 @@ -128,6 +128,6 @@ CachedTokenSource performBrowserAuth( Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()), Optional.ofNullable(tokenCache)); - return new CachedTokenSource.Builder(tokenSource).withToken(token).build(); + return new CachedTokenSource.Builder(tokenSource).setToken(token).build(); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index aa427ce62..165b504f0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -41,7 +41,7 @@ public String authType() { @Override public OAuthHeaderFactory configure(DatabricksConfig config) { CachedTokenSource cachedTokenSource = - new CachedTokenSource.Builder(tokenSource).withToken(tokenSource.getToken()).build(); + new CachedTokenSource.Builder(tokenSource).setToken(tokenSource.getToken()).build(); return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } 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 2cad04c9a..0a29f5a7c 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 @@ -67,8 +67,8 @@ public Token getToken() { CachedTokenSource source = new CachedTokenSource.Builder(tokenSource) - .withAsyncEnabled(asyncEnabled) - .withToken(initialToken) + .setAsyncEnabled(asyncEnabled) + .setToken(initialToken) .build(); Token token = source.getToken(); @@ -127,9 +127,9 @@ public Token getToken() { TestSource testSource = new TestSource(); CachedTokenSource source = new CachedTokenSource.Builder(testSource) - .withAsyncEnabled(true) - .withToken(staleToken) - .withClockSupplier(clockSupplier) + .setAsyncEnabled(true) + .setToken(staleToken) + .setClockSupplier(clockSupplier) .build(); // First call triggers async refresh, which fails diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 36fec527e..968cde75b 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -435,7 +435,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { Optional.empty()); CachedTokenSource cachedTokenSource = - new CachedTokenSource.Builder(browserAuthTokenSource).withToken(browserAuthToken).build(); + new CachedTokenSource.Builder(browserAuthTokenSource).setToken(browserAuthToken).build(); // Create config with failing HTTP client and mock token cache DatabricksConfig config = @@ -512,7 +512,7 @@ void cacheWithInvalidTokensTest() throws IOException { Optional.empty()); CachedTokenSource cachedTokenSource = - new CachedTokenSource.Builder(browserAuthTokenSource).withToken(browserAuthToken).build(); + new CachedTokenSource.Builder(browserAuthTokenSource).setToken(browserAuthToken).build(); // Create simple config DatabricksConfig config = From 3d6c39aeea2dc8d183a6457587821a4deb8b0bf2 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 26 Jun 2025 12:56:40 +0000 Subject: [PATCH 44/47] Update names --- .../sdk/core/oauth/ExternalBrowserCredentialsProvider.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 e65c9a312..0766d8a1b 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 @@ -89,10 +89,10 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } // If no cached token or refresh failed, perform browser auth - CachedTokenSource tokenSource = + CachedTokenSource cachedTokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); - tokenCache.save(tokenSource.getToken()); - return OAuthHeaderFactory.fromTokenSource(tokenSource); + tokenCache.save(cachedTokenSource.getToken()); + return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; From 28bcb03538657a639c3af9dafa720c3a7c4e2685 Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 26 Jun 2025 14:08:27 +0000 Subject: [PATCH 45/47] Enabled async by default via config --- .../sdk/core/AzureCliCredentialsProvider.java | 5 ++++- .../DatabricksCliCredentialsProvider.java | 5 ++++- .../databricks/sdk/core/DatabricksConfig.java | 4 ++-- .../AzureGithubOidcCredentialsProvider.java | 5 ++++- ...reServicePrincipalCredentialsProvider.java | 4 +++- .../sdk/core/oauth/CachedTokenSource.java | 19 +++++++++---------- .../sdk/core/oauth/DataPlaneTokenSource.java | 6 +++++- .../ExternalBrowserCredentialsProvider.java | 10 ++++++++-- .../oauth/GithubOidcCredentialsProvider.java | 5 ++++- ...2MServicePrincipalCredentialsProvider.java | 4 +++- .../oauth/TokenSourceCredentialsProvider.java | 4 +++- .../ServingEndpointsDataPlaneImpl.java | 5 ++++- .../sdk/core/oauth/CachedTokenSourceTest.java | 18 +++++++++--------- .../core/oauth/DataPlaneTokenSourceTest.java | 8 +++++--- 14 files changed, 67 insertions(+), 35 deletions(-) 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 316f40afd..6db6f1b1a 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 @@ -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 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 866d2092d..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. @@ -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); From 20ca15440fb1da49c7c2cb7a5b8531d1bec0443a Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Thu, 26 Jun 2025 16:32:48 +0000 Subject: [PATCH 46/47] revert formatter change --- .../java/com/databricks/sdk/core/DatabricksConfig.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 6db6f1b1a..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 @@ -414,17 +414,13 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) { return this; } - /** - * @deprecated Use {@link #getAzureUseMsi()} instead. - */ + /** @deprecated Use {@link #getAzureUseMsi()} instead. */ @Deprecated() public boolean getAzureUseMSI() { return azureUseMsi; } - /** - * @deprecated Use {@link #setAzureUseMsi(boolean)} instead. - */ + /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */ @Deprecated public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) { this.azureUseMsi = azureUseMsi; From 70e576c582cc3e5460252ee6f980537cde2d33dc Mon Sep 17 00:00:00 2001 From: emmyzhou-db Date: Fri, 27 Jun 2025 00:17:46 +0000 Subject: [PATCH 47/47] Update changelog --- NEXT_CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) 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