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 fe89b0a7f..b2625f82d 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) { // First try to use the cached token if available (will return null if disabled) Token cachedToken = tokenCache.load(); - if (cachedToken != null && cachedToken.getRefreshToken() != null) { + if (cachedToken != null) { LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId); try { @@ -84,9 +84,10 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { CachedTokenSource cachedTokenSource = new CachedTokenSource.Builder(tokenSource) + .setToken(cachedToken) .setAsyncDisabled(config.getDisableAsyncTokenRefresh()) .build(); - LOGGER.debug("Using cached token, will immediately refresh"); + LOGGER.debug("Using cached token, will refresh if necessary"); cachedTokenSource.getToken(); return OAuthHeaderFactory.fromTokenSource(cachedTokenSource); } catch (Exception e) { 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 f920f38d6..dbd1cba9f 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 @@ -246,25 +246,19 @@ void sessionCredentials() throws IOException { // Token caching tests @Test - void cacheWithValidTokenTest() throws IOException { - // Create mock HTTP client for token refresh + void cacheWithValidRefreshableTokenTest() throws IOException { + // Create mock HTTP client (shouldn't be called for valid token). HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - String refreshResponse = - "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; - URL url = new URL("https://test.databricks.com/"); - Mockito.doAnswer(invocation -> new Response(refreshResponse, url)) - .when(mockHttpClient) - .execute(any(Request.class)); - // Create an valid token with valid refresh token + // Create a valid token with valid refresh token (expires in 1 hour - FRESH state). 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 + // Create mock token cache that returns the valid token. TokenCache mockTokenCache = Mockito.mock(TokenCache.class); Mockito.doReturn(validToken).when(mockTokenCache).load(); - // Create config with HTTP client and mock token cache + // Create config with HTTP client and mock token cache. DatabricksConfig config = new DatabricksConfig() .setAuthType("external-browser") @@ -272,33 +266,38 @@ void cacheWithValidTokenTest() throws IOException { .setClientId("test-client-id") .setHttpClient(mockHttpClient); - // We need to provide OIDC endpoints for token refresh + // We need to provide OIDC endpoints. OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", "https://test.databricks.com/authorize"); + "https://test.databricks.com/oidc/v1/token", + "https://test.databricks.com/oidc/v1/authorize"); - // Create our provider with the mock token cache and mock the browser auth method + // Create our provider with the mock token cache. ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - // Spy on the config to inject the endpoints + // Spy on the config to inject the endpoints. DatabricksConfig spyConfig = Mockito.spy(config); Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); - // Configure provider + // Configure provider. HeaderFactory headerFactory = provider.configure(spyConfig); + assertNotNull(headerFactory, "HeaderFactory should be created"); - // Verify headers contain the refreshed token even though the cached token is valid + // Verify headers contain the CACHED valid token (no refresh needed!). Map headers = headerFactory.headers(); - assertEquals("Bearer refreshed_access_token", headers.get("Authorization")); + assertEquals( + "Bearer valid_access_token", + headers.get("Authorization"), + "Should use cached valid token without refreshing"); // Verify token was loaded from cache Mockito.verify(mockTokenCache, Mockito.times(1)).load(); - // Verify HTTP call was made to refresh the token - Mockito.verify(mockHttpClient, Mockito.times(1)).execute(any(Request.class)); + // Verify NO HTTP call was made (token is still valid, no refresh needed). + Mockito.verify(mockHttpClient, Mockito.never()).execute(any(Request.class)); - // Verify performBrowserAuth was NOT called since refresh succeeded + // Verify performBrowserAuth was NOT called since cached token is valid. Mockito.verify(provider, Mockito.never()) .performBrowserAuth( any(DatabricksConfig.class), @@ -306,23 +305,68 @@ void cacheWithValidTokenTest() throws IOException { any(String.class), any(TokenCache.class)); - // Verify token was saved back to cache - Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class)); + // Verify token was NOT saved back to cache (we're using the cached one as-is). + Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class)); + } - // Capture the token that was saved to cache to verify it's the refreshed token - ArgumentCaptor tokenCaptor = ArgumentCaptor.forClass(Token.class); - Mockito.verify(mockTokenCache).save(tokenCaptor.capture()); - Token savedToken = tokenCaptor.getValue(); + @Test + void cacheWithValidNonRefreshableTokenTest() throws IOException { + // Create mock HTTP client (shouldn't be called for valid token). + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - // Verify the saved token contains the refreshed values from the HTTP response - assertEquals( - "refreshed_access_token", - savedToken.getAccessToken(), - "Should save refreshed access token to cache"); + // Create a valid token WITHOUT refresh token (expires in 1 hour - FRESH state). + Instant futureTime = Instant.now().plusSeconds(3600); + Token validTokenNoRefresh = new Token("valid_access_token", "Bearer", null, futureTime); + + // Create mock token cache that returns the valid token. + TokenCache mockTokenCache = Mockito.mock(TokenCache.class); + Mockito.doReturn(validTokenNoRefresh).when(mockTokenCache).load(); + + // Create config with HTTP client and mock token cache. + DatabricksConfig config = + new DatabricksConfig() + .setAuthType("external-browser") + .setHost("https://test.databricks.com") + .setClientId("test-client-id") + .setHttpClient(mockHttpClient); + + // We need to provide OIDC endpoints. + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints( + "https://test.databricks.com/oidc/v1/token", + "https://test.databricks.com/oidc/v1/authorize"); + + // Create our provider with the mock token cache. + ExternalBrowserCredentialsProvider provider = + Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); + + // Spy on the config to inject the endpoints. + DatabricksConfig spyConfig = Mockito.spy(config); + Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints(); + + // Configure provider. + HeaderFactory headerFactory = provider.configure(spyConfig); + assertNotNull(headerFactory, "HeaderFactory should be created"); + + // Verify headers contain the cached token (NOT browser auth!). + Map headers = headerFactory.headers(); assertEquals( - "new_refresh_token", - savedToken.getRefreshToken(), - "Should save new refresh token to cache"); + "Bearer valid_access_token", + headers.get("Authorization"), + "Should use cached valid token even without refresh token"); + + // Verify token was loaded from cache. + Mockito.verify(mockTokenCache, Mockito.times(1)).load(); + + // Verify NO HTTP call was made (token is still valid, no refresh needed). + Mockito.verify(mockHttpClient, Mockito.never()).execute(any(Request.class)); + + // Verify performBrowserAuth was NOT called. + Mockito.verify(provider, Mockito.never()) + .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); + + // Verify no token was saved (we're using the cached one as-is). + Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class)); } @Test @@ -356,7 +400,8 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { // We need to provide OIDC endpoints for token refresh OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", "https://test.databricks.com/authorize"); + "https://test.databricks.com/oidc/v1/token", + "https://test.databricks.com/oidc/v1/authorize"); // Create our provider with the mock token cache ExternalBrowserCredentialsProvider provider = @@ -455,7 +500,8 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // We need to provide OIDC endpoints for token refresh attempt OpenIDConnectEndpoints endpoints = new OpenIDConnectEndpoints( - "https://test.databricks.com/token", "https://test.databricks.com/authorize"); + "https://test.databricks.com/oidc/v1/token", + "https://test.databricks.com/oidc/v1/authorize"); // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider =