import { WriteBatch } from '../batch/write-batch' ; import { Connection } from '../connection/connection' ; import { DocumentClientBatchTransformer } from '../transformer/document-client-batch-transformer' ; import pLimit from 'p-limit' ; import { BATCH_READ_MAX_ALLOWED_ATTEMPTS , BATCH_WRITE_CONCURRENCY_LIMIT , BATCH_WRITE_MAX_ALLOWED_ATTEMPTS , INTERNAL_ENTITY_ATTRIBUTE , MANAGER_NAME , STATS_TYPE , } from '@typedorm/common' ; import { BatchGetResponseMap , BatchWriteItemRequestMap , DocumentClient , } from 'aws-sdk/clients/dynamodb' ; import { isEmptyObject } from '../../helpers/is-empty-object' ; import { ReadBatch } from '../batch/read-batch' ; import { MetadataOptions } from '../transformer/base-transformer' ; import { getUniqueRequestId } from '../../helpers/get-unique-request-id' ; export enum REQUEST_TYPE { TRANSACT_WRITE = 'TRANSACT_WRITE' , BATCH_WRITE = 'BATCH_WRITE' , BATCH_READ = 'BATCH_READ' } /**
 * Batch manager write options
 */ export type BatchManagerWriteOptions = BatchManageBaseOptions ; /**
 * Batch manager read options
 */ export type BatchManagerReadOptions = BatchManageBaseOptions ; export interface BatchManageBaseOptions { default maxRetryAttempts ? : number ; default requestsConcurrencyLimit ? : number ; default backoffMultiplicationFactor ? : number ; } export class BatchManager { private _dcBatchTransformer : DocumentClientBatchTransformer ; private _errorQueue : { requestInput : any ; error : Error ; requestType : REQUEST_TYPE ; } [ ] ; private limit = pLimit ( BATCH_WRITE_CONCURRENCY_LIMIT ) ; constructor ( private connection ) { this . _dcBatchTransformer = new DocumentClientBatchTransformer ( connection ) ; this . resetErrorQueue ( ) ; } param batch async write ( batch , options ? , metadataOptions ? ) { this . resetErrorQueue ( ) ; const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; if ( options ?. requestsConcurrencyLimit ) { this . limit = pLimit ( options ?. requestsConcurrencyLimit ) ; } this . connection . logger . logInfo ( { requestId , scope : MANAGER_NAME . BATCH_MANAGER , log : `Running a batch write request for ${ batch . items . length } items` , } ) ; const { batchWriteRequestMapItems , lazyTransactionWriteItemListLoaderItems , transactionListItems , metadata , } = this . _dcBatchTransformer . toDynamoWriteBatchItems ( batch , { requestId , } ) ; const transactionRequests = transactionListItems . map ( ( { rawInput , transformedInput } ) => { return this . toLimited ( ( ) => this . connection . transactionManger . writeRaw ( transformedInput , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) , rawInput , REQUEST_TYPE . TRANSACT_WRITE ) ; } ) ; const lazyTransactionRequests = lazyTransactionWriteItemListLoaderItems . map ( ( { rawInput , transformedInput } ) => { return this . toLimited ( async ( ) => { const existingItem = await this . connection . entityManager . findOne ( transformedInput . entityClass , transformedInput . primaryKeyAttributes , undefined , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; const deleteTransactionItemList = transformedInput . lazyLoadTransactionWriteItems ( existingItem ) ; return this . connection . transactionManger . writeRaw ( deleteTransactionItemList , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; } , rawInput , REQUEST_TYPE . TRANSACT_WRITE ) ; } ) ; const batchRequests = batchWriteRequestMapItems . map ( batchRequestMap => { return this . toLimited ( async ( ) => this . connection . documentClient . batchWrite ( { RequestItems : { ... batchRequestMap } , ReturnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) . promise ( ) , batchRequestMap , REQUEST_TYPE . BATCH_WRITE ) ; } ) ; const allRequests = [ ... transactionRequests , ... lazyTransactionRequests , ... batchRequests , ] as Promise < DocumentClient . TransactWriteItemsOutput > [ ] | Promise < DocumentClient . BatchWriteItemOutput > [ ] ; const responses = await Promise . all ( allRequests ) ; responses . forEach ( ( response , index ) => { if ( response . ConsumedCapacity ) { this . connection . logger . logStats ( { scope : MANAGER_NAME . BATCH_MANAGER , requestId , requestSegment : index , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : response . ConsumedCapacity , } ) ; } } ) ; const unprocessedItems = await this . recursiveHandleBatchWriteItemsResponse ( responses , 0 , options , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; const transformedUnprocessedItems = unprocessedItems . flatMap ( ( unprocessedItemInput ) => this . _dcBatchTransformer . toWriteBatchInputList ( unprocessedItemInput , metadata ) ) ; const failedItemsOriginalInput = this . _errorQueue . flatMap ( item => { if ( item . requestType === REQUEST_TYPE . BATCH_WRITE ) { return this . _dcBatchTransformer . toWriteBatchInputList ( item . requestInput , metadata ) ; } else if ( item . requestType === REQUEST_TYPE . TRANSACT_WRITE ) { return item . requestInput ; } else { throw new Error ( 'Unsupported request type, if this continues please file an issue on github' ) ; } } ) ; return { unprocessedItems : transformedUnprocessedItems , failedItems : failedItemsOriginalInput , } ; } /**
     * Reads all items given in batch with default eventually consistent read type
     * _Note_: Returned items are not guaranteed to be in the same sequence as requested
     */ async read ( batch , options ? , metadataOptions ? ) { this . resetErrorQueue ( ) ; const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; if ( options ?. requestsConcurrencyLimit ) { this . limit = pLimit ( options ?. requestsConcurrencyLimit ) ; } this . connection . logger . logInfo ( { requestId , scope : MANAGER_NAME . BATCH_MANAGER , log : `Running a batch read request for ${ batch . items . length } items` , } ) ; const { batchRequestItemsList , metadata , } = this . _dcBatchTransformer . toDynamoReadBatchItems ( batch , { requestId , } ) ; const batchRequests = batchRequestItemsList . map ( batchRequestItems => { return this . toLimited ( async ( ) => this . connection . documentClient . batchGet ( { RequestItems : { ... batchRequestItems } , ReturnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) . promise ( ) , batchRequestItems , REQUEST_TYPE . BATCH_READ ) ; } ) ; const initialResponses = await Promise . all ( batchRequests ) ; initialResponses . forEach ( ( response , index ) => { if ( response . ConsumedCapacity ) { this . connection . logger . logStats ( { scope : MANAGER_NAME . BATCH_MANAGER , requestId , requestSegment : index , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : response . ConsumedCapacity , } ) ; } } ) ; const { items , unprocessedItemsList , } = await this . recursiveHandleBatchReadItemsResponse ( initialResponses , 0 , options , [ ] , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; const transformedItems = items . map ( item => { const entityPhysicalName = item [ INTERNAL_ENTITY_ATTRIBUTE . ENTITY_NAME ] ; if ( ! entityPhysicalName ) { this . connection . logger . logWarn ( { requestId , scope : MANAGER_NAME . ENTITY_MANAGER , log : `Item ${ JSON . stringify ( item ) } is not known to TypeDORM there for transform was not run` , } ) ; return item ; } const { target } = this . connection . getEntityByPhysicalName ( entityPhysicalName ) ; return this . _dcBatchTransformer . fromDynamoEntity ( target , item , { requestId , } ) ; } ) as unknown [ ] ; const unprocessedTransformedItems = unprocessedItemsList ?. flatMap ( ( item ) => this . _dcBatchTransformer . toReadBatchInputList ( item , metadata ) ) ; const failedTransformedItems = this . _errorQueue . flatMap ( item => { this . connection . logger . logError ( { requestId , scope : MANAGER_NAME . BATCH_MANAGER , log : item . error , } ) ; return this . _dcBatchTransformer . toReadBatchInputList ( item . requestInput , metadata ) ; } ) ; return { items : transformedItems ?? [ ] , unprocessedItems : unprocessedTransformedItems ?? [ ] , failedItems : failedTransformedItems ?? [ ] , } ; } param batchWriteItemOutputItems private async recursiveHandleBatchWriteItemsResponse ( batchWriteItemOutputItems , totalAttemptsSoFar , options ? , metadataOptions ? ) { const unProcessedListItems = batchWriteItemOutputItems . filter ( ( response ) => response . UnprocessedItems && ! isEmptyObject ( response . UnprocessedItems ) ) . map ( item => item . UnprocessedItems ! ) ; if ( ! unProcessedListItems . length ) { return unProcessedListItems ; } if ( totalAttemptsSoFar === ( options ?. maxRetryAttempts ?? BATCH_WRITE_MAX_ALLOWED_ATTEMPTS ) ) { this . connection . logger . logInfo ( { scope : MANAGER_NAME . BATCH_MANAGER , log : `Reached max allowed attempts ${ totalAttemptsSoFar }, aborting...` , requestId : metadataOptions ?. requestId , } ) ; return unProcessedListItems ; } await this . waitForExponentialBackoff ( totalAttemptsSoFar , undefined , metadataOptions ) ; const sortedUnprocessedItems = unProcessedListItems . reduce ( ( acc , unprocessedItems ) => { Object . entries ( unprocessedItems ! ) . forEach ( ( [ tableName , unprocessedRequests ] ) => { if ( ! acc [ tableName ] ) { acc [ tableName ] = [ ] ; } acc [ tableName ] = [ ... acc [ tableName ] , ... unprocessedRequests ] ; } ) ; return acc ; } , { } as BatchWriteItemRequestMap ) ; const batchRequestsItems = this . _dcBatchTransformer . mapTableWriteItemsToBatchWriteItems ( sortedUnprocessedItems ) ; const batchRequests = batchRequestsItems . map ( batchRequestMap => { return this . toLimited ( async ( ) => this . connection . documentClient . batchWrite ( { RequestItems : { ... batchRequestMap } , ReturnItemCollectionMetrics : metadataOptions ?. returnConsumedCapacity , } ) . promise ( ) , batchRequestMap , REQUEST_TYPE . BATCH_WRITE ) ; } ) ; const batchRequestsResponses = ( await Promise . all ( batchRequests ) ) as DocumentClient . BatchWriteItemOutput [ ] ; batchRequestsResponses . forEach ( ( response , index ) => { if ( response . ConsumedCapacity ) { this . connection . logger . logStats ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . BATCH_MANAGER , requestSegment : index , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : response . ConsumedCapacity , } ) ; } } ) ; return this . recursiveHandleBatchWriteItemsResponse ( batchRequestsResponses , ++ totalAttemptsSoFar , options , metadataOptions ) ; } private async recursiveHandleBatchReadItemsResponse ( batchReadItemOutputList , totalAttemptsSoFar , options ? , responsesStore = [ ] , metadataOptions ? ) { const batchReadResponses = batchReadItemOutputList . filter ( ( response ) => response . Responses && ! isEmptyObject ( response . Responses ) ) . map ( ( response ) => response . Responses ! ) ; if ( batchReadResponses . length ) { const mappedResponsesItemList = batchReadResponses . flatMap ( batchGetResponse => this . mapBatchGetResponseToItemList ( batchGetResponse ) ) ; responsesStore . push ( ... mappedResponsesItemList ) ; } const unprocessedItemsList = batchReadItemOutputList . filter ( ( response ) => response . UnprocessedKeys && ! isEmptyObject ( response . UnprocessedKeys ) ) ; if ( ! unprocessedItemsList . length ) { return { items : responsesStore , } ; } if ( totalAttemptsSoFar === ( options ?. maxRetryAttempts ?? BATCH_READ_MAX_ALLOWED_ATTEMPTS ) ) { this . connection . logger . logInfo ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . BATCH_MANAGER , log : `Reached max allowed attempts ${ totalAttemptsSoFar }, aborting...` , } ) ; return { items : responsesStore , unprocessedItemsList : unprocessedItemsList . map ( item => item . UnprocessedKeys ! ) , } ; } await this . waitForExponentialBackoff ( totalAttemptsSoFar , undefined , metadataOptions ) ; const sortedUnprocessedItems = unprocessedItemsList . reduce ( ( acc , { UnprocessedKeys } ) => { Object . entries ( UnprocessedKeys ! ) . forEach ( ( [ tableName , unprocessedRequests ] ) => { if ( ! acc [ tableName ] ) { acc [ tableName ] = { Keys : [ ] , } ; } acc [ tableName ] . Keys . push ( ... unprocessedRequests . Keys ) ; } ) ; return acc ; } , { } as DocumentClient . BatchGetRequestMap ) ; const batchRequestsItemsList = this . _dcBatchTransformer . mapTableReadItemsToBatchReadItems ( sortedUnprocessedItems ) ; const batchRequests = batchRequestsItemsList . map ( batchRequestMap => { return this . toLimited ( async ( ) => this . connection . documentClient . batchGet ( { RequestItems : { ... batchRequestMap } , ReturnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) . promise ( ) , batchRequestMap , REQUEST_TYPE . BATCH_READ ) ; } ) ; const batchRequestsResponses = ( await Promise . all ( batchRequests ) ) as DocumentClient . BatchGetItemOutput [ ] ; batchRequestsResponses . forEach ( ( response , index ) => { if ( response . ConsumedCapacity ) { this . connection . logger . logStats ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . BATCH_MANAGER , requestSegment : index , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : response . ConsumedCapacity , } ) ; } } ) ; return this . recursiveHandleBatchReadItemsResponse ( batchRequestsResponses , ++ totalAttemptsSoFar , options , responsesStore , metadataOptions ) ; } private resetErrorQueue ( ) { this . _errorQueue = [ ] ; } private mapBatchGetResponseToItemList ( batchGetResponse ) { return Object . entries ( batchGetResponse ) . flatMap ( ( [ , batchResponse ] ) => batchResponse ) ; } param anyPromiseFactory param requestItem param requestType private toLimited < T > ( anyPromiseFactory , requestItem , requestType ) { return this . limit ( async ( ) => { try { const response = await anyPromiseFactory ( ) ; return response ; } catch ( err ) { this . _errorQueue . push ( { requestInput : requestItem , error : err , requestType , } ) ; return err as T ; } } ) ; } private waitForExponentialBackoff ( attempts , multiplicationFactor = 1 , metadataOptions ? ) { multiplicationFactor = multiplicationFactor < 1 ? 1 : multiplicationFactor ; return new Promise ( resolve => { const backoffTime = this . exponentialBackoff ( attempts , multiplicationFactor ) ; this . connection . logger . logInfo ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . BATCH_MANAGER , log : `${ attempts } attempts so far, sleeping ${ backoffTime }ms before retrying...` , } ) ; setTimeout ( resolve , backoffTime ) ; } ) ; } param attempts private exponentialBackoff ( attempts , multiplicationFactor ) { return Math . floor ( Math . random ( ) * 10 * Math . pow ( 2 , attempts || 1 ) * multiplicationFactor ) ; } }