Support Azure AD Tokens instead of PAT tokens (Issue 121)

docs/readme.md

@@ -2,20 +2,25 @@
 
 ## Table of Contents
 
-1. [Foreword](#foreword)
-2. [Example](#example) \
-   2.1. [Error handling](#error-handling)
-3. [DBSQLSession](#dbsqlsession)
-4. [DBSQLOperation](#dbsqloperation)
-5. [Status](#status)
-6. [Finalize](#finalize)
+- [Getting started](#getting-started)
+  - [Table of Contents](#table-of-contents)
+  - [Foreword](#foreword)
+  - [Example](#example)
+    - [Error handling](#error-handling)
+  - [DBSQLSession](#dbsqlsession)
+  - [DBSQLOperation](#dbsqloperation)
+    - [Example](#example-1)
+  - [Status](#status)
+  - [Finalize](#finalize)
 
 ## Foreword
 
 The library is written using TypeScript, so the best way to get to know how it works is to look through the code [lib/](/lib/), [tests/e2e](/tests/e2e/) and [examples](/examples).
 
 If you find any mistakes, misleading or some confusion feel free to create an issue or send a pull request.
 
+By default, the token is assumed to be a PAT token. If you are using AAD tokens, pass `useAADToken: true` in the connection options. To generate an AAD token, use the command line `az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d` or use the MSAL authentication library.
+
 ## Example
 
 ```javascript
@@ -28,6 +33,7 @@ client
     host: '...',
     path: '/sql/1.0/endpoints/****************',
     token: 'dapi********************************',
+    // useAADToken: true, // If token is an AAD token
   })
   .then(async (client) => {
     const session = await client.openSession();

examples/usage.js

@@ -5,9 +5,12 @@ const client = new DBSQLClient();
 const host = '****.databricks.com';
 const path = '/sql/1.0/endpoints/****';
 const token = 'dapi********************************';
-
+const useAADToken = false;
+// For AAD tokens
+// const token = 'ey***********';
+// const useAADToken=true;
 client
-  .connect({ host, path, token })
+  .connect({ host, path, token, useAADToken })
   .then(async (client) => {
     const session = await client.openSession();
     const response = await session.getInfo(thrift.TCLIService_types.TGetInfoType.CLI_DBMS_VER);

lib/DBSQLClient.ts

@@ -92,6 +92,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
     this.authProvider = new PlainHttpAuthentication({
       username: 'token',
       password: options.token,
+      useAADToken: options.useAADToken,
       headers: {
         'User-Agent': buildUserAgentString(options.clientId),
       },

lib/connection/auth/PlainHttpAuthentication.ts

@@ -4,19 +4,23 @@ import { AuthOptions } from '../types/AuthOptions';
 
 type HttpAuthOptions = AuthOptions & {
   headers?: object;
+  useAADToken?: boolean;
 };
 
 export default class PlainHttpAuthentication implements IAuthentication {
   private username: string;
 
   private password: string;
 
+  private useAADToken: boolean;
+
   private headers: object;
 
   constructor(options: HttpAuthOptions) {
     this.username = options?.username || 'anonymous';
     this.password = options?.password !== undefined ? options?.password : 'anonymous';
     this.headers = options?.headers || {};
+    this.useAADToken = options?.useAADToken !== undefined ? options?.useAADToken : false;
   }
 
   authenticate(transport: ITransport): Promise<ITransport> {
@@ -29,6 +33,9 @@ export default class PlainHttpAuthentication implements IAuthentication {
   }
 
   private getToken(username: string, password: string): string {
+    if (this.useAADToken) {
+      return `Bearer ${password}`;
+    }
     return `Basic ${Buffer.from(`${username}:${password}`).toString('base64')}`;
   }
 }

lib/contracts/IDBSQLClient.ts

@@ -11,6 +11,7 @@ export interface ConnectionOptions {
   path: string;
   token: string;
   clientId?: string;
+  useAADToken?: boolean;
 }
 
 export interface OpenSessionRequest {

tests/e2e/batched_fetch.test.js

@@ -11,6 +11,7 @@ const openSession = async () => {
     host: config.host,
     path: config.path,
     token: config.token,
+    useAADToken: config.useAADToken,
   });
 
   return connection.openSession({
@@ -37,14 +38,14 @@ describe('Data fetching', () => {
   it('fetch chunks should return a max row set of chunkSize', async () => {
     const session = await openSession();
     const operation = await session.executeStatement(query, { runAsync: true, maxRows: null });
-    let chunkedOp = await operation.fetchChunk({ maxRows: 10 }).catch((error) => logger(error));
+    const chunkedOp = await operation.fetchChunk({ maxRows: 10 }).catch((error) => logger(error));
     expect(chunkedOp.length).to.be.equal(10);
   });
 
   it('fetch all should fetch all records', async () => {
     const session = await openSession();
     const operation = await session.executeStatement(query, { runAsync: true, maxRows: null });
-    let all = await operation.fetchAll();
+    const all = await operation.fetchAll();
     expect(all.length).to.be.equal(1000);
   });
 });

tests/e2e/data_types.test.js

@@ -11,6 +11,7 @@ const openSession = async () => {
     host: config.host,
     path: config.path,
     token: config.token,
+    useAADToken: config.useAADToken,
   });
 
   return connection.openSession({

tests/e2e/utils/config.js

@@ -16,6 +16,8 @@ module.exports = {
   path: process.env.E2E_PATH,
   // Access token: dapi********************************
   token: process.env.E2E_ACCESS_TOKEN,
+  // Use AAD token: false
+  useAADToken: process.env.E2E_USEAADTOKEN || false,
   // Catalog and database to use for testing; specify both or leave array empty to use defaults
   database: catalog || database ? [catalog, database] : [],
   // Suffix used for tables that will be created during tests

tests/unit/connection/auth/PlainHttpAuthentication.test.js

@@ -17,6 +17,7 @@ describe('PlainHttpAuthentication', () => {
 
     expect(auth.username).to.be.eq('user');
     expect(auth.password).to.be.eq('pass');
+    expect(auth.useAADToken).to.be.eq(false);
   });
 
   it('empty password must be set', () => {
@@ -27,9 +28,10 @@ describe('PlainHttpAuthentication', () => {
 
     expect(auth.username).to.be.eq('user');
     expect(auth.password).to.be.eq('');
+    expect(auth.useAADToken).to.be.eq(false);
   });
 
-  it('auth token must be set to header', () => {
+  it('auth (PAT) token must be set to header', () => {
     const auth = new PlainHttpAuthentication();
     const transportMock = {
       setOptions(name, value) {
@@ -41,4 +43,21 @@ describe('PlainHttpAuthentication', () => {
       expect(transport).to.be.eq(transportMock);
     });
   });
+
+  it('auth (Azure) token must be set to header', () => {
+    const auth = new PlainHttpAuthentication({
+      username: 'user',
+      password: 'azureadtoken',
+      useAADToken: true,
+    });
+    const transportMock = {
+      setOptions(name, value) {
+        expect(name).to.be.eq('headers');
+        expect(value.Authorization).to.be.eq('Bearer azureadtoken');
+      },
+    };
+    return auth.authenticate(transportMock).then((transport) => {
+      expect(transport).to.be.eq(transportMock);
+    });
+  });
 });