Skip to content

Commit ab64d0f

Browse files
vikrantpuppalavikrantpuppala
committed
update logic used to get client id
1 parent 174c3c2 commit ab64d0f

File tree

5 files changed

+200
-30
lines changed

5 files changed

+200
-30
lines changed

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

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -48,31 +48,33 @@ public HeaderFactory configure(DatabricksConfig config) {
4848
if (config.getHost() == null || !Objects.equals(config.getAuthType(), "external-browser")) {
4949
return null;
5050
}
51-
if (config.getClientId() == null && config.getAzureClientId() == null) {
52-
config.setClientId("databricks-cli");
53-
}
51+
52+
// Use the utility class to resolve client ID and client secret
53+
String[] clientCreds = OAuthClientUtils.resolveClientCredentials(config);
54+
String clientId = clientCreds[0];
55+
String clientSecret = clientCreds[1];
56+
5457
try {
5558
if (tokenCache == null) {
5659
// Create a default FileTokenCache based on config
5760
Path cachePath =
58-
TokenCacheUtils.getCacheFilePath(
59-
config.getHost(), config.getClientId(), config.getScopes());
61+
TokenCacheUtils.getCacheFilePath(config.getHost(), clientId, config.getScopes());
6062
tokenCache = new FileTokenCache(cachePath);
6163
}
6264

6365
// First try to use the cached token if available (will return null if disabled)
6466
Token cachedToken = tokenCache.load();
6567
if (cachedToken != null && cachedToken.getRefreshToken() != null) {
66-
LOGGER.debug("Found cached token for {}:{}", config.getHost(), config.getClientId());
68+
LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId);
6769

6870
try {
6971
// Create SessionCredentials with the cached token and try to refresh if needed
7072
SessionCredentials cachedCreds =
7173
new SessionCredentials.Builder()
7274
.withToken(cachedToken)
7375
.withHttpClient(config.getHttpClient())
74-
.withClientId(config.getClientId())
75-
.withClientSecret(config.getClientSecret())
76+
.withClientId(clientId)
77+
.withClientSecret(clientSecret)
7678
.withTokenUrl(config.getOidcEndpoints().getTokenEndpoint())
7779
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
7880
.withTokenCache(tokenCache)
@@ -88,7 +90,8 @@ public HeaderFactory configure(DatabricksConfig config) {
8890
}
8991

9092
// If no cached token or refresh failed, perform browser auth
91-
SessionCredentials credentials = performBrowserAuth(config, tokenCache);
93+
SessionCredentials credentials =
94+
performBrowserAuth(config, clientId, clientSecret, tokenCache);
9295
tokenCache.save(credentials.getToken());
9396
return credentials.configure(config);
9497
} catch (IOException | DatabricksException e) {
@@ -97,10 +100,19 @@ public HeaderFactory configure(DatabricksConfig config) {
97100
}
98101
}
99102

100-
SessionCredentials performBrowserAuth(DatabricksConfig config, TokenCache tokenCache)
103+
SessionCredentials performBrowserAuth(
104+
DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache)
101105
throws IOException {
102106
LOGGER.debug("Performing browser authentication");
103-
OAuthClient client = new OAuthClient(config);
107+
OAuthClient client =
108+
new OAuthClient.Builder()
109+
.withHttpClient(config.getHttpClient())
110+
.withClientId(clientId)
111+
.withClientSecret(clientSecret)
112+
.withHost(config.getHost())
113+
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
114+
.withScopes(config.getScopes())
115+
.build();
104116
Consent consent = client.initiateConsent();
105117

106118
// Use the existing browser flow to get credentials

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

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -85,17 +85,6 @@ public OAuthClient build() throws IOException {
8585
private final boolean isAws;
8686
private final boolean isAzure;
8787

88-
public OAuthClient(DatabricksConfig config) throws IOException {
89-
this(
90-
new Builder()
91-
.withHttpClient(config.getHttpClient())
92-
.withClientId(config.getClientId())
93-
.withClientSecret(config.getClientSecret())
94-
.withHost(config.getHost())
95-
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
96-
.withScopes(config.getScopes()));
97-
}
98-
9988
private OAuthClient(Builder b) throws IOException {
10089
this.clientId = Objects.requireNonNull(b.clientId);
10190
this.clientSecret = b.clientSecret;

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksConfig;
4+
5+
/**
6+
* Utility methods for OAuth client credentials resolution.
7+
*/
8+
public class OAuthClientUtils {
9+
10+
/**
11+
* Default client ID to use when no client ID is specified.
12+
*/
13+
private static final String DEFAULT_CLIENT_ID = "databricks-cli";
14+
15+
/**
16+
* Resolves the OAuth client ID from the configuration.
17+
* Prioritizes regular OAuth client ID, then Azure client ID, and falls back to default client ID.
18+
*
19+
* @param config The Databricks configuration
20+
* @return The resolved client ID
21+
*/
22+
public static String resolveClientId(DatabricksConfig config) {
23+
if (config.getClientId() != null) {
24+
return config.getClientId();
25+
} else if (config.getAzureClientId() != null) {
26+
return config.getAzureClientId();
27+
}
28+
return DEFAULT_CLIENT_ID;
29+
}
30+
31+
/**
32+
* Resolves the OAuth client secret from the configuration.
33+
* Prioritizes regular OAuth client secret, then Azure client secret.
34+
*
35+
* @param config The Databricks configuration
36+
* @return The resolved client secret, or null if not present
37+
*/
38+
public static String resolveClientSecret(DatabricksConfig config) {
39+
if (config.getClientSecret() != null) {
40+
return config.getClientSecret();
41+
} else if (config.getAzureClientSecret() != null) {
42+
return config.getAzureClientSecret();
43+
}
44+
return null;
45+
}
46+
47+
/**
48+
* Resolves both client ID and client secret from the configuration.
49+
*
50+
* @param config The Databricks configuration
51+
* @return An array containing the client ID and client secret (may be null)
52+
*/
53+
public static String[] resolveClientCredentials(DatabricksConfig config) {
54+
String clientId = resolveClientId(config);
55+
String clientSecret = resolveClientSecret(config);
56+
return new String[]{clientId, clientSecret};
57+
}
58+
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,14 @@ void clientAndConsentTest() throws IOException {
4242

4343
assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint());
4444

45-
OAuthClient testClient = new OAuthClient(config);
45+
OAuthClient testClient = new OAuthClient.Builder()
46+
.withHttpClient(config.getHttpClient())
47+
.withClientId(config.getClientId())
48+
.withClientSecret(config.getClientSecret())
49+
.withHost(config.getHost())
50+
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
51+
.withScopes(config.getScopes())
52+
.build();
4653
assertEquals("test-client-id", testClient.getClientId());
4754

4855
Consent testConsent = testClient.initiateConsent();
@@ -79,7 +86,14 @@ void clientAndConsentTestWithCustomRedirectUrl() throws IOException {
7986

8087
assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint());
8188

82-
OAuthClient testClient = new OAuthClient(config);
89+
OAuthClient testClient = new OAuthClient.Builder()
90+
.withHttpClient(config.getHttpClient())
91+
.withClientId(config.getClientId())
92+
.withClientSecret(config.getClientSecret())
93+
.withHost(config.getHost())
94+
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
95+
.withScopes(config.getScopes())
96+
.build();
8397
assertEquals("test-client-id", testClient.getClientId());
8498

8599
Consent testConsent = testClient.initiateConsent();
@@ -261,7 +275,7 @@ void cacheWithValidTokenTest() throws IOException {
261275

262276
// Verify performBrowserAuth was NOT called since refresh succeeded
263277
Mockito.verify(provider, Mockito.never())
264-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
278+
.performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class));
265279

266280
// Verify token was saved back to cache
267281
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -338,7 +352,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException {
338352

339353
// Verify performBrowserAuth was NOT called since refresh succeeded
340354
Mockito.verify(provider, Mockito.never())
341-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
355+
.performBrowserAuth(any(DatabricksConfig.class), any(String.class), any(String.class), any(TokenCache.class));
342356

343357
// Verify token was saved back to cache
344358
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -409,7 +423,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
409423
Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache));
410424
Mockito.doReturn(browserAuthCreds)
411425
.when(provider)
412-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
426+
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
413427

414428
// Spy on the config to inject the endpoints
415429
DatabricksConfig spyConfig = Mockito.spy(config);
@@ -427,7 +441,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
427441

428442
// Verify performBrowserAuth was called since refresh failed
429443
Mockito.verify(provider, Mockito.times(1))
430-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
444+
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
431445

432446
// Verify token was saved after browser auth (for the new token)
433447
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -470,7 +484,7 @@ void cacheWithInvalidTokensTest() throws IOException {
470484
Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache));
471485
Mockito.doReturn(browserAuthCreds)
472486
.when(provider)
473-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
487+
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
474488

475489
// Configure provider
476490
HeaderFactory headerFactory = provider.configure(config);
@@ -483,7 +497,7 @@ void cacheWithInvalidTokensTest() throws IOException {
483497

484498
// Verify performBrowserAuth was called since we had an invalid token
485499
Mockito.verify(provider, Mockito.times(1))
486-
.performBrowserAuth(any(DatabricksConfig.class), any(TokenCache.class));
500+
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
487501

488502
// Verify token was saved after browser auth (for the new token)
489503
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/OAuthClientUtilsTest.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNull;
5+
6+
import com.databricks.sdk.core.DatabricksConfig;
7+
import org.junit.jupiter.api.Test;
8+
9+
public class OAuthClientUtilsTest {
10+
11+
@Test
12+
void resolveClientIdTest() {
13+
// Test with regular client ID
14+
DatabricksConfig config = new DatabricksConfig()
15+
.setClientId("test-client-id")
16+
.setAzureClientId("azure-client-id");
17+
assertEquals("test-client-id", OAuthClientUtils.resolveClientId(config));
18+
19+
// Test with only Azure client ID
20+
config = new DatabricksConfig()
21+
.setClientId(null)
22+
.setAzureClientId("azure-client-id");
23+
assertEquals("azure-client-id", OAuthClientUtils.resolveClientId(config));
24+
25+
// Test with no client IDs
26+
config = new DatabricksConfig()
27+
.setClientId(null)
28+
.setAzureClientId(null);
29+
assertEquals("databricks-cli", OAuthClientUtils.resolveClientId(config));
30+
}
31+
32+
@Test
33+
void resolveClientSecretTest() {
34+
// Test with regular client secret
35+
DatabricksConfig config = new DatabricksConfig()
36+
.setClientSecret("test-client-secret")
37+
.setAzureClientSecret("azure-client-secret");
38+
assertEquals("test-client-secret", OAuthClientUtils.resolveClientSecret(config));
39+
40+
// Test with only Azure client secret
41+
config = new DatabricksConfig()
42+
.setClientSecret(null)
43+
.setAzureClientSecret("azure-client-secret");
44+
assertEquals("azure-client-secret", OAuthClientUtils.resolveClientSecret(config));
45+
46+
// Test with no client secrets
47+
config = new DatabricksConfig()
48+
.setClientSecret(null)
49+
.setAzureClientSecret(null);
50+
assertNull(OAuthClientUtils.resolveClientSecret(config));
51+
}
52+
53+
@Test
54+
void resolveClientCredentialsTest() {
55+
// Test with both client ID and secret
56+
DatabricksConfig config = new DatabricksConfig()
57+
.setClientId("test-client-id")
58+
.setClientSecret("test-client-secret");
59+
String[] credentials = OAuthClientUtils.resolveClientCredentials(config);
60+
assertEquals("test-client-id", credentials[0]);
61+
assertEquals("test-client-secret", credentials[1]);
62+
63+
// Test with only client ID
64+
config = new DatabricksConfig()
65+
.setClientId("test-client-id")
66+
.setClientSecret(null);
67+
credentials = OAuthClientUtils.resolveClientCredentials(config);
68+
assertEquals("test-client-id", credentials[0]);
69+
assertNull(credentials[1]);
70+
71+
// Test with Azure credentials
72+
config = new DatabricksConfig()
73+
.setClientId(null)
74+
.setClientSecret(null)
75+
.setAzureClientId("azure-client-id")
76+
.setAzureClientSecret("azure-client-secret");
77+
credentials = OAuthClientUtils.resolveClientCredentials(config);
78+
assertEquals("azure-client-id", credentials[0]);
79+
assertEquals("azure-client-secret", credentials[1]);
80+
81+
// Test with no credentials
82+
config = new DatabricksConfig();
83+
credentials = OAuthClientUtils.resolveClientCredentials(config);
84+
assertEquals("databricks-cli", credentials[0]);
85+
assertNull(credentials[1]);
86+
87+
// Test mixed credentials preference
88+
config = new DatabricksConfig()
89+
.setClientId("test-client-id")
90+
.setClientSecret(null)
91+
.setAzureClientId("azure-client-id")
92+
.setAzureClientSecret("azure-client-secret");
93+
credentials = OAuthClientUtils.resolveClientCredentials(config);
94+
assertEquals("test-client-id", credentials[0]);
95+
assertEquals("azure-client-secret", credentials[1]);
96+
}
97+
}

0 commit comments

Comments
 (0)