Skip to content

Commit f06c006

Browse files
sreekanth-dbsreekanth-db
committed
addressing review comments
Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
1 parent 312c32a commit f06c006

3 files changed

Lines changed: 42 additions & 18 deletions

File tree

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

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@
2323
public class DatabricksConfig {
2424

2525
private static final Logger logger = LoggerFactory.getLogger(DatabricksConfig.class);
26+
27+
/** Azure authentication endpoint for tenant ID discovery */
28+
private static final String AZURE_AUTH_ENDPOINT = "/aad/auth";
29+
2630
private CredentialsProvider credentialsProvider = new DefaultCredentialsProvider();
2731

2832
@ConfigAttribute(env = "DATABRICKS_HOST")
@@ -731,7 +735,7 @@ private DatabricksConfig clone(Set<String> fieldsToSkip) {
731735
}
732736

733737
public DatabricksConfig clone() {
734-
return clone(new HashSet<>(Collections.singletonList("logger")));
738+
return clone(new HashSet<>(Arrays.asList("logger", "AZURE_AUTH_ENDPOINT")));
735739
}
736740

737741
public DatabricksConfig newWithWorkspaceHost(String host) {
@@ -742,6 +746,7 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
742746
// ID, and also omits
743747
// the account ID.
744748
"logger",
749+
"AZURE_AUTH_ENDPOINT",
745750
"host",
746751
"accountId",
747752
"azureWorkspaceResourceId",
@@ -765,33 +770,40 @@ public String getEffectiveOAuthRedirectUrl() {
765770
/**
766771
* [Internal] Load the Azure tenant ID from the Azure Databricks login page. If the tenant ID is
767772
* already set, this method does nothing.
773+
*
774+
* @return true if tenant ID is available (either was already set or successfully loaded), false otherwise
768775
*/
769-
public void loadAzureTenantId() {
776+
public boolean loadAzureTenantId() {
770777

771-
if (!isAzure() || azureTenantId != null || host == null) {
772-
return;
778+
if (azureTenantId != null) {
779+
return true; // Tenant ID already available - success
780+
}
781+
782+
if (!isAzure() || host == null) {
783+
return false; // Configuration issue - can't perform operation
773784
}
774785

775-
final String azureAuthEndpoint = "/aad/auth";
776-
String loginUrl = host + azureAuthEndpoint;
786+
String loginUrl = host + AZURE_AUTH_ENDPOINT;
777787
logger.debug("Loading tenant ID from {}", loginUrl);
778788

779789
try {
780790
String redirectLocation = getRedirectLocation(loginUrl);
781791
if (redirectLocation == null) {
782-
return;
792+
return false; // Failed to get redirect location
783793
}
784794

785795
String extractedTenantId = extractTenantIdFromUrl(redirectLocation);
786796
if (extractedTenantId == null) {
787-
return;
797+
return false; // Failed to extract tenant ID
788798
}
789799

790800
this.azureTenantId = extractedTenantId;
791801
logger.debug("Loaded tenant ID: {}", this.azureTenantId);
802+
return true; // Successfully loaded
792803

793804
} catch (Exception e) {
794805
logger.warn("Failed to load tenant ID: {}", e.getMessage());
806+
return false;
795807
}
796808
}
797809

@@ -802,9 +814,9 @@ private String getRedirectLocation(String loginUrl) throws IOException {
802814
Response response = getHttpClient().execute(request);
803815
int statusCode = response.getStatusCode();
804816

805-
if (statusCode / 100 != 3) {
817+
if (statusCode != 302) {
806818
logger.warn(
807-
"Failed to get tenant ID from {}: expected status code 3xx, got {}",
819+
"Failed to get tenant ID from {}: expected status code 302, got {}",
808820
loginUrl,
809821
statusCode);
810822
return null;

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,12 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
2727
}
2828
AzureUtils.ensureHostPresent(
2929
config, mapper, AzureServicePrincipalCredentialsProvider::tokenSourceFor);
30-
config.loadAzureTenantId();
30+
31+
boolean tenantIdLoaded = config.loadAzureTenantId();
32+
if (!tenantIdLoaded) {
33+
return null;
34+
}
35+
3136
CachedTokenSource inner = tokenSourceFor(config, config.getEffectiveAzureLoginAppId());
3237
CachedTokenSource cloud =
3338
tokenSourceFor(config, config.getAzureEnvironment().getServiceManagementEndpoint());

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,8 @@ public void testLoadAzureTenantId404() throws IOException {
257257
DatabricksConfig config = new DatabricksConfig();
258258
config.setHost(server.getUrl());
259259
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
260-
config.loadAzureTenantId();
260+
boolean result = config.loadAzureTenantId();
261+
assertFalse(result);
261262
assertNull(config.getAzureTenantId());
262263
}
263264
}
@@ -268,7 +269,8 @@ public void testLoadAzureTenantIdNoLocationHeader() throws IOException {
268269
DatabricksConfig config = new DatabricksConfig();
269270
config.setHost(server.getUrl());
270271
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
271-
config.loadAzureTenantId();
272+
boolean result = config.loadAzureTenantId();
273+
assertFalse(result);
272274
assertNull(config.getAzureTenantId());
273275
}
274276
}
@@ -286,7 +288,8 @@ public void testLoadAzureTenantIdUnparsableLocationHeader() throws IOException {
286288
DatabricksConfig config = new DatabricksConfig();
287289
config.setHost(server.getUrl());
288290
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
289-
config.loadAzureTenantId();
291+
boolean result = config.loadAzureTenantId();
292+
assertFalse(result);
290293
assertNull(config.getAzureTenantId());
291294
}
292295
}
@@ -306,7 +309,8 @@ public void testLoadAzureTenantIdHappyPath() throws IOException {
306309
config.setAzureWorkspaceResourceId(
307310
"/subscriptions/123/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/ws");
308311
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
309-
config.loadAzureTenantId();
312+
boolean result = config.loadAzureTenantId();
313+
assertTrue(result);
310314
assertEquals("test-tenant-id", config.getAzureTenantId());
311315
}
312316
}
@@ -316,7 +320,8 @@ public void testLoadAzureTenantIdSkipsWhenNotAzure() throws IOException {
316320
DatabricksConfig config = new DatabricksConfig();
317321
config.setHost("https://my-workspace.cloud.databricks.com"); // non-azure host
318322
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
319-
config.loadAzureTenantId();
323+
boolean result = config.loadAzureTenantId();
324+
assertFalse(result);
320325
assertNull(config.getAzureTenantId());
321326
}
322327

@@ -326,15 +331,17 @@ public void testLoadAzureTenantIdSkipsWhenAlreadySet() throws IOException {
326331
config.setHost("https://adb-123.0.azuredatabricks.net");
327332
config.setAzureTenantId("existing-tenant-id");
328333
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
329-
config.loadAzureTenantId();
334+
boolean result = config.loadAzureTenantId();
335+
assertTrue(result);
330336
assertEquals("existing-tenant-id", config.getAzureTenantId());
331337
}
332338

333339
@Test
334340
public void testLoadAzureTenantIdSkipsWhenNoHost() throws IOException {
335341
DatabricksConfig config = new DatabricksConfig();
336342
config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
337-
config.loadAzureTenantId();
343+
boolean result = config.loadAzureTenantId();
344+
assertFalse(result);
338345
assertNull(config.getAzureTenantId());
339346
}
340347
}

0 commit comments

Comments
 (0)