@naturalcycles/db-lib/dist/pipeline/dbPipelineRestore.js

Lowest Common Denominator API to supported Databases

import { _hb, _mapValues, _passthroughMapper, ErrorMode, localTime, pMap, } from '@naturalcycles/js-lib';
import { _pipeline, boldWhite, dimWhite, fs2, grey, NDJsonStats, transformChunk, transformFilterSync, transformLogProgress, transformMap, transformTap, writableForEach, yellow, } from '@naturalcycles/nodejs-lib';
/**
 * Pipeline from NDJSON files in a folder (optionally gzipped) to CommonDB.
 * Allows to define a mapper and a predicate to map/filter objects between input and output.
 * Handles backpressure.
 *
 * Optionally you can provide mapperPerTable and @param transformMapOptions (one for all mappers) - it will run for each table.
 */
export async function dbPipelineRestore(opt) {
    const { db, concurrency = 16, chunkSize = 100, limit, sinceUpdated, inputDirPath, mapperPerTable = {}, saveOptionsPerTable = {}, transformMapOptions, errorMode = ErrorMode.SUPPRESS, recreateTables = false, } = opt;
    const onlyTables = opt.tables && new Set(opt.tables);
    const sinceUpdatedStr = sinceUpdated ? ' since ' + grey(localTime(sinceUpdated).toPretty()) : '';
    console.log(`>> ${dimWhite('dbPipelineRestore')} started in ${grey(inputDirPath)}...${sinceUpdatedStr}`);
    fs2.ensureDir(inputDirPath);
    const tablesToGzip = new Set();
    const sizeByTable = {};
    const statsPerTable = {};
    const tables = [];
    fs2.readdir(inputDirPath).forEach(f => {
        let table;
        let gzip = false;
        if (f.endsWith('.ndjson')) {
            table = f.slice(0, f.length - '.ndjson'.length);
        }
        else if (f.endsWith('.ndjson.gz')) {
            table = f.slice(0, f.length - '.ndjson.gz'.length);
            gzip = true;
        }
        else {
            return;
        }
        if (onlyTables && !onlyTables.has(table))
            return; // skip table
        tables.push(table);
        if (gzip)
            tablesToGzip.add(table);
        sizeByTable[table] = fs2.stat(`${inputDirPath}/${f}`).size;
    });
    const sizeStrByTable = _mapValues(sizeByTable, (_k, b) => _hb(b));
    console.log(`${yellow(tables.length)} ${boldWhite('table(s)')}:\n`, sizeStrByTable);
    // const schemaByTable: Record<string, CommonSchema> = {}
    if (recreateTables) {
        await pMap(tables, async (table) => {
            const schemaFilePath = `${inputDirPath}/${table}.schema.json`;
            if (!fs2.pathExists(schemaFilePath)) {
                console.warn(`${schemaFilePath} does not exist!`);
                return;
            }
            const schema = await fs2.readJsonAsync(schemaFilePath);
            await db.createTable(table, schema, { dropIfExists: true });
        });
    }
    await pMap(tables, async (table) => {
        const gzip = tablesToGzip.has(table);
        const filePath = `${inputDirPath}/${table}.ndjson` + (gzip ? '.gz' : '');
        const saveOptions = saveOptionsPerTable[table] || {};
        const started = Date.now();
        let rows = 0;
        const sizeBytes = sizeByTable[table];
        console.log(`<< ${grey(filePath)} ${dimWhite(_hb(sizeBytes))} started...`);
        await _pipeline([
            fs2.createReadStreamAsNDJSON(filePath).take(limit || Number.POSITIVE_INFINITY),
            transformTap(() => rows++),
            transformLogProgress({
                logEvery: 1000,
                ...opt,
                metric: table,
            }),
            ...(sinceUpdated
                ? [transformFilterSync(r => r.updated >= sinceUpdated)]
                : []),
            transformMap(mapperPerTable[table] || _passthroughMapper, {
                errorMode,
                flattenArrayOutput: true,
                ...transformMapOptions,
                metric: table,
            }),
            transformChunk({ chunkSize }),
            writableForEach(async (dbms) => {
                await db.saveBatch(table, dbms, saveOptions);
            }),
        ]);
        const stats = NDJsonStats.create({
            tookMillis: Date.now() - started,
            rows,
            sizeBytes,
        });
        console.log(`<< ${grey(filePath)}\n` + stats.toPretty());
        statsPerTable[table] = stats;
    }, { concurrency, errorMode });
    const statsTotal = NDJsonStats.createCombined(Object.values(statsPerTable));
    console.log(statsTotal.toPretty('total'));
    return statsTotal;
}