Skip to content

Commit a4f4fdb

Browse files
authored
Use same logic as Go SDK for DatabricksConfig.isAzure() (#174)
## Changes The current isAzure() logic only checks if ".azuredatabricks." is contained in the hostname, which fails for workspaces in other Azure environments. This PR copies the logic directly from the Go SDK: https://github.com/databricks/databricks-sdk-go/blob/main/config/config.go#L142 ## Tests - [x] Some unit tests ensuring that isAzure() returns the right values for various hosts, including AWS, Azure public, Azure US Gov, and Azure China.
1 parent bf686df commit a4f4fdb

File tree

4 files changed

+29
-3
lines changed

4 files changed

+29
-3
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -463,7 +463,9 @@ public boolean isAzure() {
463463
if (host == null) {
464464
return false;
465465
}
466-
return host.contains(".azuredatabricks.");
466+
return host.contains(".azuredatabricks.net")
467+
|| host.contains("databricks.azure.cn")
468+
|| host.contains(".databricks.azure.us");
467469
}
468470

469471
public synchronized void authenticate(HttpMessage request) {

databricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ void testFallbackWhenTailsToGetTokenForSubscription() {
8686
void testGetTokenWithoutWorkspaceResourceID() {
8787
AzureCliCredentialsProvider provider = getAzureCliCredentialsProvider(mockTokenSource());
8888
DatabricksConfig config =
89-
new DatabricksConfig().setHost(".azuredatabricks.").setCredentialsProvider(provider);
89+
new DatabricksConfig().setHost(".azuredatabricks.net").setCredentialsProvider(provider);
9090

9191
ArgumentCaptor<List<String>> argument = ArgumentCaptor.forClass(List.class);
9292

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,28 @@ public void testIsAccountHostWorkspace() {
2525
assertFalse(
2626
new DatabricksConfig().setHost("https://westeurope.azuredatabricks.net").isAccountClient());
2727
}
28+
29+
@Test
30+
public void testIsAzureForAwsHost() {
31+
assertFalse(
32+
new DatabricksConfig().setHost("https://my-workspace.cloud.databricks.com/").isAzure());
33+
}
34+
35+
@Test
36+
public void testIsAzurePublic() {
37+
assertTrue(
38+
new DatabricksConfig().setHost("https://adb-1234567890.0.azuredatabricks.net/").isAzure());
39+
}
40+
41+
@Test
42+
public void testIsAzureMooncake() {
43+
assertTrue(
44+
new DatabricksConfig().setHost("https://adb-1234567890.0.databricks.azure.cn/").isAzure());
45+
}
46+
47+
@Test
48+
public void testIsAzureUsGov() {
49+
assertTrue(
50+
new DatabricksConfig().setHost("https://adb-1234567890.0.databricks.azure.us/").isAzure());
51+
}
2852
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ void testGetToken() {
4444
getAzureServicePrincipalCredentialsProvider(mockTokenSource());
4545
DatabricksConfig config =
4646
new DatabricksConfig()
47-
.setHost(".azuredatabricks.")
47+
.setHost(".azuredatabricks.net")
4848
.setCredentialsProvider(provider)
4949
.setAzureClientId("clientID")
5050
.setAzureClientSecret("clientSecret")

0 commit comments

Comments
 (0)