Adjust useDatabricksOAuthInAzure behavior

Signed-off-by: Levko Kravets <levko.ne@gmail.com>
databricks · Feb 1, 2024 · f53b654 · f53b654
1 parent f3c53a5
commit f53b654
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts
@@ -198,19 +198,19 @@ export default abstract class OAuthManager {
     }
 
     if (options.useDatabricksOAuthInAzure) {
-      const domains = ['.azuredatabricks.net'];
-      const isSupportedDomain = domains.some((domain) => host.endsWith(domain));
-      if (isSupportedDomain) {
+      const azureDomains = ['.azuredatabricks.net'];
+      const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain));
+      if (isAzureDomain) {
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
         return new DatabricksOAuthManager(options);
       }
-    }
-
-    const azureDomains = ['.azuredatabricks.net', '.databricks.azure.us', '.databricks.azure.cn'];
-    const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain));
-    if (isAzureDomain) {
-      // eslint-disable-next-line @typescript-eslint/no-use-before-define
-      return new AzureOAuthManager(options);
+    } else {
+      const azureDomains = ['.azuredatabricks.net', '.databricks.azure.us', '.databricks.azure.cn'];
+      const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain));
+      if (isAzureDomain) {
+        // eslint-disable-next-line @typescript-eslint/no-use-before-define
+        return new AzureOAuthManager(options);
+      }
     }
 
     throw new Error(`OAuth is not supported for ${options.host}`);

diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js
@@ -366,6 +366,9 @@ describe('DBSQLClient.initAuthProvider', () => {
   it('should use Databricks InHouse OAuth method (Azure)', () => {
     const client = new DBSQLClient();
 
+    // When `useDatabricksOAuthInAzure = true`, it should use Databricks OAuth method
+    // only for supported Azure hosts, and fail for others
+
     case1: {
       const provider = client.initAuthProvider({
         authType: 'databricks-oauth',
@@ -379,15 +382,14 @@ describe('DBSQLClient.initAuthProvider', () => {
     }
 
     case2: {
-      const provider = client.initAuthProvider({
-        authType: 'databricks-oauth',
-        // host is used when creating OAuth manager, so make it look like a real Azure instance
-        host: 'example.databricks.azure.us',
-        useDatabricksOAuthInAzure: true,
-      });
-
-      expect(provider).to.be.instanceOf(DatabricksOAuth);
-      expect(provider.manager).to.be.instanceOf(AzureOAuthManager);
+      expect(() => {
+        const provider = client.initAuthProvider({
+          authType: 'databricks-oauth',
+          // host is used when creating OAuth manager, so make it look like a real Azure instance
+          host: 'example.databricks.azure.us',
+          useDatabricksOAuthInAzure: true,
+        });
+      }).to.throw();
     }
   });