Merge branch 'main' into adjust-useDatabricksOAuthInAzure-behaviour

databricks · Feb 2, 2024 · 4aec809 · 4aec809
2 parents f53b654 + 4d4078b
commit 4aec809
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts
@@ -197,6 +197,13 @@ export default abstract class OAuthManager {
       return new DatabricksOAuthManager(options);
     }
 
+    const gcpDomains = ['.gcp.databricks.com'];
+    const isGCPDomain = gcpDomains.some((domain) => host.endsWith(domain));
+    if (isGCPDomain) {
+      // eslint-disable-next-line @typescript-eslint/no-use-before-define
+      return new DatabricksOAuthManager(options);
+    }
+
     if (options.useDatabricksOAuthInAzure) {
       const azureDomains = ['.azuredatabricks.net'];
       const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain));

diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js
@@ -363,6 +363,19 @@ describe('DBSQLClient.initAuthProvider', () => {
     expect(provider.manager).to.be.instanceOf(AzureOAuthManager);
   });
 
+  it('should use Databricks OAuth method (GCP)', () => {
+    const client = new DBSQLClient();
+
+    const provider = client.initAuthProvider({
+      authType: 'databricks-oauth',
+      // host is used when creating OAuth manager, so make it look like a real AWS instance
+      host: 'example.gcp.databricks.com',
+    });
+
+    expect(provider).to.be.instanceOf(DatabricksOAuth);
+    expect(provider.manager).to.be.instanceOf(DatabricksOAuthManager);
+  });
+
   it('should use Databricks InHouse OAuth method (Azure)', () => {
     const client = new DBSQLClient();