Async Generator of Batches

query-batches-generator.js

const F = require('functional-pipelines');

async function* queryBatches(table, {keyConditionExpression, expressionAttributeValues, strict = false, limit, flatten = false} = {}, options = {}, params = {}) {
    let LastEvaluatedKey;
    let Count;
    let ScannedCount;
    let ConsumedCapacity;
    let itemsBatch = [];
    let resultsCount = 0;
    let yieldCount = 0;
    let continuation = {};
    do {
        ({
            Items: itemsBatch,
            LastEvaluatedKey,
            Count,
            ScannedCount,
            ConsumedCapacity
        } = await query(table, {keyConditionExpression, expressionAttributeValues}, options, {...params, ...continuation}));
        continuation = {ExclusiveStartKey: LastEvaluatedKey};
        const metadata = F.lazy({Count, ScannedCount, ConsumedCapacity, LastEvaluatedKey});
        const batchSize = itemsBatch.length;
        resultsCount += batchSize;
        if (strict && resultsCount >= limit) {
            const lastBatchSize = limit - yieldCount;
            if (flatten) {
                yield* itemsBatch.slice(0, lastBatchSize);
            } else {
                yield F.iterator(itemsBatch.slice(0, lastBatchSize), {metadata});
            }
            yieldCount += lastBatchSize;
            return yieldCount; // TIP: the return value from a generator is ignored by for..of loop, only retrievable by calling .next() manually to get the last {value: returnValue, done: true}, or `const result = yield* gen1` inside gen2 function, result here would get the return value
        }
        yieldCount += batchSize;
        if (flatten) {
            yield* itemsBatch;
        } else {
            yield F.iterator(itemsBatch, {metadata});
        }
    } while (LastEvaluatedKey && (limit ? resultsCount < limit : true));
}