import { inject , injectable } from 'inversify' ; import { client , CosmosContainerClient , cosmosContainerClientTypes , CosmosOperationResponse } from 'azure-services' ; import { flatMap , groupBy } from 'lodash' ; import { ItemType , OnDemandPageScanResult } from 'storage-documents' ; import * as cosmos from '@azure/cosmos' ; import pLimit from 'p-limit' ; import { PartitionKeyFactory } from '../factories/partition-key-factory' ; import { OperationResult } from './operation-result' ; @ injectable ( ) export class OnDemandPageScanRunResultProvider { public maxConcurrencyLimit = 5 ; constructor ( @ inject ( cosmosContainerClientTypes . OnDemandScanRunsCosmosContainerClient ) private readonly cosmosContainerClient , @ inject ( PartitionKeyFactory ) private readonly partitionKeyFactory ) { } public async readScanRun ( scanId ) { return ( await this . cosmosContainerClient . readDocument < OnDemandPageScanResult > ( scanId , this . getPartitionKey ( scanId ) ) ) . item ; } public async readScanRuns ( scanIds ) { const maxItemsPerQuery = 1000 ; if ( scanIds . length > maxItemsPerQuery ) { throw new Error ( `Can't read more than ${ maxItemsPerQuery } scan documents per query.` ) ; } const scanIdsByPartition = groupBy ( scanIds , ( scanId ) => { return this . getPartitionKey ( scanId ) ; } ) ; const limit = pLimit ( this . maxConcurrencyLimit ) ; const response = await Promise . all ( Object . keys ( scanIdsByPartition ) . map ( async ( pKey ) => { return limit ( async ( ) => this . cosmosContainerClient . executeQueryWithContinuationToken < OnDemandPageScanResult > ( async ( token ) => { return this . cosmosContainerClient . queryDocuments < OnDemandPageScanResult > ( this . getQuery ( scanIdsByPartition [ pKey ] , pKey ) , token ) ; } ) ) ; } ) ) ; return flatMap ( response ) ; } param pageScanResult public async tryUpdateScanRun ( pageScanResult ) { const operationResponse = await this . updateScanRunImpl ( pageScanResult , false ) ; if ( client . isSuccessStatusCode ( operationResponse ) ) { return { succeeded : true , result : operationResponse . item } ; } else if ( operationResponse . statusCode === 412 ) { return { succeeded : false } ; } else { throw new Error ( `Scan result document operation failed. Scan Id: ${ pageScanResult . id } Response status code: ${ operationResponse . statusCode } Response: ${ operationResponse . response }` ) ; } } param pageScanResult public async updateScanRun ( pageScanResult ) { const operationResponse = await this . updateScanRunImpl ( pageScanResult , true ) ; return operationResponse . item ; } param scanRuns public async writeScanRuns ( scanRuns ) { const scanRunsByPartition = groupBy ( scanRuns , ( scanRun ) => { this . setSystemProperties ( scanRun ) ; return scanRun . partitionKey ; } ) ; const limit = pLimit ( this . maxConcurrencyLimit ) ; await Promise . all ( Object . keys ( scanRunsByPartition ) . map ( async ( pKey ) => { return limit ( async ( ) => this . cosmosContainerClient . writeDocuments ( scanRunsByPartition [ pKey ] , pKey ) ) ; } ) ) ; } private async updateScanRunImpl ( pageScanResult , throwIfNotSuccess ) { if ( pageScanResult . id === undefined ) { throw new Error ( `Cannot update scan run using partial scan run without id: ${ JSON . stringify ( pageScanResult ) }` ) ; } const persistedResult = pageScanResult as OnDemandPageScanResult ; this . setSystemProperties ( persistedResult ) ; return this . cosmosContainerClient . mergeOrWriteDocument ( persistedResult , undefined , throwIfNotSuccess ) ; } private getQuery ( scanIds , partitionKey ) { return { query : 'SELECT * FROM c WHERE c.partitionKey = @partitionKey AND ARRAY_CONTAINS(@scanIds, c.id)' , parameters : [ { name : '@partitionKey' , value : partitionKey , } , { name : '@scanIds' , value : scanIds , } , ] , } ; } private setSystemProperties ( pageScanResult ) { pageScanResult . itemType = ItemType . onDemandPageScanRunResult ; pageScanResult . partitionKey = this . getPartitionKey ( pageScanResult . id ) ; } private getPartitionKey ( scanId ) { return this . partitionKeyFactory . createPartitionKeyForDocument ( ItemType . onDemandPageScanRunResult , scanId ) ; } }