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

[divyavanmahajan]

Issue #121
Support Azure AD Tokens instead of PAT tokens

Summary:

PAT vs Azure AD tokens - PAT tokens are workspace specific and usually have a long expiry. Azure AD tokens expire within a few hours and the same token will work across all workspaces with SSO enabled.
Python driver supports Azure tokens but the NodeJS SQL driver does not.
This pull request is to support Azure AD tokens.

Description of change:

• Introduces useAADToken in connection options.
• When useAADToken=true, the PlainHTTPAuthentication class will send a Bearer authorization header with the Azure AD token, instead of the Basic authorization header when using a PAT token.
• Unit tests extended for PlainHTTPAuthentication for Azure tokens.
• E2E tests extended to support AAD Tokens

Workaround:

To use AAD token - you can tweak (@databricks/sql 1.1.1) as follows to use Azure AD tokens.

const { DBSQLClient, auth } = require('@databricks/sql');
// Patch to Databricks SQL driver to send Azure AD token as Bearer token.
Reflect.defineProperty(auth.PlainHttpAuthentication.prototype, 'getToken', {
  value: function () {
    return `Bearer ${this.password}`
  }
})

// Use normally - replacing token with your AAD token.

Submitted by:
Divya van Mahajan [email protected]

[superdupershant]

Wow awesome thanks for the submission @divyavanmahajan! we'll take a look.

[kravets-levko]

LGTM! 🚀

[kravets-levko]

@divyavanmahajan could you please run npm run prettier:fix + npm run lint over the codebase so lint check passes?

[andrefurlan-db]

I totally agree with the intent. It is unfortunately to wrong solution to the right issue.

Azure AD tokens work in all other drivers because Personal Access Token is implemented using Bearer tokens, and not basic auth. The fact that PAT works with basic auth is a Databricks idiosyncrasy.

So, the correct solution would be to set any token as Bearer, and not to add yet another flag to the driver interface.

Add the bearer Auth and change the code in DBSQLClient.ts

  async connect(options: ConnectionOptions): Promise<IDBSQLClient> {
    this.authProvider = new PlainHttpAuthentication({
      username: 'token',
      password: options.token, <-- root of the problem
      headers: {
        'User-Agent': buildUserAgentString(options.clientId),
      },
    });

[divyavanmahajan]

I totally agree with the intent. It is unfortunately to wrong solution to the right issue.

Azure AD tokens work in all other drivers because Personal Access Token is implemented using Bearer tokens, and not basic auth. The fact that PAT works with basic auth is a Databricks idiosyncrasy.

So, the correct solution would be to set any token as Bearer, and not to add yet another flag to the driver interface.

Add the bearer Auth and change the code in DBSQLClient.ts

  async connect(options: ConnectionOptions): Promise<IDBSQLClient> {
    this.authProvider = new PlainHttpAuthentication({
      username: 'token',
      password: options.token, <-- root of the problem
      headers: {
        'User-Agent': buildUserAgentString(options.clientId),
      },
    });

Hi Andre,
The logic to generate the Authorization header is buried deep inside the authProvider (PlainHTTPAuthentication class).
A user of the library can only pass parameters through the DBSQLClient.
So there were two design choices

1. Allow a parameter to override the authProvider and replace PlainHTTPAuthentication with user provided authProvider. (This was requested by fjakobs in Make auth provider configurable #120 ).
2. Allow a parameter that can be passed through to the authProvider.

In either case, you need to change the driver.
While selecting the design, I saw that #120 is not accepted so I dropped the first option. The second option is easier for a developer and had minimal changes to existing code.

Regards,
Divya

[kravets-levko]

Hi @divyavanmahajan,

I think Andre means that you don't need to add an option to specify token type, because Databricks one can work as a Bearer token. So basically you can remove the whole branch related to Basic auth, and keep Bearer one passing both Databricks PAT and AAD tokens, and it should work:

--- a/lib/connection/auth/PlainHttpAuthentication.ts
+++ b/lib/connection/auth/PlainHttpAuthentication.ts
@@ -29,6 +29,6 @@ export default class PlainHttpAuthentication implements IAuthentication {
   }
 
   private getToken(username: string, password: string): string {
-    return `Basic ${Buffer.from(`${username}:${password}`).toString('base64')}`;
+    return `Bearer ${password}`;
   }
 }

[divyavanmahajan]

Based on the suggestions from Andre - I am closing this pull request - and will create a new (smaller) pull request.
Thanks!