node-es-transformer 1.0.0-beta1 → 1.0.0-beta3

package/dist/node-es-transformer.esm.js

@@ -1,10 +1,20 @@
 import fs from 'fs';
 import es from 'event-stream';
 import glob from 'glob';
+import split from 'split2';
+import { Readable } from 'stream';
 import cliProgress from 'cli-progress';
 import elasticsearch from '@elastic/elasticsearch';
 
-var DEFAULT_BUFFER_SIZE = 1000;
+// In earlier versions this was used to set the number of docs to index in a
+// single bulk request. Since we switched to use the helpers.bulk() method from
+// the ES client, this now translates to the `flushBytes` option of the helper.
+// However, for kind of a backwards compability with the old values, this uses
+// KBytes instead of Bytes. It will be multiplied by 1024 in the index queue.
+var DEFAULT_BUFFER_SIZE = 5120;
+
+// The default number of docs to fetch in a single search request when reindexing.
+var DEFAULT_SEARCH_SIZE = 1000;
 
 function createMappingFactory(ref) {
   var sourceClient = ref.sourceClient;
@@ -15,6 +25,7 @@ function createMappingFactory(ref) {
   var mappingsOverride = ref.mappingsOverride;
   var indexMappingTotalFieldsLimit = ref.indexMappingTotalFieldsLimit;
   var verbose = ref.verbose;
+  var deleteIndex = ref.deleteIndex;
 
   return async function () {
     var targetMappings = mappingsOverride ? undefined : mappings;
@@ -24,7 +35,14 @@ function createMappingFactory(ref) {
         var mapping = await sourceClient.indices.getMapping({
           index: sourceIndexName,
         });
-        targetMappings = mapping[sourceIndexName].mappings;
+        if (mapping[sourceIndexName]) {
+          targetMappings = mapping[sourceIndexName].mappings;
+        } else {
+          var allMappings = Object.values(mapping);
+          if (allMappings.length > 0) {
+            targetMappings = Object.values(mapping)[0].mappings;
+          }
+        }
       } catch (err) {
         console.log('Error reading source mapping', err);
         return;
@@ -39,18 +57,28 @@ function createMappingFactory(ref) {
       }
 
       try {
-        var resp = await targetClient.indices.create({
-          index: targetIndexName,
-          body: Object.assign({}, {mappings: targetMappings},
-            (indexMappingTotalFieldsLimit !== undefined
-              ? {
-                  settings: {
-                    'index.mapping.total_fields.limit': indexMappingTotalFieldsLimit,
-                  },
-                }
-              : {})),
-        });
-        if (verbose) { console.log('Created target mapping', resp); }
+        var indexExists = await targetClient.indices.exists({ index: targetIndexName });
+
+        if (indexExists === true && deleteIndex === true) {
+          await targetClient.indices.delete({ index: targetIndexName });
+        }
+
+        if (indexExists === false || deleteIndex === true) {
+          var resp = await targetClient.indices.create({
+            index: targetIndexName,
+            body: Object.assign({}, {mappings: targetMappings},
+              (indexMappingTotalFieldsLimit !== undefined
+                ? {
+                    settings: {
+                      'index.mapping.total_fields.limit': indexMappingTotalFieldsLimit,
+                      'index.number_of_shards': 1,
+                      'index.number_of_replicas': 0,
+                    },
+                  }
+                : {})),
+          });
+          if (verbose) { console.log('Created target mapping', resp); }
+        }
       } catch (err) {
         console.log('Error creating target mapping', err);
       }
@@ -58,22 +86,28 @@ function createMappingFactory(ref) {
   };
 }
 
-var MAX_QUEUE_SIZE = 15;
-
 function fileReaderFactory(indexer, fileName, transform, splitRegex, verbose) {
   function startIndex(files) {
-    var ingestQueueSize = 0;
     var finished = false;
 
     var file = files.shift();
     var s = fs
       .createReadStream(file)
-      .pipe(es.split(splitRegex))
+      .pipe(split(splitRegex))
       .pipe(
         es
           .mapSync(function (line) {
             try {
-              var doc = typeof transform === 'function' ? transform(line) : line;
+              // skip empty lines
+              if (line === '') {
+                return;
+              }
+
+              var doc =
+                typeof transform === 'function'
+                  ? JSON.stringify(transform(JSON.parse(line)))
+                  : line;
+
               // if doc is undefined we'll skip indexing it
               if (typeof doc === 'undefined') {
                 s.resume();
@@ -107,20 +141,13 @@ function fileReaderFactory(indexer, fileName, transform, splitRegex, verbose) {
           })
       );
 
-    indexer.queueEmitter.on('queue-size', async function (size) {
+    indexer.queueEmitter.on('pause', function () {
       if (finished) { return; }
-      ingestQueueSize = size;
-
-      if (ingestQueueSize < MAX_QUEUE_SIZE) {
-        s.resume();
-      } else {
-        s.pause();
-      }
+      s.pause();
     });
 
     indexer.queueEmitter.on('resume', function () {
       if (finished) { return; }
-      ingestQueueSize = 0;
       s.resume();
     });
   }
@@ -136,7 +163,7 @@ var EventEmitter = require('events');
 
 var queueEmitter = new EventEmitter();
 
-var parallelCalls = 1;
+var parallelCalls = 5;
 
 // a simple helper queue to bulk index documents
 function indexQueueFactory(ref) {
@@ -146,78 +173,76 @@ function indexQueueFactory(ref) {
   var skipHeader = ref.skipHeader; if ( skipHeader === void 0 ) skipHeader = false;
   var verbose = ref.verbose; if ( verbose === void 0 ) verbose = true;
 
-  var buffer = [];
-  var queue = [];
-  var ingesting = 0;
-  var ingestTimes = [];
-  var finished = false;
+  var flushBytes = bufferSize * 1024; // Convert KB to Bytes
+  var highWaterMark = flushBytes * parallelCalls;
 
-  var ingest = function (b) {
-    if (typeof b !== 'undefined') {
-      queue.push(b);
-      queueEmitter.emit('queue-size', queue.length);
-    }
+  // Create a Readable stream
+  var stream = new Readable({
+    read: function read() {}, // Implement read but we manage pushing manually
+    highWaterMark: highWaterMark, // Buffer size for backpressure management
+  });
 
-    if (ingestTimes.length > 5) { ingestTimes = ingestTimes.slice(-5); }
+  async function* ndjsonStreamIterator(readableStream) {
+    var buffer = ''; // To hold the incomplete data
+    var skippedHeader = false;
 
-    if (ingesting < parallelCalls) {
-      var docs = queue.shift();
+    // Iterate over the stream using async iteration
+    for await (var chunk of readableStream) {
+      buffer += chunk.toString(); // Accumulate the chunk data in the buffer
 
-      queueEmitter.emit('queue-size', queue.length);
-      if (queue.length <= 5) {
-        queueEmitter.emit('resume');
-      }
+      // Split the buffer into lines (NDJSON items)
+      var lines = buffer.split('\n');
 
-      ingesting += 1;
-
-      if (verbose)
-        { console.log(("bulk ingest docs: " + (docs.length / 2) + ", queue length: " + (queue.length))); }
-
-      var start = Date.now();
-      client
-        .bulk({ body: docs })
-        .then(function () {
-          var end = Date.now();
-          var delta = end - start;
-          ingestTimes.push(delta);
-          ingesting -= 1;
-
-          var ingestTimesMovingAverage =
-            ingestTimes.length > 0
-              ? ingestTimes.reduce(function (p, c) { return p + c; }, 0) / ingestTimes.length
-              : 0;
-          var ingestTimesMovingAverageSeconds = Math.floor(ingestTimesMovingAverage / 1000);
-
-          if (
-            ingestTimes.length > 0 &&
-            ingestTimesMovingAverageSeconds < 30 &&
-            parallelCalls < 10
-          ) {
-            parallelCalls += 1;
-          } else if (
-            ingestTimes.length > 0 &&
-            ingestTimesMovingAverageSeconds >= 30 &&
-            parallelCalls > 1
-          ) {
-            parallelCalls -= 1;
-          }
+      // The last line might be incomplete, so hold it back in the buffer
+      buffer = lines.pop();
 
-          if (queue.length > 0) {
-            ingest();
-          } else if (queue.length === 0 && finished) {
-            queueEmitter.emit('finish');
-          }
-        })
-        .catch(function (error) {
-          console.error(error);
-          ingesting -= 1;
-          parallelCalls = 1;
-          if (queue.length > 0) {
-            ingest();
+      // Yield each complete JSON object
+      for (var line of lines) {
+        if (line.trim()) {
+          try {
+            if (!skipHeader || (skipHeader && !skippedHeader)) {
+              yield JSON.parse(line); // Parse and yield the JSON object
+              skippedHeader = true;
+            }
+          } catch (err) {
+            // Handle JSON parse errors if necessary
+            console.error('Failed to parse JSON:', err);
           }
-        });
+        }
+      }
     }
-  };
+
+    // Handle any remaining data in the buffer after the stream ends
+    if (buffer.trim()) {
+      try {
+        yield JSON.parse(buffer);
+      } catch (err) {
+        console.error('Failed to parse final JSON:', err);
+      }
+    }
+  }
+
+  var finished = false;
+
+  // Async IIFE to start bulk indexing
+  (async function () {
+    console.log('START BULK INDEXING');
+    await client.helpers.bulk({
+      concurrency: parallelCalls,
+      flushBytes: flushBytes,
+      flushInterval: 1000,
+      refreshOnCompletion: true,
+      datasource: ndjsonStreamIterator(stream),
+      onDocument: function onDocument(doc) {
+        return {
+          index: { _index: targetIndexName },
+        };
+      },
+    });
+    console.log('FINISHED BULK INDEXING');
+
+    queueEmitter.emit('finish');
+  })();
 
   return {
     add: function (doc) {
@@ -225,37 +250,22 @@ function indexQueueFactory(ref) {
         throw new Error('Unexpected doc added after indexer should finish.');
       }
 
-      if (!skipHeader) {
-        var header = { index: { _index: targetIndexName } };
-        buffer.push(header);
-      }
-      buffer.push(doc);
-
-      if (queue.length === 0) {
-        queueEmitter.emit('resume');
-      }
-
-      if (buffer.length >= bufferSize * 2) {
-        ingest(buffer);
-        buffer = [];
+      var canContinue = stream.push(((JSON.stringify(doc)) + "\n"));
+      if (!canContinue) {
+        queueEmitter.emit('pause');
+        stream.once('drain', function () {
+          queueEmitter.emit('resume');
+        });
       }
     },
     finish: function () {
       finished = true;
-
-      if (buffer.length > 0) {
-        ingest(buffer);
-        buffer = [];
-      } else if (queue.length === 0 && ingesting === 0) {
-        queueEmitter.emit('finish');
-      }
+      stream.push(null);
     },
     queueEmitter: queueEmitter,
   };
 }
 
-var MAX_QUEUE_SIZE$1 = 15;
-
 // create a new progress bar instance and use shades_classic theme
 var progressBar = new cliProgress.SingleBar({}, cliProgress.Presets.shades_classic);
 
@@ -265,40 +275,69 @@ function indexReaderFactory(
   transform,
   client,
   query,
-  bufferSize
+  searchSize,
+  populatedFields
 ) {
-  if ( bufferSize === void 0 ) bufferSize = DEFAULT_BUFFER_SIZE;
+  if ( searchSize === void 0 ) searchSize = DEFAULT_SEARCH_SIZE;
+  if ( populatedFields === void 0 ) populatedFields = false;
 
   return async function indexReader() {
-    var responseQueue = [];
     var docsNum = 0;
+    var scrollId;
+    var finished = false;
+    var readActive = false;
+    var backPressurePause = false;
 
-    function search() {
-      return client.search({
-        index: sourceIndexName,
-        scroll: '30s',
-        size: bufferSize,
-        query: query,
-      });
+    async function fetchPopulatedFields() {
+      try {
+        var response = await client.search({
+          index: sourceIndexName,
+          size: searchSize,
+          query: {
+            function_score: {
+              query: query,
+              random_score: {},
+            },
+          },
+        });
+
+        // Get all field names for each returned doc and flatten it
+        // to a list of unique field names used across all docs.
+        return Array.from(new Set(response.hits.hits.map(function (d) { return Object.keys(d._source); }).flat(1)));
+      } catch (e) {
+        console.log('error', e);
+      }
+    }
+
+    function search(fields) {
+      return client.search(Object.assign({}, {index: sourceIndexName,
+        scroll: '600s',
+        size: searchSize,
+        query: query},
+        (fields ? { _source: fields } : {})));
     }
 
     function scroll(id) {
       return client.scroll({
         scroll_id: id,
-        scroll: '30s',
+        scroll: '600s',
       });
     }
 
-    // start things off by searching, setting a scroll timeout, and pushing
-    // our first response into the queue to be processed
-    var se = await search();
-    responseQueue.push(se);
-    progressBar.start(se.hits.total.value, 0);
+    var fieldsWithData;
+
+    // identify populated fields
+    if (populatedFields) {
+      fieldsWithData = await fetchPopulatedFields();
+    }
+
+    await fetchNextResponse();
 
     function processHit(hit) {
       docsNum += 1;
       try {
         var doc = typeof transform === 'function' ? transform(hit._source) : hit._source; // eslint-disable-line no-underscore-dangle
+
         // if doc is undefined we'll skip indexing it
         if (typeof doc === 'undefined') {
           return;
@@ -317,68 +356,117 @@ function indexReaderFactory(
       }
     }
 
-    var ingestQueueSize = 0;
-    var scrollId = se._scroll_id; // eslint-disable-line no-underscore-dangle
-    var readActive = false;
+    async function fetchNextResponse() {
+      readActive = true;
 
-    async function processResponseQueue() {
-      while (responseQueue.length) {
-        readActive = true;
-        var response = responseQueue.shift();
+      var sc = scrollId ? await scroll(scrollId) : await search(fieldsWithData);
 
-        // collect the docs from this response
-        response.hits.hits.forEach(processHit);
+      if (!scrollId) {
+        progressBar.start(sc.hits.total.value, 0);
+      }
 
-        progressBar.update(docsNum);
+      scrollId = sc._scroll_id;
+      readActive = false;
 
-        // check to see if we have collected all of the docs
-        if (response.hits.total.value === docsNum) {
-          indexer.finish();
-          break;
-        }
-
-        if (ingestQueueSize < MAX_QUEUE_SIZE$1) {
-          // get the next response if there are more docs to fetch
-          var sc = await scroll(response._scroll_id); // eslint-disable-line no-await-in-loop,no-underscore-dangle,max-len
-          scrollId = sc._scroll_id; // eslint-disable-line no-underscore-dangle
-          responseQueue.push(sc);
-        } else {
-          readActive = false;
-        }
-      }
+      processResponse(sc);
     }
 
-    indexer.queueEmitter.on('queue-size', async function (size) {
-      ingestQueueSize = size;
+    async function processResponse(response) {
+      // collect the docs from this response
+      response.hits.hits.forEach(processHit);
 
-      if (!readActive && ingestQueueSize < MAX_QUEUE_SIZE$1) {
-        // get the next response if there are more docs to fetch
-        var sc = await scroll(scrollId); // eslint-disable-line no-await-in-loop,no-underscore-dangle,max-len
-        scrollId = sc._scroll_id; // eslint-disable-line no-underscore-dangle
-        responseQueue.push(sc);
-        processResponseQueue();
+      progressBar.update(docsNum);
+
+      // check to see if we have collected all of the docs
+      if (response.hits.total.value === docsNum) {
+        indexer.finish();
+        return;
       }
+
+      if (!backPressurePause) {
+        await fetchNextResponse();
+      }
+    }
+
+    indexer.queueEmitter.on('pause', async function () {
+      backPressurePause = true;
     });
 
     indexer.queueEmitter.on('resume', async function () {
-      ingestQueueSize = 0;
+      backPressurePause = false;
 
-      if (readActive) {
+      if (readActive || finished) {
         return;
       }
 
-      // get the next response if there are more docs to fetch
-      var sc = await scroll(scrollId); // eslint-disable-line no-await-in-loop,no-underscore-dangle,max-len
-      scrollId = sc._scroll_id; // eslint-disable-line no-underscore-dangle
-      responseQueue.push(sc);
-      processResponseQueue();
+      await fetchNextResponse();
     });
 
     indexer.queueEmitter.on('finish', function () {
+      finished = true;
       progressBar.stop();
     });
+  };
+}
 
-    processResponseQueue();
+function streamReaderFactory(indexer, stream, transform, splitRegex, verbose) {
+  function startIndex() {
+    console.log('START INDEX', splitRegex);
+    var finished = false;
+
+    var s = stream.pipe(split(splitRegex)).pipe(
+      es
+        .mapSync(function (line) {
+          try {
+            // skip empty lines
+            if (line === '') {
+              return;
+            }
+
+            var doc =
+              typeof transform === 'function' ? JSON.stringify(transform(JSON.parse(line))) : line;
+
+            // if doc is undefined we'll skip indexing it
+            if (typeof doc === 'undefined') {
+              s.resume();
+              return;
+            }
+
+            // the transform callback may return an array of docs so we can emit
+            // multiple docs from a single line
+            if (Array.isArray(doc)) {
+              doc.forEach(function (d) { return indexer.add(d); });
+              return;
+            }
+
+            indexer.add(doc);
+          } catch (e) {
+            console.log('error', e);
+          }
+        })
+        .on('error', function (err) {
+          console.log('Error while reading file.', err);
+        })
+        .on('end', function () {
+          if (verbose) { console.log('Read entire stream.'); }
+          indexer.finish();
+          finished = true;
+        })
+    );
+
+    indexer.queueEmitter.on('pause', function () {
+      if (finished) { return; }
+      s.pause();
+    });
+
+    indexer.queueEmitter.on('resume', function () {
+      if (finished) { return; }
+      s.resume();
+    });
+  }
+
+  return function () {
+    startIndex();
   };
 }
 
@@ -387,6 +475,8 @@ async function transformer(ref) {
   var sourceClientConfig = ref.sourceClientConfig;
   var targetClientConfig = ref.targetClientConfig;
   var bufferSize = ref.bufferSize; if ( bufferSize === void 0 ) bufferSize = DEFAULT_BUFFER_SIZE;
+  var searchSize = ref.searchSize; if ( searchSize === void 0 ) searchSize = DEFAULT_SEARCH_SIZE;
+  var stream = ref.stream;
   var fileName = ref.fileName;
   var splitRegex = ref.splitRegex; if ( splitRegex === void 0 ) splitRegex = /\n/;
   var sourceIndexName = ref.sourceIndexName;
@@ -394,11 +484,13 @@ async function transformer(ref) {
   var mappings = ref.mappings;
   var mappingsOverride = ref.mappingsOverride; if ( mappingsOverride === void 0 ) mappingsOverride = false;
   var indexMappingTotalFieldsLimit = ref.indexMappingTotalFieldsLimit;
+  var populatedFields = ref.populatedFields; if ( populatedFields === void 0 ) populatedFields = false;
   var query = ref.query;
   var skipHeader = ref.skipHeader; if ( skipHeader === void 0 ) skipHeader = false;
   var transform = ref.transform;
   var verbose = ref.verbose; if ( verbose === void 0 ) verbose = true;
 
+  console.log('TRANSFORMER');
   if (typeof targetIndexName === 'undefined') {
     throw Error('targetIndexName must be specified.');
   }
@@ -421,6 +513,7 @@ async function transformer(ref) {
     mappingsOverride: mappingsOverride,
     indexMappingTotalFieldsLimit: indexMappingTotalFieldsLimit,
     verbose: verbose,
+    deleteIndex: deleteIndex,
   });
   var indexer = indexQueueFactory({
     targetClient: targetClient,
@@ -435,8 +528,12 @@ async function transformer(ref) {
       throw Error('Only either one of fileName or sourceIndexName can be specified.');
     }
 
-    if (typeof fileName === 'undefined' && typeof sourceIndexName === 'undefined') {
-      throw Error('Either fileName or sourceIndexName must be specified.');
+    if (
+      (typeof fileName !== 'undefined' && typeof sourceIndexName !== 'undefined') ||
+      (typeof fileName !== 'undefined' && typeof stream !== 'undefined') ||
+      (typeof sourceIndexName !== 'undefined' && typeof stream !== 'undefined')
+    ) {
+      throw Error('Only one of fileName, sourceIndexName, or stream can be specified.');
     }
 
     if (typeof fileName !== 'undefined') {
@@ -450,17 +547,25 @@ async function transformer(ref) {
         transform,
         sourceClient,
         query,
-        bufferSize
+        searchSize,
+        populatedFields
       );
     }
 
+    if (typeof stream !== 'undefined') {
+      console.log('STREAM READER');
+      return streamReaderFactory(indexer, stream, transform, splitRegex, verbose);
+    }
+
     return null;
   }
 
   var reader = getReader();
+  console.log('READER INITIALIZED');
 
   try {
     var indexExists = await targetClient.indices.exists({ index: targetIndexName });
+    console.log('INDEX EXISTS', indexExists);
 
     if (indexExists === false) {
       await createMapping();