Re-create thrift client when auth credentials (e.g. oauth token) change

Signed-off-by: Levko Kravets <[email protected]>

lib/DBSQLClient.ts

@@ -3,7 +3,7 @@ import thrift, { HttpHeaders } from 'thrift';
 import { EventEmitter } from 'events';
 import TCLIService from '../thrift/TCLIService';
 import { TProtocolVersion } from '../thrift/TCLIService_types';
-import IDBSQLClient, { ConnectionOptions, OpenSessionRequest, ClientOptions } from './contracts/IDBSQLClient';
+import IDBSQLClient, { ClientOptions, ConnectionOptions, OpenSessionRequest } from './contracts/IDBSQLClient';
 import HiveDriver from './hive/HiveDriver';
 import { Int64 } from './hive/Types';
 import DBSQLSession from './DBSQLSession';
@@ -13,7 +13,7 @@ import HttpConnection from './connection/connections/HttpConnection';
 import IConnectionOptions from './connection/contracts/IConnectionOptions';
 import Status from './dto/Status';
 import HiveDriverError from './errors/HiveDriverError';
-import { buildUserAgentString, definedOrError } from './utils';
+import { areHeadersEqual, buildUserAgentString, definedOrError } from './utils';
 import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication';
 import DatabricksOAuth from './connection/auth/DatabricksOAuth';
 import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger';
@@ -46,6 +46,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
 
   private connectionOptions: ConnectionOptions | null = null;
 
+  private additionalHeaders: HttpHeaders = {};
+
   private readonly logger: IDBSQLLogger;
 
   private readonly thrift = thrift;
@@ -153,9 +155,13 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient {
       throw new HiveDriverError('DBSQLClient: not connected');
     }
 
-    if (!this.client) {
-      const authHeaders = await this.authProvider.authenticate();
-      const connectionOptions = this.getConnectionOptions(this.connectionOptions, authHeaders);
+    const authHeaders = await this.authProvider.authenticate();
+    // When auth headers change - recreate client. Thrift library does not provide API for updating
+    // changed options, therefore we have to recreate both connection and client to apply new headers
+    if (!this.client || !areHeadersEqual(this.additionalHeaders, authHeaders)) {
+      this.logger.log(LogLevel.info, 'DBSQLClient: initializing thrift client');
+      this.additionalHeaders = authHeaders;
+      const connectionOptions = this.getConnectionOptions(this.connectionOptions, this.additionalHeaders);
 
       const connection = await this.createConnection(connectionOptions);
       this.client = this.thrift.createClient(TCLIService, connection.getConnection());

lib/utils/areHeadersEqual.ts

@@ -0,0 +1,54 @@
+import { HttpHeaders } from 'thrift';
+
+function areArraysEqual<T>(a: Array<T>, b: Array<T>): boolean {
+  // If they're the same object - they're equal
+  if (a === b) {
+    return true;
+  }
+
+  // If they have a different size - they're definitely not equal
+  if (a.length !== b.length) {
+    return false;
+  }
+
+  // Here we have arrays of same size. Compare elements - if any pair is different
+  // then arrays are different
+  for (let i = 0; i < a.length; i += 1) {
+    if (a[i] !== b[i]) {
+      return false;
+    }
+  }
+
+  // If all corresponding elements in both arrays are equal - arrays are equal too
+  return true;
+}
+
+export default function areHeadersEqual(a: HttpHeaders, b: HttpHeaders): boolean {
+  // If they're the same object - they're equal
+  if (a === b) {
+    return true;
+  }
+
+  // If both objects have different keys - they're not equal
+  const keysOfA = Object.keys(a);
+  const keysOfB = Object.keys(b);
+  if (!areArraysEqual(keysOfA, keysOfB)) {
+    return false;
+  }
+
+  // Compare corresponding properties of both objects. If any pair is different - objects are different
+  for (const key of keysOfA) {
+    const propA = a[key];
+    const propB = b[key];
+
+    if (Array.isArray(propA) && Array.isArray(propB)) {
+      if (!areArraysEqual(propA, propB)) {
+        return false;
+      }
+    } else if (propA !== propB) {
+      return false;
+    }
+  }
+
+  return true;
+}

lib/utils/index.ts

@@ -1,5 +1,6 @@
+import areHeadersEqual from './areHeadersEqual';
 import definedOrError from './definedOrError';
 import buildUserAgentString from './buildUserAgentString';
 import formatProgress, { ProgressUpdateTransformer } from './formatProgress';
 
-export { definedOrError, buildUserAgentString, formatProgress, ProgressUpdateTransformer };
+export { areHeadersEqual, definedOrError, buildUserAgentString, formatProgress, ProgressUpdateTransformer };