import { Subject } from 'rxjs' ; import { omit , defaults } from 'lodash' ; import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey' ; import { SavedObject , SavedObjectsSerializer , SavedObjectsRawDoc , ISavedObjectsRepository , SavedObjectsUpdateResponse , ElasticsearchClient , } from '../../../../src/core/server' ; import { asOk , asErr , Result } from './lib/result_type' ; import { ConcreteTaskInstance , TaskInstance , TaskLifecycle , TaskLifecycleResult , SerializedConcreteTaskInstance , } from './task' ; import { TaskTypeDictionary } from './task_type_dictionary' ; export interface StoreOpts { esClient : ElasticsearchClient ; index : string ; taskManagerId : string ; definitions : TaskTypeDictionary ; savedObjectsRepository : ISavedObjectsRepository ; serializer : SavedObjectsSerializer ; } export interface SearchOpts { search_after ? : Array < number | string > ; size ? : number ; sort ? : estypes . Sort ; query ? : estypes . QueryDslQueryContainer ; seq_no_primary_term ? : boolean ; } export interface AggregationOpts { aggs : Record < string , estypes . AggregationsAggregationContainer > ; query ? : estypes . QueryDslQueryContainer ; size ? : number ; } export interface UpdateByQuerySearchOpts extends SearchOpts { script ? : estypes . Script ; } export interface UpdateByQueryOpts extends SearchOpts { max_docs ? : number ; } export interface FetchResult { docs : ConcreteTaskInstance [ ] ; } export type BulkUpdateResult = Result < ConcreteTaskInstance , { entity : ConcreteTaskInstance ; error : Error ; } > ; export interface UpdateByQueryResult { updated : number ; version_conflicts : number ; total : number ; } /**
 * Wraps an elasticsearch connection and provides a task manager-specific
 * interface into the index.
 */ export class TaskStore { public readonly index : string ; public readonly taskManagerId : string ; public readonly errors$ = new Subject < Error > ( ) ; private esClient : ElasticsearchClient ; private definitions : TaskTypeDictionary ; private savedObjectsRepository : ISavedObjectsRepository ; private serializer : SavedObjectsSerializer ; param { StoreOpts } opts prop prop prop prop prop constructor ( opts ) { this . esClient = opts . esClient ; this . index = opts . index ; this . taskManagerId = opts . taskManagerId ; this . definitions = opts . definitions ; this . serializer = opts . serializer ; this . savedObjectsRepository = opts . savedObjectsRepository ; } param tasks public convertToSavedObjectIds ( taskIds ) { return taskIds . map ( ( id ) => this . serializer . generateRawId ( undefined , 'task' , id ) ) ; } param task public async schedule ( taskInstance ) { this . definitions . ensureHas ( taskInstance . taskType ) ; let savedObject ; try { savedObject = await this . savedObjectsRepository . create < SerializedConcreteTaskInstance > ( 'task' , taskInstanceToAttributes ( taskInstance ) , { id : taskInstance . id , refresh : false } ) ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } return savedObjectToConcreteTaskInstance ( savedObject ) ; } param opts public async fetch ( { sort = [ { 'task.runAt' : 'asc' } ] , ... opts } = { } ) { return this . search ( { ... opts , sort , } ) ; } param { TaskDoc } doc returns { Promise < TaskDoc > } public async update ( doc ) { const attributes = taskInstanceToAttributes ( doc ) ; let updatedSavedObject ; try { updatedSavedObject = await this . savedObjectsRepository . update < SerializedConcreteTaskInstance > ( 'task' , doc . id , attributes , { refresh : false , version : doc . version , } ) ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } return savedObjectToConcreteTaskInstance ( { ... updatedSavedObject , attributes : defaults ( updatedSavedObject . attributes , attributes ) } ) ; } param { Array < TaskDoc > } docs returns { Promise < Array < TaskDoc > > } public async bulkUpdate ( docs ) { const attributesByDocId = docs . reduce ( ( attrsById , doc ) => { attrsById . set ( doc . id , taskInstanceToAttributes ( doc ) ) ; return attrsById ; } , new Map ( ) ) ; let updatedSavedObjects ; try { ( { saved_objects : updatedSavedObjects } = await this . savedObjectsRepository . bulkUpdate < SerializedConcreteTaskInstance > ( docs . map ( ( doc ) => ( { type : 'task' , id : doc . id , options : { version : doc . version } , attributes : attributesByDocId . get ( doc . id ) ! , } ) ) , { refresh : false , } ) ) ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } return updatedSavedObjects . map < BulkUpdateResult > ( ( updatedSavedObject , index ) => isSavedObjectsUpdateResponse ( updatedSavedObject ) ? asOk ( savedObjectToConcreteTaskInstance ( { ... updatedSavedObject , attributes : defaults ( updatedSavedObject . attributes , attributesByDocId . get ( updatedSavedObject . id ) ! ) , } ) ) : asErr ( { entity : docs [ index ] , error : updatedSavedObject , } ) ) ; } param { string } id returns { Promise < void > } public async remove ( id ) { try { await this . savedObjectsRepository . delete ( 'task' , id ) ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } } param { string } id returns { Promise < void > } public async get ( id ) { let result ; try { result = await this . savedObjectsRepository . get < SerializedConcreteTaskInstance > ( 'task' , id ) ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } return savedObjectToConcreteTaskInstance ( result ) ; } param { string } id returns { Promise < void > } public async getLifecycle ( id ) { try { const task = await this . get ( id ) ; return task . status ; } catch ( err ) { if ( err . output && err . output . statusCode === 404 ) { return TaskLifecycleResult . NotFound ; } throw err ; } } private async search ( opts = { } ) { const { query } = ensureQueryOnlyReturnsTaskObjects ( opts ) ; try { const { body : { hits : { hits : tasks } , } , } = await this . esClient . search < SavedObjectsRawDoc [ '_source' ] > ( { index : this . index , ignore_unavailable : true , body : { ... opts , query , } , } ) ; return { docs : tasks . filter ( ( doc ) => this . serializer . isRawSavedObject ( doc ) ) . map ( ( doc ) => this . serializer . rawToSavedObject ( doc ) ) . map ( ( doc ) => omit ( doc , 'namespace' ) as SavedObject < SerializedConcreteTaskInstance > ) . map ( savedObjectToConcreteTaskInstance ) , } ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } } public async aggregate < TSearchRequest extends AggregationOpts > ( { aggs , query , size = 0 , } ) { const { body } = await this . esClient . search < ConcreteTaskInstance , Record < string , estypes . AggregationsAggregate > > ( { index : this . index , ignore_unavailable : true , track_total_hits : true , body : ensureAggregationOnlyReturnsTaskObjects ( { query , aggs , size , } ) , } ) ; return body ; } public async updateByQuery ( opts = { } , { max_docs : max_docs } = { } ) { const { query } = ensureQueryOnlyReturnsTaskObjects ( opts ) ; try { const { body : { total , updated , version_conflicts } , } = await this . esClient . updateByQuery ( { index : this . index , ignore_unavailable : true , refresh : true , conflicts : 'proceed' , body : { ... opts , max_docs , query , } , } ) ; const conflictsCorrectedForContinuation = correctVersionConflictsForContinuation ( updated , version_conflicts , max_docs ) ; return { total : total || 0 , updated : updated || 0 , version_conflicts : conflictsCorrectedForContinuation , } ; } catch ( e ) { this . errors$ . next ( e ) ; throw e ; } } } /**
 * When we run updateByQuery with conflicts='proceed', it's possible for the `version_conflicts`
 * to count against the specified `max_docs`, as per https://github.com/elastic/elasticsearch/issues/63671
 * In order to correct for that happening, we only count `version_conflicts` if we haven't updated as
 * many docs as we could have.
 * This is still no more than an estimation, as there might have been less docuemnt to update that the
 * `max_docs`, but we bias in favour of over zealous `version_conflicts` as that's the best indicator we
 * have for an unhealthy cluster distribution of Task Manager polling intervals
 */ export function correctVersionConflictsForContinuation ( updated , versionConflicts , maxDocs ? ) { return maxDocs && versionConflicts + updated > maxDocs ? maxDocs - updated : versionConflicts ; } function taskInstanceToAttributes ( doc ) { return { ... omit ( doc , 'id' , 'version' ) , params : JSON . stringify ( doc . params || { } ) , state : JSON . stringify ( doc . state || { } ) , attempts : ( doc as ConcreteTaskInstance ) . attempts || 0 , scheduledAt : ( doc . scheduledAt || new Date ( ) ) . toISOString ( ) , startedAt : ( doc . startedAt && doc . startedAt . toISOString ( ) ) || null , retryAt : ( doc . retryAt && doc . retryAt . toISOString ( ) ) || null , runAt : ( doc . runAt || new Date ( ) ) . toISOString ( ) , status : ( doc as ConcreteTaskInstance ) . status || 'idle' , } as SerializedConcreteTaskInstance ; } export function savedObjectToConcreteTaskInstance ( savedObject ) { return { ... savedObject . attributes , id : savedObject . id , version : savedObject . version , scheduledAt : new Date ( savedObject . attributes . scheduledAt ) , runAt : new Date ( savedObject . attributes . runAt ) , startedAt : savedObject . attributes . startedAt ? new Date ( savedObject . attributes . startedAt ) : null , retryAt : savedObject . attributes . retryAt ? new Date ( savedObject . attributes . retryAt ) : null , state : parseJSONField ( savedObject . attributes . state , 'state' , savedObject . id ) , params : parseJSONField ( savedObject . attributes . params , 'params' , savedObject . id ) , } ; } function parseJSONField ( json , fieldName , id ) { try { return json ? JSON . parse ( json ) : { } ; } catch ( error ) { throw new Error ( `Task "${ id }"'s ${ fieldName } field has invalid JSON: ${ json }` ) ; } } function ensureQueryOnlyReturnsTaskObjects ( opts ) { const originalQuery = opts . query ; const queryOnlyTasks = { term : { type : 'task' } } ; const query = originalQuery ? { bool : { must : [ queryOnlyTasks , originalQuery ] } } : queryOnlyTasks ; return { ... opts , query , } ; } function ensureAggregationOnlyReturnsTaskObjects ( opts ) { const originalQuery = opts . query ; const filterToOnlyTasks = { bool : { filter : [ { term : { type : 'task' } } ] , } , } ; const query = originalQuery ? { bool : { must : [ filterToOnlyTasks , originalQuery ] } } : filterToOnlyTasks ; return { ... opts , query , } ; } function isSavedObjectsUpdateResponse ( result ) { return result && typeof ( result as SavedObjectsUpdateResponse ) . id === 'string' ; }