From d1ce0c75911eacedb7f9842dde41e242c75cc95e 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      | 24 ++++++++++---------
 tests/unit/DBSQLClient.test.js                | 20 +++++++++-------
 2 files changed, 24 insertions(+), 20 deletions(-)

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