From 593861d59e172856f8ce9ad8b5f882c709b208ff Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Wed, 24 Jan 2024 14:51:33 -0800 Subject: [PATCH 1/7] Support Databricks InHouse OAuth in Azure Signed-off-by: Jacky Hu --- lib/DBSQLClient.ts | 1 + .../auth/DatabricksOAuth/OAuthManager.ts | 23 +++++++++++++++++-- lib/contracts/IDBSQLClient.ts | 1 + tests/unit/DBSQLClient.test.js | 14 +++++++++++ 4 files changed, 37 insertions(+), 2 deletions(-) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 5c25d540..c8bc1e10 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -128,6 +128,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I azureTenantId: options.azureTenantId, clientId: options.oauthClientId, clientSecret: options.oauthClientSecret, + useDatabricksOAuthInAzure: options.useDatabricksOAuthInAzure, context: this, }); case 'custom': diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index de805d1a..a2ee7286 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -13,6 +13,7 @@ export interface OAuthManagerOptions { clientId?: string; azureTenantId?: string; clientSecret?: string; + useDatabricksOAuthInAzure?: boolean; context: IClientContext; } @@ -38,6 +39,10 @@ export default abstract class OAuthManager { this.context = options.context; } + protected canUse(): boolean { + return true; + } + protected abstract getOIDCConfigUrl(): string; protected abstract getAuthorizationUrl(): string; @@ -195,7 +200,10 @@ export default abstract class OAuthManager { for (const OAuthManagerClass of managers) { for (const domain of OAuthManagerClass.domains) { if (host.endsWith(domain)) { - return new OAuthManagerClass(options); + const manager: OAuthManager = new OAuthManagerClass(options); + if (manager.canUse()) { + return manager; + } } } } @@ -204,13 +212,21 @@ export default abstract class OAuthManager { } } + // Databricks InHouse OAuth Manager export class AWSOAuthManager extends OAuthManager { - public static domains = ['.cloud.databricks.com', '.dev.databricks.com']; + public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; + + public static domains = ['.cloud.databricks.com', '.dev.databricks.com'].concat(this.azureDomains); public static defaultClientId = 'databricks-sql-connector'; public static defaultCallbackPorts = [8030]; + protected canUse(): boolean { + const isAzureDomain = AWSOAuthManager.azureDomains.some((domain) => this.options.host.endsWith(domain)); + return !isAzureDomain || !!this.options.useDatabricksOAuthInAzure; + } + protected getOIDCConfigUrl(): string { return `${getDatabricksOIDCUrl(this.options.host)}/.well-known/oauth-authorization-server`; } @@ -231,6 +247,9 @@ export class AWSOAuthManager extends OAuthManager { export class AzureOAuthManager extends OAuthManager { public static domains = ['.azuredatabricks.net', '.databricks.azure.cn', '.databricks.azure.us']; + public static canBeUsed = (options: OAuthManagerOptions) => + this.domains.some((domain) => options.host.endsWith(domain)); + public static defaultClientId = '96eecda7-19ea-49cc-abb5-240097d554f5'; public static defaultCallbackPorts = [8030]; diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts index 0a14f435..0c0cee89 100644 --- a/lib/contracts/IDBSQLClient.ts +++ b/lib/contracts/IDBSQLClient.ts @@ -19,6 +19,7 @@ type AuthOptions = azureTenantId?: string; oauthClientId?: string; oauthClientSecret?: string; + useDatabricksOAuthInAzure?: boolean; } | { authType: 'custom'; diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index b1a1f3f2..11c33da2 100644 --- a/tests/unit/DBSQLClient.test.js +++ b/tests/unit/DBSQLClient.test.js @@ -359,6 +359,20 @@ describe('DBSQLClient.initAuthProvider', () => { expect(provider.manager).to.be.instanceOf(AzureOAuthManager); }); + it('should use Databricks InHouse OAuth method (Azure)', () => { + 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 Azure instance + host: 'example.databricks.azure.us', + useDatabricksOAuthInAzure: true + }); + + expect(provider).to.be.instanceOf(DatabricksOAuth); + expect(provider.manager).to.be.instanceOf(AWSOAuthManager); + }); + it('should throw error when OAuth not supported for host', () => { const client = new DBSQLClient(); From e37600b4e96a99c6df662020337208cb327d800d Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Wed, 24 Jan 2024 17:48:48 -0800 Subject: [PATCH 2/7] Rename AWSOAuthManager to DatabricksOAuthManager Signed-off-by: Jacky Hu --- lib/connection/auth/DatabricksOAuth/OAuthManager.ts | 12 ++++++------ tests/unit/DBSQLClient.test.js | 12 ++++++++---- .../auth/DatabricksOAuth/OAuthManager.test.js | 4 ++-- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index a2ee7286..4f2e1a79 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -195,7 +195,7 @@ export default abstract class OAuthManager { const host = options.host.toLowerCase().replace('https://', '').split('/')[0]; // eslint-disable-next-line @typescript-eslint/no-use-before-define - const managers = [AWSOAuthManager, AzureOAuthManager]; + const managers = [DatabricksOAuthManager, AzureOAuthManager]; for (const OAuthManagerClass of managers) { for (const domain of OAuthManagerClass.domains) { @@ -212,8 +212,8 @@ export default abstract class OAuthManager { } } - // Databricks InHouse OAuth Manager -export class AWSOAuthManager extends OAuthManager { +// Databricks InHouse OAuth Manager +export class DatabricksOAuthManager extends OAuthManager { public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; public static domains = ['.cloud.databricks.com', '.dev.databricks.com'].concat(this.azureDomains); @@ -223,7 +223,7 @@ export class AWSOAuthManager extends OAuthManager { public static defaultCallbackPorts = [8030]; protected canUse(): boolean { - const isAzureDomain = AWSOAuthManager.azureDomains.some((domain) => this.options.host.endsWith(domain)); + const isAzureDomain = DatabricksOAuthManager.azureDomains.some((domain) => this.options.host.endsWith(domain)); return !isAzureDomain || !!this.options.useDatabricksOAuthInAzure; } @@ -236,11 +236,11 @@ export class AWSOAuthManager extends OAuthManager { } protected getClientId(): string { - return this.options.clientId ?? AWSOAuthManager.defaultClientId; + return this.options.clientId ?? DatabricksOAuthManager.defaultClientId; } protected getCallbackPorts(): Array { - return this.options.callbackPorts ?? AWSOAuthManager.defaultCallbackPorts; + return this.options.callbackPorts ?? DatabricksOAuthManager.defaultCallbackPorts; } } diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index 11c33da2..5d9d9c6f 100644 --- a/tests/unit/DBSQLClient.test.js +++ b/tests/unit/DBSQLClient.test.js @@ -5,9 +5,13 @@ const DBSQLSession = require('../../dist/DBSQLSession').default; const PlainHttpAuthentication = require('../../dist/connection/auth/PlainHttpAuthentication').default; const DatabricksOAuth = require('../../dist/connection/auth/DatabricksOAuth').default; -const { AWSOAuthManager, AzureOAuthManager } = require('../../dist/connection/auth/DatabricksOAuth/OAuthManager'); +const { + DatabricksOAuthManager, + AzureOAuthManager, +} = require('../../dist/connection/auth/DatabricksOAuth/OAuthManager'); const HttpConnectionModule = require('../../dist/connection/connections/HttpConnection'); + const { default: HttpConnection } = HttpConnectionModule; class AuthProviderMock { @@ -343,7 +347,7 @@ describe('DBSQLClient.initAuthProvider', () => { }); expect(provider).to.be.instanceOf(DatabricksOAuth); - expect(provider.manager).to.be.instanceOf(AWSOAuthManager); + expect(provider.manager).to.be.instanceOf(DatabricksOAuthManager); }); it('should use Databricks OAuth method (Azure)', () => { @@ -366,11 +370,11 @@ describe('DBSQLClient.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 + useDatabricksOAuthInAzure: true, }); expect(provider).to.be.instanceOf(DatabricksOAuth); - expect(provider.manager).to.be.instanceOf(AWSOAuthManager); + expect(provider.manager).to.be.instanceOf(DatabricksOAuthManager); }); it('should throw error when OAuth not supported for host', () => { diff --git a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js index 814f3faa..aab82975 100644 --- a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js +++ b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js @@ -3,7 +3,7 @@ const sinon = require('sinon'); const openidClientLib = require('openid-client'); const { DBSQLLogger, LogLevel } = require('../../../../../dist'); const { - AWSOAuthManager, + DatabricksOAuthManager, AzureOAuthManager, } = require('../../../../../dist/connection/auth/DatabricksOAuth/OAuthManager'); const OAuthToken = require('../../../../../dist/connection/auth/DatabricksOAuth/OAuthToken').default; @@ -110,7 +110,7 @@ class OAuthClientMock { } } -[AWSOAuthManager, AzureOAuthManager].forEach((OAuthManagerClass) => { +[DatabricksOAuthManager, AzureOAuthManager].forEach((OAuthManagerClass) => { function prepareTestInstances(options) { const oauthClient = new OAuthClientMock(); sinon.stub(oauthClient, 'grant').callThrough(); From 767c8a3928126c22d52d697ffe25f4e2fc997b25 Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Wed, 24 Jan 2024 18:16:13 -0800 Subject: [PATCH 3/7] Remove dead code Signed-off-by: Jacky Hu --- lib/connection/auth/DatabricksOAuth/OAuthManager.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 4f2e1a79..7d756151 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -216,7 +216,7 @@ export default abstract class OAuthManager { export class DatabricksOAuthManager extends OAuthManager { public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; - public static domains = ['.cloud.databricks.com', '.dev.databricks.com'].concat(this.azureDomains); + public static domains = ['.cloud.databricks.com', '.dev.databricks.com','.gcp.databricks.com'].concat(this.azureDomains); public static defaultClientId = 'databricks-sql-connector'; @@ -247,9 +247,6 @@ export class DatabricksOAuthManager extends OAuthManager { export class AzureOAuthManager extends OAuthManager { public static domains = ['.azuredatabricks.net', '.databricks.azure.cn', '.databricks.azure.us']; - public static canBeUsed = (options: OAuthManagerOptions) => - this.domains.some((domain) => options.host.endsWith(domain)); - public static defaultClientId = '96eecda7-19ea-49cc-abb5-240097d554f5'; public static defaultCallbackPorts = [8030]; From 3c33dbd7fe8bb7ec943dbee206c7b1f58323f5ff Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Thu, 25 Jan 2024 10:05:27 -0800 Subject: [PATCH 4/7] Format fix Signed-off-by: Jacky Hu --- lib/connection/auth/DatabricksOAuth/OAuthManager.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 7d756151..ecf1edb6 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -216,7 +216,9 @@ export default abstract class OAuthManager { export class DatabricksOAuthManager extends OAuthManager { public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; - public static domains = ['.cloud.databricks.com', '.dev.databricks.com','.gcp.databricks.com'].concat(this.azureDomains); + public static domains = ['.cloud.databricks.com', '.dev.databricks.com', '.gcp.databricks.com'].concat( + this.azureDomains, + ); public static defaultClientId = 'databricks-sql-connector'; From c993f1ed221881cb85967d85559f7a98da963b14 Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Thu, 25 Jan 2024 10:29:46 -0800 Subject: [PATCH 5/7] Fix domain list Signed-off-by: Jacky Hu --- lib/connection/auth/DatabricksOAuth/OAuthManager.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index ecf1edb6..81914b6d 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -216,9 +216,7 @@ export default abstract class OAuthManager { export class DatabricksOAuthManager extends OAuthManager { public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; - public static domains = ['.cloud.databricks.com', '.dev.databricks.com', '.gcp.databricks.com'].concat( - this.azureDomains, - ); + public static domains = ['.cloud.databricks.com', '.dev.databricks.com'].concat(this.azureDomains); public static defaultClientId = 'databricks-sql-connector'; From b4f5aa66d4060f1c6449d347b36f496410d48421 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 29 Jan 2024 18:37:32 +0200 Subject: [PATCH 6/7] Refine code Signed-off-by: Levko Kravets --- .../auth/DatabricksOAuth/OAuthManager.ts | 43 ++++++++----------- 1 file changed, 17 insertions(+), 26 deletions(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 81914b6d..1e7baa14 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -39,10 +39,6 @@ export default abstract class OAuthManager { this.context = options.context; } - protected canUse(): boolean { - return true; - } - protected abstract getOIDCConfigUrl(): string; protected abstract getAuthorizationUrl(): string; @@ -194,18 +190,24 @@ export default abstract class OAuthManager { // normalize const host = options.host.toLowerCase().replace('https://', '').split('/')[0]; - // eslint-disable-next-line @typescript-eslint/no-use-before-define - const managers = [DatabricksOAuthManager, AzureOAuthManager]; - - for (const OAuthManagerClass of managers) { - for (const domain of OAuthManagerClass.domains) { - if (host.endsWith(domain)) { - const manager: OAuthManager = new OAuthManagerClass(options); - if (manager.canUse()) { - return manager; - } - } + const awsDomains = ['.cloud.databricks.com', '.dev.databricks.com']; + const isAWSDomain = awsDomains.some((domain) => host.endsWith(domain)); + + if (isAWSDomain) { + // eslint-disable-next-line @typescript-eslint/no-use-before-define + return new DatabricksOAuthManager(options); + } + + const azureDomains = ['.azuredatabricks.net', '.databricks.azure.cn', '.databricks.azure.us']; + const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain)); + + if (isAzureDomain) { + if (options.useDatabricksOAuthInAzure) { + // eslint-disable-next-line @typescript-eslint/no-use-before-define + return new DatabricksOAuthManager(options); } + // eslint-disable-next-line @typescript-eslint/no-use-before-define + return new AzureOAuthManager(options); } throw new Error(`OAuth is not supported for ${options.host}`); @@ -214,19 +216,10 @@ export default abstract class OAuthManager { // Databricks InHouse OAuth Manager export class DatabricksOAuthManager extends OAuthManager { - public static azureDomains = ['.azuredatabricks.net', '.databricks.azure.us']; - - public static domains = ['.cloud.databricks.com', '.dev.databricks.com'].concat(this.azureDomains); - public static defaultClientId = 'databricks-sql-connector'; public static defaultCallbackPorts = [8030]; - protected canUse(): boolean { - const isAzureDomain = DatabricksOAuthManager.azureDomains.some((domain) => this.options.host.endsWith(domain)); - return !isAzureDomain || !!this.options.useDatabricksOAuthInAzure; - } - protected getOIDCConfigUrl(): string { return `${getDatabricksOIDCUrl(this.options.host)}/.well-known/oauth-authorization-server`; } @@ -245,8 +238,6 @@ export class DatabricksOAuthManager extends OAuthManager { } export class AzureOAuthManager extends OAuthManager { - public static domains = ['.azuredatabricks.net', '.databricks.azure.cn', '.databricks.azure.us']; - public static defaultClientId = '96eecda7-19ea-49cc-abb5-240097d554f5'; public static defaultCallbackPorts = [8030]; From 3f4ef7c4ca9a33d80dc0a436ee3cd522b4166920 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 29 Jan 2024 20:06:40 +0200 Subject: [PATCH 7/7] Refine code Signed-off-by: Levko Kravets --- .../auth/DatabricksOAuth/OAuthManager.ts | 15 ++++++---- tests/unit/DBSQLClient.test.js | 30 ++++++++++++++----- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 1e7baa14..07c6a0c9 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -192,20 +192,23 @@ export default abstract class OAuthManager { const awsDomains = ['.cloud.databricks.com', '.dev.databricks.com']; const isAWSDomain = awsDomains.some((domain) => host.endsWith(domain)); - if (isAWSDomain) { // eslint-disable-next-line @typescript-eslint/no-use-before-define return new DatabricksOAuthManager(options); } - const azureDomains = ['.azuredatabricks.net', '.databricks.azure.cn', '.databricks.azure.us']; - const isAzureDomain = azureDomains.some((domain) => host.endsWith(domain)); - - if (isAzureDomain) { - if (options.useDatabricksOAuthInAzure) { + 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); } diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index 5d9d9c6f..2ec1ca29 100644 --- a/tests/unit/DBSQLClient.test.js +++ b/tests/unit/DBSQLClient.test.js @@ -366,15 +366,29 @@ describe('DBSQLClient.initAuthProvider', () => { it('should use Databricks InHouse OAuth method (Azure)', () => { 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 Azure instance - host: 'example.databricks.azure.us', - useDatabricksOAuthInAzure: true, - }); + case1: { + 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.azuredatabricks.net', + useDatabricksOAuthInAzure: true, + }); - expect(provider).to.be.instanceOf(DatabricksOAuth); - expect(provider.manager).to.be.instanceOf(DatabricksOAuthManager); + expect(provider).to.be.instanceOf(DatabricksOAuth); + expect(provider.manager).to.be.instanceOf(DatabricksOAuthManager); + } + + 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); + } }); it('should throw error when OAuth not supported for host', () => {