Skip to content

Commit c4433de

Browse files
vikrantpuppalavikrantpuppala
committed
address comments
1 parent c9fc008 commit c4433de

File tree

5 files changed

+86
-125
lines changed

5 files changed

+86
-125
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java‎

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,13 @@
44
import com.databricks.sdk.core.http.HttpClient;
55
import com.databricks.sdk.core.http.Request;
66
import com.databricks.sdk.core.http.Response;
7-
import com.databricks.sdk.core.oauth.FileTokenCache;
87
import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
9-
import com.databricks.sdk.core.oauth.TokenCache;
10-
import com.databricks.sdk.core.oauth.TokenCacheUtils;
118
import com.databricks.sdk.core.utils.Cloud;
129
import com.databricks.sdk.core.utils.Environment;
1310
import com.fasterxml.jackson.databind.ObjectMapper;
1411
import java.io.File;
1512
import java.io.IOException;
1613
import java.lang.reflect.Field;
17-
import java.nio.file.Path;
1814
import java.util.*;
1915
import org.apache.http.HttpMessage;
2016

@@ -145,9 +141,6 @@ public class DatabricksConfig {
145141

146142
private DatabricksEnvironment databricksEnvironment;
147143

148-
// Lazily initialized OAuth token cache
149-
private transient TokenCache tokenCache;
150-
151144
public Environment getEnv() {
152145
return env;
153146
}
@@ -677,17 +670,6 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
677670
return clone(fieldsToSkip).setHost(host);
678671
}
679672

680-
/**
681-
* Sets a custom TokenCache implementation.
682-
*
683-
* @param tokenCache the TokenCache implementation to use
684-
* @return this config instance
685-
*/
686-
public DatabricksConfig setTokenCache(TokenCache tokenCache) {
687-
this.tokenCache = tokenCache;
688-
return this;
689-
}
690-
691673
/**
692674
* Gets the default OAuth redirect URL. If one is not provided explicitly, uses
693675
* http://localhost:8080/callback
@@ -697,20 +679,4 @@ public DatabricksConfig setTokenCache(TokenCache tokenCache) {
697679
public String getEffectiveOAuthRedirectUrl() {
698680
return redirectUrl != null ? redirectUrl : "http://localhost:8080/callback";
699681
}
700-
701-
/**
702-
* Gets the TokenCache instance for the current configuration.
703-
*
704-
* <p>If a custom TokenCache has been set, it will be returned. Otherwise, a SimpleFileTokenCache
705-
* will be created based on the configuration properties.
706-
*
707-
* @return A TokenCache instance
708-
*/
709-
public synchronized TokenCache getTokenCache() {
710-
if (tokenCache == null) {
711-
Path cachePath = TokenCacheUtils.getCacheFilePath(getHost(), getClientId(), getScopes());
712-
tokenCache = new FileTokenCache(cachePath);
713-
}
714-
return tokenCache;
715-
}
716682
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java‎

Lines changed: 48 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,34 +5,56 @@
55
import com.databricks.sdk.core.DatabricksException;
66
import com.databricks.sdk.core.HeaderFactory;
77
import java.io.IOException;
8+
import java.nio.file.Path;
89
import org.slf4j.Logger;
910
import org.slf4j.LoggerFactory;
1011

1112
/**
1213
* A {@code CredentialsProvider} which implements the Authorization Code + PKCE flow by opening a
13-
* browser for the user to authorize the application. Uses the cache from {@link
14-
* DatabricksConfig#getTokenCache()}, tokens will be cached to avoid repeated authentication.
14+
* browser for the user to authorize the application. Uses a specified TokenCache or creates a
15+
* default one if none is provided.
1516
*/
1617
public class ExternalBrowserCredentialsProvider implements CredentialsProvider {
1718
private static final Logger LOGGER =
1819
LoggerFactory.getLogger(ExternalBrowserCredentialsProvider.class);
1920

21+
private TokenCache tokenCache;
22+
23+
/**
24+
* Creates a new ExternalBrowserCredentialsProvider with the specified TokenCache.
25+
*
26+
* @param tokenCache the TokenCache to use for caching tokens
27+
*/
28+
public ExternalBrowserCredentialsProvider(TokenCache tokenCache) {
29+
this.tokenCache = tokenCache;
30+
}
31+
32+
/**
33+
* Creates a new ExternalBrowserCredentialsProvider with a default TokenCache. A FileTokenCache
34+
* will be created when credentials are configured.
35+
*/
36+
public ExternalBrowserCredentialsProvider() {
37+
this(null);
38+
}
39+
2040
@Override
2141
public String authType() {
2242
return "external-browser";
2343
}
2444

2545
@Override
2646
public HeaderFactory configure(DatabricksConfig config) {
27-
if (config.getHost() == null
28-
|| config.getClientId() == null
29-
|| !config.getAuthType().equals("external-browser")) {
47+
if (config.getHost() == null || config.getAuthType() != "external-browser") {
3048
return null;
3149
}
32-
3350
try {
34-
// Get the token cache from config
35-
TokenCache tokenCache = config.getTokenCache();
51+
if (tokenCache == null) {
52+
// Create a default FileTokenCache based on config
53+
Path cachePath =
54+
TokenCacheUtils.getCacheFilePath(
55+
config.getHost(), config.getClientId(), config.getScopes());
56+
tokenCache = new FileTokenCache(cachePath);
57+
}
3658

3759
// First try to use the cached token if available (will return null if disabled)
3860
Token cachedToken = tokenCache.load();
@@ -49,11 +71,11 @@ public HeaderFactory configure(DatabricksConfig config) {
4971
.withClientSecret(config.getClientSecret())
5072
.withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
5173
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
74+
.withTokenCache(tokenCache)
5275
.build();
5376

5477
LOGGER.debug("Using cached token, will immediately refresh");
5578
cachedCreds.token = cachedCreds.refresh();
56-
tokenCache.save(cachedCreds.getToken());
5779
return cachedCreds.configure(config);
5880
} catch (Exception e) {
5981
// If token refresh fails, log and continue to browser auth
@@ -62,7 +84,7 @@ public HeaderFactory configure(DatabricksConfig config) {
6284
}
6385

6486
// If no cached token or refresh failed, perform browser auth
65-
SessionCredentials credentials = performBrowserAuth(config);
87+
SessionCredentials credentials = performBrowserAuth(config, tokenCache);
6688
tokenCache.save(credentials.getToken());
6789
return credentials.configure(config);
6890
} catch (IOException | DatabricksException e) {
@@ -71,10 +93,24 @@ public HeaderFactory configure(DatabricksConfig config) {
7193
}
7294
}
7395

74-
SessionCredentials performBrowserAuth(DatabricksConfig config) throws IOException {
96+
SessionCredentials performBrowserAuth(DatabricksConfig config, TokenCache tokenCache)
97+
throws IOException {
7598
LOGGER.debug("Performing browser authentication");
7699
OAuthClient client = new OAuthClient(config);
77100
Consent consent = client.initiateConsent();
78-
return consent.launchExternalBrowser();
101+
102+
// Use the existing browser flow to get credentials
103+
SessionCredentials credentials = consent.launchExternalBrowser();
104+
105+
// Create a new SessionCredentials with the same token but with our token cache
106+
return new SessionCredentials.Builder()
107+
.withToken(credentials.getToken())
108+
.withHttpClient(config.getHttpClient())
109+
.withClientId(config.getClientId())
110+
.withClientSecret(config.getClientSecret())
111+
.withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
112+
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
113+
.withTokenCache(tokenCache)
114+
.build();
79115
}
80116
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,6 @@ public String authType() {
3131

3232
@Override
3333
public HeaderFactory configure(DatabricksConfig config) {
34-
this.tokenCache = config.getTokenCache();
35-
3634
return () -> {
3735
Map<String, String> headers = new HashMap<>();
3836
headers.put(
@@ -48,6 +46,7 @@ static class Builder {
4846
private String redirectUrl;
4947
private String clientId;
5048
private String clientSecret;
49+
private TokenCache tokenCache;
5150

5251
public Builder withHttpClient(HttpClient hc) {
5352
this.hc = hc;
@@ -79,6 +78,11 @@ public Builder withClientSecret(String clientSecret) {
7978
return this;
8079
}
8180

81+
public Builder withTokenCache(TokenCache tokenCache) {
82+
this.tokenCache = tokenCache;
83+
return this;
84+
}
85+
8286
public SessionCredentials build() {
8387
return new SessionCredentials(this);
8488
}
@@ -89,7 +93,7 @@ public SessionCredentials build() {
8993
private final String redirectUrl;
9094
private final String clientId;
9195
private final String clientSecret;
92-
private transient TokenCache tokenCache;
96+
private final TokenCache tokenCache;
9397

9498
private SessionCredentials(Builder b) {
9599
super(b.token);
@@ -98,6 +102,7 @@ private SessionCredentials(Builder b) {
98102
this.redirectUrl = b.redirectUrl;
99103
this.clientId = b.clientId;
100104
this.clientSecret = b.clientSecret;
105+
this.tokenCache = b.tokenCache;
101106
}
102107

103108
@Override

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java‎

Lines changed: 0 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,14 @@
33
import static org.junit.jupiter.api.Assertions.*;
44

55
import com.databricks.sdk.core.commons.CommonsHttpClient;
6-
import com.databricks.sdk.core.oauth.FileTokenCache;
76
import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
8-
import com.databricks.sdk.core.oauth.TokenCache;
97
import com.databricks.sdk.core.utils.Environment;
108
import java.io.IOException;
119
import java.util.ArrayList;
12-
import java.util.Arrays;
1310
import java.util.HashMap;
1411
import java.util.List;
1512
import java.util.Map;
1613
import org.junit.jupiter.api.Test;
17-
import org.mockito.Mockito;
1814

1915
public class DatabricksConfigTest {
2016
@Test
@@ -199,43 +195,4 @@ public void testClone() {
199195
assert newWorkspaceConfig.getClientId().equals("my-client-id");
200196
assert newWorkspaceConfig.getClientSecret().equals("my-client-secret");
201197
}
202-
203-
@Test
204-
public void testGetTokenCache() {
205-
// Test that getTokenCache returns a FileTokenCache by default
206-
String testHost = "https://test-host.cloud.databricks.com";
207-
String testClientId = "test-client-id";
208-
List<String> testScopes = Arrays.asList("offline_access", "clusters", "sql");
209-
210-
DatabricksConfig config =
211-
new DatabricksConfig().setHost(testHost).setClientId(testClientId).setScopes(testScopes);
212-
213-
TokenCache cache = config.getTokenCache();
214-
215-
assertNotNull(cache, "TokenCache from config should not be null");
216-
assertInstanceOf(
217-
FileTokenCache.class, cache, "Default cache should be a FileTokenCache instance");
218-
}
219-
220-
@Test
221-
public void testSetTokenCache() {
222-
// Test that a custom token cache can be set and is returned by getTokenCache
223-
String testHost = "https://test-host.cloud.databricks.com";
224-
String testClientId = "test-client-id";
225-
List<String> testScopes = Arrays.asList("offline_access", "clusters", "sql");
226-
227-
// Create a mock token cache
228-
TokenCache mockCache = Mockito.mock(TokenCache.class);
229-
230-
DatabricksConfig config =
231-
new DatabricksConfig()
232-
.setHost(testHost)
233-
.setClientId(testClientId)
234-
.setScopes(testScopes)
235-
.setTokenCache(mockCache);
236-
237-
// Verify the custom cache is returned
238-
TokenCache returnedCache = config.getTokenCache();
239-
assertSame(mockCache, returnedCache, "Should return the custom token cache");
240-
}
241198
}

0 commit comments

Comments
 (0)