Adjust useDatabricksOAuthInAzure behavior

Signed-off-by: Levko Kravets <[email protected]>
databricks · Feb 1, 2024 · d1ce0c7 · d1ce0c7
1 parent f3c53a5
commit d1ce0c7
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 20 deletions.
diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts
@@ -197,20 +197,22 @@ export default abstract class OAuthManager {
       return new DatabricksOAuthManager(options);
     }
 
-    if (options.useDatabricksOAuthInAzure) {
-      const domains = ['.azuredatabricks.net'];
-      const isSupportedDomain = domains.some((domain) => host.endsWith(domain));
-      if (isSupportedDomain) {
-        // 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);
+      // When `useDatabricksOAuthInAzure = true`, it should use Databricks OAuth method
+      // only for supported Azure hosts, and fail for others
+      if (options.useDatabricksOAuthInAzure) {
+        const domains = ['.azuredatabricks.net'];
+        const isSupportedDomain = domains.some((domain) => host.endsWith(domain));
+        if (isSupportedDomain) {
+          // eslint-disable-next-line @typescript-eslint/no-use-before-define
+          return new DatabricksOAuthManager(options);
+        }
+      } else {
+        // 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();
     }
   });