From 37cd8914608921661962a9a30a94d9b709a91f1f Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 2 Oct 2023 13:59:13 +0300 Subject: [PATCH 1/4] Introduce client context Signed-off-by: Levko Kravets --- lib/DBSQLClient.ts | 52 ++-- lib/DBSQLOperation/index.ts | 47 +-- lib/DBSQLSession.ts | 33 +- .../auth/DatabricksOAuth/AuthorizationCode.ts | 15 +- .../auth/DatabricksOAuth/OAuthManager.ts | 24 +- lib/connection/auth/DatabricksOAuth/index.ts | 9 +- .../auth/PlainHttpAuthentication.ts | 5 + lib/contracts/IClientContext.ts | 11 + lib/hive/HiveDriver.ts | 54 ++-- tests/unit/DBSQLClient.test.js | 16 +- tests/unit/DBSQLOperation.test.js | 289 ++++++++++++------ tests/unit/DBSQLSession.test.js | 8 +- .../DatabricksOAuth/AuthorizationCode.test.js | 6 + .../auth/DatabricksOAuth/OAuthManager.test.js | 6 + .../auth/PlainHttpAuthentication.test.js | 4 +- tests/unit/hive/HiveDriver.test.js | 10 +- 16 files changed, 378 insertions(+), 211 deletions(-) create mode 100644 lib/contracts/IClientContext.ts diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index e8188d2e..31edbf47 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -4,6 +4,7 @@ import { EventEmitter } from 'events'; import TCLIService from '../thrift/TCLIService'; import { TProtocolVersion } from '../thrift/TCLIService_types'; import IDBSQLClient, { ClientOptions, ConnectionOptions, OpenSessionRequest } from './contracts/IDBSQLClient'; +import IClientContext from './contracts/IClientContext'; import HiveDriver from './hive/HiveDriver'; import { Int64 } from './hive/Types'; import DBSQLSession from './DBSQLSession'; @@ -41,7 +42,7 @@ function getInitialNamespaceOptions(catalogName?: string, schemaName?: string) { }; } -export default class DBSQLClient extends EventEmitter implements IDBSQLClient { +export default class DBSQLClient extends EventEmitter implements IDBSQLClient, IClientContext { private connectionProvider?: IConnectionProvider; private authProvider?: IAuthentication; @@ -73,7 +74,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { }; } - private getAuthProvider(options: ConnectionOptions, authProvider?: IAuthentication): IAuthentication { + private initAuthProvider(options: ConnectionOptions, authProvider?: IAuthentication): IAuthentication { if (authProvider) { return authProvider; } @@ -84,15 +85,16 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { return new PlainHttpAuthentication({ username: 'token', password: options.token, + context: this, }); case 'databricks-oauth': return new DatabricksOAuth({ host: options.host, - logger: this.logger, persistence: options.persistence, azureTenantId: options.azureTenantId, clientId: options.oauthClientId, clientSecret: options.oauthClientSecret, + context: this, }); case 'custom': return options.provider; @@ -110,7 +112,7 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { * const session = client.connect({host, path, token}); */ public async connect(options: ConnectionOptions, authProvider?: IAuthentication): Promise { - this.authProvider = this.getAuthProvider(options, authProvider); + this.authProvider = this.initAuthProvider(options, authProvider); this.connectionProvider = new HttpConnection(this.getConnectionOptions(options)); @@ -156,7 +158,9 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { * const session = await client.openSession(); */ public async openSession(request: OpenSessionRequest = {}): Promise { - const driver = new HiveDriver(() => this.getClient()); + const driver = new HiveDriver({ + context: this, + }); const response = await driver.openSession({ client_protocol_i64: new Int64(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8), @@ -164,36 +168,48 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient { }); Status.assert(response.status); - const session = new DBSQLSession(driver, definedOrError(response.sessionHandle), { - logger: this.logger, + const session = new DBSQLSession({ + driver, + handle: definedOrError(response.sessionHandle), + context: this, }); this.sessions.add(session); return session; } - private async getClient() { + public async close(): Promise { + await this.sessions.closeAll(); + + this.client = undefined; + this.connectionProvider = undefined; + this.authProvider = undefined; + } + + public getLogger(): IDBSQLLogger { + return this.logger; + } + + public async getConnectionProvider(): Promise { if (!this.connectionProvider) { throw new HiveDriverError('DBSQLClient: not connected'); } + return this.connectionProvider; + } + + public async getClient(): Promise { + const connectionProvider = await this.getConnectionProvider(); + if (!this.client) { this.logger.log(LogLevel.info, 'DBSQLClient: initializing thrift client'); - this.client = this.thrift.createClient(TCLIService, await this.connectionProvider.getThriftConnection()); + this.client = this.thrift.createClient(TCLIService, await connectionProvider.getThriftConnection()); } if (this.authProvider) { const authHeaders = await this.authProvider.authenticate(); - this.connectionProvider.setHeaders(authHeaders); + connectionProvider.setHeaders(authHeaders); } return this.client; } - - public async close(): Promise { - await this.sessions.closeAll(); - - this.client = undefined; - this.connectionProvider = undefined; - this.authProvider = undefined; - } } diff --git a/lib/DBSQLOperation/index.ts b/lib/DBSQLOperation/index.ts index 31690224..27087c41 100644 --- a/lib/DBSQLOperation/index.ts +++ b/lib/DBSQLOperation/index.ts @@ -18,7 +18,7 @@ import { } from '../../thrift/TCLIService_types'; import Status from '../dto/Status'; import FetchResultsHelper from './FetchResultsHelper'; -import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger'; +import { LogLevel } from '../contracts/IDBSQLLogger'; import OperationStateError, { OperationStateErrorCode } from '../errors/OperationStateError'; import IOperationResult from '../result/IOperationResult'; import JsonResult from '../result/JsonResult'; @@ -26,11 +26,15 @@ import ArrowResult from '../result/ArrowResult'; import CloudFetchResult from '../result/CloudFetchResult'; import { definedOrError } from '../utils'; import HiveDriverError from '../errors/HiveDriverError'; +import IClientContext from '../contracts/IClientContext'; const defaultMaxRows = 100000; interface DBSQLOperationConstructorOptions { - logger: IDBSQLLogger; + handle: TOperationHandle; + driver: HiveDriver; + directResults?: TSparkDirectResults; + context: IClientContext; } async function delay(ms?: number): Promise { @@ -42,12 +46,12 @@ async function delay(ms?: number): Promise { } export default class DBSQLOperation implements IOperation { + private readonly context: IClientContext; + private readonly driver: HiveDriver; private readonly operationHandle: TOperationHandle; - private readonly logger: IDBSQLLogger; - public onClose?: () => void; private readonly _data: FetchResultsHelper; @@ -70,19 +74,14 @@ export default class DBSQLOperation implements IOperation { private resultHandler?: IOperationResult; - constructor( - driver: HiveDriver, - operationHandle: TOperationHandle, - { logger }: DBSQLOperationConstructorOptions, - directResults?: TSparkDirectResults, - ) { + constructor({ driver, handle, directResults, context }: DBSQLOperationConstructorOptions) { this.driver = driver; - this.operationHandle = operationHandle; - this.logger = logger; + this.operationHandle = handle; + this.context = context; const useOnlyPrefetchedResults = Boolean(directResults?.closeOperation); - this.hasResultSet = operationHandle.hasResultSet; + this.hasResultSet = this.operationHandle.hasResultSet; if (directResults?.operationStatus) { this.processOperationStatusResponse(directResults.operationStatus); } @@ -95,7 +94,7 @@ export default class DBSQLOperation implements IOperation { useOnlyPrefetchedResults, ); this.closeOperation = directResults?.closeOperation; - this.logger.log(LogLevel.debug, `Operation created with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Operation created with id: ${this.getId()}`); } public getId() { @@ -118,7 +117,7 @@ export default class DBSQLOperation implements IOperation { const chunk = await this.fetchChunk(options); data.push(chunk); } while (await this.hasMoreRows()); // eslint-disable-line no-await-in-loop - this.logger?.log(LogLevel.debug, `Fetched all data from operation with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Fetched all data from operation with id: ${this.getId()}`); return data.flat(); } @@ -149,10 +148,12 @@ export default class DBSQLOperation implements IOperation { await this.failIfClosed(); const result = await resultHandler.getValue(data ? [data] : []); - this.logger?.log( - LogLevel.debug, - `Fetched chunk of size: ${options?.maxRows || defaultMaxRows} from operation with id: ${this.getId()}`, - ); + this.context + .getLogger() + .log( + LogLevel.debug, + `Fetched chunk of size: ${options?.maxRows || defaultMaxRows} from operation with id: ${this.getId()}`, + ); return result; } @@ -163,7 +164,7 @@ export default class DBSQLOperation implements IOperation { */ public async status(progress: boolean = false): Promise { await this.failIfClosed(); - this.logger?.log(LogLevel.debug, `Fetching status for operation with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Fetching status for operation with id: ${this.getId()}`); if (this.operationStatus) { return this.operationStatus; @@ -186,7 +187,7 @@ export default class DBSQLOperation implements IOperation { return Status.success(); } - this.logger?.log(LogLevel.debug, `Cancelling operation with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Cancelling operation with id: ${this.getId()}`); const response = await this.driver.cancelOperation({ operationHandle: this.operationHandle, @@ -209,7 +210,7 @@ export default class DBSQLOperation implements IOperation { return Status.success(); } - this.logger?.log(LogLevel.debug, `Closing operation with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Closing operation with id: ${this.getId()}`); const response = this.closeOperation ?? @@ -254,7 +255,7 @@ export default class DBSQLOperation implements IOperation { await this.waitUntilReady(options); - this.logger?.log(LogLevel.debug, `Fetching schema for operation with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Fetching schema for operation with id: ${this.getId()}`); const metadata = await this.fetchMetadata(); return metadata.schema ?? null; } diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 5bb2b4a2..781e94f3 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -32,12 +32,13 @@ import Status from './dto/Status'; import InfoValue from './dto/InfoValue'; import { definedOrError } from './utils'; import CloseableCollection from './utils/CloseableCollection'; -import IDBSQLLogger, { LogLevel } from './contracts/IDBSQLLogger'; +import { LogLevel } from './contracts/IDBSQLLogger'; import HiveDriverError from './errors/HiveDriverError'; import globalConfig from './globalConfig'; import StagingError from './errors/StagingError'; import { DBSQLParameter, DBSQLParameterValue } from './DBSQLParameter'; import ParameterError from './errors/ParameterError'; +import IClientContext from './contracts/IClientContext'; const defaultMaxRows = 100000; @@ -130,27 +131,29 @@ function getQueryParameters( } interface DBSQLSessionConstructorOptions { - logger: IDBSQLLogger; + handle: TSessionHandle; + driver: HiveDriver; + context: IClientContext; } export default class DBSQLSession implements IDBSQLSession { + private readonly context: IClientContext; + private readonly driver: HiveDriver; private readonly sessionHandle: TSessionHandle; - private readonly logger: IDBSQLLogger; - private isOpen = true; public onClose?: () => void; private operations = new CloseableCollection(); - constructor(driver: HiveDriver, sessionHandle: TSessionHandle, { logger }: DBSQLSessionConstructorOptions) { + constructor({ handle, driver, context }: DBSQLSessionConstructorOptions) { this.driver = driver; - this.sessionHandle = sessionHandle; - this.logger = logger; - this.logger.log(LogLevel.debug, `Session created with id: ${this.getId()}`); + this.sessionHandle = handle; + this.context = context; + this.context.getLogger().log(LogLevel.debug, `Session created with id: ${this.getId()}`); } public getId() { @@ -496,21 +499,19 @@ export default class DBSQLSession implements IDBSQLSession { this.onClose?.(); this.isOpen = false; - this.logger.log(LogLevel.debug, `Session closed with id: ${this.getId()}`); + this.context.getLogger().log(LogLevel.debug, `Session closed with id: ${this.getId()}`); return new Status(response.status); } private createOperation(response: OperationResponseShape): DBSQLOperation { Status.assert(response.status); const handle = definedOrError(response.operationHandle); - const operation = new DBSQLOperation( - this.driver, + const operation = new DBSQLOperation({ + driver: this.driver, handle, - { - logger: this.logger, - }, - response.directResults, - ); + directResults: response.directResults, + context: this.context, + }); this.operations.add(operation); diff --git a/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts b/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts index aee4bc17..03845c74 100644 --- a/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts +++ b/lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts @@ -1,13 +1,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 { LogLevel } from '../../../contracts/IDBSQLLogger'; import { OAuthScopes, scopeDelimiter } from './OAuthScope'; +import IClientContext from '../../../contracts/IClientContext'; export interface AuthorizationCodeOptions { client: BaseClient; ports: Array; - logger?: IDBSQLLogger; + context: IClientContext; } async function startServer( @@ -57,18 +58,18 @@ export interface AuthorizationCodeFetchResult { } export default class AuthorizationCode { + private readonly context: IClientContext; + private readonly client: BaseClient; private readonly host: string = 'localhost'; private readonly ports: Array; - private readonly logger?: IDBSQLLogger; - constructor(options: AuthorizationCodeOptions) { this.client = options.client; this.ports = options.ports; - this.logger = options.logger; + this.context = options.context; } private async openUrl(url: string) { @@ -125,7 +126,7 @@ export default class AuthorizationCode { const host = this.host; // eslint-disable-line prefer-destructuring try { const server = await startServer(host, port, requestHandler); // eslint-disable-line no-await-in-loop - this.logger?.log(LogLevel.info, `Listening for OAuth authorization callback at ${host}:${port}`); + this.context.getLogger().log(LogLevel.info, `Listening for OAuth authorization callback at ${host}:${port}`); let resolveStopped: () => void; let rejectStopped: (reason?: any) => void; @@ -144,7 +145,7 @@ export default class AuthorizationCode { } catch (error) { // if port already in use - try another one, otherwise re-throw an exception if (error instanceof Error && 'code' in error && error.code === 'EADDRINUSE') { - this.logger?.log(LogLevel.debug, `Failed to start server at ${host}:${port}: ${error.code}`); + this.context.getLogger().log(LogLevel.debug, `Failed to start server at ${host}:${port}: ${error.code}`); } else { throw error; } diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 64b1af0e..fde9ea30 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -1,9 +1,10 @@ import { Issuer, BaseClient } from 'openid-client'; import HiveDriverError from '../../../errors/HiveDriverError'; -import IDBSQLLogger, { LogLevel } from '../../../contracts/IDBSQLLogger'; +import { LogLevel } from '../../../contracts/IDBSQLLogger'; import OAuthToken from './OAuthToken'; import AuthorizationCode from './AuthorizationCode'; import { OAuthScope, OAuthScopes } from './OAuthScope'; +import IClientContext from '../../../contracts/IClientContext'; export interface OAuthManagerOptions { host: string; @@ -11,7 +12,7 @@ export interface OAuthManagerOptions { clientId?: string; azureTenantId?: string; clientSecret?: string; - logger?: IDBSQLLogger; + context: IClientContext; } function getDatabricksOIDCUrl(host: string): string { @@ -21,9 +22,9 @@ function getDatabricksOIDCUrl(host: string): string { } export default abstract class OAuthManager { - protected readonly options: OAuthManagerOptions; + protected readonly context: IClientContext; - protected readonly logger?: IDBSQLLogger; + protected readonly options: OAuthManagerOptions; protected issuer?: Issuer; @@ -31,7 +32,7 @@ export default abstract class OAuthManager { constructor(options: OAuthManagerOptions) { this.options = options; - this.logger = options.logger; + this.context = options.context; } protected abstract getOIDCConfigUrl(): string; @@ -71,15 +72,14 @@ export default abstract class OAuthManager { private async refreshAccessTokenU2M(token: OAuthToken): Promise { if (!token.refreshToken) { const message = `OAuth access token expired on ${token.expirationTime}.`; - this.logger?.log(LogLevel.error, message); + this.context.getLogger().log(LogLevel.error, message); throw new HiveDriverError(message); } // Try to refresh using the refresh token - this.logger?.log( - LogLevel.debug, - `Attempting to refresh OAuth access token that expired on ${token.expirationTime}`, - ); + this.context + .getLogger() + .log(LogLevel.debug, `Attempting to refresh OAuth access token that expired on ${token.expirationTime}`); const client = await this.getClient(); const { access_token: accessToken, refresh_token: refreshToken } = await client.refresh(token.refreshToken); @@ -106,7 +106,7 @@ export default abstract class OAuthManager { return token; } } catch (error) { - this.logger?.log(LogLevel.error, `${error}`); + this.context.getLogger().log(LogLevel.error, `${error}`); throw error; } @@ -119,7 +119,7 @@ export default abstract class OAuthManager { const authCode = new AuthorizationCode({ client, ports: this.getCallbackPorts(), - logger: this.logger, + context: this.context, }); const mappedScopes = this.getScopes(scopes); diff --git a/lib/connection/auth/DatabricksOAuth/index.ts b/lib/connection/auth/DatabricksOAuth/index.ts index 29243b8f..37489974 100644 --- a/lib/connection/auth/DatabricksOAuth/index.ts +++ b/lib/connection/auth/DatabricksOAuth/index.ts @@ -1,29 +1,28 @@ import { HeadersInit } from 'node-fetch'; import IAuthentication from '../../contracts/IAuthentication'; -import IDBSQLLogger from '../../../contracts/IDBSQLLogger'; import OAuthPersistence, { OAuthPersistenceCache } from './OAuthPersistence'; import OAuthManager, { OAuthManagerOptions } from './OAuthManager'; import { OAuthScopes, defaultOAuthScopes } from './OAuthScope'; +import IClientContext from '../../../contracts/IClientContext'; interface DatabricksOAuthOptions extends OAuthManagerOptions { scopes?: OAuthScopes; - logger?: IDBSQLLogger; persistence?: OAuthPersistence; headers?: HeadersInit; } export default class DatabricksOAuth implements IAuthentication { - private readonly options: DatabricksOAuthOptions; + private readonly context: IClientContext; - private readonly logger?: IDBSQLLogger; + private readonly options: DatabricksOAuthOptions; private readonly manager: OAuthManager; private readonly defaultPersistence = new OAuthPersistenceCache(); constructor(options: DatabricksOAuthOptions) { + this.context = options.context; this.options = options; - this.logger = options.logger; this.manager = OAuthManager.getManager(this.options); } diff --git a/lib/connection/auth/PlainHttpAuthentication.ts b/lib/connection/auth/PlainHttpAuthentication.ts index 1115c745..977f1cc5 100644 --- a/lib/connection/auth/PlainHttpAuthentication.ts +++ b/lib/connection/auth/PlainHttpAuthentication.ts @@ -1,13 +1,17 @@ import { HeadersInit } from 'node-fetch'; import IAuthentication from '../contracts/IAuthentication'; +import IClientContext from '../../contracts/IClientContext'; interface PlainHttpAuthenticationOptions { username?: string; password?: string; headers?: HeadersInit; + context: IClientContext; } export default class PlainHttpAuthentication implements IAuthentication { + private readonly context: IClientContext; + private readonly username: string; private readonly password: string; @@ -15,6 +19,7 @@ export default class PlainHttpAuthentication implements IAuthentication { private readonly headers: HeadersInit; constructor(options: PlainHttpAuthenticationOptions) { + this.context = options.context; this.username = options?.username || 'anonymous'; this.password = options?.password ?? 'anonymous'; this.headers = options?.headers || {}; diff --git a/lib/contracts/IClientContext.ts b/lib/contracts/IClientContext.ts new file mode 100644 index 00000000..0e6d5ed3 --- /dev/null +++ b/lib/contracts/IClientContext.ts @@ -0,0 +1,11 @@ +import IDBSQLLogger from './IDBSQLLogger'; +import IConnectionProvider from '../connection/contracts/IConnectionProvider'; +import TCLIService from '../../thrift/TCLIService'; + +export default interface IClientContext { + getLogger(): IDBSQLLogger; + + getConnectionProvider(): Promise; + + getClient(): Promise; +} diff --git a/lib/hive/HiveDriver.ts b/lib/hive/HiveDriver.ts index f4199bac..9c34d743 100644 --- a/lib/hive/HiveDriver.ts +++ b/lib/hive/HiveDriver.ts @@ -1,4 +1,3 @@ -import TCLIService from '../../thrift/TCLIService'; import { TOpenSessionReq, TCloseSessionReq, @@ -43,138 +42,141 @@ import CloseOperationCommand from './Commands/CloseOperationCommand'; import GetDelegationTokenCommand from './Commands/GetDelegationTokenCommand'; import CancelDelegationTokenCommand from './Commands/CancelDelegationTokenCommand'; import RenewDelegationTokenCommand from './Commands/RenewDelegationTokenCommand'; +import IClientContext from '../contracts/IClientContext'; -export type ClientFactory = () => Promise; +export interface HiveDriverOptions { + context: IClientContext; +} export default class HiveDriver { - private readonly clientFactory: ClientFactory; + private readonly context: IClientContext; - constructor(clientFactory: ClientFactory) { - this.clientFactory = clientFactory; + constructor(options: HiveDriverOptions) { + this.context = options.context; } async openSession(request: TOpenSessionReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const action = new OpenSessionCommand(client); return action.execute(request); } async closeSession(request: TCloseSessionReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new CloseSessionCommand(client); return command.execute(request); } async executeStatement(request: TExecuteStatementReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new ExecuteStatementCommand(client); return command.execute(request); } async getResultSetMetadata(request: TGetResultSetMetadataReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetResultSetMetadataCommand(client); return command.execute(request); } async fetchResults(request: TFetchResultsReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new FetchResultsCommand(client); return command.execute(request); } async getInfo(request: TGetInfoReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetInfoCommand(client); return command.execute(request); } async getTypeInfo(request: TGetTypeInfoReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetTypeInfoCommand(client); return command.execute(request); } async getCatalogs(request: TGetCatalogsReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetCatalogsCommand(client); return command.execute(request); } async getSchemas(request: TGetSchemasReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetSchemasCommand(client); return command.execute(request); } async getTables(request: TGetTablesReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetTablesCommand(client); return command.execute(request); } async getTableTypes(request: TGetTableTypesReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetTableTypesCommand(client); return command.execute(request); } async getColumns(request: TGetColumnsReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetColumnsCommand(client); return command.execute(request); } async getFunctions(request: TGetFunctionsReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetFunctionsCommand(client); return command.execute(request); } async getPrimaryKeys(request: TGetPrimaryKeysReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetPrimaryKeysCommand(client); return command.execute(request); } async getCrossReference(request: TGetCrossReferenceReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetCrossReferenceCommand(client); return command.execute(request); } async getOperationStatus(request: TGetOperationStatusReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetOperationStatusCommand(client); return command.execute(request); } async cancelOperation(request: TCancelOperationReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new CancelOperationCommand(client); return command.execute(request); } async closeOperation(request: TCloseOperationReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new CloseOperationCommand(client); return command.execute(request); } async getDelegationToken(request: TGetDelegationTokenReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new GetDelegationTokenCommand(client); return command.execute(request); } async cancelDelegationToken(request: TCancelDelegationTokenReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new CancelDelegationTokenCommand(client); return command.execute(request); } async renewDelegationToken(request: TRenewDelegationTokenReq) { - const client = await this.clientFactory(); + const client = await this.context.getClient(); const command = new RenewDelegationTokenCommand(client); return command.execute(request); } diff --git a/tests/unit/DBSQLClient.test.js b/tests/unit/DBSQLClient.test.js index a9a1d304..55231707 100644 --- a/tests/unit/DBSQLClient.test.js +++ b/tests/unit/DBSQLClient.test.js @@ -306,12 +306,12 @@ describe('DBSQLClient.close', () => { }); }); -describe('DBSQLClient.getAuthProvider', () => { +describe('DBSQLClient.initAuthProvider', () => { it('should use access token auth method', () => { const client = new DBSQLClient(); const testAccessToken = 'token'; - const provider = client.getAuthProvider({ + const provider = client.initAuthProvider({ authType: 'access-token', token: testAccessToken, }); @@ -324,7 +324,7 @@ describe('DBSQLClient.getAuthProvider', () => { const client = new DBSQLClient(); const testAccessToken = 'token'; - const provider = client.getAuthProvider({ + const provider = client.initAuthProvider({ // note: no `authType` provided token: testAccessToken, }); @@ -336,7 +336,7 @@ describe('DBSQLClient.getAuthProvider', () => { it('should use Databricks OAuth method (AWS)', () => { const client = new DBSQLClient(); - const provider = client.getAuthProvider({ + const provider = client.initAuthProvider({ authType: 'databricks-oauth', // host is used when creating OAuth manager, so make it look like a real AWS instance host: 'example.dev.databricks.com', @@ -349,7 +349,7 @@ describe('DBSQLClient.getAuthProvider', () => { it('should use Databricks OAuth method (Azure)', () => { const client = new DBSQLClient(); - const provider = client.getAuthProvider({ + const provider = client.initAuthProvider({ authType: 'databricks-oauth', // host is used when creating OAuth manager, so make it look like a real Azure instance host: 'example.databricks.azure.us', @@ -363,7 +363,7 @@ describe('DBSQLClient.getAuthProvider', () => { const client = new DBSQLClient(); expect(() => { - client.getAuthProvider({ + client.initAuthProvider({ authType: 'databricks-oauth', // use host which is not supported for sure host: 'example.com', @@ -376,7 +376,7 @@ describe('DBSQLClient.getAuthProvider', () => { const customProvider = {}; - const provider = client.getAuthProvider({ + const provider = client.initAuthProvider({ authType: 'custom', provider: customProvider, }); @@ -389,7 +389,7 @@ describe('DBSQLClient.getAuthProvider', () => { const customProvider = {}; - const provider = client.getAuthProvider( + const provider = client.initAuthProvider( // custom provider from second arg should be used no matter what's specified in config { authType: 'access-token', token: 'token' }, customProvider, diff --git a/tests/unit/DBSQLOperation.test.js b/tests/unit/DBSQLOperation.test.js index b0654170..4b45f240 100644 --- a/tests/unit/DBSQLOperation.test.js +++ b/tests/unit/DBSQLOperation.test.js @@ -107,25 +107,33 @@ class DriverMock { describe('DBSQLOperation', () => { describe('status', () => { it('should pick up state from operation handle', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; const driver = new DriverMock(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.true; }); it('should pick up state from directResults', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); - const operation = new DBSQLOperation(driver, handle, logger, { - operationStatus: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - operationState: TOperationState.FINISHED_STATE, - hasResultSet: true, + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + operationStatus: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.FINISHED_STATE, + hasResultSet: true, + }, }, }); @@ -134,6 +142,8 @@ describe('DBSQLOperation', () => { }); it('should fetch status and update internal state', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -142,7 +152,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.false; @@ -156,6 +166,8 @@ describe('DBSQLOperation', () => { }); it('should request progress', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -163,7 +175,7 @@ describe('DBSQLOperation', () => { sinon.spy(driver, 'getOperationStatus'); driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.status(true); expect(driver.getOperationStatus.called).to.be.true; @@ -172,6 +184,8 @@ describe('DBSQLOperation', () => { }); it('should not fetch status once operation is finished', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -179,7 +193,7 @@ describe('DBSQLOperation', () => { sinon.spy(driver, 'getOperationStatus'); driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.false; @@ -210,6 +224,8 @@ describe('DBSQLOperation', () => { }); it('should fetch status if directResults status is not finished', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -218,11 +234,16 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation(driver, handle, logger, { - operationStatus: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - operationState: TOperationState.RUNNING_STATE, - hasResultSet: false, + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + operationStatus: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.RUNNING_STATE, + hasResultSet: false, + }, }, }); @@ -238,6 +259,8 @@ describe('DBSQLOperation', () => { }); it('should not fetch status if directResults status is finished', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -246,11 +269,16 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.RUNNING_STATE; driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation(driver, handle, logger, { - operationStatus: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - operationState: TOperationState.FINISHED_STATE, - hasResultSet: false, + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + operationStatus: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.FINISHED_STATE, + hasResultSet: false, + }, }, }); @@ -266,10 +294,11 @@ describe('DBSQLOperation', () => { }); it('should throw an error in case of a status error', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); try { await operation.status(false); @@ -285,10 +314,11 @@ describe('DBSQLOperation', () => { describe('cancel', () => { it('should cancel operation and update state', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -301,10 +331,11 @@ describe('DBSQLOperation', () => { }); it('should return immediately if already cancelled', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -321,11 +352,12 @@ describe('DBSQLOperation', () => { }); it('should return immediately if already closed', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'cancelOperation'); sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -342,10 +374,11 @@ describe('DBSQLOperation', () => { }); it('should throw an error in case of a status error and keep state', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.cancelOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -364,9 +397,10 @@ describe('DBSQLOperation', () => { }); it('should reject all methods once cancelled', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.cancel(); expect(operation.cancelled).to.be.true; @@ -381,10 +415,11 @@ describe('DBSQLOperation', () => { describe('close', () => { it('should close operation and update state', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -397,10 +432,11 @@ describe('DBSQLOperation', () => { }); it('should return immediately if already closed', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -417,11 +453,12 @@ describe('DBSQLOperation', () => { }); it('should return immediately if already cancelled', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'closeOperation'); sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -438,12 +475,18 @@ describe('DBSQLOperation', () => { }); it('should initialize from directResults', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation(driver, handle, logger, { - closeOperation: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + closeOperation: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + }, }, }); @@ -459,10 +502,11 @@ describe('DBSQLOperation', () => { }); it('should throw an error in case of a status error and keep state', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -481,9 +525,10 @@ describe('DBSQLOperation', () => { }); it('should reject all methods once closed', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.close(); expect(operation.closed).to.be.true; @@ -502,6 +547,7 @@ describe('DBSQLOperation', () => { it(`should wait for finished state starting from TOperationState.${TOperationState[operationState]}`, async () => { const attemptsUntilFinished = 3; + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = operationState; @@ -514,7 +560,7 @@ describe('DBSQLOperation', () => { return driver.getOperationStatus.wrappedMethod.apply(driver, args); }); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); @@ -527,6 +573,7 @@ describe('DBSQLOperation', () => { ); it('should request progress', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -539,7 +586,7 @@ describe('DBSQLOperation', () => { return driver.getOperationStatus.wrappedMethod.apply(driver, args); }); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.finished({ progress: true }); expect(driver.getOperationStatus.called).to.be.true; @@ -550,6 +597,7 @@ describe('DBSQLOperation', () => { it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -562,7 +610,7 @@ describe('DBSQLOperation', () => { return driver.getOperationStatus.wrappedMethod.apply(driver, args); }); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const callback = sinon.stub(); @@ -573,17 +621,23 @@ describe('DBSQLOperation', () => { }); it('should pick up finished state from directResults', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); sinon.spy(driver, 'getOperationStatus'); driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - const operation = new DBSQLOperation(driver, handle, logger, { - operationStatus: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - operationState: TOperationState.FINISHED_STATE, - hasResultSet: true, + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + operationStatus: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + operationState: TOperationState.FINISHED_STATE, + hasResultSet: true, + }, }, }); @@ -594,11 +648,12 @@ describe('DBSQLOperation', () => { }); it('should throw an error in case of a status error', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); try { await operation.finished(); @@ -619,11 +674,12 @@ describe('DBSQLOperation', () => { TOperationState.TIMEDOUT_STATE, ].forEach((operationState) => { it(`should throw an error in case of a TOperationState.${TOperationState[operationState]}`, async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; driver.getOperationStatusResp.operationState = operationState; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); try { await operation.finished(); @@ -640,6 +696,8 @@ describe('DBSQLOperation', () => { describe('getSchema', () => { it('should return immediately if operation has no results', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; @@ -647,7 +705,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = false; sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const schema = await operation.getSchema(); @@ -656,6 +714,7 @@ describe('DBSQLOperation', () => { }); it('should wait for operation to complete', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -670,7 +729,7 @@ describe('DBSQLOperation', () => { driver.getResultSetMetadataResp.schema = { columns: [] }; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const schema = await operation.getSchema(); @@ -680,6 +739,7 @@ describe('DBSQLOperation', () => { }); it('should request progress', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -692,7 +752,7 @@ describe('DBSQLOperation', () => { return driver.getOperationStatus.wrappedMethod.apply(driver, args); }); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.getSchema({ progress: true }); expect(driver.getOperationStatus.called).to.be.true; @@ -703,6 +763,7 @@ describe('DBSQLOperation', () => { it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -715,7 +776,7 @@ describe('DBSQLOperation', () => { return driver.getOperationStatus.wrappedMethod.apply(driver, args); }); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const callback = sinon.stub(); @@ -726,6 +787,8 @@ describe('DBSQLOperation', () => { }); it('should fetch schema if operation has data', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -733,7 +796,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const schema = await operation.getSchema(); @@ -742,6 +805,8 @@ describe('DBSQLOperation', () => { }); it('should return cached schema on subsequent calls', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -749,7 +814,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const schema1 = await operation.getSchema(); expect(schema1).to.deep.equal(driver.getResultSetMetadataResp.schema); @@ -761,6 +826,8 @@ describe('DBSQLOperation', () => { }); it('should use schema from directResults', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -777,7 +844,7 @@ describe('DBSQLOperation', () => { }, }, }; - const operation = new DBSQLOperation(driver, handle, logger, directResults); + const operation = new DBSQLOperation({ driver, handle, context, directResults }); const schema = await operation.getSchema(); @@ -786,6 +853,8 @@ describe('DBSQLOperation', () => { }); it('should throw an error in case of a status error', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -793,7 +862,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.getResultSetMetadataResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); try { await operation.getSchema(); @@ -807,6 +876,8 @@ describe('DBSQLOperation', () => { }); it('should use appropriate result handler', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -819,7 +890,7 @@ describe('DBSQLOperation', () => { driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.COLUMN_BASED_SET; driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const resultHandler = await operation.getResultHandler(); expect(driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(JsonResult); @@ -829,7 +900,7 @@ describe('DBSQLOperation', () => { driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ARROW_BASED_SET; driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const resultHandler = await operation.getResultHandler(); expect(driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(ArrowResult); @@ -839,7 +910,7 @@ describe('DBSQLOperation', () => { driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.URL_BASED_SET; driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const resultHandler = await operation.getResultHandler(); expect(driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(CloudFetchResult); @@ -849,13 +920,15 @@ describe('DBSQLOperation', () => { describe('fetchChunk', () => { it('should return immediately if operation has no results', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = false; const driver = new DriverMock(); sinon.spy(driver, 'getResultSetMetadata'); sinon.spy(driver, 'fetchResults'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const results = await operation.fetchChunk(); @@ -865,6 +938,7 @@ describe('DBSQLOperation', () => { }); it('should wait for operation to complete', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -881,7 +955,7 @@ describe('DBSQLOperation', () => { driver.fetchResultsResp.hasMoreRows = false; driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const results = await operation.fetchChunk(); @@ -891,6 +965,7 @@ describe('DBSQLOperation', () => { }); it('should request progress', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -907,7 +982,7 @@ describe('DBSQLOperation', () => { driver.fetchResultsResp.hasMoreRows = false; driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); await operation.fetchChunk({ progress: true }); expect(driver.getOperationStatus.called).to.be.true; @@ -918,6 +993,7 @@ describe('DBSQLOperation', () => { it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; @@ -934,7 +1010,7 @@ describe('DBSQLOperation', () => { driver.fetchResultsResp.hasMoreRows = false; driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const callback = sinon.stub(); @@ -945,6 +1021,8 @@ describe('DBSQLOperation', () => { }); it('should fetch schema and data and return array of records', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -953,7 +1031,7 @@ describe('DBSQLOperation', () => { sinon.spy(driver, 'getResultSetMetadata'); sinon.spy(driver, 'fetchResults'); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const results = await operation.fetchChunk(); @@ -963,6 +1041,8 @@ describe('DBSQLOperation', () => { }); it('should return data from directResults (all the data in directResults)', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -971,18 +1051,23 @@ describe('DBSQLOperation', () => { sinon.spy(driver, 'getResultSetMetadata'); sinon.spy(driver, 'fetchResults'); - const operation = new DBSQLOperation(driver, handle, logger, { - resultSet: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - hasMoreRows: false, - results: { - columns: [ - { - i32Val: { - values: [5, 6], + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + resultSet: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + hasMoreRows: false, + results: { + columns: [ + { + i32Val: { + values: [5, 6], + }, }, - }, - ], + ], + }, }, }, }); @@ -995,6 +1080,8 @@ describe('DBSQLOperation', () => { }); it('should return data from directResults (first chunk in directResults, next chunk fetched)', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1004,18 +1091,23 @@ describe('DBSQLOperation', () => { sinon.spy(driver, 'getResultSetMetadata'); sinon.spy(driver, 'fetchResults'); - const operation = new DBSQLOperation(driver, handle, logger, { - resultSet: { - status: { statusCode: TStatusCode.SUCCESS_STATUS }, - hasMoreRows: true, - results: { - columns: [ - { - i32Val: { - values: [5, 6], + const operation = new DBSQLOperation({ + driver, + handle, + context, + directResults: { + resultSet: { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + hasMoreRows: true, + results: { + columns: [ + { + i32Val: { + values: [5, 6], + }, }, - }, - ], + ], + }, }, }, }); @@ -1034,6 +1126,8 @@ describe('DBSQLOperation', () => { }); it('should fail on unsupported result format', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1043,7 +1137,7 @@ describe('DBSQLOperation', () => { driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ROW_BASED_SET; driver.getResultSetMetadataResp.schema = { columns: [] }; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); try { await operation.fetchChunk(); @@ -1059,9 +1153,10 @@ describe('DBSQLOperation', () => { describe('fetchAll', () => { it('should fetch data while available and return it all', async () => { + const context = { getLogger: () => logger }; const handle = new OperationHandleMock(); const driver = new DriverMock(); - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); const originalData = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]; @@ -1090,6 +1185,8 @@ describe('DBSQLOperation', () => { describe('hasMoreRows', () => { it('should return False until first chunk of data fetched', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1097,7 +1194,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1105,6 +1202,8 @@ describe('DBSQLOperation', () => { }); it('should return False if operation was closed', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1112,7 +1211,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1122,6 +1221,8 @@ describe('DBSQLOperation', () => { }); it('should return False if operation was cancelled', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1129,7 +1230,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1139,6 +1240,8 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag was set in response', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1146,7 +1249,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1154,6 +1257,8 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag is False but there is actual data', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1161,7 +1266,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = false; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1169,6 +1274,8 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag is unset but there is actual data', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1176,7 +1283,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = undefined; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1184,6 +1291,8 @@ describe('DBSQLOperation', () => { }); it('should return False if hasMoreRows flag is False and there is no data', async () => { + const context = { getLogger: () => logger }; + const handle = new OperationHandleMock(); handle.hasResultSet = true; @@ -1192,7 +1301,7 @@ describe('DBSQLOperation', () => { driver.getOperationStatusResp.hasResultSet = true; driver.fetchResultsResp.hasMoreRows = false; driver.fetchResultsResp.results = undefined; - const operation = new DBSQLOperation(driver, handle, logger); + const operation = new DBSQLOperation({ driver, handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); diff --git a/tests/unit/DBSQLSession.test.js b/tests/unit/DBSQLSession.test.js index 8a439b87..2c76caae 100644 --- a/tests/unit/DBSQLSession.test.js +++ b/tests/unit/DBSQLSession.test.js @@ -34,7 +34,13 @@ function createDriverMock(customMethodHandler) { function createSession(customMethodHandler) { const driver = createDriverMock(customMethodHandler); - return new DBSQLSession(driver, { sessionId: 'id' }, logger); + return new DBSQLSession({ + driver, + handle: { sessionId: 'id' }, + context: { + getLogger: () => logger, + }, + }); } async function expectFailure(fn) { diff --git a/tests/unit/connection/auth/DatabricksOAuth/AuthorizationCode.test.js b/tests/unit/connection/auth/DatabricksOAuth/AuthorizationCode.test.js index bc83e559..a9e61f4d 100644 --- a/tests/unit/connection/auth/DatabricksOAuth/AuthorizationCode.test.js +++ b/tests/unit/connection/auth/DatabricksOAuth/AuthorizationCode.test.js @@ -2,8 +2,11 @@ const { expect, AssertionError } = require('chai'); const { EventEmitter } = require('events'); const sinon = require('sinon'); const http = require('http'); +const { DBSQLLogger, LogLevel } = require('../../../../../dist'); const AuthorizationCode = require('../../../../../dist/connection/auth/DatabricksOAuth/AuthorizationCode').default; +const logger = new DBSQLLogger({ level: LogLevel.error }); + class HttpServerMock extends EventEmitter { constructor() { super(); @@ -66,6 +69,9 @@ function prepareTestInstances(options) { const authCode = new AuthorizationCode({ client: oauthClient, ...options, + context: { + getLogger: () => logger, + }, }); sinon.stub(http, 'createServer').callsFake((requestHandler) => { diff --git a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js index 3f4a0f6b..30dbd05c 100644 --- a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js +++ b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.js @@ -1,6 +1,7 @@ const { expect, AssertionError } = require('chai'); const sinon = require('sinon'); const openidClientLib = require('openid-client'); +const { DBSQLLogger, LogLevel } = require('../../../../../dist'); const { AWSOAuthManager, AzureOAuthManager, @@ -10,6 +11,8 @@ const AuthorizationCodeModule = require('../../../../../dist/connection/auth/Dat const { createValidAccessToken, createExpiredAccessToken } = require('./utils'); +const logger = new DBSQLLogger({ level: LogLevel.error }); + class AuthorizationCodeMock { constructor() { this.fetchResult = undefined; @@ -129,6 +132,9 @@ class OAuthClientMock { const oauthManager = new OAuthManagerClass({ host: 'https://example.com', ...options, + context: { + getLogger: () => logger, + }, }); const authCode = new AuthorizationCodeMock(); diff --git a/tests/unit/connection/auth/PlainHttpAuthentication.test.js b/tests/unit/connection/auth/PlainHttpAuthentication.test.js index 9a53c76b..c799b071 100644 --- a/tests/unit/connection/auth/PlainHttpAuthentication.test.js +++ b/tests/unit/connection/auth/PlainHttpAuthentication.test.js @@ -3,7 +3,7 @@ const PlainHttpAuthentication = require('../../../../dist/connection/auth/PlainH describe('PlainHttpAuthentication', () => { it('username and password must be anonymous if nothing passed', () => { - const auth = new PlainHttpAuthentication(); + const auth = new PlainHttpAuthentication({}); expect(auth.username).to.be.eq('anonymous'); expect(auth.password).to.be.eq('anonymous'); @@ -30,7 +30,7 @@ describe('PlainHttpAuthentication', () => { }); it('auth token must be set to header', async () => { - const auth = new PlainHttpAuthentication(); + const auth = new PlainHttpAuthentication({}); const transportMock = { updateHeaders(headers) { expect(headers).to.deep.equal({ diff --git a/tests/unit/hive/HiveDriver.test.js b/tests/unit/hive/HiveDriver.test.js index 59c665f8..254969f5 100644 --- a/tests/unit/hive/HiveDriver.test.js +++ b/tests/unit/hive/HiveDriver.test.js @@ -7,8 +7,12 @@ const toTitleCase = (str) => str[0].toUpperCase() + str.slice(1); const testCommand = async (command, request) => { const client = {}; - const clientFactory = sinon.stub().returns(Promise.resolve(client)); - const driver = new HiveDriver(clientFactory); + const clientContext = { + getClient: sinon.stub().returns(Promise.resolve(client)), + }; + const driver = new HiveDriver({ + context: clientContext, + }); const response = { response: 'value' }; client[toTitleCase(command)] = function (req, cb) { @@ -18,7 +22,7 @@ const testCommand = async (command, request) => { const resp = await driver[command](request); expect(resp).to.be.deep.eq(response); - expect(clientFactory.called).to.be.true; + expect(clientContext.getClient.called).to.be.true; }; describe('HiveDriver', () => { From 2fa8c1bbf7daf63c4d16b2b1b622f7b05c312c14 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 2 Oct 2023 14:20:00 +0300 Subject: [PATCH 2/4] Pass context to all relevant classes Signed-off-by: Levko Kravets --- lib/DBSQLOperation/FetchResultsHelper.ts | 5 +++ lib/DBSQLOperation/index.ts | 7 +++-- lib/result/ArrowResult.ts | 6 +++- lib/result/CloudFetchResult.ts | 5 +-- lib/result/JsonResult.ts | 6 +++- tests/unit/result/ArrowResult.test.js | 12 +++++--- tests/unit/result/CloudFetchResult.test.js | 19 +++++++++--- tests/unit/result/JsonResult.test.js | 36 ++++++++++++++++------ tests/unit/result/compatibility.test.js | 9 ++++-- 9 files changed, 77 insertions(+), 28 deletions(-) diff --git a/lib/DBSQLOperation/FetchResultsHelper.ts b/lib/DBSQLOperation/FetchResultsHelper.ts index e1d638ad..a6cc2c22 100644 --- a/lib/DBSQLOperation/FetchResultsHelper.ts +++ b/lib/DBSQLOperation/FetchResultsHelper.ts @@ -8,6 +8,7 @@ import { import { ColumnCode, FetchType, Int64 } from '../hive/Types'; import HiveDriver from '../hive/HiveDriver'; import Status from '../dto/Status'; +import IClientContext from '../contracts/IClientContext'; function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { if (response.hasMoreRows) { @@ -36,6 +37,8 @@ function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { } export default class FetchResultsHelper { + private readonly context: IClientContext; + private readonly driver: HiveDriver; private readonly operationHandle: TOperationHandle; @@ -49,11 +52,13 @@ export default class FetchResultsHelper { public hasMoreRows: boolean = false; constructor( + context: IClientContext, driver: HiveDriver, operationHandle: TOperationHandle, prefetchedResults: Array, returnOnlyPrefetchedResults: boolean, ) { + this.context = context; this.driver = driver; this.operationHandle = operationHandle; prefetchedResults.forEach((item) => { diff --git a/lib/DBSQLOperation/index.ts b/lib/DBSQLOperation/index.ts index 27087c41..5fee4cf0 100644 --- a/lib/DBSQLOperation/index.ts +++ b/lib/DBSQLOperation/index.ts @@ -88,6 +88,7 @@ export default class DBSQLOperation implements IOperation { this.metadata = directResults?.resultSetMetadata; this._data = new FetchResultsHelper( + this.context, this.driver, this.operationHandle, [directResults?.resultSet], @@ -350,13 +351,13 @@ export default class DBSQLOperation implements IOperation { if (!this.resultHandler) { switch (resultFormat) { case TSparkRowSetType.COLUMN_BASED_SET: - this.resultHandler = new JsonResult(metadata.schema); + this.resultHandler = new JsonResult(this.context, metadata.schema); break; case TSparkRowSetType.ARROW_BASED_SET: - this.resultHandler = new ArrowResult(metadata.schema, metadata.arrowSchema); + this.resultHandler = new ArrowResult(this.context, metadata.schema, metadata.arrowSchema); break; case TSparkRowSetType.URL_BASED_SET: - this.resultHandler = new CloudFetchResult(metadata.schema); + this.resultHandler = new CloudFetchResult(this.context, metadata.schema); break; default: this.resultHandler = undefined; diff --git a/lib/result/ArrowResult.ts b/lib/result/ArrowResult.ts index 0aea9786..2d2e9756 100644 --- a/lib/result/ArrowResult.ts +++ b/lib/result/ArrowResult.ts @@ -12,6 +12,7 @@ import { util as arrowUtils, } from 'apache-arrow'; import { TRowSet, TTableSchema, TColumnDesc } from '../../thrift/TCLIService_types'; +import IClientContext from '../contracts/IClientContext'; import IOperationResult from './IOperationResult'; import { getSchemaColumns, convertThriftValue } from './utils'; @@ -21,11 +22,14 @@ type ArrowSchema = Schema; type ArrowSchemaField = Field>; export default class ArrowResult implements IOperationResult { + protected readonly context: IClientContext; + private readonly schema: Array; private readonly arrowSchema?: Buffer; - constructor(schema?: TTableSchema, arrowSchema?: Buffer) { + constructor(context: IClientContext, schema?: TTableSchema, arrowSchema?: Buffer) { + this.context = context; this.schema = getSchemaColumns(schema); this.arrowSchema = arrowSchema; } diff --git a/lib/result/CloudFetchResult.ts b/lib/result/CloudFetchResult.ts index a3cdce56..1f4a14c1 100644 --- a/lib/result/CloudFetchResult.ts +++ b/lib/result/CloudFetchResult.ts @@ -1,6 +1,7 @@ import { Buffer } from 'buffer'; import fetch, { RequestInfo, RequestInit } from 'node-fetch'; import { TRowSet, TSparkArrowResultLink, TTableSchema } from '../../thrift/TCLIService_types'; +import IClientContext from '../contracts/IClientContext'; import ArrowResult from './ArrowResult'; import globalConfig from '../globalConfig'; @@ -9,10 +10,10 @@ export default class CloudFetchResult extends ArrowResult { private downloadedBatches: Array = []; - constructor(schema?: TTableSchema) { + constructor(context: IClientContext, schema?: TTableSchema) { // Arrow schema returned in metadata is not needed for CloudFetch results: // each batch already contains schema and could be decoded as is - super(schema, Buffer.alloc(0)); + super(context, schema, Buffer.alloc(0)); } async hasPendingData() { diff --git a/lib/result/JsonResult.ts b/lib/result/JsonResult.ts index e2bbeb3f..0c7daefa 100644 --- a/lib/result/JsonResult.ts +++ b/lib/result/JsonResult.ts @@ -1,12 +1,16 @@ import { ColumnCode } from '../hive/Types'; import { TRowSet, TTableSchema, TColumn, TColumnDesc } from '../../thrift/TCLIService_types'; +import IClientContext from '../contracts/IClientContext'; import IOperationResult from './IOperationResult'; import { getSchemaColumns, convertThriftValue } from './utils'; export default class JsonResult implements IOperationResult { + private readonly context: IClientContext; + private readonly schema: Array; - constructor(schema?: TTableSchema) { + constructor(context: IClientContext, schema?: TTableSchema) { + this.context = context; this.schema = getSchemaColumns(schema); } diff --git a/tests/unit/result/ArrowResult.test.js b/tests/unit/result/ArrowResult.test.js index 905e5268..631838f9 100644 --- a/tests/unit/result/ArrowResult.test.js +++ b/tests/unit/result/ArrowResult.test.js @@ -68,13 +68,15 @@ const sampleRowSet4 = { describe('ArrowResult', () => { it('should not buffer any data', async () => { - const result = new ArrowResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + const result = new ArrowResult(context, sampleThriftSchema, sampleArrowSchema); await result.getValue([sampleRowSet1]); expect(await result.hasPendingData()).to.be.false; }); it('should convert data', async () => { - const result = new ArrowResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + const result = new ArrowResult(context, sampleThriftSchema, sampleArrowSchema); expect(await result.getValue([sampleRowSet1])).to.be.deep.eq([]); expect(await result.getValue([sampleRowSet2])).to.be.deep.eq([]); expect(await result.getValue([sampleRowSet3])).to.be.deep.eq([]); @@ -82,13 +84,15 @@ describe('ArrowResult', () => { }); it('should return empty array if no data to process', async () => { - const result = new ArrowResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + const result = new ArrowResult(context, sampleThriftSchema, sampleArrowSchema); expect(await result.getValue()).to.be.deep.eq([]); expect(await result.getValue([])).to.be.deep.eq([]); }); it('should return empty array if no schema available', async () => { - const result = new ArrowResult(); + const context = {}; + const result = new ArrowResult(context); expect(await result.getValue([sampleRowSet4])).to.be.deep.eq([]); }); }); diff --git a/tests/unit/result/CloudFetchResult.test.js b/tests/unit/result/CloudFetchResult.test.js index bbf7e0aa..20451093 100644 --- a/tests/unit/result/CloudFetchResult.test.js +++ b/tests/unit/result/CloudFetchResult.test.js @@ -106,7 +106,8 @@ describe('CloudFetchResult', () => { }); it('should report pending data if there are any', async () => { - const result = new CloudFetchResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + const result = new CloudFetchResult({}, sampleThriftSchema, sampleArrowSchema); case1: { result.pendingLinks = []; @@ -130,7 +131,9 @@ describe('CloudFetchResult', () => { it('should extract links from row sets', async () => { globalConfig.cloudFetchConcurrentDownloads = 0; // this will prevent it from downloading batches - const result = new CloudFetchResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + + const result = new CloudFetchResult({}, sampleThriftSchema, sampleArrowSchema); sinon.stub(result, 'fetch').returns( Promise.resolve({ @@ -153,7 +156,9 @@ describe('CloudFetchResult', () => { it('should download batches according to settings', async () => { globalConfig.cloudFetchConcurrentDownloads = 2; - const result = new CloudFetchResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + + const result = new CloudFetchResult({}, sampleThriftSchema, sampleArrowSchema); sinon.stub(result, 'fetch').returns( Promise.resolve({ @@ -199,7 +204,9 @@ describe('CloudFetchResult', () => { it('should handle HTTP errors', async () => { globalConfig.cloudFetchConcurrentDownloads = 1; - const result = new CloudFetchResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + + const result = new CloudFetchResult({}, sampleThriftSchema, sampleArrowSchema); sinon.stub(result, 'fetch').returns( Promise.resolve({ @@ -225,7 +232,9 @@ describe('CloudFetchResult', () => { }); it('should handle expired links', async () => { - const result = new CloudFetchResult(sampleThriftSchema, sampleArrowSchema); + const context = {}; + + const result = new CloudFetchResult(context, sampleThriftSchema, sampleArrowSchema); sinon.stub(result, 'fetch').returns( Promise.resolve({ diff --git a/tests/unit/result/JsonResult.test.js b/tests/unit/result/JsonResult.test.js index e44a7e8a..f7e90259 100644 --- a/tests/unit/result/JsonResult.test.js +++ b/tests/unit/result/JsonResult.test.js @@ -38,7 +38,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); await result.getValue(data); expect(await result.hasPendingData()).to.be.false; }); @@ -124,7 +126,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue(data)).to.be.deep.eq([ { @@ -194,7 +198,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue(data)).to.be.deep.eq([ { @@ -234,7 +240,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue(data)).to.be.deep.eq([ { 'table.id': '0' }, @@ -245,7 +253,9 @@ describe('JsonResult', () => { }); it('should detect nulls', () => { - const result = new JsonResult(null); + const context = {}; + + const result = new JsonResult(context, null); const buf = Buffer.from([0x55, 0xaa, 0xc3]); [ @@ -357,7 +367,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue(data)).to.be.deep.eq([ { @@ -386,7 +398,9 @@ describe('JsonResult', () => { columns: [getColumnSchema('table.id', TCLIService_types.TTypeId.STRING_TYPE, 1)], }; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue()).to.be.deep.eq([]); expect(await result.getValue([])).to.be.deep.eq([]); @@ -403,7 +417,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(); + const context = {}; + + const result = new JsonResult(context); expect(await result.getValue(data)).to.be.deep.eq([]); }); @@ -436,7 +452,9 @@ describe('JsonResult', () => { }, ]; - const result = new JsonResult(schema); + const context = {}; + + const result = new JsonResult(context, schema); expect(await result.getValue(data)).to.be.deep.eq([ { diff --git a/tests/unit/result/compatibility.test.js b/tests/unit/result/compatibility.test.js index 8df3e6a8..5b27d39e 100644 --- a/tests/unit/result/compatibility.test.js +++ b/tests/unit/result/compatibility.test.js @@ -9,19 +9,22 @@ const fixtureArrowNT = require('../../fixtures/compatibility/arrow_native_types' describe('Result handlers compatibility tests', () => { it('colum-based data', async () => { - const result = new JsonResult(fixtureColumn.schema); + const context = {}; + const result = new JsonResult(context, fixtureColumn.schema); const rows = await result.getValue(fixtureColumn.rowSets); expect(rows).to.deep.equal(fixtureColumn.expected); }); it('arrow-based data without native types', async () => { - const result = new ArrowResult(fixtureArrow.schema, fixtureArrow.arrowSchema); + const context = {}; + const result = new ArrowResult(context, fixtureArrow.schema, fixtureArrow.arrowSchema); const rows = await result.getValue(fixtureArrow.rowSets); expect(fixArrowResult(rows)).to.deep.equal(fixtureArrow.expected); }); it('arrow-based data with native types', async () => { - const result = new ArrowResult(fixtureArrowNT.schema, fixtureArrowNT.arrowSchema); + const context = {}; + const result = new ArrowResult(context, fixtureArrowNT.schema, fixtureArrowNT.arrowSchema); const rows = await result.getValue(fixtureArrowNT.rowSets); expect(fixArrowResult(rows)).to.deep.equal(fixtureArrowNT.expected); }); From 92d874a3bb304ee298b2efa6d177bcca9ee27374 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 2 Oct 2023 15:23:16 +0300 Subject: [PATCH 3/4] Make driver a part of context Signed-off-by: Levko Kravets --- lib/DBSQLClient.ts | 16 +- lib/DBSQLOperation/FetchResultsHelper.ts | 8 +- lib/DBSQLOperation/index.ts | 20 +- lib/DBSQLSession.ts | 44 +- lib/contracts/IClientContext.ts | 3 + lib/contracts/IDriver.ts | 88 +++ lib/hive/HiveDriver.ts | 3 +- tests/unit/DBSQLOperation.test.js | 683 +++++++++++------------ tests/unit/DBSQLSession.test.js | 17 +- 9 files changed, 479 insertions(+), 403 deletions(-) create mode 100644 lib/contracts/IDriver.ts diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 31edbf47..f943c211 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -4,6 +4,7 @@ import { EventEmitter } from 'events'; import TCLIService from '../thrift/TCLIService'; import { TProtocolVersion } from '../thrift/TCLIService_types'; import IDBSQLClient, { ClientOptions, ConnectionOptions, OpenSessionRequest } from './contracts/IDBSQLClient'; +import IDriver from './contracts/IDriver'; import IClientContext from './contracts/IClientContext'; import HiveDriver from './hive/HiveDriver'; import { Int64 } from './hive/Types'; @@ -49,6 +50,10 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I private client?: TCLIService.Client; + private readonly driver = new HiveDriver({ + context: this, + }); + private readonly logger: IDBSQLLogger; private readonly thrift = thrift; @@ -158,18 +163,13 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I * const session = await client.openSession(); */ public async openSession(request: OpenSessionRequest = {}): Promise { - const driver = new HiveDriver({ - context: this, - }); - - const response = await driver.openSession({ + const response = await this.driver.openSession({ client_protocol_i64: new Int64(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8), ...getInitialNamespaceOptions(request.initialCatalog, request.initialSchema), }); Status.assert(response.status); const session = new DBSQLSession({ - driver, handle: definedOrError(response.sessionHandle), context: this, }); @@ -212,4 +212,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I return this.client; } + + public async getDriver(): Promise { + return this.driver; + } } diff --git a/lib/DBSQLOperation/FetchResultsHelper.ts b/lib/DBSQLOperation/FetchResultsHelper.ts index a6cc2c22..79f82603 100644 --- a/lib/DBSQLOperation/FetchResultsHelper.ts +++ b/lib/DBSQLOperation/FetchResultsHelper.ts @@ -6,7 +6,6 @@ import { TRowSet, } from '../../thrift/TCLIService_types'; import { ColumnCode, FetchType, Int64 } from '../hive/Types'; -import HiveDriver from '../hive/HiveDriver'; import Status from '../dto/Status'; import IClientContext from '../contracts/IClientContext'; @@ -39,8 +38,6 @@ function checkIfOperationHasMoreRows(response: TFetchResultsResp): boolean { export default class FetchResultsHelper { private readonly context: IClientContext; - private readonly driver: HiveDriver; - private readonly operationHandle: TOperationHandle; private fetchOrientation: TFetchOrientation = TFetchOrientation.FETCH_FIRST; @@ -53,13 +50,11 @@ export default class FetchResultsHelper { constructor( context: IClientContext, - driver: HiveDriver, operationHandle: TOperationHandle, prefetchedResults: Array, returnOnlyPrefetchedResults: boolean, ) { this.context = context; - this.driver = driver; this.operationHandle = operationHandle; prefetchedResults.forEach((item) => { if (item) { @@ -90,7 +85,8 @@ export default class FetchResultsHelper { return this.processFetchResponse(prefetchedResponse); } - const response = await this.driver.fetchResults({ + const driver = await this.context.getDriver(); + const response = await driver.fetchResults({ operationHandle: this.operationHandle, orientation: this.fetchOrientation, maxRows: new Int64(maxRows), diff --git a/lib/DBSQLOperation/index.ts b/lib/DBSQLOperation/index.ts index 5fee4cf0..5198d726 100644 --- a/lib/DBSQLOperation/index.ts +++ b/lib/DBSQLOperation/index.ts @@ -5,7 +5,6 @@ import IOperation, { GetSchemaOptions, WaitUntilReadyOptions, } from '../contracts/IOperation'; -import HiveDriver from '../hive/HiveDriver'; import { TGetOperationStatusResp, TOperationHandle, @@ -32,7 +31,6 @@ const defaultMaxRows = 100000; interface DBSQLOperationConstructorOptions { handle: TOperationHandle; - driver: HiveDriver; directResults?: TSparkDirectResults; context: IClientContext; } @@ -48,8 +46,6 @@ async function delay(ms?: number): Promise { export default class DBSQLOperation implements IOperation { private readonly context: IClientContext; - private readonly driver: HiveDriver; - private readonly operationHandle: TOperationHandle; public onClose?: () => void; @@ -74,8 +70,7 @@ export default class DBSQLOperation implements IOperation { private resultHandler?: IOperationResult; - constructor({ driver, handle, directResults, context }: DBSQLOperationConstructorOptions) { - this.driver = driver; + constructor({ handle, directResults, context }: DBSQLOperationConstructorOptions) { this.operationHandle = handle; this.context = context; @@ -89,7 +84,6 @@ export default class DBSQLOperation implements IOperation { this.metadata = directResults?.resultSetMetadata; this._data = new FetchResultsHelper( this.context, - this.driver, this.operationHandle, [directResults?.resultSet], useOnlyPrefetchedResults, @@ -171,7 +165,8 @@ export default class DBSQLOperation implements IOperation { return this.operationStatus; } - const response = await this.driver.getOperationStatus({ + const driver = await this.context.getDriver(); + const response = await driver.getOperationStatus({ operationHandle: this.operationHandle, getProgressUpdate: progress, }); @@ -190,7 +185,8 @@ export default class DBSQLOperation implements IOperation { this.context.getLogger().log(LogLevel.debug, `Cancelling operation with id: ${this.getId()}`); - const response = await this.driver.cancelOperation({ + const driver = await this.context.getDriver(); + const response = await driver.cancelOperation({ operationHandle: this.operationHandle, }); Status.assert(response.status); @@ -213,9 +209,10 @@ export default class DBSQLOperation implements IOperation { this.context.getLogger().log(LogLevel.debug, `Closing operation with id: ${this.getId()}`); + const driver = await this.context.getDriver(); const response = this.closeOperation ?? - (await this.driver.closeOperation({ + (await driver.closeOperation({ operationHandle: this.operationHandle, })); Status.assert(response.status); @@ -334,7 +331,8 @@ export default class DBSQLOperation implements IOperation { private async fetchMetadata() { if (!this.metadata) { - const metadata = await this.driver.getResultSetMetadata({ + const driver = await this.context.getDriver(); + const metadata = await driver.getResultSetMetadata({ operationHandle: this.operationHandle, }); Status.assert(metadata.status); diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 781e94f3..e64c1ac2 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -12,7 +12,6 @@ import { TSparkParameter, TProtocolVersion, } from '../thrift/TCLIService_types'; -import HiveDriver from './hive/HiveDriver'; import { Int64 } from './hive/Types'; import IDBSQLSession, { ExecuteStatementOptions, @@ -132,15 +131,12 @@ function getQueryParameters( interface DBSQLSessionConstructorOptions { handle: TSessionHandle; - driver: HiveDriver; context: IClientContext; } export default class DBSQLSession implements IDBSQLSession { private readonly context: IClientContext; - private readonly driver: HiveDriver; - private readonly sessionHandle: TSessionHandle; private isOpen = true; @@ -149,8 +145,7 @@ export default class DBSQLSession implements IDBSQLSession { private operations = new CloseableCollection(); - constructor({ handle, driver, context }: DBSQLSessionConstructorOptions) { - this.driver = driver; + constructor({ handle, context }: DBSQLSessionConstructorOptions) { this.sessionHandle = handle; this.context = context; this.context.getLogger().log(LogLevel.debug, `Session created with id: ${this.getId()}`); @@ -170,7 +165,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getInfo(infoType: number): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getInfo({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getInfo({ sessionHandle: this.sessionHandle, infoType, }); @@ -190,7 +186,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async executeStatement(statement: string, options: ExecuteStatementOptions = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.executeStatement({ + const driver = await this.context.getDriver(); + const operationPromise = driver.executeStatement({ sessionHandle: this.sessionHandle, statement, queryTimeout: options.queryTimeout, @@ -315,7 +312,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getTypeInfo(request: TypeInfoRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getTypeInfo({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getTypeInfo({ sessionHandle: this.sessionHandle, runAsync: true, ...getDirectResultsOptions(request.maxRows), @@ -332,7 +330,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getCatalogs(request: CatalogsRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getCatalogs({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getCatalogs({ sessionHandle: this.sessionHandle, runAsync: true, ...getDirectResultsOptions(request.maxRows), @@ -349,7 +348,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getSchemas(request: SchemasRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getSchemas({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getSchemas({ sessionHandle: this.sessionHandle, catalogName: request.catalogName, schemaName: request.schemaName, @@ -368,7 +368,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getTables(request: TablesRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getTables({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getTables({ sessionHandle: this.sessionHandle, catalogName: request.catalogName, schemaName: request.schemaName, @@ -389,7 +390,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getTableTypes(request: TableTypesRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getTableTypes({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getTableTypes({ sessionHandle: this.sessionHandle, runAsync: true, ...getDirectResultsOptions(request.maxRows), @@ -406,7 +408,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getColumns(request: ColumnsRequest = {}): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getColumns({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getColumns({ sessionHandle: this.sessionHandle, catalogName: request.catalogName, schemaName: request.schemaName, @@ -427,7 +430,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getFunctions(request: FunctionsRequest): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getFunctions({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getFunctions({ sessionHandle: this.sessionHandle, catalogName: request.catalogName, schemaName: request.schemaName, @@ -441,7 +445,8 @@ export default class DBSQLSession implements IDBSQLSession { public async getPrimaryKeys(request: PrimaryKeysRequest): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getPrimaryKeys({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getPrimaryKeys({ sessionHandle: this.sessionHandle, catalogName: request.catalogName, schemaName: request.schemaName, @@ -461,7 +466,8 @@ export default class DBSQLSession implements IDBSQLSession { */ public async getCrossReference(request: CrossReferenceRequest): Promise { await this.failIfClosed(); - const operationPromise = this.driver.getCrossReference({ + const driver = await this.context.getDriver(); + const operationPromise = driver.getCrossReference({ sessionHandle: this.sessionHandle, parentCatalogName: request.parentCatalogName, parentSchemaName: request.parentSchemaName, @@ -489,7 +495,8 @@ export default class DBSQLSession implements IDBSQLSession { // Close owned operations one by one, removing successfully closed ones from the list await this.operations.closeAll(); - const response = await this.driver.closeSession({ + const driver = await this.context.getDriver(); + const response = await driver.closeSession({ sessionHandle: this.sessionHandle, }); // check status for being successful @@ -507,7 +514,6 @@ export default class DBSQLSession implements IDBSQLSession { Status.assert(response.status); const handle = definedOrError(response.operationHandle); const operation = new DBSQLOperation({ - driver: this.driver, handle, directResults: response.directResults, context: this.context, diff --git a/lib/contracts/IClientContext.ts b/lib/contracts/IClientContext.ts index 0e6d5ed3..8df0f7e9 100644 --- a/lib/contracts/IClientContext.ts +++ b/lib/contracts/IClientContext.ts @@ -1,4 +1,5 @@ import IDBSQLLogger from './IDBSQLLogger'; +import IDriver from './IDriver'; import IConnectionProvider from '../connection/contracts/IConnectionProvider'; import TCLIService from '../../thrift/TCLIService'; @@ -8,4 +9,6 @@ export default interface IClientContext { getConnectionProvider(): Promise; getClient(): Promise; + + getDriver(): Promise; } diff --git a/lib/contracts/IDriver.ts b/lib/contracts/IDriver.ts new file mode 100644 index 00000000..4ee755dc --- /dev/null +++ b/lib/contracts/IDriver.ts @@ -0,0 +1,88 @@ +import { + TOpenSessionReq, + TOpenSessionResp, + TCloseSessionReq, + TCloseSessionResp, + TExecuteStatementReq, + TExecuteStatementResp, + TGetResultSetMetadataReq, + TGetResultSetMetadataResp, + TFetchResultsReq, + TFetchResultsResp, + TGetInfoReq, + TGetInfoResp, + TGetTypeInfoReq, + TGetTypeInfoResp, + TGetCatalogsReq, + TGetCatalogsResp, + TGetSchemasReq, + TGetSchemasResp, + TGetTablesReq, + TGetTablesResp, + TGetTableTypesReq, + TGetTableTypesResp, + TGetColumnsReq, + TGetColumnsResp, + TGetFunctionsReq, + TGetFunctionsResp, + TGetPrimaryKeysReq, + TGetPrimaryKeysResp, + TGetCrossReferenceReq, + TGetCrossReferenceResp, + TGetOperationStatusReq, + TGetOperationStatusResp, + TCancelOperationReq, + TCancelOperationResp, + TCloseOperationReq, + TCloseOperationResp, + TGetDelegationTokenReq, + TGetDelegationTokenResp, + TCancelDelegationTokenReq, + TCancelDelegationTokenResp, + TRenewDelegationTokenReq, + TRenewDelegationTokenResp, +} from '../../thrift/TCLIService_types'; + +export default interface IDriver { + openSession(request: TOpenSessionReq): Promise; + + closeSession(request: TCloseSessionReq): Promise; + + executeStatement(request: TExecuteStatementReq): Promise; + + getResultSetMetadata(request: TGetResultSetMetadataReq): Promise; + + fetchResults(request: TFetchResultsReq): Promise; + + getInfo(request: TGetInfoReq): Promise; + + getTypeInfo(request: TGetTypeInfoReq): Promise; + + getCatalogs(request: TGetCatalogsReq): Promise; + + getSchemas(request: TGetSchemasReq): Promise; + + getTables(request: TGetTablesReq): Promise; + + getTableTypes(request: TGetTableTypesReq): Promise; + + getColumns(request: TGetColumnsReq): Promise; + + getFunctions(request: TGetFunctionsReq): Promise; + + getPrimaryKeys(request: TGetPrimaryKeysReq): Promise; + + getCrossReference(request: TGetCrossReferenceReq): Promise; + + getOperationStatus(request: TGetOperationStatusReq): Promise; + + cancelOperation(request: TCancelOperationReq): Promise; + + closeOperation(request: TCloseOperationReq): Promise; + + getDelegationToken(request: TGetDelegationTokenReq): Promise; + + cancelDelegationToken(request: TCancelDelegationTokenReq): Promise; + + renewDelegationToken(request: TRenewDelegationTokenReq): Promise; +} diff --git a/lib/hive/HiveDriver.ts b/lib/hive/HiveDriver.ts index 9c34d743..9b7384b1 100644 --- a/lib/hive/HiveDriver.ts +++ b/lib/hive/HiveDriver.ts @@ -42,13 +42,14 @@ import CloseOperationCommand from './Commands/CloseOperationCommand'; import GetDelegationTokenCommand from './Commands/GetDelegationTokenCommand'; import CancelDelegationTokenCommand from './Commands/CancelDelegationTokenCommand'; import RenewDelegationTokenCommand from './Commands/RenewDelegationTokenCommand'; +import IDriver from '../contracts/IDriver'; import IClientContext from '../contracts/IClientContext'; export interface HiveDriverOptions { context: IClientContext; } -export default class HiveDriver { +export default class HiveDriver implements IDriver { private readonly context: IClientContext; constructor(options: HiveDriverOptions) { diff --git a/tests/unit/DBSQLOperation.test.js b/tests/unit/DBSQLOperation.test.js index 4b45f240..94834baf 100644 --- a/tests/unit/DBSQLOperation.test.js +++ b/tests/unit/DBSQLOperation.test.js @@ -10,10 +10,6 @@ const JsonResult = require('../../dist/result/JsonResult').default; const ArrowResult = require('../../dist/result/ArrowResult').default; const CloudFetchResult = require('../../dist/result/CloudFetchResult').default; -// Create logger that won't emit -// -const logger = new DBSQLLogger({ level: LogLevel.error }); - class OperationHandleMock { constructor(hasResultSet = true) { this.operationId = 1; @@ -104,28 +100,41 @@ class DriverMock { } } +class ClientContextMock { + constructor(props) { + // Create logger that won't emit + this.logger = new DBSQLLogger({ level: LogLevel.error }); + this.driver = new DriverMock(); + } + + getLogger() { + return this.logger; + } + + async getDriver() { + return this.driver; + } +} + describe('DBSQLOperation', () => { describe('status', () => { it('should pick up state from operation handle', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.true; }); it('should pick up state from directResults', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -142,100 +151,95 @@ describe('DBSQLOperation', () => { }); it('should fetch status and update internal state', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.false; const status = await operation.status(); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(status.operationState).to.equal(TOperationState.FINISHED_STATE); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); expect(operation.hasResultSet).to.be.true; }); it('should request progress', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.status(true); - expect(driver.getOperationStatus.called).to.be.true; - const request = driver.getOperationStatus.getCall(0).args[0]; + expect(context.driver.getOperationStatus.called).to.be.true; + const request = context.driver.getOperationStatus.getCall(0).args[0]; expect(request.getProgressUpdate).to.be.true; }); it('should not fetch status once operation is finished', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.hasResultSet = true; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); expect(operation.hasResultSet).to.be.false; // First call - should fetch data and cache - driver.getOperationStatusResp = { - ...driver.getOperationStatusResp, + context.driver.getOperationStatusResp = { + ...context.driver.getOperationStatusResp, operationState: TOperationState.FINISHED_STATE, }; const status1 = await operation.status(); - expect(driver.getOperationStatus.callCount).to.equal(1); + expect(context.driver.getOperationStatus.callCount).to.equal(1); expect(status1.operationState).to.equal(TOperationState.FINISHED_STATE); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); expect(operation.hasResultSet).to.be.true; // Second call - should return cached data - driver.getOperationStatusResp = { - ...driver.getOperationStatusResp, + context.driver.getOperationStatusResp = { + ...context.driver.getOperationStatusResp, operationState: TOperationState.RUNNING_STATE, }; const status2 = await operation.status(); - expect(driver.getOperationStatus.callCount).to.equal(1); + expect(context.driver.getOperationStatus.callCount).to.equal(1); expect(status2.operationState).to.equal(TOperationState.FINISHED_STATE); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); expect(operation.hasResultSet).to.be.true; }); it('should fetch status if directResults status is not finished', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -252,25 +256,23 @@ describe('DBSQLOperation', () => { const status = await operation.status(false); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(status.operationState).to.equal(TOperationState.FINISHED_STATE); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); expect(operation.hasResultSet).to.be.true; }); it('should not fetch status if directResults status is finished', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.operationState = TOperationState.RUNNING_STATE; - driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.operationState = TOperationState.RUNNING_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -287,18 +289,18 @@ describe('DBSQLOperation', () => { const status = await operation.status(false); - expect(driver.getOperationStatus.called).to.be.false; + expect(context.driver.getOperationStatus.called).to.be.false; expect(status.operationState).to.equal(TOperationState.FINISHED_STATE); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); expect(operation.hasResultSet).to.be.false; }); it('should throw an error in case of a status error', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation({ driver, handle, context }); + + context.driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; + const operation = new DBSQLOperation({ handle, context }); try { await operation.status(false); @@ -314,71 +316,71 @@ describe('DBSQLOperation', () => { describe('cancel', () => { it('should cancel operation and update state', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'cancelOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.cancel(); - expect(driver.cancelOperation.called).to.be.true; + expect(context.driver.cancelOperation.called).to.be.true; expect(operation.cancelled).to.be.true; expect(operation.closed).to.be.false; }); it('should return immediately if already cancelled', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'cancelOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.cancel(); - expect(driver.cancelOperation.callCount).to.be.equal(1); + expect(context.driver.cancelOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.true; expect(operation.closed).to.be.false; await operation.cancel(); - expect(driver.cancelOperation.callCount).to.be.equal(1); + expect(context.driver.cancelOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.true; expect(operation.closed).to.be.false; }); it('should return immediately if already closed', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'cancelOperation'); - sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'cancelOperation'); + sinon.spy(context.driver, 'closeOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.close(); - expect(driver.closeOperation.callCount).to.be.equal(1); + expect(context.driver.closeOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; await operation.cancel(); - expect(driver.cancelOperation.callCount).to.be.equal(0); + expect(context.driver.cancelOperation.callCount).to.be.equal(0); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; }); it('should throw an error in case of a status error and keep state', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.cancelOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation({ driver, handle, context }); + + context.driver.cancelOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -397,10 +399,9 @@ describe('DBSQLOperation', () => { }); it('should reject all methods once cancelled', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.cancel(); expect(operation.cancelled).to.be.true; @@ -415,72 +416,71 @@ describe('DBSQLOperation', () => { describe('close', () => { it('should close operation and update state', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'closeOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.close(); - expect(driver.closeOperation.called).to.be.true; + expect(context.driver.closeOperation.called).to.be.true; expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; }); it('should return immediately if already closed', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'closeOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'closeOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.close(); - expect(driver.closeOperation.callCount).to.be.equal(1); + expect(context.driver.closeOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; await operation.close(); - expect(driver.closeOperation.callCount).to.be.equal(1); + expect(context.driver.closeOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; }); it('should return immediately if already cancelled', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'closeOperation'); - sinon.spy(driver, 'cancelOperation'); - const operation = new DBSQLOperation({ driver, handle, context }); + + sinon.spy(context.driver, 'closeOperation'); + sinon.spy(context.driver, 'cancelOperation'); + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; await operation.cancel(); - expect(driver.cancelOperation.callCount).to.be.equal(1); + expect(context.driver.cancelOperation.callCount).to.be.equal(1); expect(operation.cancelled).to.be.true; expect(operation.closed).to.be.false; await operation.close(); - expect(driver.closeOperation.callCount).to.be.equal(0); + expect(context.driver.closeOperation.callCount).to.be.equal(0); expect(operation.cancelled).to.be.true; expect(operation.closed).to.be.false; }); it('should initialize from directResults', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'closeOperation'); + + sinon.spy(context.driver, 'closeOperation'); const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -495,18 +495,18 @@ describe('DBSQLOperation', () => { await operation.close(); - expect(driver.closeOperation.called).to.be.false; + expect(context.driver.closeOperation.called).to.be.false; expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.true; - expect(driver.closeOperation.callCount).to.be.equal(0); + expect(context.driver.closeOperation.callCount).to.be.equal(0); }); it('should throw an error in case of a status error and keep state', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation({ driver, handle, context }); + + context.driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS; + const operation = new DBSQLOperation({ handle, context }); expect(operation.cancelled).to.be.false; expect(operation.closed).to.be.false; @@ -525,10 +525,9 @@ describe('DBSQLOperation', () => { }); it('should reject all methods once closed', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.close(); expect(operation.closed).to.be.true; @@ -547,89 +546,88 @@ describe('DBSQLOperation', () => { it(`should wait for finished state starting from TOperationState.${TOperationState[operationState]}`, async () => { const attemptsUntilFinished = 3; - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = operationState; + + context.driver.getOperationStatusResp.operationState = operationState; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onCall(attemptsUntilFinished - 1) // count is zero-based .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); expect(operation.state).to.equal(TOperationState.INITIALIZED_STATE); await operation.finished(); - expect(driver.getOperationStatus.callCount).to.be.equal(attemptsUntilFinished); + expect(context.driver.getOperationStatus.callCount).to.be.equal(attemptsUntilFinished); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); }); }, ); it('should request progress', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onSecondCall() .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.finished({ progress: true }); - expect(driver.getOperationStatus.called).to.be.true; - const request = driver.getOperationStatus.getCall(0).args[0]; + expect(context.driver.getOperationStatus.called).to.be.true; + const request = context.driver.getOperationStatus.getCall(0).args[0]; expect(request.getProgressUpdate).to.be.true; }); it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onCall(attemptsUntilFinished - 1) // count is zero-based .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const callback = sinon.stub(); await operation.finished({ callback }); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(callback.callCount).to.be.equal(attemptsUntilFinished); }); it('should pick up finished state from directResults', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - sinon.spy(driver, 'getOperationStatus'); - driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + + sinon.spy(context.driver, 'getOperationStatus'); + context.driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -644,16 +642,16 @@ describe('DBSQLOperation', () => { await operation.finished(); // Once operation is finished - no need to fetch status again - expect(driver.getOperationStatus.called).to.be.false; + expect(context.driver.getOperationStatus.called).to.be.false; }); it('should throw an error in case of a status error', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - const operation = new DBSQLOperation({ driver, handle, context }); + + context.driver.getOperationStatusResp.status.statusCode = TStatusCode.ERROR_STATUS; + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + const operation = new DBSQLOperation({ handle, context }); try { await operation.finished(); @@ -674,12 +672,12 @@ describe('DBSQLOperation', () => { TOperationState.TIMEDOUT_STATE, ].forEach((operationState) => { it(`should throw an error in case of a TOperationState.${TOperationState[operationState]}`, async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; - driver.getOperationStatusResp.operationState = operationState; - const operation = new DBSQLOperation({ driver, handle, context }); + + context.driver.getOperationStatusResp.status.statusCode = TStatusCode.SUCCESS_STATUS; + context.driver.getOperationStatusResp.operationState = operationState; + const operation = new DBSQLOperation({ handle, context }); try { await operation.finished(); @@ -696,145 +694,141 @@ describe('DBSQLOperation', () => { describe('getSchema', () => { it('should return immediately if operation has no results', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = false; - sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = false; + sinon.spy(context.driver, 'getResultSetMetadata'); + const operation = new DBSQLOperation({ handle, context }); const schema = await operation.getSchema(); expect(schema).to.be.null; - expect(driver.getResultSetMetadata.called).to.be.false; + expect(context.driver.getResultSetMetadata.called).to.be.false; }); it('should wait for operation to complete', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onSecondCall() .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - driver.getResultSetMetadataResp.schema = { columns: [] }; + context.driver.getResultSetMetadataResp.schema = { columns: [] }; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const schema = await operation.getSchema(); - expect(driver.getOperationStatus.called).to.be.true; - expect(schema).to.deep.equal(driver.getResultSetMetadataResp.schema); + expect(context.driver.getOperationStatus.called).to.be.true; + expect(schema).to.deep.equal(context.driver.getResultSetMetadataResp.schema); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); }); it('should request progress', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onSecondCall() .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.getSchema({ progress: true }); - expect(driver.getOperationStatus.called).to.be.true; - const request = driver.getOperationStatus.getCall(0).args[0]; + expect(context.driver.getOperationStatus.called).to.be.true; + const request = context.driver.getOperationStatus.getCall(0).args[0]; expect(request.getProgressUpdate).to.be.true; }); it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onCall(attemptsUntilFinished - 1) // count is zero-based .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const callback = sinon.stub(); await operation.getSchema({ callback }); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(callback.callCount).to.be.equal(attemptsUntilFinished); }); it('should fetch schema if operation has data', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getResultSetMetadata'); + const operation = new DBSQLOperation({ handle, context }); const schema = await operation.getSchema(); - expect(schema).to.deep.equal(driver.getResultSetMetadataResp.schema); - expect(driver.getResultSetMetadata.called).to.be.true; + expect(schema).to.deep.equal(context.driver.getResultSetMetadataResp.schema); + expect(context.driver.getResultSetMetadata.called).to.be.true; }); it('should return cached schema on subsequent calls', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - sinon.spy(driver, 'getResultSetMetadata'); - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getResultSetMetadata'); + const operation = new DBSQLOperation({ handle, context }); const schema1 = await operation.getSchema(); - expect(schema1).to.deep.equal(driver.getResultSetMetadataResp.schema); - expect(driver.getResultSetMetadata.callCount).to.equal(1); + expect(schema1).to.deep.equal(context.driver.getResultSetMetadataResp.schema); + expect(context.driver.getResultSetMetadata.callCount).to.equal(1); const schema2 = await operation.getSchema(); - expect(schema2).to.deep.equal(driver.getResultSetMetadataResp.schema); - expect(driver.getResultSetMetadata.callCount).to.equal(1); // no additional requests + expect(schema2).to.deep.equal(context.driver.getResultSetMetadataResp.schema); + expect(context.driver.getResultSetMetadata.callCount).to.equal(1); // no additional requests }); it('should use schema from directResults', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - sinon.spy(driver, 'getResultSetMetadata'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getResultSetMetadata'); const directResults = { resultSetMetadata: { @@ -844,25 +838,24 @@ describe('DBSQLOperation', () => { }, }, }; - const operation = new DBSQLOperation({ driver, handle, context, directResults }); + const operation = new DBSQLOperation({ handle, context, directResults }); const schema = await operation.getSchema(); expect(schema).to.deep.equal(directResults.resultSetMetadata.schema); - expect(driver.getResultSetMetadata.called).to.be.false; + expect(context.driver.getResultSetMetadata.called).to.be.false; }); it('should throw an error in case of a status error', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.getResultSetMetadataResp.status.statusCode = TStatusCode.ERROR_STATUS; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.getResultSetMetadataResp.status.statusCode = TStatusCode.ERROR_STATUS; + const operation = new DBSQLOperation({ handle, context }); try { await operation.getSchema(); @@ -876,43 +869,42 @@ describe('DBSQLOperation', () => { }); it('should use appropriate result handler', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - sinon.spy(driver, 'getResultSetMetadata'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getResultSetMetadata'); jsonHandler: { - driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.COLUMN_BASED_SET; - driver.getResultSetMetadata.resetHistory(); + context.driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.COLUMN_BASED_SET; + context.driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const resultHandler = await operation.getResultHandler(); - expect(driver.getResultSetMetadata.called).to.be.true; + expect(context.driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(JsonResult); } arrowHandler: { - driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ARROW_BASED_SET; - driver.getResultSetMetadata.resetHistory(); + context.driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ARROW_BASED_SET; + context.driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const resultHandler = await operation.getResultHandler(); - expect(driver.getResultSetMetadata.called).to.be.true; + expect(context.driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(ArrowResult); } cloudFetchHandler: { - driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.URL_BASED_SET; - driver.getResultSetMetadata.resetHistory(); + context.driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.URL_BASED_SET; + context.driver.getResultSetMetadata.resetHistory(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const resultHandler = await operation.getResultHandler(); - expect(driver.getResultSetMetadata.called).to.be.true; + expect(context.driver.getResultSetMetadata.called).to.be.true; expect(resultHandler).to.be.instanceOf(CloudFetchResult); } }); @@ -920,139 +912,135 @@ describe('DBSQLOperation', () => { describe('fetchChunk', () => { it('should return immediately if operation has no results', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = false; - const driver = new DriverMock(); - sinon.spy(driver, 'getResultSetMetadata'); - sinon.spy(driver, 'fetchResults'); - const operation = new DBSQLOperation({ driver, handle, context }); + sinon.spy(context.driver, 'getResultSetMetadata'); + sinon.spy(context.driver, 'fetchResults'); + const operation = new DBSQLOperation({ handle, context }); const results = await operation.fetchChunk(); expect(results).to.deep.equal([]); - expect(driver.getResultSetMetadata.called).to.be.false; - expect(driver.fetchResults.called).to.be.false; + expect(context.driver.getResultSetMetadata.called).to.be.false; + expect(context.driver.fetchResults.called).to.be.false; }); it('should wait for operation to complete', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onSecondCall() .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - driver.getResultSetMetadataResp.schema = { columns: [] }; - driver.fetchResultsResp.hasMoreRows = false; - driver.fetchResultsResp.results.columns = []; + context.driver.getResultSetMetadataResp.schema = { columns: [] }; + context.driver.fetchResultsResp.hasMoreRows = false; + context.driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const results = await operation.fetchChunk(); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(results).to.deep.equal([]); expect(operation.state).to.equal(TOperationState.FINISHED_STATE); }); it('should request progress', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onSecondCall() .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - driver.getResultSetMetadataResp.schema = { columns: [] }; - driver.fetchResultsResp.hasMoreRows = false; - driver.fetchResultsResp.results.columns = []; + context.driver.getResultSetMetadataResp.schema = { columns: [] }; + context.driver.fetchResultsResp.hasMoreRows = false; + context.driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); await operation.fetchChunk({ progress: true }); - expect(driver.getOperationStatus.called).to.be.true; - const request = driver.getOperationStatus.getCall(0).args[0]; + expect(context.driver.getOperationStatus.called).to.be.true; + const request = context.driver.getOperationStatus.getCall(0).args[0]; expect(request.getProgressUpdate).to.be.true; }); it('should invoke progress callback', async () => { const attemptsUntilFinished = 3; - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; + + context.driver.getOperationStatusResp.operationState = TOperationState.INITIALIZED_STATE; sinon - .stub(driver, 'getOperationStatus') + .stub(context.driver, 'getOperationStatus') .callThrough() .onCall(attemptsUntilFinished - 1) // count is zero-based .callsFake((...args) => { - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - return driver.getOperationStatus.wrappedMethod.apply(driver, args); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + return context.driver.getOperationStatus.wrappedMethod.apply(context.driver, args); }); - driver.getResultSetMetadataResp.schema = { columns: [] }; - driver.fetchResultsResp.hasMoreRows = false; - driver.fetchResultsResp.results.columns = []; + context.driver.getResultSetMetadataResp.schema = { columns: [] }; + context.driver.fetchResultsResp.hasMoreRows = false; + context.driver.fetchResultsResp.results.columns = []; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const callback = sinon.stub(); await operation.fetchChunk({ callback }); - expect(driver.getOperationStatus.called).to.be.true; + expect(context.driver.getOperationStatus.called).to.be.true; expect(callback.callCount).to.be.equal(attemptsUntilFinished); }); it('should fetch schema and data and return array of records', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - sinon.spy(driver, 'getResultSetMetadata'); - sinon.spy(driver, 'fetchResults'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + sinon.spy(context.driver, 'getResultSetMetadata'); + sinon.spy(context.driver, 'fetchResults'); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const results = await operation.fetchChunk(); expect(results).to.deep.equal([{ test: 1 }, { test: 2 }, { test: 3 }]); - expect(driver.getResultSetMetadata.called).to.be.true; - expect(driver.fetchResults.called).to.be.true; + expect(context.driver.getResultSetMetadata.called).to.be.true; + expect(context.driver.fetchResults.called).to.be.true; }); it('should return data from directResults (all the data in directResults)', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - sinon.spy(driver, 'getResultSetMetadata'); - sinon.spy(driver, 'fetchResults'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + sinon.spy(context.driver, 'getResultSetMetadata'); + sinon.spy(context.driver, 'fetchResults'); const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -1075,24 +1063,22 @@ describe('DBSQLOperation', () => { const results = await operation.fetchChunk(); expect(results).to.deep.equal([{ test: 5 }, { test: 6 }]); - expect(driver.getResultSetMetadata.called).to.be.true; - expect(driver.fetchResults.called).to.be.false; + expect(context.driver.getResultSetMetadata.called).to.be.true; + expect(context.driver.fetchResults.called).to.be.false; }); it('should return data from directResults (first chunk in directResults, next chunk fetched)', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - sinon.spy(driver, 'getResultSetMetadata'); - sinon.spy(driver, 'fetchResults'); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + sinon.spy(context.driver, 'getResultSetMetadata'); + sinon.spy(context.driver, 'fetchResults'); const operation = new DBSQLOperation({ - driver, handle, context, directResults: { @@ -1115,29 +1101,28 @@ describe('DBSQLOperation', () => { const results1 = await operation.fetchChunk(); expect(results1).to.deep.equal([{ test: 5 }, { test: 6 }]); - expect(driver.getResultSetMetadata.callCount).to.be.eq(1); - expect(driver.fetchResults.callCount).to.be.eq(0); + expect(context.driver.getResultSetMetadata.callCount).to.be.eq(1); + expect(context.driver.fetchResults.callCount).to.be.eq(0); const results2 = await operation.fetchChunk(); expect(results2).to.deep.equal([{ test: 1 }, { test: 2 }, { test: 3 }]); - expect(driver.getResultSetMetadata.callCount).to.be.eq(1); - expect(driver.fetchResults.callCount).to.be.eq(1); + expect(context.driver.getResultSetMetadata.callCount).to.be.eq(1); + expect(context.driver.fetchResults.callCount).to.be.eq(1); }); it('should fail on unsupported result format', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ROW_BASED_SET; - driver.getResultSetMetadataResp.schema = { columns: [] }; + context.driver.getResultSetMetadataResp.resultFormat = TSparkRowSetType.ROW_BASED_SET; + context.driver.getResultSetMetadataResp.schema = { columns: [] }; - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); try { await operation.fetchChunk(); @@ -1153,10 +1138,9 @@ describe('DBSQLOperation', () => { describe('fetchAll', () => { it('should fetch data while available and return it all', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); - const driver = new DriverMock(); - const operation = new DBSQLOperation({ driver, handle, context }); + const operation = new DBSQLOperation({ handle, context }); const originalData = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]; @@ -1185,16 +1169,15 @@ describe('DBSQLOperation', () => { describe('hasMoreRows', () => { it('should return False until first chunk of data fetched', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = true; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1202,16 +1185,15 @@ describe('DBSQLOperation', () => { }); it('should return False if operation was closed', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = true; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1221,16 +1203,15 @@ describe('DBSQLOperation', () => { }); it('should return False if operation was cancelled', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = true; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1240,16 +1221,15 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag was set in response', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = true; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = true; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1257,16 +1237,15 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag is False but there is actual data', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = false; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = false; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1274,16 +1253,15 @@ describe('DBSQLOperation', () => { }); it('should return True if hasMoreRows flag is unset but there is actual data', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = undefined; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = undefined; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); @@ -1291,17 +1269,16 @@ describe('DBSQLOperation', () => { }); it('should return False if hasMoreRows flag is False and there is no data', async () => { - const context = { getLogger: () => logger }; + const context = new ClientContextMock(); const handle = new OperationHandleMock(); handle.hasResultSet = true; - const driver = new DriverMock(); - driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; - driver.getOperationStatusResp.hasResultSet = true; - driver.fetchResultsResp.hasMoreRows = false; - driver.fetchResultsResp.results = undefined; - const operation = new DBSQLOperation({ driver, handle, context }); + context.driver.getOperationStatusResp.operationState = TOperationState.FINISHED_STATE; + context.driver.getOperationStatusResp.hasResultSet = true; + context.driver.fetchResultsResp.hasMoreRows = false; + context.driver.fetchResultsResp.results = undefined; + const operation = new DBSQLOperation({ handle, context }); expect(await operation.hasMoreRows()).to.be.false; await operation.fetchChunk(); diff --git a/tests/unit/DBSQLSession.test.js b/tests/unit/DBSQLSession.test.js index 2c76caae..9ff5172b 100644 --- a/tests/unit/DBSQLSession.test.js +++ b/tests/unit/DBSQLSession.test.js @@ -5,6 +5,7 @@ const DBSQLSession = require('../../dist/DBSQLSession').default; const InfoValue = require('../../dist/dto/InfoValue').default; const Status = require('../../dist/dto/Status').default; const DBSQLOperation = require('../../dist/DBSQLOperation').default; +const HiveDriver = require('../../dist/hive/HiveDriver').default; // Create logger that won't emit // @@ -13,10 +14,12 @@ const logger = new DBSQLLogger({ level: LogLevel.error }); function createDriverMock(customMethodHandler) { customMethodHandler = customMethodHandler || ((methodName, value) => value); - return new Proxy( - {}, - { - get: function (target, prop) { + const driver = new HiveDriver({}); + + return new Proxy(driver, { + get: function (target, prop) { + // Mock only methods of driver + if (typeof target[prop] === 'function') { return () => Promise.resolve( customMethodHandler(prop, { @@ -27,18 +30,18 @@ function createDriverMock(customMethodHandler) { infoValue: {}, }), ); - }, + } }, - ); + }); } function createSession(customMethodHandler) { const driver = createDriverMock(customMethodHandler); return new DBSQLSession({ - driver, handle: { sessionId: 'id' }, context: { getLogger: () => logger, + getDriver: async () => driver, }, }); } From 64d21931578c0432de17cbd612f1438b282c8ba3 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 2 Oct 2023 16:01:33 +0300 Subject: [PATCH 4/4] Fix tests Signed-off-by: Levko Kravets --- tests/e2e/batched_fetch.test.js | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/e2e/batched_fetch.test.js b/tests/e2e/batched_fetch.test.js index 272b9492..5218088e 100644 --- a/tests/e2e/batched_fetch.test.js +++ b/tests/e2e/batched_fetch.test.js @@ -37,14 +37,14 @@ describe('Data fetching', () => { it('fetch chunks should return a max row set of chunkSize', async () => { const session = await openSession(); - sinon.spy(session.driver, 'fetchResults'); + sinon.spy(session.context.driver, 'fetchResults'); try { // set `maxRows` to null to disable direct results so all the data are fetched through `driver.fetchResults` const operation = await session.executeStatement(query, { maxRows: null }); let chunkedOp = await operation.fetchChunk({ maxRows: 10 }).catch((error) => logger(error)); expect(chunkedOp.length).to.be.equal(10); // we explicitly requested only one chunk - expect(session.driver.fetchResults.callCount).to.equal(1); + expect(session.context.driver.fetchResults.callCount).to.equal(1); } finally { await session.close(); } @@ -52,14 +52,14 @@ describe('Data fetching', () => { it('fetch all should fetch all records', async () => { const session = await openSession(); - sinon.spy(session.driver, 'fetchResults'); + sinon.spy(session.context.driver, 'fetchResults'); try { // set `maxRows` to null to disable direct results so all the data are fetched through `driver.fetchResults` const operation = await session.executeStatement(query, { maxRows: null }); let all = await operation.fetchAll({ maxRows: 200 }); expect(all.length).to.be.equal(1000); // 1000/200 = 5 chunks + one extra request to ensure that there's no more data - expect(session.driver.fetchResults.callCount).to.equal(6); + expect(session.context.driver.fetchResults.callCount).to.equal(6); } finally { await session.close(); } @@ -67,14 +67,14 @@ describe('Data fetching', () => { it('should fetch all records if they fit within directResults response', async () => { const session = await openSession(); - sinon.spy(session.driver, 'fetchResults'); + sinon.spy(session.context.driver, 'fetchResults'); try { // here `maxRows` enables direct results with limit of the first batch const operation = await session.executeStatement(query, { maxRows: 1000 }); let all = await operation.fetchAll(); expect(all.length).to.be.equal(1000); // all the data returned immediately from direct results, so no additional requests - expect(session.driver.fetchResults.callCount).to.equal(0); + expect(session.context.driver.fetchResults.callCount).to.equal(0); } finally { await session.close(); } @@ -82,7 +82,7 @@ describe('Data fetching', () => { it('should fetch all records if only part of them fit within directResults response', async () => { const session = await openSession(); - sinon.spy(session.driver, 'fetchResults'); + sinon.spy(session.context.driver, 'fetchResults'); try { // here `maxRows` enables direct results with limit of the first batch const operation = await session.executeStatement(query, { maxRows: 200 }); @@ -91,7 +91,7 @@ describe('Data fetching', () => { expect(all.length).to.be.equal(1000); // 1 chunk returned immediately from direct results + 4 remaining chunks + one extra chunk to ensure // that there's no more data - expect(session.driver.fetchResults.callCount).to.equal(5); + expect(session.context.driver.fetchResults.callCount).to.equal(5); } finally { await session.close(); }