From f53b654d5f9a38dc03435086ad2fbc7675fabdc6 Mon Sep 17 00:00:00 2001
From: Levko Kravets <levko.ne@gmail.com>
Date: Thu, 1 Feb 2024 16:46:20 +0200
Subject: [PATCH] Adjust `useDatabricksOAuthInAzure` behavior

Signed-off-by: Levko Kravets <levko.ne@gmail.com>
---
 .../auth/DatabricksOAuth/OAuthManager.ts      | 20 +++++++++----------
 tests/unit/DBSQLClient.test.js                | 20 ++++++++++---------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts
index 07c6a0c9..f5235cb8 100644
--- 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
index 2ec1ca29..89eac9f9 100644
--- 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();
     }
   });