From 9f3bd40dc82584a14d8c102c9e9a88959cf970db Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Wed, 31 Jan 2024 12:58:31 +0200 Subject: [PATCH] [PECO-1433] Support Databricks OAuth on GCP Signed-off-by: Levko Kravets --- lib/connection/auth/DatabricksOAuth/OAuthManager.ts | 7 +++++++ tests/unit/DBSQLClient.test.js | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 07c6a0c9..40bf7480 100644 --- 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 domains = ['.azuredatabricks.net']; const isSupportedDomain = domains.some((domain) => host.endsWith(domain)); diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index 2ec1ca29..f527aa26 100644 --- 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();