import { Athena } from 'aws-sdk' ; import * as csv from 'csv-parser' ; import { Transform } from 'stream' ; import { setTimeout } from 'timers' ; import { AthenaRequest , AthenaRequestConfig } from './request' ; import * as util from './util' ; export interface AthenaExecutionResult < T > { records : T [ ] ; queryExecution : Athena . QueryExecution ; } export interface AthenaExecutionSelect < T > { toPromise : ( ) => Promise < AthenaExecutionResult < T > > ; toStream : ( ) => Transform ; } export interface AthenaClientConfig extends AthenaRequestConfig { pollingInterval ? : number ; queryTimeout ? : number ; concurrentExecMax ? : number ; execRightCheckInterval ? : number ; skipFetchResult ? : boolean ; } const defaultPollingInterval = 1000 ; const defaultQueryTimeout = 0 ; const defaultExecRightCheckInterval = 100 ; let concurrentExecMax = 5 ; let concurrentExecNum = 0 ; export function setConcurrentExecMax ( val ) { concurrentExecMax = val ; } export class AthenaClient { private config : AthenaClientConfig ; private concurrentExecNum : number ; private request : AthenaRequest ; constructor ( request , config ) { this . request = request ; this . config = config ; if ( config . concurrentExecMax ) { console . warn ( `[WARN] please use 'athena.setConcurrentExecMax()' instead 'clientConfig.concurrentExecMax'` ) ; concurrentExecMax = config . concurrentExecMax ; } if ( ( config . encryptionOption === 'SSE_KMS' || config . encryptionOption === 'CSE_KMS' ) && config . encryptionKmsKey === undefined ) { throw new Error ( 'KMS key required' ) ; } } public execute < T > ( query ) ; public execute < T > ( query , callback ) ; public execute < T > ( query , callback ? ) { const currentConfig = { ... this . config } ; const csvTransform = new csv ( ) ; this . _execute ( query , csvTransform , currentConfig ) ; if ( callback !== undefined ) { let isEnd = false ; const records = [ ] ; let queryExecution ; csvTransform . on ( 'data' , ( record ) => { records . push ( record ) ; } ) ; csvTransform . on ( 'query_end' , ( q ) => { queryExecution = q ; } ) ; csvTransform . on ( 'end' , ( record ) => { if ( isEnd ) { return ; } const result = { records , queryExecution , } ; callback ( undefined , result ) ; } ) ; csvTransform . on ( 'error' , ( err ) => { isEnd = true ; callback ( err ) ; } ) ; return ; } else { return { toPromise : ( ) => { return new Promise < AthenaExecutionResult < T > > ( ( resolve , reject ) => { const records = [ ] ; let queryExecution ; csvTransform . on ( 'data' , ( record ) => { records . push ( record ) ; } ) ; csvTransform . on ( 'query_end' , ( q ) => { queryExecution = q ; } ) ; csvTransform . on ( 'end' , ( record ) => { const result = { records , queryExecution , } ; return resolve ( result ) ; } ) ; csvTransform . on ( 'error' , ( err ) => { return reject ( err ) ; } ) ; } ) ; } , toStream : ( ) : Transform => { return csvTransform ; } , } ; } } private async _execute ( query , csvTransform , config ) { while ( ! this . canStartQuery ( ) ) { await util . sleep ( config . execRightCheckInterval || defaultExecRightCheckInterval ) ; } let queryExecution ; try { this . startQuery ( ) ; const queryId = await this . request . startQuery ( query , config ) ; let isTimeout = false ; if ( ( config . queryTimeout || defaultQueryTimeout ) !== 0 ) { setTimeout ( ( ) => { isTimeout = true ; } , config . queryTimeout || defaultQueryTimeout ) ; } while ( ! isTimeout && ! ( await this . request . checkQuery ( queryId , config ) ) ) { await util . sleep ( config . pollingInterval || defaultPollingInterval ) ; } if ( isTimeout ) { await this . request . stopQuery ( queryId , config ) ; throw new Error ( 'query timeout' ) ; } queryExecution = await this . request . getQueryExecution ( queryId , config ) ; csvTransform . emit ( 'query_end' , queryExecution ) ; this . endQuery ( ) ; } catch ( err ) { this . endQuery ( ) ; csvTransform . emit ( 'error' , err ) ; return ; } if ( config . skipFetchResult ) { csvTransform . end ( ) ; return ; } else { try { if ( ! queryExecution . ResultConfiguration || ! queryExecution . ResultConfiguration . OutputLocation ) { throw new Error ( 'query outputlocation is empty' ) ; } const resultsStream = this . request . getResultsStream ( queryExecution . ResultConfiguration . OutputLocation ) ; resultsStream . pipe ( csvTransform ) ; } catch ( err ) { csvTransform . emit ( 'error' , err ) ; return ; } } } private canStartQuery ( ) { return concurrentExecNum < concurrentExecMax ; } private startQuery ( ) { concurrentExecNum = Math . min ( ++ concurrentExecNum , concurrentExecMax ) ; } private endQuery ( ) { concurrentExecNum = Math . max ( -- concurrentExecNum , 0 ) ; } }