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.
databricks · Oct 25, 2023 · a4f4fdb · a4f4fdb
1 parent bf686df
commit a4f4fdb
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 3 deletions.
diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
@@ -463,7 +463,9 @@ public boolean isAzure() {
     if (host == null) {
       return false;
     }
-    return host.contains(".azuredatabricks.");
+    return host.contains(".azuredatabricks.net")
+        || host.contains("databricks.azure.cn")
+        || host.contains(".databricks.azure.us");
   }
 
   public synchronized void authenticate(HttpMessage request) {

diff --git a/...ricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java b/...ricks-sdk-java/src/test/java/com/databricks/sdk/core/AzureCliCredentialsProviderTest.java
@@ -86,7 +86,7 @@ void testFallbackWhenTailsToGetTokenForSubscription() {
   void testGetTokenWithoutWorkspaceResourceID() {
     AzureCliCredentialsProvider provider = getAzureCliCredentialsProvider(mockTokenSource());
     DatabricksConfig config =
-        new DatabricksConfig().setHost(".azuredatabricks.").setCredentialsProvider(provider);
+        new DatabricksConfig().setHost(".azuredatabricks.net").setCredentialsProvider(provider);
 
     ArgumentCaptor<List<String>> argument = ArgumentCaptor.forClass(List.class);
 

diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
@@ -25,4 +25,28 @@ public void testIsAccountHostWorkspace() {
     assertFalse(
         new DatabricksConfig().setHost("https://westeurope.azuredatabricks.net").isAccountClient());
   }
+
+  @Test
+  public void testIsAzureForAwsHost() {
+    assertFalse(
+        new DatabricksConfig().setHost("https://my-workspace.cloud.databricks.com/").isAzure());
+  }
+
+  @Test
+  public void testIsAzurePublic() {
+    assertTrue(
+        new DatabricksConfig().setHost("https://adb-1234567890.0.azuredatabricks.net/").isAzure());
+  }
+
+  @Test
+  public void testIsAzureMooncake() {
+    assertTrue(
+        new DatabricksConfig().setHost("https://adb-1234567890.0.databricks.azure.cn/").isAzure());
+  }
+
+  @Test
+  public void testIsAzureUsGov() {
+    assertTrue(
+        new DatabricksConfig().setHost("https://adb-1234567890.0.databricks.azure.us/").isAzure());
+  }
 }
diff --git a/...test/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProviderTest.java b/...test/java/com/databricks/sdk/core/oauth/AzureServicePrincipalCredentialsProviderTest.java
@@ -44,7 +44,7 @@ void testGetToken() {
         getAzureServicePrincipalCredentialsProvider(mockTokenSource());
     DatabricksConfig config =
         new DatabricksConfig()
-            .setHost(".azuredatabricks.")
+            .setHost(".azuredatabricks.net")
             .setCredentialsProvider(provider)
             .setAzureClientId("clientID")
             .setAzureClientSecret("clientSecret")