import { CoreStart , PluginInitializerContext , CoreSetup , Plugin , Logger , OpenSearchDashboardsRequest , RouteMethod , RequestHandler , } from 'src/core/server' ; import { schema } from '@osd/config-schema' ; import { Subject } from 'rxjs' ; import { StreamingResponseHandler , BatchRequestData , BatchResponseItem , ErrorLike , removeLeadingSlash , normalizeError , } from '../common' ; import { StreamingRequestHandler } from './types' ; import { createNDJSONStream } from './streaming' ; export interface BfetchServerSetupDependencies { } export interface BfetchServerStartDependencies { } export interface BatchProcessingRouteParams < BatchItemData , BatchItemResult > { onBatchItem : ( data : BatchItemData ) => Promise < BatchItemResult > ; } public export interface BfetchServerSetup { addBatchProcessingRoute : < BatchItemData extends object , BatchItemResult extends object > ( path : string , handler : ( request : OpenSearchDashboardsRequest ) => BatchProcessingRouteParams < BatchItemData , BatchItemResult > ) => void ; addStreamingResponseRoute : < Payload , Response > ( path : string , params : ( request : OpenSearchDashboardsRequest ) => StreamingResponseHandler < Payload , Response > ) => void ; example param streamHandler createStreamingRequestHandler : < Response , P , Q , B , Method extends RouteMethod = any > ( streamHandler : StreamingRequestHandler < Response , P , Q , B , Method > ) => RequestHandler < P , Q , B , Method > ; } export interface BfetchServerStart { } const streamingHeaders = { 'Content-Type' : 'application/x-ndjson' , Connection : 'keep-alive' , 'Transfer-Encoding' : 'chunked' , } ; export class BfetchServerPlugin implements Plugin < BfetchServerSetup , BfetchServerStart , BfetchServerSetupDependencies , BfetchServerStartDependencies > { constructor ( private readonly initializerContext ) { } public setup ( core , plugins ) { const logger = this . initializerContext . logger . get ( ) ; const router = core . http . createRouter ( ) ; const addStreamingResponseRoute = this . addStreamingResponseRoute ( { router , logger } ) ; const addBatchProcessingRoute = this . addBatchProcessingRoute ( addStreamingResponseRoute ) ; const createStreamingRequestHandler = this . createStreamingRequestHandler ( { logger } ) ; return { addBatchProcessingRoute , addStreamingResponseRoute , createStreamingRequestHandler , } ; } public start ( core , plugins ) { return { } ; } public stop ( ) { } private addStreamingResponseRoute = ( { router , logger , } ) : BfetchServerSetup [ 'addStreamingResponseRoute' ] => ( path , handler ) => { router . post ( { path : `/${ removeLeadingSlash ( path ) }` , validate : { body : schema . any ( ) , } , } , async ( context , request , response ) => { const handlerInstance = handler ( request ) ; const data = request . body ; return response . ok ( { headers : streamingHeaders , body : createNDJSONStream ( handlerInstance . getResponseStream ( data ) , logger ) , } ) ; } ) ; } ; private createStreamingRequestHandler = ( { logger , } ) : BfetchServerSetup [ 'createStreamingRequestHandler' ] => ( streamHandler ) => async ( context , request , response ) => { const response$ = await streamHandler ( context , request ) ; return response . ok ( { headers : streamingHeaders , body : createNDJSONStream ( response$ , logger ) , } ) ; } ; private addBatchProcessingRoute = ( addStreamingResponseRoute ) : BfetchServerSetup [ 'addBatchProcessingRoute' ] => < BatchItemData extends object , BatchItemResult extends object , E extends ErrorLike = ErrorLike > ( path , handler ) => { addStreamingResponseRoute < BatchRequestData < BatchItemData > , BatchResponseItem < BatchItemResult , E > > ( path , ( request ) => { const handlerInstance = handler ( request ) ; return { getResponseStream : ( { batch } ) => { const subject = new Subject < BatchResponseItem < BatchItemResult , E > > ( ) ; let cnt = batch . length ; batch . forEach ( async ( batchItem , id ) => { try { const result = await handlerInstance . onBatchItem ( batchItem ) ; subject . next ( { id , result } ) ; } catch ( err ) { const error = normalizeError < E > ( err ) ; subject . next ( { id , error } ) ; } finally { cnt -- ; if ( ! cnt ) subject . complete ( ) ; } } ) ; return subject ; } , } ; } ) ; } ; }