import { EntityTarget , INTERNAL_ENTITY_ATTRIBUTE , InvalidParallelScanLimitOptionError , MANAGER_NAME , PARALLEL_SCAN_CONCURRENCY_LIMIT , STATS_TYPE , } from '@typedorm/common' ; import { DynamoDB } from 'aws-sdk' ; import pLimit from 'p-limit' ; import { getUniqueRequestId } from '../../helpers/get-unique-request-id' ; import { Connection } from '../connection/connection' ; import { FilterOptions } from '../expression/filter-options-type' ; import { ProjectionKeys } from '../expression/projection-keys-options-type' ; import { MetadataOptions } from '../transformer/base-transformer' ; import { DocumentClientScanTransformer } from '../transformer/document-client-scan-transformer' ; interface ScanManageBaseOptions < Entity , PartitionKey > { default scanIndex ? : string ; default limit ? : number ; default cursor ? : DynamoDB . DocumentClient . Key ; default where ? : FilterOptions < Entity , PartitionKey > ; default select ? : ProjectionKeys < Entity > ; } export interface ScanManagerFindOptions < Entity > extends ScanManageBaseOptions < Entity , { } > { default totalSegments ? : number ; default limitPerSegment ? : number ; default cursor ? : Record < number , ScanManageBaseOptions < Entity , { } > [ 'cursor' ] > | ScanManageBaseOptions < Entity , { } > [ 'cursor' ] ; default requestsConcurrencyLimit ? : number ; } export type ScanManagerCountOptions < Entity > = Pick < ScanManageBaseOptions < Entity , { } > , 'scanIndex' | 'where' > ; export interface ScanManagerParallelScanOptions extends ScanManageBaseOptions < any , any > { /**
     * Total number of segments to divide this scan in
     */ totalSegments : number ; default limitPerSegment ? : number ; default entity ? : EntityTarget < any > ; default cursor ? : Record < number , ScanManagerScanOptions [ 'cursor' ] > ; default requestsConcurrencyLimit ? : number ; } export interface ScanManagerScanOptions extends ScanManageBaseOptions < any , any > { default entity ? : EntityTarget < any > ; default segment ? : number ; default totalSegments ? : number ; default limitPerSegment ? : number ; } export class ScanManager { private itemsFetchedSoFarTotalParallelCount : number ; private limit = pLimit ( PARALLEL_SCAN_CONCURRENCY_LIMIT ) ; private _dcScanTransformer : DocumentClientScanTransformer ; constructor ( private connection ) { this . _dcScanTransformer = new DocumentClientScanTransformer ( connection ) ; this . itemsFetchedSoFarTotalParallelCount = 0 ; } param entityClass param findOptions param metadataOptions async find < Entity > ( entityClass , findOptions ? , metadataOptions ? ) { const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; let response ; if ( findOptions ?. totalSegments ) { ( response = await this . parallelScan < Entity > ( { ... findOptions , entity : entityClass , } as ScanManagerParallelScanOptions ) ) , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ; } else { response = await this . scan < Entity > ( { ... findOptions , entity : entityClass } as ScanManagerScanOptions , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; } if ( response . unknownItems ?. length ) { this . connection . logger . logWarn ( { requestId , scope : MANAGER_NAME . SCAN_MANAGER , log : `"${ response . unknownItems . length }" items were skipped from the response because TypDORM failed to resolve them.` , } ) ; } return { items : response . items , cursor : response . cursor , } ; } param entityClass param scanOptions param metadataOptions async count < Entity > ( entityClass , scanOptions ? , metadataOptions ? ) { const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; const dynamoScanInput = this . _dcScanTransformer . toDynamoScanItem ( { ... scanOptions , entity : entityClass , onlyCount : true , select : undefined } , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; const count = await this . _internalRecursiveCount ( { scanInput : dynamoScanInput , metadataOptions : { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } , } ) ; return count ; } param scanOptions param metadataOptions async parallelScan < Entity > ( scanOptions , metadataOptions ? ) { this . itemsFetchedSoFarTotalParallelCount = 0 ; const concurrencyLimit = PARALLEL_SCAN_CONCURRENCY_LIMIT || scanOptions . requestsConcurrencyLimit ; const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; if ( scanOptions . requestsConcurrencyLimit ) { this . limit = pLimit ( scanOptions . requestsConcurrencyLimit ) ; } const parallelScanOptions = [ ] ; if ( scanOptions ?. limit && scanOptions ?. limitPerSegment && scanOptions ?. limit < scanOptions ?. limitPerSegment ) { throw new InvalidParallelScanLimitOptionError ( scanOptions ?. limit , scanOptions ?. limitPerSegment ) ; } for ( let index = 0 ; index < scanOptions . totalSegments ; index ++ ) { const cursorForSegment = scanOptions . cursor ? scanOptions . cursor [ index ] : undefined ; parallelScanOptions . push ( { ... scanOptions , cursor : cursorForSegment , segment : index , } ) ; } this . connection . logger . logInfo ( { requestId , scope : MANAGER_NAME . SCAN_MANAGER , log : `Running scan in parallel with ${ scanOptions . totalSegments } segments.` , } ) ; if ( concurrencyLimit < scanOptions . totalSegments ) { this . connection . logger . logInfo ( { requestId , scope : MANAGER_NAME . SCAN_MANAGER , log : `Current request concurrency limit ${ concurrencyLimit } is lower than requested segments count ${ scanOptions . totalSegments }
        So requests will be run in a batch of ${ concurrencyLimit } at a time until all segments ${ scanOptions . totalSegments } have processed.` , } ) ; } const allPromisesResponse = await Promise . all ( parallelScanOptions . map ( options => this . toLimited ( this . scan < Entity > ( options , metadataOptions ) ) ) ) ; const response = allPromisesResponse . reduce ( ( acc , current , index ) => { if ( current . items ?. length ) { if ( ! acc . items ) { acc . items = [ ] ; } acc . items = [ ... acc . items , ... current . items ] ; } if ( current . unknownItems ?. length ) { if ( ! acc . unknownItems ) { acc . unknownItems = [ ] ; } acc . unknownItems = [ ... acc . unknownItems , ... current . unknownItems ] ; } if ( current . cursor ) { if ( ! acc . cursor ) { acc . cursor = { } ; } acc . cursor = { ... acc . cursor , [ index ] : current . cursor , } ; } return acc ; } , { } as { items : Entity [ ] ; unknownItems : DynamoDB . DocumentClient . AttributeMap [ ] ; cursor : Record < number , DynamoDB . DocumentClient . Key > ; } ) ; return response ; } param scanOptions param metadataOptions async scan < Entity > ( scanOptions ? , metadataOptions ? ) { const requestId = getUniqueRequestId ( metadataOptions ?. requestId ) ; const dynamoScanInput = this . _dcScanTransformer . toDynamoScanItem ( { ... scanOptions , limit : scanOptions ?. totalSegments ? scanOptions ?. limitPerSegment : scanOptions ?. limit , } , { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } ) ; const response = await this . _internalRecursiveScan ( { scanInput : dynamoScanInput , limit : scanOptions ?. limit , cursor : scanOptions ?. cursor , metadataOptions : { requestId , returnConsumedCapacity : metadataOptions ?. returnConsumedCapacity , } , } ) ; const entities = this . _dcScanTransformer . fromDynamoScanResponseItemList < Entity > ( response . items ) ; if ( scanOptions ?. entity && entities . unknownItems ) { this . connection . logger . logWarn ( { requestId , scope : MANAGER_NAME . SCAN_MANAGER , log : `
        There were some items that looked like ${ scanOptions ?. entity . name } but TypeDORM was unable to convert it back to entity type,
        This can happen when there are items in the table with "${ INTERNAL_ENTITY_ATTRIBUTE . ENTITY_NAME } but was not created by TypeDORM.
        You should remove them or update it to something different."` , } ) ; } return { items : entities . items ?. length ? entities . items : undefined , unknownItems : entities . unknownItems ?. length ? entities . unknownItems : undefined , cursor : response . cursor , } ; } /**
     * Recursively scans table with given options
     */ private async _internalRecursiveScan ( { scanInput , limit , cursor , itemsFetched = [ ] , metadataOptions , } ) { if ( limit && this . itemsFetchedSoFarTotalParallelCount >= limit ) { return { items : itemsFetched , cursor , } ; } const { LastEvaluatedKey , Items = [ ] , ConsumedCapacity , } = await this . connection . documentClient . scan ( { ... scanInput , ExclusiveStartKey : cursor } ) . promise ( ) ; if ( ConsumedCapacity ) { this . connection . logger . logStats ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . SCAN_MANAGER , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : ConsumedCapacity , } ) ; } if ( limit && this . itemsFetchedSoFarTotalParallelCount >= limit ) { return { items : itemsFetched , cursor , } ; } itemsFetched = [ ... itemsFetched , ... Items ] ; this . itemsFetchedSoFarTotalParallelCount += Items . length ; if ( LastEvaluatedKey ) { return this . _internalRecursiveScan ( { scanInput , limit , cursor : LastEvaluatedKey , itemsFetched , metadataOptions , } ) ; } return { items : itemsFetched , cursor : LastEvaluatedKey , } ; } /**
     * Recursively counts items form table with given options
     */ private async _internalRecursiveCount ( { scanInput , cursor , currentCount = 0 , metadataOptions , } ) { const { Count , LastEvaluatedKey , ConsumedCapacity , } = await this . connection . documentClient . scan ( { ... scanInput , ExclusiveStartKey : cursor } ) . promise ( ) ; if ( ConsumedCapacity ) { this . connection . logger . logStats ( { requestId : metadataOptions ?. requestId , scope : MANAGER_NAME . SCAN_MANAGER , statsType : STATS_TYPE . CONSUMED_CAPACITY , consumedCapacityData : ConsumedCapacity , } ) ; } currentCount += Count || 0 ; if ( LastEvaluatedKey ) { return this . _internalRecursiveCount ( { scanInput , cursor : LastEvaluatedKey , currentCount , metadataOptions , } ) ; } return currentCount ; } param promise returns private toLimited < T > ( promise ) { return this . limit ( ( ) => promise ) ; } }