import { BATCH_READ_ITEMS_LIMIT , BATCH_WRITE_ITEMS_LIMIT , InvalidBatchWriteItemError , TRANSFORM_BATCH_TYPE , } from '@typedorm/common' ; import { DocumentClient } from 'aws-sdk/clients/dynamodb' ; import { v4 } from 'uuid' ; import { getHashedIdForInput } from '../../helpers/get-hashed-id-for-input' ; import { ReadBatch , ReadBatchItem } from '../batch/read-batch' ; import { isBatchAddCreateItem , isBatchAddDeleteItem } from '../batch/type-guards' ; import { WriteBatchItem , WriteBatch } from '../batch/write-batch' ; import { Connection } from '../connection/connection' ; import { isWriteTransactionItemList } from '../transaction/type-guards' ; import { MetadataOptions } from './base-transformer' ; import { isLazyTransactionWriteItemListLoader , LazyTransactionWriteItemListLoader , } from './is-lazy-transaction-write-item-list-loader' ; import { LowOrderTransformers } from './low-order-transformers' ; export type WriteRequestWithMeta = { tableName : string ; writeRequest : DocumentClient . WriteRequest ; } ; export type ReadRequestWithMeta = { tableName : string ; readRequest : DocumentClient . Key ; } ; export type BatchWriteItemTransform < Transformed > = { rawInput : any ; transformedInput : Transformed ; } ; export type BatchWriteItemRequestMapTransform < Transformed > = { [ key ] : BatchWriteItemTransform < Transformed > [ ] ; } ; export class DocumentClientBatchTransformer extends LowOrderTransformers { constructor ( connection ) { super ( connection ) ; } toDynamoWriteBatchItems ( writeBatch , metadataOptions ? ) { const { items } = writeBatch ; this . connection . logger . logTransformBatch ( { requestId : metadataOptions ?. requestId , operation : TRANSFORM_BATCH_TYPE . BATCH_WRITE , prefix : 'Before' , body : items , } ) ; const { lazyTransactionWriteItemListLoaderItems , simpleBatchRequestItems , transactionListItems , metadata , } = this . innerTransformBatchWriteItems ( items , metadataOptions ) ; const sorted = this . getWriteRequestsSortedByTable ( simpleBatchRequestItems ) ; const batchWriteRequestItems = this . mapTableWriteItemsToBatchWriteItems ( sorted ) ; const transformed = { batchWriteRequestMapItems : batchWriteRequestItems , transactionListItems , lazyTransactionWriteItemListLoaderItems , metadata , } ; this . connection . logger . logTransformBatch ( { requestId : metadataOptions ?. requestId , operation : TRANSFORM_BATCH_TYPE . BATCH_WRITE , prefix : 'After' , body : transformed , } ) ; return transformed ; } toDynamoReadBatchItems ( readBatch , metadataOptions ? ) { const { items } = readBatch ; this . connection . logger . logTransformBatch ( { requestId : metadataOptions ?. requestId , operation : TRANSFORM_BATCH_TYPE . BATCH_READ , prefix : 'Before' , body : items , } ) ; const { metadata , batchReadRequestItems } = this . transformBatchReadItems ( items ) ; const sortedByTableName = this . getReadRequestsSortedByTable ( batchReadRequestItems ) ; const batchRequestItemsList = this . mapTableReadItemsToBatchReadItems ( sortedByTableName ) ; const transformed = { batchRequestItemsList , metadata , } ; this . connection . logger . logTransformBatch ( { requestId : metadataOptions ?. requestId , operation : TRANSFORM_BATCH_TYPE . BATCH_READ , prefix : 'After' , body : transformed , } ) ; return transformed ; } mapTableWriteItemsToBatchWriteItems ( requestsSortedByTable ) { let currentFillingIndex = 0 ; let totalItemsAtCurrentFillingIndex = 0 ; const multiBatchItems = Object . entries ( requestsSortedByTable ) . reduce ( ( acc , [ tableName , perTableItems ] ) => { while ( perTableItems . length ) { if ( ! acc [ currentFillingIndex ] ) { acc [ currentFillingIndex ] = { } ; } if ( ! acc [ currentFillingIndex ] [ tableName ] ) { acc [ currentFillingIndex ] [ tableName ] = [ ] ; } acc [ currentFillingIndex ] [ tableName ] . push ( perTableItems . shift ( ) ! ) ; ++ totalItemsAtCurrentFillingIndex ; if ( totalItemsAtCurrentFillingIndex === BATCH_WRITE_ITEMS_LIMIT ) { ++ currentFillingIndex ; totalItemsAtCurrentFillingIndex = 0 ; } } return acc ; } , [ ] as DocumentClient . BatchWriteItemRequestMap [ ] ) ; return multiBatchItems ; } mapTableReadItemsToBatchReadItems ( requestsSortedByTable ) { let currentFillingIndex = 0 ; let totalItemsAtCurrentFillingIndex = 0 ; const multiBatchItems = Object . entries ( requestsSortedByTable ) . reduce ( ( acc , [ tableName , perTableItems ] ) => { while ( perTableItems . Keys . length ) { if ( ! acc [ currentFillingIndex ] ) { acc [ currentFillingIndex ] = { } ; } if ( ! acc [ currentFillingIndex ] [ tableName ] ) { acc [ currentFillingIndex ] [ tableName ] = { Keys : [ ] , } ; } acc [ currentFillingIndex ] [ tableName ] . Keys . push ( perTableItems . Keys . shift ( ) ! ) ; ++ totalItemsAtCurrentFillingIndex ; if ( totalItemsAtCurrentFillingIndex === BATCH_READ_ITEMS_LIMIT ) { ++ currentFillingIndex ; totalItemsAtCurrentFillingIndex = 0 ; } } return acc ; } , [ ] as DocumentClient . BatchGetRequestMap [ ] ) ; return multiBatchItems ; } /**
     * Reverse transforms batch write item request to write batch item
     */ toWriteBatchInputList ( requestMap , { namespaceId , itemTransformHashMap , } ) { return Object . entries ( requestMap ) . flatMap ( ( [ , writeRequests ] ) => { return this . toRawBatchInputItem < DocumentClient . WriteRequest , WriteBatchItem < any , any > > ( writeRequests , { namespaceId , itemTransformHashMap , } ) ; } ) ; } /**
     * Reverse transforms batch read item request to read batch item
     */ toReadBatchInputList ( requestMap , { namespaceId , itemTransformHashMap , } ) { return Object . entries ( requestMap ) . flatMap ( ( [ , readRequests ] ) => { return this . toRawBatchInputItem < DocumentClient . Key , ReadBatchItem < any , any > > ( readRequests . Keys , { namespaceId , itemTransformHashMap , } ) ; } ) ; } /**
     * Converts batch item input to pre transformed item input
     */ private toRawBatchInputItem < Input , Output > ( transformedItems , { namespaceId , itemTransformHashMap , } ) { return transformedItems . map ( transformedItem => { const hashId = getHashedIdForInput ( namespaceId , transformedItem ) ; const originalItem = itemTransformHashMap . get ( hashId ) ; return originalItem ! ; } ) ; } /**
     * Parse each item in the request to be in one of the following collections
     * - simpleBatchRequestItems: simple items that can be processed in batch (i.e no uniques)
     * - transactionListItems: items that must be processed in a transaction instead of batch (i.e items with unique attributes)
     * - LazyTransactionWriteItemListLoaderItems: items that are must be processed in transaction but also requires other requests to be made first (i.e delete of unique items)
     */ private innerTransformBatchWriteItems ( batchItems , metadataOptions ? ) { const namespaceId = v4 ( ) ; const itemTransformHashMap = new Map < string , WriteBatchItem < any , any > > ( ) ; return batchItems . reduce ( ( acc , batchItem ) => { if ( isBatchAddCreateItem ( batchItem ) ) { const dynamoPutItem = this . toDynamoPutItem ( batchItem . create . item , undefined , metadataOptions ) ; if ( ! isWriteTransactionItemList ( dynamoPutItem ) ) { const transformedWriteRequest = { PutRequest : { Item : dynamoPutItem . Item , } , } ; acc . simpleBatchRequestItems . push ( { writeRequest : transformedWriteRequest , tableName : dynamoPutItem . TableName , } ) ; const itemHashId = getHashedIdForInput ( namespaceId , transformedWriteRequest ) ; itemTransformHashMap . set ( itemHashId , batchItem ) ; } else { acc . transactionListItems . push ( { rawInput : batchItem , transformedInput : dynamoPutItem , } ) ; } } else if ( isBatchAddDeleteItem ( batchItem ) ) { const { delete : { item , primaryKey } , } = batchItem ; const itemToRemove = this . toDynamoDeleteItem ( item , primaryKey , undefined , metadataOptions ) ; if ( ! isLazyTransactionWriteItemListLoader ( itemToRemove ) ) { const transformedItemRequest = { DeleteRequest : { Key : itemToRemove . Key , } , } ; acc . simpleBatchRequestItems . push ( { writeRequest : transformedItemRequest , tableName : itemToRemove . TableName , } ) ; const itemHashId = getHashedIdForInput ( namespaceId , transformedItemRequest ) ; itemTransformHashMap . set ( itemHashId , batchItem ) ; } else { acc . lazyTransactionWriteItemListLoaderItems . push ( { rawInput : batchItem , transformedInput : itemToRemove , } ) ; } } else { throw new InvalidBatchWriteItemError ( batchItem ) ; } return acc ; } , { simpleBatchRequestItems : [ ] as WriteRequestWithMeta [ ] , transactionListItems : [ ] as BatchWriteItemTransform < DocumentClient . TransactWriteItemList > [ ] , lazyTransactionWriteItemListLoaderItems : [ ] as BatchWriteItemTransform < LazyTransactionWriteItemListLoader > [ ] , metadata : { namespaceId , itemTransformHashMap , } , } ) ; } private transformBatchReadItems ( batchItems , metadataOptions ? ) { const namespaceId = v4 ( ) ; const itemTransformHashMap = new Map < string , ReadBatchItem < any , any > > ( ) ; return batchItems . reduce ( ( acc , batchItem ) => { const { item , primaryKey } = batchItem ; const itemToGet = this . toDynamoGetItem ( item , primaryKey , undefined , metadataOptions ) ; acc . batchReadRequestItems . push ( { readRequest : itemToGet . Key , tableName : itemToGet . TableName , } ) ; const transformedItemHashId = getHashedIdForInput ( namespaceId , itemToGet . Key ) ; itemTransformHashMap . set ( transformedItemHashId , batchItem ) ; return acc ; } , { batchReadRequestItems : [ ] as ReadRequestWithMeta [ ] , metadata : { namespaceId , itemTransformHashMap , } , } ) ; } private getWriteRequestsSortedByTable ( allRequestItems ) { return allRequestItems . reduce ( ( acc , requestItemWithMeta ) => { if ( ! acc [ requestItemWithMeta . tableName ] ) { acc [ requestItemWithMeta . tableName ] = [ ] ; } acc [ requestItemWithMeta . tableName ] . push ( requestItemWithMeta . writeRequest ) ; return acc ; } , { } as DocumentClient . BatchWriteItemRequestMap ) ; } private getReadRequestsSortedByTable ( allRequestItems ) { return allRequestItems . reduce ( ( acc , requestItemWithMeta ) => { if ( ! acc [ requestItemWithMeta . tableName ] ) { acc [ requestItemWithMeta . tableName ] = { Keys : [ ] , } ; } acc [ requestItemWithMeta . tableName ] . Keys . push ( requestItemWithMeta . readRequest ) ; return acc ; } , { } as DocumentClient . BatchGetRequestMap ) ; } }