Skip to content

Commit 8aad2be

Browse files
authored
Change DatabricksOAuthTokenSource to be refreshable (#446)
## What changes are proposed in this pull request? ## Make DatabricksOAuthTokenSource refreshable This change extends DatabricksOAuthTokenSource from RefreshableTokenSource to enable automatic token refresh functionality while maintaining the existing OAuth token exchange flow ## How is this tested? Unit testing NO_CHANGELOG=true
1 parent c092a76 commit 8aad2be

File tree

3 files changed

+9
-6
lines changed

3 files changed

+9
-6
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
* Implementation of TokenSource that handles OAuth token exchange for Databricks authentication.
1818
* This class manages the OAuth token exchange flow using ID tokens to obtain access tokens.
1919
*/
20-
public class DatabricksOAuthTokenSource implements TokenSource {
20+
public class DatabricksOAuthTokenSource extends RefreshableTokenSource {
2121
private static final Logger LOG = LoggerFactory.getLogger(DatabricksOAuthTokenSource.class);
2222

2323
/** OAuth client ID used for token exchange. */
@@ -154,7 +154,7 @@ private static void validate(Object value, String fieldName) {
154154
* parameters are missing.
155155
*/
156156
@Override
157-
public Token getToken() {
157+
public Token refresh() {
158158
// Validate all required parameters
159159
validate(clientId, "ClientID");
160160
validate(host, "Host");

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,14 @@ public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType)
4141
@Override
4242
public HeaderFactory configure(DatabricksConfig config) {
4343
try {
44-
// Validate that we can get a token before returning the HeaderFactory
45-
String accessToken = tokenSource.getToken().getAccessToken();
44+
// Validate that we can get a token before returning a HeaderFactory
45+
tokenSource.getToken().getAccessToken();
4646

4747
return () -> {
4848
Map<String, String> headers = new HashMap<>();
49-
headers.put("Authorization", "Bearer " + accessToken);
49+
// Some TokenSource implementations cache tokens internally, so an additional getToken()
50+
// call is not costly
51+
headers.put("Authorization", "Bearer " + tokenSource.getToken().getAccessToken());
5052
return headers;
5153
};
5254
} catch (Exception e) {

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ void testTokenScenarios(
4040
Map<String, String> headers = headerFactory.headers();
4141
assertEquals(expectedAuthHeader, headers.get("Authorization"));
4242
}
43-
verify(mockTokenSource).getToken();
43+
44+
verify(mockTokenSource, atLeastOnce()).getToken();
4445
assertEquals(TEST_AUTH_TYPE, provider.authType());
4546
}
4647

0 commit comments

Comments
 (0)