diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 36250f95..e8188d2e 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -1,4 +1,4 @@ -import thrift, { HttpHeaders } from 'thrift'; +import thrift from 'thrift'; import { EventEmitter } from 'events'; import TCLIService from '../thrift/TCLIService'; @@ -13,12 +13,13 @@ import HttpConnection from './connection/connections/HttpConnection'; import IConnectionOptions from './connection/contracts/IConnectionOptions'; import Status from './dto/Status'; import HiveDriverError from './errors/HiveDriverError'; -import { areHeadersEqual, buildUserAgentString, definedOrError } from './utils'; +import { buildUserAgentString, definedOrError } from './utils'; import PlainHttpAuthentication from './connection/auth/PlainHttpAuthentication'; import DatabricksOAuth from './connection/auth/DatabricksOAuth'; import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger'; import DBSQLLogger from './DBSQLLogger'; import CloseableCollection from './utils/CloseableCollection'; +import IConnectionProvider from './connection/contracts/IConnectionProvider'; function prependSlash(str: string): string { if (str.length > 0 && str.charAt(0) !== '/') { @@ -41,13 +42,11 @@ function getInitialNamespaceOptions(catalogName?: string, schemaName?: string) { } export default class DBSQLClient extends EventEmitter implements IDBSQLClient { - private client: TCLIService.Client | null = null; + private connectionProvider?: IConnectionProvider; - private authProvider: IAuthentication | null = null; + private authProvider?: IAuthentication; - private connectionOptions: ConnectionOptions | null = null; - - private additionalHeaders: HttpHeaders = {}; + private client?: TCLIService.Client; private readonly logger: IDBSQLLogger; @@ -61,30 +60,14 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { this.logger.log(LogLevel.info, 'Created DBSQLClient'); } - private getConnectionOptions(options: ConnectionOptions, headers: HttpHeaders): IConnectionOptions { - const { - host, - port, - path, - clientId, - authType, - // @ts-expect-error TS2339: Property 'token' does not exist on type 'ConnectionOptions' - token, - // @ts-expect-error TS2339: Property 'persistence' does not exist on type 'ConnectionOptions' - persistence, - // @ts-expect-error TS2339: Property 'provider' does not exist on type 'ConnectionOptions' - provider, - ...otherOptions - } = options; - + private getConnectionOptions(options: ConnectionOptions): IConnectionOptions { return { - host, - port: port || 443, - path: prependSlash(path), + host: options.host, + port: options.port || 443, + path: prependSlash(options.path), https: true, - ...otherOptions, + socketTimeout: options.socketTimeout, headers: { - ...headers, 'User-Agent': buildUserAgentString(options.clientId), }, }; @@ -128,7 +111,38 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { */ public async connect(options: ConnectionOptions, authProvider?: IAuthentication): Promise { this.authProvider = this.getAuthProvider(options, authProvider); - this.connectionOptions = options; + + this.connectionProvider = new HttpConnection(this.getConnectionOptions(options)); + + const thriftConnection = await this.connectionProvider.getThriftConnection(); + + thriftConnection.on('error', (error: Error) => { + // Error.stack already contains error type and message, so log stack if available, + // otherwise fall back to just error type + message + this.logger.log(LogLevel.error, error.stack || `${error.name}: ${error.message}`); + try { + this.emit('error', error); + } catch (e) { + // EventEmitter will throw unhandled error when emitting 'error' event. + // Since we already logged it few lines above, just suppress this behaviour + } + }); + + thriftConnection.on('reconnecting', (params: { delay: number; attempt: number }) => { + this.logger.log(LogLevel.debug, `Reconnecting, params: ${JSON.stringify(params)}`); + this.emit('reconnecting', params); + }); + + thriftConnection.on('close', () => { + this.logger.log(LogLevel.debug, 'Closing connection.'); + this.emit('close'); + }); + + thriftConnection.on('timeout', () => { + this.logger.log(LogLevel.debug, 'Connection timed out.'); + this.emit('timeout'); + }); + return this; } @@ -158,65 +172,28 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { } private async getClient() { - if (!this.connectionOptions || !this.authProvider) { + if (!this.connectionProvider) { throw new HiveDriverError('DBSQLClient: not connected'); } - 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)) { + if (!this.client) { this.logger.log(LogLevel.info, 'DBSQLClient: initializing thrift client'); - this.additionalHeaders = authHeaders; - const connectionOptions = this.getConnectionOptions(this.connectionOptions, this.additionalHeaders); + this.client = this.thrift.createClient(TCLIService, await this.connectionProvider.getThriftConnection()); + } - const connection = await this.createConnection(connectionOptions); - this.client = this.thrift.createClient(TCLIService, connection.getConnection()); + if (this.authProvider) { + const authHeaders = await this.authProvider.authenticate(); + this.connectionProvider.setHeaders(authHeaders); } return this.client; } - private async createConnection(options: IConnectionOptions) { - const connectionProvider = new HttpConnection(); - const connection = await connectionProvider.connect(options); - const thriftConnection = connection.getConnection(); - - thriftConnection.on('error', (error: Error) => { - // Error.stack already contains error type and message, so log stack if available, - // otherwise fall back to just error type + message - this.logger.log(LogLevel.error, error.stack || `${error.name}: ${error.message}`); - try { - this.emit('error', error); - } catch (e) { - // EventEmitter will throw unhandled error when emitting 'error' event. - // Since we already logged it few lines above, just suppress this behaviour - } - }); - - thriftConnection.on('reconnecting', (params: { delay: number; attempt: number }) => { - this.logger.log(LogLevel.debug, `Reconnecting, params: ${JSON.stringify(params)}`); - this.emit('reconnecting', params); - }); - - thriftConnection.on('close', () => { - this.logger.log(LogLevel.debug, 'Closing connection.'); - this.emit('close'); - }); - - thriftConnection.on('timeout', () => { - this.logger.log(LogLevel.debug, 'Connection timed out.'); - this.emit('timeout'); - }); - - return connection; - } - public async close(): Promise { await this.sessions.closeAll(); - this.client = null; - this.authProvider = null; - this.connectionOptions = null; + this.client = undefined; + this.connectionProvider = undefined; + this.authProvider = undefined; } } diff --git a/lib/connection/auth/DatabricksOAuth/index.ts b/lib/connection/auth/DatabricksOAuth/index.ts index 20e29e16..29243b8f 100644 --- a/lib/connection/auth/DatabricksOAuth/index.ts +++ b/lib/connection/auth/DatabricksOAuth/index.ts @@ -1,4 +1,4 @@ -import { HttpHeaders } from 'thrift'; +import { HeadersInit } from 'node-fetch'; import IAuthentication from '../../contracts/IAuthentication'; import IDBSQLLogger from '../../../contracts/IDBSQLLogger'; import OAuthPersistence, { OAuthPersistenceCache } from './OAuthPersistence'; @@ -9,7 +9,7 @@ interface DatabricksOAuthOptions extends OAuthManagerOptions { scopes?: OAuthScopes; logger?: IDBSQLLogger; persistence?: OAuthPersistence; - headers?: HttpHeaders; + headers?: HeadersInit; } export default class DatabricksOAuth implements IAuthentication { @@ -27,7 +27,7 @@ export default class DatabricksOAuth implements IAuthentication { this.manager = OAuthManager.getManager(this.options); } - public async authenticate(): Promise { + public async authenticate(): Promise { const { host, scopes, headers } = this.options; const persistence = this.options.persistence ?? this.defaultPersistence; diff --git a/lib/connection/auth/PlainHttpAuthentication.ts b/lib/connection/auth/PlainHttpAuthentication.ts index ab0283e3..1115c745 100644 --- a/lib/connection/auth/PlainHttpAuthentication.ts +++ b/lib/connection/auth/PlainHttpAuthentication.ts @@ -1,10 +1,10 @@ -import { HttpHeaders } from 'thrift'; +import { HeadersInit } from 'node-fetch'; import IAuthentication from '../contracts/IAuthentication'; interface PlainHttpAuthenticationOptions { username?: string; password?: string; - headers?: HttpHeaders; + headers?: HeadersInit; } export default class PlainHttpAuthentication implements IAuthentication { @@ -12,7 +12,7 @@ export default class PlainHttpAuthentication implements IAuthentication { private readonly password: string; - private readonly headers: HttpHeaders; + private readonly headers: HeadersInit; constructor(options: PlainHttpAuthenticationOptions) { this.username = options?.username || 'anonymous'; @@ -20,7 +20,7 @@ export default class PlainHttpAuthentication implements IAuthentication { this.headers = options?.headers || {}; } - public async authenticate(): Promise { + public async authenticate(): Promise { return { ...this.headers, Authorization: `Bearer ${this.password}`, diff --git a/lib/connection/connections/HttpConnection.ts b/lib/connection/connections/HttpConnection.ts index 98d64202..e9e132d1 100644 --- a/lib/connection/connections/HttpConnection.ts +++ b/lib/connection/connections/HttpConnection.ts @@ -1,18 +1,36 @@ import thrift from 'thrift'; import https from 'https'; import http from 'http'; +import { HeadersInit } from 'node-fetch'; -import IThriftConnection from '../contracts/IThriftConnection'; import IConnectionProvider from '../contracts/IConnectionProvider'; import IConnectionOptions from '../contracts/IConnectionOptions'; import globalConfig from '../../globalConfig'; import ThriftHttpConnection from './ThriftHttpConnection'; -export default class HttpConnection implements IConnectionProvider, IThriftConnection { - private connection: any; +export default class HttpConnection implements IConnectionProvider { + private readonly options: IConnectionOptions; + + private headers: HeadersInit = {}; + + private connection?: ThriftHttpConnection; + + constructor(options: IConnectionOptions) { + this.options = options; + } + + public setHeaders(headers: HeadersInit) { + this.headers = headers; + this.connection?.setHeaders({ + ...this.options.headers, + ...this.headers, + }); + } + + private async getAgent(): Promise { + const { options } = this; - private async getAgent(options: IConnectionOptions): Promise { const httpAgentOptions: http.AgentOptions = { keepAlive: true, maxSockets: 5, @@ -32,30 +50,28 @@ export default class HttpConnection implements IConnectionProvider, IThriftConne return options.https ? new https.Agent(httpsAgentOptions) : new http.Agent(httpAgentOptions); } - async connect(options: IConnectionOptions): Promise { - const agent = await this.getAgent(options); - - this.connection = new ThriftHttpConnection( - { - url: `${options.https ? 'https' : 'http'}://${options.host}:${options.port}${options.path ?? '/'}`, - transport: thrift.TBufferedTransport, - protocol: thrift.TBinaryProtocol, - }, - { - agent, - timeout: options.socketTimeout ?? globalConfig.socketTimeout, - headers: options.headers, - }, - ); - - return this; - } + public async getThriftConnection(): Promise { + if (!this.connection) { + const { options } = this; + const agent = await this.getAgent(); - getConnection() { - return this.connection; - } + this.connection = new ThriftHttpConnection( + { + url: `${options.https ? 'https' : 'http'}://${options.host}:${options.port}${options.path ?? '/'}`, + transport: thrift.TBufferedTransport, + protocol: thrift.TBinaryProtocol, + }, + { + agent, + timeout: options.socketTimeout ?? globalConfig.socketTimeout, + headers: { + ...options.headers, + ...this.headers, + }, + }, + ); + } - isConnected(): boolean { - return !!this.connection; + return this.connection; } } diff --git a/lib/connection/connections/ThriftHttpConnection.ts b/lib/connection/connections/ThriftHttpConnection.ts index 83be7a07..1864b957 100644 --- a/lib/connection/connections/ThriftHttpConnection.ts +++ b/lib/connection/connections/ThriftHttpConnection.ts @@ -6,7 +6,7 @@ import { EventEmitter } from 'events'; import { TBinaryProtocol, TBufferedTransport, Thrift, TProtocol, TProtocolConstructor, TTransport } from 'thrift'; -import fetch, { RequestInit, Response, FetchError } from 'node-fetch'; +import fetch, { RequestInit, HeadersInit, Response, FetchError } from 'node-fetch'; // @ts-expect-error TS7016: Could not find a declaration file for module import InputBufferUnderrunError from 'thrift/lib/nodejs/lib/thrift/input_buffer_underrun_error'; @@ -48,7 +48,7 @@ type ThriftClient = { export default class ThriftHttpConnection extends EventEmitter { private readonly url: string; - private readonly config: RequestInit; + private config: RequestInit; // This field is used by Thrift internally, so name and type are important private readonly transport: TTransportType; @@ -67,6 +67,13 @@ export default class ThriftHttpConnection extends EventEmitter { this.protocol = options.protocol ?? TBinaryProtocol; } + public setHeaders(headers: HeadersInit) { + this.config = { + ...this.config, + headers, + }; + } + public write(data: Buffer, seqId: number) { const requestConfig: RequestInit = { ...this.config, diff --git a/lib/connection/contracts/IAuthentication.ts b/lib/connection/contracts/IAuthentication.ts index 03bdbe1b..e75b0b21 100644 --- a/lib/connection/contracts/IAuthentication.ts +++ b/lib/connection/contracts/IAuthentication.ts @@ -1,5 +1,5 @@ -import { HttpHeaders } from 'thrift'; +import { HeadersInit } from 'node-fetch'; export default interface IAuthentication { - authenticate(): Promise; + authenticate(): Promise; } diff --git a/lib/connection/contracts/IConnectionProvider.ts b/lib/connection/contracts/IConnectionProvider.ts index 0d25ccb6..f0efd739 100644 --- a/lib/connection/contracts/IConnectionProvider.ts +++ b/lib/connection/contracts/IConnectionProvider.ts @@ -1,6 +1,7 @@ -import IConnectionOptions from './IConnectionOptions'; -import IThriftConnection from './IThriftConnection'; +import { HeadersInit } from 'node-fetch'; export default interface IConnectionProvider { - connect(options: IConnectionOptions): Promise; + getThriftConnection(): Promise; + + setHeaders(headers: HeadersInit): void; } diff --git a/lib/connection/contracts/IThriftConnection.ts b/lib/connection/contracts/IThriftConnection.ts deleted file mode 100644 index fe074e7d..00000000 --- a/lib/connection/contracts/IThriftConnection.ts +++ /dev/null @@ -1,5 +0,0 @@ -export default interface IThriftConnection { - getConnection(): any; - - isConnected(): boolean; -} diff --git a/lib/utils/areHeadersEqual.ts b/lib/utils/areHeadersEqual.ts deleted file mode 100644 index f17433d6..00000000 --- a/lib/utils/areHeadersEqual.ts +++ /dev/null @@ -1,54 +0,0 @@ -import { HttpHeaders } from 'thrift'; - -function areArraysEqual(a: Array, b: Array): 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; -} diff --git a/lib/utils/index.ts b/lib/utils/index.ts index f89f71ea..4603277a 100644 --- a/lib/utils/index.ts +++ b/lib/utils/index.ts @@ -1,6 +1,5 @@ -import areHeadersEqual from './areHeadersEqual'; import definedOrError from './definedOrError'; import buildUserAgentString from './buildUserAgentString'; import formatProgress, { ProgressUpdateTransformer } from './formatProgress'; -export { areHeadersEqual, definedOrError, buildUserAgentString, formatProgress, ProgressUpdateTransformer }; +export { definedOrError, buildUserAgentString, formatProgress, ProgressUpdateTransformer }; diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index a7e9c4e3..a9a1d304 100644 --- a/tests/unit/DBSQLClient.test.js +++ b/tests/unit/DBSQLClient.test.js @@ -8,6 +8,7 @@ const DatabricksOAuth = require('../../dist/connection/auth/DatabricksOAuth').de const { AWSOAuthManager, AzureOAuthManager } = require('../../dist/connection/auth/DatabricksOAuth/OAuthManager'); const HttpConnectionModule = require('../../dist/connection/connections/HttpConnection'); +const { default: HttpConnection } = HttpConnectionModule; class AuthProviderMock { constructor() { @@ -26,6 +27,10 @@ describe('DBSQLClient.connect', () => { token: 'dapi********************************', }; + afterEach(() => { + HttpConnectionModule.default.restore?.(); + }); + it('should prepend "/" to path if it is missing', async () => { const client = new DBSQLClient(); @@ -47,15 +52,31 @@ describe('DBSQLClient.connect', () => { it('should initialize connection state', async () => { const client = new DBSQLClient(); - expect(client.client).to.be.null; - expect(client.authProvider).to.be.null; - expect(client.connectionOptions).to.be.null; + expect(client.client).to.be.undefined; + expect(client.authProvider).to.be.undefined; + expect(client.connectionProvider).to.be.undefined; await client.connect(options); - expect(client.client).to.be.null; // it should not be initialized at this point + expect(client.client).to.be.undefined; // it should not be initialized at this point expect(client.authProvider).to.be.instanceOf(PlainHttpAuthentication); - expect(client.connectionOptions).to.be.deep.equal(options); + expect(client.connectionProvider).to.be.instanceOf(HttpConnection); + }); + + it('should listen for Thrift connection events', async () => { + const client = new DBSQLClient(); + + const thriftConnectionMock = { + on: sinon.stub(), + }; + + sinon.stub(HttpConnectionModule, 'default').returns({ + getThriftConnection: () => Promise.resolve(thriftConnectionMock), + }); + + await client.connect(options); + + expect(thriftConnectionMock.on.called).to.be.true; }); }); @@ -163,39 +184,40 @@ describe('DBSQLClient.getClient', () => { const thriftClient = {}; client.authProvider = new AuthProviderMock(); - client.connectionOptions = { ...options }; + client.connectionProvider = new HttpConnection({ ...options }); client.thrift = { createClient: sinon.stub().returns(thriftClient), }; - sinon.stub(client, 'createConnection').returns({ - getConnection: () => null, - }); const result = await client.getClient(); expect(client.thrift.createClient.called).to.be.true; - expect(client.createConnection.called).to.be.true; expect(result).to.be.equal(thriftClient); }); - it('should re-create client if auth credentials change', async () => { + it('should update auth credentials each time when client is requested', async () => { const client = new DBSQLClient(); const thriftClient = {}; - client.authProvider = new AuthProviderMock(); - client.connectionOptions = { ...options }; + client.connectionProvider = new HttpConnection({ ...options }); client.thrift = { createClient: sinon.stub().returns(thriftClient), }; - sinon.stub(client, 'createConnection').returns({ - getConnection: () => null, - }); + + sinon.stub(client.connectionProvider, 'setHeaders').callThrough(); + + // just a sanity check - authProvider should be initialized by this time, but if not it should not be used + expect(client.connectionProvider.setHeaders.callCount).to.be.equal(0); + await client.getClient(); + expect(client.connectionProvider.setHeaders.callCount).to.be.equal(0); + + client.authProvider = new AuthProviderMock(); // initialize client firstCall: { const result = await client.getClient(); expect(client.thrift.createClient.callCount).to.be.equal(1); - expect(client.createConnection.callCount).to.be.equal(1); + expect(client.connectionProvider.setHeaders.callCount).to.be.equal(1); expect(result).to.be.equal(thriftClient); } @@ -203,7 +225,7 @@ describe('DBSQLClient.getClient', () => { secondCall: { const result = await client.getClient(); expect(client.thrift.createClient.callCount).to.be.equal(1); - expect(client.createConnection.callCount).to.be.equal(1); + expect(client.connectionProvider.setHeaders.callCount).to.be.equal(2); expect(result).to.be.equal(thriftClient); } @@ -212,54 +234,24 @@ describe('DBSQLClient.getClient', () => { client.authProvider.authResult = { b: 2 }; const result = await client.getClient(); - expect(client.thrift.createClient.callCount).to.be.equal(2); - expect(client.createConnection.callCount).to.be.equal(2); + expect(client.thrift.createClient.callCount).to.be.equal(1); + expect(client.connectionProvider.setHeaders.callCount).to.be.equal(3); expect(result).to.be.equal(thriftClient); } }); }); -describe('DBSQLClient.createConnection', () => { - afterEach(() => { - HttpConnectionModule.default.restore?.(); - }); - - it('should create connection', async () => { - const thriftConnection = { - on: sinon.stub(), - }; - - const connectionMock = { - getConnection: sinon.stub().returns(thriftConnection), - }; - - const connectionProviderMock = { - connect: sinon.stub().returns(Promise.resolve(connectionMock)), - }; - - sinon.stub(HttpConnectionModule, 'default').returns(connectionProviderMock); - - const client = new DBSQLClient(); - - const result = await client.createConnection({}); - expect(result).to.be.equal(connectionMock); - expect(connectionProviderMock.connect.called).to.be.true; - expect(connectionMock.getConnection.called).to.be.true; - expect(thriftConnection.on.called).to.be.true; - }); -}); - describe('DBSQLClient.close', () => { it('should close the connection if it was initiated', async () => { const client = new DBSQLClient(); client.client = {}; + client.connectionProvider = {}; client.authProvider = {}; - client.connectionOptions = {}; await client.close(); - expect(client.client).to.be.null; - expect(client.authProvider).to.be.null; - expect(client.connectionOptions).to.be.null; + expect(client.client).to.be.undefined; + expect(client.connectionProvider).to.be.undefined; + expect(client.authProvider).to.be.undefined; // No additional asserts needed - it should just reach this point }); @@ -267,9 +259,9 @@ describe('DBSQLClient.close', () => { const client = new DBSQLClient(); await client.close(); - expect(client.client).to.be.null; - expect(client.authProvider).to.be.null; - expect(client.connectionOptions).to.be.null; + expect(client.client).to.be.undefined; + expect(client.connectionProvider).to.be.undefined; + expect(client.authProvider).to.be.undefined; // No additional asserts needed - it should just reach this point }); diff --git a/tests/unit/connection/connections/HttpConnection.test.js b/tests/unit/connection/connections/HttpConnection.test.js index 5875d670..a9a21136 100644 --- a/tests/unit/connection/connections/HttpConnection.test.js +++ b/tests/unit/connection/connections/HttpConnection.test.js @@ -3,37 +3,26 @@ const { expect } = require('chai'); const HttpConnection = require('../../../../dist/connection/connections/HttpConnection').default; const ThriftHttpConnection = require('../../../../dist/connection/connections/ThriftHttpConnection').default; -const thriftMock = (connection) => ({ - createHttpConnection(host, port, options) { - this.host = host; - this.port = port; - this.options = options; - this.executed = true; - return connection; - }, -}); - describe('HttpConnection.connect', () => { - it('should successfully connect', async () => { - const connection = new HttpConnection(); - - expect(connection.isConnected()).to.be.false; - - await connection.connect({ + it('should create Thrift connection', async () => { + const connection = new HttpConnection({ host: 'localhost', port: 10001, path: '/hive', }); - expect(connection.connection.url).to.be.eq('http://localhost:10001/hive'); - expect(connection.getConnection()).to.be.instanceOf(ThriftHttpConnection); - expect(connection.isConnected()).to.be.true; + const thriftConnection = await connection.getThriftConnection(); + + expect(thriftConnection).to.be.instanceOf(ThriftHttpConnection); + expect(thriftConnection.url).to.be.equal('http://localhost:10001/hive'); + + // We expect that connection will be cached + const anotherConnection = await connection.getThriftConnection(); + expect(anotherConnection).to.eq(thriftConnection); }); it('should set SSL certificates and disable rejectUnauthorized', async () => { - const connection = new HttpConnection(); - - await connection.connect({ + const connection = new HttpConnection({ host: 'localhost', port: 10001, path: '/hive', @@ -43,22 +32,82 @@ describe('HttpConnection.connect', () => { key: 'key', }); - expect(connection.connection.config.agent.options.rejectUnauthorized).to.be.false; - expect(connection.connection.config.agent.options.ca).to.be.eq('ca'); - expect(connection.connection.config.agent.options.cert).to.be.eq('cert'); - expect(connection.connection.config.agent.options.key).to.be.eq('key'); + const thriftConnection = await connection.getThriftConnection(); + + expect(thriftConnection.config.agent.options.rejectUnauthorized).to.be.false; + expect(thriftConnection.config.agent.options.ca).to.be.eq('ca'); + expect(thriftConnection.config.agent.options.cert).to.be.eq('cert'); + expect(thriftConnection.config.agent.options.key).to.be.eq('key'); }); it('should initialize http agents', async () => { - const connection = new HttpConnection(); - - await connection.connect({ + const connection = new HttpConnection({ host: 'localhost', port: 10001, https: false, path: '/hive', }); - expect(connection.connection.config.agent).to.be.instanceOf(http.Agent); + const thriftConnection = await connection.getThriftConnection(); + + expect(thriftConnection.config.agent).to.be.instanceOf(http.Agent); + }); + + it('should update headers (case 1: Thrift connection not initialized)', async () => { + const initialHeaders = { + a: 'test header A', + b: 'test header B', + }; + + const connection = new HttpConnection({ + host: 'localhost', + port: 10001, + path: '/hive', + headers: initialHeaders, + }); + + const extraHeaders = { + b: 'new header B', + c: 'test header C', + }; + connection.setHeaders(extraHeaders); + expect(connection.headers).to.deep.equal(extraHeaders); + + const thriftConnection = await connection.getThriftConnection(); + + expect(thriftConnection.config.headers).to.deep.equal({ + ...initialHeaders, + ...extraHeaders, + }); + }); + + it('should update headers (case 2: Thrift connection initialized)', async () => { + const initialHeaders = { + a: 'test header A', + b: 'test header B', + }; + + const connection = new HttpConnection({ + host: 'localhost', + port: 10001, + path: '/hive', + headers: initialHeaders, + }); + + const thriftConnection = await connection.getThriftConnection(); + + expect(connection.headers).to.deep.equal({}); + expect(thriftConnection.config.headers).to.deep.equal(initialHeaders); + + const extraHeaders = { + b: 'new header B', + c: 'test header C', + }; + connection.setHeaders(extraHeaders); + expect(connection.headers).to.deep.equal(extraHeaders); + expect(thriftConnection.config.headers).to.deep.equal({ + ...initialHeaders, + ...extraHeaders, + }); }); }); diff --git a/tests/unit/utils.test.js b/tests/unit/utils.test.js index 0458ec7b..5e07ed16 100644 --- a/tests/unit/utils.test.js +++ b/tests/unit/utils.test.js @@ -1,12 +1,6 @@ const { expect, AssertionError } = require('chai'); -const { - areHeadersEqual, - buildUserAgentString, - definedOrError, - formatProgress, - ProgressUpdateTransformer, -} = require('../../dist/utils'); +const { buildUserAgentString, definedOrError, formatProgress, ProgressUpdateTransformer } = require('../../dist/utils'); const CloseableCollection = require('../../dist/utils/CloseableCollection').default; describe('buildUserAgentString', () => { @@ -99,103 +93,6 @@ describe('definedOrError', () => { }); }); -describe('areHeadersEqual', () => { - it('should return true for same objects', () => { - const a = {}; - expect(areHeadersEqual(a, a)).to.be.true; - }); - - it('should return false for objects with different keys', () => { - const a = { a: 1, x: 2 }; - const b = { b: 3, x: 4 }; - const c = { c: 5 }; - - expect(areHeadersEqual(a, b)).to.be.false; - expect(areHeadersEqual(b, a)).to.be.false; - expect(areHeadersEqual(a, c)).to.be.false; - expect(areHeadersEqual(c, a)).to.be.false; - }); - - it('should compare different types of properties', () => { - case1: { - expect( - areHeadersEqual( - { - a: 1, - b: 'b', - c: ['x', 'y', 'z'], - }, - { - a: 1, - b: 'b', - c: ['x', 'y', 'z'], - }, - ), - ).to.be.true; - } - - case2: { - const arr = ['a', 'b']; - - expect( - areHeadersEqual( - { - a: 1, - b: 'b', - c: arr, - }, - { - a: 1, - b: 'b', - c: arr, - }, - ), - ).to.be.true; - } - - case3: { - expect( - areHeadersEqual( - { - arr: ['a'], - }, - { - arr: ['b'], - }, - ), - ).to.be.false; - } - - case4: { - expect( - areHeadersEqual( - { - arr: ['a'], - }, - { - arr: ['a', 'b'], - }, - ), - ).to.be.false; - } - - case5: { - expect( - areHeadersEqual( - { - arr: ['a'], - prop: 'x', - }, - { - arr: ['a'], - prop: 1, - }, - ), - ).to.be.false; - } - }); -}); - describe('CloseableCollection', () => { it('should add item if not already added', () => { const collection = new CloseableCollection();