kourou 0.19.4 → 0.20.2

package/lib/support/dump-collection.js

@@ -6,72 +6,222 @@ const fs_1 = (0, tslib_1.__importDefault)(require("fs"));
 const path_1 = (0, tslib_1.__importDefault)(require("path"));
 const cli_ux_1 = (0, tslib_1.__importDefault)(require("cli-ux"));
 const ndjson_1 = (0, tslib_1.__importDefault)(require("ndjson"));
-async function dumpCollectionData(sdk, index, collection, batchSize, destPath, query = {}, format = 'jsonl') {
-    const collectionDir = path_1.default.join(destPath, collection);
-    const filename = path_1.default.join(collectionDir, (format.toLowerCase() === 'jsonl' ? 'documents.jsonl' : 'documents.json'));
-    const writeStream = fs_1.default.createWriteStream(filename);
-    const waitWrite = new Promise(resolve => writeStream.on('finish', resolve));
-    const ndjsonStream = ndjson_1.default.stringify();
-    const options = {
-        scroll: '1m',
-        size: batchSize
-    };
-    const writeLine = (content) => {
+const common_1 = require("../common");
+/**
+ * Flatten an object transform:
+ * {
+ *  title: "kuzzle",
+ *  info : {
+ *    tag: "news"
+ *  }
+ * }
+ *
+ * Into an object like:
+ * {
+ *  title: "kuzzle",
+ *  info.tag: news
+ * }
+ *
+ * @param {Object} target the object we have to flatten
+ * @returns {Object} the flattened object
+ */
+function flattenObject(target) {
+    const output = {};
+    flattenStep(output, target);
+    return output;
+}
+function flattenStep(output, object, prev = null) {
+    const keys = Object.keys(object);
+    for (let i = 0; i < keys.length; i++) {
+        const key = keys[i];
+        const value = object[key];
+        const newKey = prev ? prev + '.' + key : key;
+        if (Object.prototype.toString.call(value) === '[object Object]') {
+            output[newKey] = value;
+            flattenStep(output, value, newKey);
+        }
+        output[newKey] = value;
+    }
+}
+class AbstractDumper {
+    constructor(sdk, index, collection, batchSize, destPath, query = {}) {
+        this.sdk = sdk;
+        this.index = index;
+        this.collection = collection;
+        this.batchSize = batchSize;
+        this.destPath = destPath;
+        this.query = query;
+        this.collectionDir = path_1.default.join(this.destPath, this.collection);
+        this.options = {
+            scroll: '2s',
+            size: batchSize
+        };
+    }
+    /**
+     * One-shot call before the dump. Can be used to
+     * perform setup operations before dumping.
+     *
+     * @returns void
+     */
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    async setup() { }
+    /**
+     * One-shot call before iterating over the data. Can be
+     * used to write the header of the dumped output.
+     */
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    async writeHeader() { }
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    async tearDown() { }
+    /**
+     * The loop that iterates over the documents of the collection and
+     * calls all the other hooks.
+     *
+     * @returns a promise resolving when the dump is finished.
+     */
+    async dump() {
+        this.filename = path_1.default.join(this.collectionDir, `documents.${this.fileExtension}`);
+        this.writeStream = fs_1.default.createWriteStream(this.filename);
+        const waitWrite = new Promise((resolve, reject) => this.writeStream ? this.writeStream.on('finish', resolve) : reject());
+        this.writeStream.on('error', error => {
+            throw error;
+        });
+        await this.setup();
+        await this.writeHeader();
+        let results = await this.sdk.document.search(this.index, this.collection, { query: this.query }, this.options);
+        const progressBar = cli_ux_1.default.progress({
+            format: `Dumping ${this.collection} |{bar}| {percentage}% || {value}/{total} documents`
+        });
+        progressBar.start(results.total, 0);
+        do {
+            progressBar.update(results.fetched);
+            for (const hit of results.hits) {
+                await this.onResult({
+                    _id: hit._id,
+                    _source: hit._source
+                });
+            }
+        } while ((results = await results.next()));
+        await this.tearDown();
+        progressBar.stop();
+        this.writeStream.end();
+        return waitWrite;
+    }
+}
+class JSONLDumper extends AbstractDumper {
+    constructor() {
+        super(...arguments);
+        this.ndjsonStream = ndjson_1.default.stringify();
+    }
+    async setup() {
+        this.ndjsonStream.on('data', (line) => {
+            if (!this.writeStream) {
+                throw new Error('Cannot write data: WriteStream is not initialized.');
+            }
+            this.writeStream.write(line);
+        });
+        return;
+    }
+    async writeHeader() {
+        await this.writeLine({ type: 'collection', index: this.index, collection: this.collection });
+    }
+    writeLine(content) {
         return new Promise(resolve => {
-            if (ndjsonStream.write(content)) {
+            if (this.ndjsonStream.write(content)) {
                 resolve(undefined);
             }
             else {
-                ndjsonStream.once('drain', resolve);
+                this.ndjsonStream.once('drain', resolve);
             }
         });
-    };
-    writeStream.on('error', error => {
-        throw error;
-    });
-    ndjsonStream.on('data', (line) => writeStream.write(line));
-    fs_1.default.mkdirSync(collectionDir, { recursive: true });
-    if (format.toLowerCase() === 'jsonl') {
-        await writeLine({ type: 'collection', index, collection });
-    }
-    let results = await sdk.document.search(index, collection, { query }, options);
-    const progressBar = cli_ux_1.default.progress({
-        format: `Dumping ${collection} |{bar}| {percentage}% || {value}/{total} documents`
-    });
-    progressBar.start(results.total, 0);
-    const rawDocuments = {
-        [index]: {
-            [collection]: []
-        }
-    };
-    do {
-        progressBar.update(results.fetched);
-        for (const hit of results.hits) {
-            let document = null;
-            if (format.toLocaleLowerCase() === 'jsonl') {
-                document = {
-                    _id: hit._id,
-                    body: hit._source
-                };
-                await writeLine(document);
+    }
+    onResult(document) {
+        return this.writeLine({ _id: document._id, body: document._source });
+    }
+    get fileExtension() {
+        return 'jsonl';
+    }
+}
+class KuzzleDumper extends JSONLDumper {
+    constructor() {
+        super(...arguments);
+        this.rawDocuments = {
+            [this.index]: {
+                [this.collection]: []
             }
-            else {
-                rawDocuments[index][collection].push({
-                    index: {
-                        _id: hit._id
-                    }
-                });
-                rawDocuments[index][collection].push(hit._source);
+        };
+    }
+    get fileExtension() {
+        return 'json';
+    }
+    async onResult(document) {
+        this.rawDocuments[this.index][this.collection].push({
+            index: {
+                _id: document._id
+            }
+        });
+        this.rawDocuments[this.index][this.collection].push(document._source);
+        return;
+    }
+    tearDown() {
+        return this.writeLine(this.rawDocuments);
+    }
+}
+class CSVDumper extends AbstractDumper {
+    constructor(sdk, index, collection, batchSize, destPath, query = {}, fields, separator = ',') {
+        super(sdk, index, collection, batchSize, destPath, query);
+        this.fields = fields;
+        this.separator = separator;
+    }
+    get fileExtension() {
+        return 'csv';
+    }
+    async setup() {
+        if (!this.fields.length) {
+            // If no field has been selected, then all fields are selected.
+            const mappings = await this.sdk.collection.getMapping(this.index, this.collection);
+            if (!mappings.properties) {
+                return;
             }
+            this.fields = Object.keys(flattenObject(mappings.properties));
+        }
+        else {
+            // Delete '_id' from the selected fields, since IDs are
+            // _always_ exported.
+            if (this.fields.includes('_id')) {
+                this.fields.splice(this.fields.indexOf('_id'), 1);
+            }
+        }
+    }
+    writeHeader() {
+        return this.writeLine(['_id', ...this.fields].join(this.separator));
+    }
+    async writeLine(line) {
+        if (!this.writeStream) {
+            throw new Error('Cannot write data: WriteStream is not initialized.');
         }
-    } while ((results = await results.next()));
-    if (format.toLocaleLowerCase() === 'kuzzle') {
-        await writeLine(rawDocuments);
-    }
-    progressBar.stop();
-    ndjsonStream.end();
-    writeStream.end();
-    return waitWrite;
+        this.writeStream.write(line);
+        this.writeStream.write('\n');
+        return;
+    }
+    onResult(document) {
+        const values = [document._id, ...(0, common_1.pickValues)(document._source, this.fields)];
+        return this.writeLine(values.join(this.separator));
+    }
+}
+async function dumpCollectionData(sdk, index, collection, batchSize, destPath, query = {}, format = 'jsonl', fields = []) {
+    let dumper;
+    switch (format.toLowerCase()) {
+        case 'jsonl':
+            dumper = new JSONLDumper(sdk, index, collection, batchSize, destPath, query);
+            return dumper.dump();
+        case 'csv':
+            dumper = new CSVDumper(sdk, index, collection, batchSize, destPath, query, fields);
+            return dumper.dump();
+        default:
+            dumper = new KuzzleDumper(sdk, index, collection, batchSize, destPath, query);
+            return dumper.dump();
+    }
 }
 exports.dumpCollectionData = dumpCollectionData;
 async function dumpCollectionMappings(sdk, index, collection, destPath, format = 'jsonl') {