wkr-util/esm/streams.js

Utility library for wkr project.

import {Transform} from 'stream'
import {StringDecoder} from 'string_decoder'

// copied from https://www.ramielcreations.com/the-hidden-power-of-node-js-streams/
export const mapStream = (fn, options = {}) => new Transform({
    objectMode: true,
    ...options,
    transform: async (chunk, encoding, callback) => {
        try {
            const result = fn(chunk)

            if(result.then) {
                callback(null, await result)
            }
            else {
                callback(null, result)
            }
        }
        catch(e) {
            return callback(e)
        }
    },
})

// Run an async function on each chunk of a stream. Each chunk is assumed to be an array.
// f needs to return a promise.
// only one promise is run at a time, each are awaited. this is just to reduce complexity of implementation.
// returns a promise resolving to undefined.
// will reject on first exception. no more chunks will be processed.
export const doOnEachRecord = (f, stream) => {
    return new Promise((resolve, reject) => {
        // if we catch an exception while running f, need a way to tell all the async stuff that we're actually done
        let done = false
        // if the stream ends while we're still running f, need a way to tell the loop that
        // it needs to resolve instead of leaving it to the 'end' callback.
        let shouldEnd = false
        let inReadable = false

        const handleReadable = async () => {
            inReadable = true
            // we only want one readable running at the same time
            stream.removeListener('readable', handleReadable)

            if(done) return // just make sure we don't continue running

            let records

            while((records = stream.read())) {
                try {
                    for(let i = 0; i < records.length; i++) {
                        await f(records[i])
                    }
                }
                catch(err) {
                    done = true
                    reject(err)
                    return
                }
            }

            if(!done && shouldEnd) {
                done = true
                resolve()
                return
            }

            inReadable = false

            // start listening again
            stream.on('readable', handleReadable)
        }

        stream.on('readable', handleReadable)

        stream.on('end', () => {
            if(!done && !inReadable) {
                done = true
                resolve()
            }
            else {
                shouldEnd = true
            }
        })

        stream.on('error', err => {
            if(!done) {
                done = true
                reject(err)
            }
        })
    })
}

export const LineToRecordTransform = transformFn => {
    const xtream = new Transform({objectMode: true})

    let partialLine = null

    const decoder = new StringDecoder('utf-8')

    xtream._transform = function(chunk, encoding, done) {
        // Convert the Buffer chunks to String
        chunk = decoder.write(chunk)

        if(partialLine) {
            chunk = partialLine + chunk
        }

        const lines = chunk.split('\n')
        partialLine = lines.splice(lines.length - 1, 1)[0]
        this.push(lines.map(transformFn))

        done()
    }

    xtream._flush = function(done) {
        if(partialLine) {
            this.push([transformFn(partialLine)])
        }

        partialLine = null
        done()
    }

    return xtream
}