import * as EventEmitter from 'events' ; import * as http2 from 'http2' ; export interface QueryStreamArgs { sql : string ; push ? : boolean ; properties ? : Record < string , any > ; } export interface CloseQueryArgs { queryId : { id : string ; } ; } export interface InsertsStreamArgs { target : string ; requiresAcks : boolean ; properties : Record < string , any > ; } export interface ErrorResponse { status : string ; errorCode : number ; message : string ; } export interface InsertAck { status : 'ok' ; } export interface QueryResponseMetadata { queryId : string ; columnNames : string [ ] ; columnTypes : string [ ] ; rowCount : number | null ; } export type QueryResponseRow = any [ ] ; type IteratorSignature = { [ nameKey ] : { [ key ] : any ; } ; } ; export const DELIMITED_CONTENT_TYPE = 'application/vnd.ksqlapi.delimited.v1' ; export abstract class QueryStream < T > { closed = false ; constructor ( protected session , protected queryStreamArgs , protected mapRow ) { this . session = session ; this . queryStreamArgs = queryStreamArgs ; this . mapRow = mapRow ; } headers ( ) { return { [ http2 . constants . HTTP2_HEADER_PATH ] : '/query-stream' , [ http2 . constants . HTTP2_HEADER_METHOD ] : 'POST' , [ http2 . constants . HTTP2_HEADER_CONTENT_TYPE ] : DELIMITED_CONTENT_TYPE , } ; } parseChunk ( buf ) { return buf . toString ( ) . split ( '\n' ) . filter ( str => str ) ; } initQueryStream ( ) { const stream = this . session . request ( this . headers ( ) ) ; const reqPayload = Buffer . from ( JSON . stringify ( this . queryStreamArgs ) ) ; stream . end ( reqPayload ) ; return stream ; } } interface QueryStreamEventEmitter extends EventEmitter { emit ( event ) : boolean ; emit ( event , err ) : boolean ; emit ( event , buf ) : boolean ; on ( event , listener ) : this ; on ( event , listener ) : this ; on ( event , listener ) : this ; } export class EventEmitterQueryStream < T > extends QueryStream < T > { queryEvents ( ) { return null ; } } export class AsyncIteratorQueryStream < T > extends QueryStream < T > { [ Symbol . asyncIterator ] ( ) { const promised = [ ] ; const received = [ ] ; const stream = this . initQueryStream ( ) ; const destroyStream = ( err ? ) : void => { for ( const [ resolve , reject ] of promised ) { if ( err ) { reject ( err ) ; } else { const [ error ] = received ; try { const parsedError = error && JSON . parse ( error ) ; if ( parsedError ?. status === 'error' ) { reject ( new Error ( parsedError . message ) ) ; } else { return resolve ( { value : received , done : true } ) ; } } catch ( e ) { reject ( new Error ( e . message ) ) ; } } } stream . destroy ( ) ; this . closed = true ; } ; stream . on ( 'error' , error => destroyStream ( error ) ) . on ( 'close' , ( ) => destroyStream ( ) ) . on ( 'abort' , ( ) => destroyStream ( new Error ( 'abort' ) ) ) . on ( 'timeout' , ( ) => destroyStream ( new Error ( 'timeout' ) ) ) ; const doWork = ( ) : void => { if ( this . closed ) { return ; } while ( promised . length && received . length ) { const [ resolve ] = promised . shift ( ) ; const rawJson = received . shift ( ) ; resolve ( rawJson ) ; } if ( promised . length ) { const next = stream . read ( ) ; if ( next != null ) { const nextLines = this . parseChunk ( next ) ; received . push ( ... nextLines ) ; } setImmediate ( ( ) => doWork ( ) ) ; } } ; const nextPromise = ( ) : Promise < string > => new Promise ( ( resolve , reject ) => promised . push ( [ resolve , reject ] ) ) ; const getRowKeys = nextPromise ( ) . then ( ( rawMeta ) => { try { let meta = { } ; if ( typeof rawMeta !== 'string' && rawMeta . value ) { meta = JSON . parse ( rawMeta . value [ 0 ] ) ; } else if ( typeof rawMeta === 'string' ) { meta = JSON . parse ( rawMeta ) ; } if ( meta . status === 'error' ) { const err = meta as Error ; destroyStream ( err ) ; throw err ; } return meta . columnNames ; } catch ( e ) { destroyStream ( new Error ( e . message ) ) ; } } ) ; doWork ( ) ; return { next : ( ) : any => { if ( this . closed ) { return Promise . resolve ( { value : undefined , done : true } ) ; } return getRowKeys . then ( ks => { const enqueued = nextPromise ( ) . then ( rawRow => { const value = this . mapRow ( rawRow , ks ) ; return { value , done : false } ; } ) ; doWork ( ) ; return enqueued ; } ) ; } , return : ( ) : Promise < any > => { destroyStream ( ) ; return Promise . resolve ( { value : undefined , done : true } ) ; } , } ; } } export const asyncIteratorQueryStream = ( session , queryStreamArgs , nameKey ) : AsyncIteratorQueryStream < IteratorSignature > => { const mapRow = ( rawRow , ks ) : IteratorSignature => { let row = [ ] ; try { row = JSON . parse ( rawRow ) ; } catch ( e ) { row = [ rawRow ] ; } return ks . reduce ( ( acc , k , i ) => { acc [ nameKey ] [ k ] = row [ i ] ; return acc ; } , { [ nameKey ] : { } } ) as Record < string , IteratorSignature > ; } ; return new AsyncIteratorQueryStream ( session , queryStreamArgs , mapRow ) ; } ;