Support OAuth flow for Databricks Azure (#154)

kravets-levko · web-flow · commit 253ab0ff8acb · 2023-07-25T10:20:28.000+03:00
* [PECO-833] Support Azure OAuth flow

Signed-off-by: Levko Kravets &lt;levko.ne@gmail.com&gt;

* Fix tests

Signed-off-by: Levko Kravets &lt;levko.ne@gmail.com&gt;

* Add tests

Signed-off-by: Levko Kravets &lt;levko.ne@gmail.com&gt;

* Use custom authorization endpoint for Azure OAuth flow

Signed-off-by: Levko Kravets &lt;levko.ne@gmail.com&gt;

---------

Signed-off-by: Levko Kravets &lt;levko.ne@gmail.com&gt;
diff --git a/.eslintrc b/.eslintrc
@@ -14,7 +14,8 @@
         "no-bitwise": "off",
         "@typescript-eslint/no-throw-literal": "off",
         "no-restricted-syntax": "off",
-        "no-case-declarations": "off"
+        "no-case-declarations": "off",
+        "max-classes-per-file": "off"
       }
     }
   ]
diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts
@@ -108,6 +108,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
           host: options.host,
           logger: this.logger,
           persistence: options.persistence,
+          azureTenantId: options.azureTenantId,
         });
       case 'custom':
         return options.provider;
diff --git a/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts b/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts
@@ -2,15 +2,14 @@ import http, { IncomingMessage, Server, ServerResponse } from 'http';
 import { BaseClient, CallbackParamsType, generators } from 'openid-client';
 import open from 'open';
 import IDBSQLLogger, { LogLevel } from '../../../contracts/IDBSQLLogger';
+import { OAuthScopes, scopeDelimiter } from './OAuthScope';
 
 export interface AuthorizationCodeOptions {
   client: BaseClient;
   ports: Array<number>;
   logger?: IDBSQLLogger;
 }
 
-const scopeDelimiter = ' ';
-
 async function startServer(
   host: string,
   port: number,
@@ -76,7 +75,7 @@ export default class AuthorizationCode {
     return open(url);
   }
 
-  public async fetch(scopes: Array<string>): Promise<AuthorizationCodeFetchResult> {
+  public async fetch(scopes: OAuthScopes): Promise<AuthorizationCodeFetchResult> {
     const verifierString = generators.codeVerifier(32);
     const challengeString = generators.codeChallenge(verifierString);
     const state = generators.state(16);
@@ -96,7 +95,7 @@ export default class AuthorizationCode {
       }
     });
 
-    const redirectUri = `http://${server.host}:${server.port}/`;
+    const redirectUri = `http://${server.host}:${server.port}`;
     const authUrl = this.client.authorizationUrl({
       response_type: 'code',
       response_mode: 'query',
diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts
@@ -3,41 +3,62 @@ import HiveDriverError from '../../../errors/HiveDriverError';
 import IDBSQLLogger, { LogLevel } from '../../../contracts/IDBSQLLogger';
 import OAuthToken from './OAuthToken';
 import AuthorizationCode from './AuthorizationCode';
-
-const oidcConfigPath = 'oidc/.well-known/oauth-authorization-server';
+import { OAuthScope, OAuthScopes } from './OAuthScope';
 
 export interface OAuthManagerOptions {
   host: string;
-  callbackPorts: Array<number>;
-  clientId: string;
+  callbackPorts?: Array<number>;
+  clientId?: string;
+  azureTenantId?: string;
   logger?: IDBSQLLogger;
 }
 
-export default class OAuthManager {
-  private readonly options: OAuthManagerOptions;
+function getDatabricksOIDCUrl(host: string): string {
+  const schema = host.startsWith('https://') ? '' : 'https://';
+  const trailingSlash = host.endsWith('/') ? '' : '/';
+  return `${schema}${host}${trailingSlash}oidc`;
+}
+
+export default abstract class OAuthManager {
+  protected readonly options: OAuthManagerOptions;
 
-  private readonly logger?: IDBSQLLogger;
+  protected readonly logger?: IDBSQLLogger;
 
-  private issuer?: Issuer;
+  protected issuer?: Issuer;
 
-  private client?: BaseClient;
+  protected client?: BaseClient;
 
   constructor(options: OAuthManagerOptions) {
     this.options = options;
     this.logger = options.logger;
   }
 
-  private async getClient(): Promise<BaseClient> {
+  protected abstract getOIDCConfigUrl(): string;
+
+  protected abstract getAuthorizationUrl(): string;
+
+  protected abstract getClientId(): string;
+
+  protected abstract getCallbackPorts(): Array<number>;
+
+  protected getScopes(requestedScopes: OAuthScopes): OAuthScopes {
+    return requestedScopes;
+  }
+
+  protected async getClient(): Promise<BaseClient> {
     if (!this.issuer) {
-      const { host } = this.options;
-      const schema = host.startsWith('https://') ? '' : 'https://';
-      const trailingSlash = host.endsWith('/') ? '' : '/';
-      this.issuer = await Issuer.discover(`${schema}${host}${trailingSlash}${oidcConfigPath}`);
+      const issuer = await Issuer.discover(this.getOIDCConfigUrl());
+      // Overwrite `authorization_endpoint` in default config (specifically needed for Azure flow
+      // where this URL has to be different)
+      this.issuer = new Issuer({
+        ...issuer.metadata,
+        authorization_endpoint: this.getAuthorizationUrl(),
+      });
     }
 
     if (!this.client) {
       this.client = new this.issuer.Client({
-        client_id: this.options.clientId,
+        client_id: this.getClientId(),
         token_endpoint_auth_method: 'none',
       });
     }
@@ -76,15 +97,17 @@ export default class OAuthManager {
     return new OAuthToken(accessToken, refreshToken);
   }
 
-  public async getToken(scopes: Array<string>): Promise<OAuthToken> {
+  public async getToken(scopes: OAuthScopes): Promise<OAuthToken> {
     const client = await this.getClient();
     const authCode = new AuthorizationCode({
       client,
-      ports: this.options.callbackPorts,
+      ports: this.getCallbackPorts(),
       logger: this.logger,
     });
 
-    const { code, verifier, redirectUri } = await authCode.fetch(scopes);
+    const mappedScopes = this.getScopes(scopes);
+
+    const { code, verifier, redirectUri } = await authCode.fetch(mappedScopes);
 
     const { access_token: accessToken, refresh_token: refreshToken } = await client.grant({
       grant_type: 'authorization_code',
@@ -99,4 +122,84 @@ export default class OAuthManager {
 
     return new OAuthToken(accessToken, refreshToken);
   }
+
+  public static getManager(options: OAuthManagerOptions): OAuthManager {
+    // normalize
+    const host = options.host.toLowerCase().replace('https://', '').split('/')[0];
+
+    // eslint-disable-next-line @typescript-eslint/no-use-before-define
+    const managers = [AWSOAuthManager, AzureOAuthManager];
+
+    for (const OAuthManagerClass of managers) {
+      for (const domain of OAuthManagerClass.domains) {
+        if (host.endsWith(domain)) {
+          return new OAuthManagerClass(options);
+        }
+      }
+    }
+
+    throw new Error(`OAuth is not supported for ${options.host}`);
+  }
+}
+
+export class AWSOAuthManager extends OAuthManager {
+  public static domains = ['.cloud.databricks.com', '.dev.databricks.com'];
+
+  public static defaultClientId = 'databricks-sql-connector';
+
+  public static defaultCallbackPorts = [8030];
+
+  protected getOIDCConfigUrl(): string {
+    return `${getDatabricksOIDCUrl(this.options.host)}/.well-known/oauth-authorization-server`;
+  }
+
+  protected getAuthorizationUrl(): string {
+    return `${getDatabricksOIDCUrl(this.options.host)}/oauth2/v2.0/authorize`;
+  }
+
+  protected getClientId(): string {
+    return this.options.clientId ?? AWSOAuthManager.defaultClientId;
+  }
+
+  protected getCallbackPorts(): Array<number> {
+    return this.options.callbackPorts ?? AWSOAuthManager.defaultCallbackPorts;
+  }
+}
+
+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];
+
+  public static datatricksAzureApp = '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d';
+
+  protected getOIDCConfigUrl(): string {
+    return 'https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration';
+  }
+
+  protected getAuthorizationUrl(): string {
+    return `${getDatabricksOIDCUrl(this.options.host)}/oauth2/v2.0/authorize`;
+  }
+
+  protected getClientId(): string {
+    return this.options.clientId ?? AzureOAuthManager.defaultClientId;
+  }
+
+  protected getCallbackPorts(): Array<number> {
+    return this.options.callbackPorts ?? AzureOAuthManager.defaultCallbackPorts;
+  }
+
+  protected getScopes(requestedScopes: OAuthScopes): OAuthScopes {
+    // There is no corresponding scopes in Azure, instead, access control will be delegated to Databricks
+    const tenantId = this.options.azureTenantId ?? AzureOAuthManager.datatricksAzureApp;
+    const azureScopes = [`${tenantId}/user_impersonation`];
+
+    if (requestedScopes.includes(OAuthScope.offlineAccess)) {
+      azureScopes.push(OAuthScope.offlineAccess);
+    }
+
+    return azureScopes;
+  }
 }
diff --git a/lib/connection/auth/DatabricksOAuth/OAuthScope.ts b/lib/connection/auth/DatabricksOAuth/OAuthScope.ts
@@ -0,0 +1,10 @@
+export enum OAuthScope {
+  offlineAccess = 'offline_access',
+  SQL = 'sql',
+}
+
+export type OAuthScopes = Array<string>;
+
+export const defaultOAuthScopes: OAuthScopes = [OAuthScope.SQL, OAuthScope.offlineAccess];
+
+export const scopeDelimiter = ' ';
diff --git a/lib/connection/auth/DatabricksOAuth/index.ts b/lib/connection/auth/DatabricksOAuth/index.ts
@@ -3,69 +3,42 @@ import IAuthentication from '../../contracts/IAuthentication';
 import HttpTransport from '../../transports/HttpTransport';
 import IDBSQLLogger from '../../../contracts/IDBSQLLogger';
 import OAuthPersistence from './OAuthPersistence';
-import OAuthManager from './OAuthManager';
+import OAuthManager, { OAuthManagerOptions } from './OAuthManager';
+import { OAuthScopes, defaultOAuthScopes } from './OAuthScope';
 
-interface DatabricksOAuthOptions {
-  host: string;
-  redirectPorts?: Array<number>;
-  clientId?: string;
-  scopes?: Array<string>;
+interface DatabricksOAuthOptions extends OAuthManagerOptions {
+  scopes?: OAuthScopes;
   logger?: IDBSQLLogger;
   persistence?: OAuthPersistence;
   headers?: HttpHeaders;
 }
 
-const defaultOAuthOptions = {
-  clientId: 'databricks-sql-connector',
-  redirectPorts: [8030],
-  scopes: ['sql', 'offline_access'],
-} satisfies Partial<DatabricksOAuthOptions>;
-
 export default class DatabricksOAuth implements IAuthentication {
-  private readonly host: string;
-
-  private readonly redirectPorts: Array<number>;
-
-  private readonly clientId: string;
-
-  private readonly scopes: Array<string>;
+  private readonly options: DatabricksOAuthOptions;
 
   private readonly logger?: IDBSQLLogger;
 
-  private readonly persistence?: OAuthPersistence;
-
-  private readonly headers?: HttpHeaders;
-
   private readonly manager: OAuthManager;
 
   constructor(options: DatabricksOAuthOptions) {
-    this.host = options.host;
-    this.redirectPorts = options.redirectPorts || defaultOAuthOptions.redirectPorts;
-    this.clientId = options.clientId || defaultOAuthOptions.clientId;
-    this.scopes = options.scopes || defaultOAuthOptions.scopes;
+    this.options = options;
     this.logger = options.logger;
-    this.persistence = options.persistence;
-    this.headers = options.headers;
-
-    this.manager = new OAuthManager({
-      host: this.host,
-      callbackPorts: this.redirectPorts,
-      clientId: this.clientId,
-      logger: this.logger,
-    });
+    this.manager = OAuthManager.getManager(this.options);
   }
 
   public async authenticate(transport: HttpTransport): Promise<void> {
-    let token = await this.persistence?.read(this.host);
+    const { host, scopes, headers, persistence } = this.options;
+
+    let token = await persistence?.read(host);
     if (!token) {
-      token = await this.manager.getToken(this.scopes);
+      token = await this.manager.getToken(scopes ?? defaultOAuthScopes);
     }
 
     token = await this.manager.refreshAccessToken(token);
-    await this.persistence?.persist(this.host, token);
+    await persistence?.persist(host, token);
 
     transport.updateHeaders({
-      ...this.headers,
+      ...headers,
       Authorization: `Bearer ${token.accessToken}`,
     });
   }
diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts
@@ -15,6 +15,7 @@ type AuthOptions =
   | {
       authType: 'databricks-oauth';
       persistence?: OAuthPersistence;
+      azureTenantId?: string;
     }
   | {
       authType: 'custom';
diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js
@@ -5,6 +5,7 @@ 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 HttpConnection = require('../../dist/connection/connections/HttpConnection').default;
 
 const ConnectionProviderMock = (connection) => ({
@@ -255,14 +256,42 @@ describe('DBSQLClient.getAuthProvider', () => {
     expect(provider.password).to.be.equal(testAccessToken);
   });
 
-  it('should use Databricks OAuth method', () => {
+  it('should use Databricks OAuth method (AWS)', () => {
     const client = new DBSQLClient();
 
     const provider = client.getAuthProvider({
       authType: 'databricks-oauth',
+      // host is used when creating OAuth manager, so make it look like a real AWS instance
+      host: 'example.dev.databricks.com',
     });
 
     expect(provider).to.be.instanceOf(DatabricksOAuth);
+    expect(provider.manager).to.be.instanceOf(AWSOAuthManager);
+  });
+
+  it('should use Databricks OAuth method (Azure)', () => {
+    const client = new DBSQLClient();
+
+    const provider = client.getAuthProvider({
+      authType: 'databricks-oauth',
+      // host is used when creating OAuth manager, so make it look like a real Azure instance
+      host: 'example.databricks.azure.us',
+    });
+
+    expect(provider).to.be.instanceOf(DatabricksOAuth);
+    expect(provider.manager).to.be.instanceOf(AzureOAuthManager);
+  });
+
+  it('should throw error when OAuth not supported for host', () => {
+    const client = new DBSQLClient();
+
+    expect(() => {
+      client.getAuthProvider({
+        authType: 'databricks-oauth',
+        // use host which is not supported for sure
+        host: 'example.com',
+      });
+    }).to.throw();
   });
 
   it('should use custom auth method', () => {
diff --git a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js
diff --git a/tests/unit/connection/auth/DatabricksOAuth/index.test.js b/tests/unit/connection/auth/DatabricksOAuth/index.test.js

Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,8 @@`
`14`	`14`	`"no-bitwise": "off",`
`15`	`15`	`"@typescript-eslint/no-throw-literal": "off",`
`16`	`16`	`"no-restricted-syntax": "off",`
`17`		`- "no-case-declarations": "off"`
	`17`	`+ "no-case-declarations": "off",`
	`18`	`+ "max-classes-per-file": "off"`
`18`	`19`	`}`
`19`	`20`	`}`
`20`	`21`	`]`
Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@ type AuthOptions =`
`15`	`15`	`\| {`
`16`	`16`	`authType: 'databricks-oauth';`
`17`	`17`	`persistence?: OAuthPersistence;`
	`18`	`+ azureTenantId?: string;`
`18`	`19`	`}`
`19`	`20`	`\| {`
`20`	`21`	`authType: 'custom';`