Add multi-get support as a transform stream

Author: hmalphettes

README.md

@@ -27,7 +27,7 @@ var bulkExec = function(bulkCmds, callback) {
 var ws = new WritableBulk(bulkExec);
 
 // stream 42 random records into ES
-require('random-document-stream')(42).pipe(ws);
+require('random-document-stream')(42).pipe(ws).on('finish', done);
 ```
 
 ## Stream search results from Elasticsearch
@@ -53,7 +53,7 @@ ws._write = function(chunk, enc, next) {
   next();
 };
 
-rs.pipe(ws);
+rs.pipe(ws).on('finish', done);
 ```
 
 If we want to start the stream at an offset and define a limit:
@@ -103,13 +103,37 @@ rs = new ReadableSearch(scrollExec);
 
 ## Stream IDs into Elasticsearch multi-get and get documents out.
 ```
-# TODO a duplex stream
+var mgetExec = function(docs, callback) {
+  client.mget({
+    index: 'myindex',
+    type: 'mytype',
+    body: {
+      docs: { ids: docs }
+    }
+  }, callback);
+};
+ts = new PipableDocs(mgetExec, 4);
+
+// Naive read stream of 12 ids that are numbers
+var rs = new Readable({objectMode: true});
+rs._read = function() {
+  for (var i = 0; i < 12; i++) {
+    rs.push(i);
+  }
+  rs.push(null);
+};
+
+var ws = new require('stream').Writable({objectMode:true});
+ws._write = function(chunk, enc, next) {
+  console.log(hit._id + ' found: ' + hit._found, hit);
+  next();
+};
+
+rs.pipe(ts).pipe(ws).on('finish', onFinish);
 ```
 
 ## TODO
 ### Short term
-* Document more
-* Multi-get as a duplex
 * Bulk document errors should be emitted as errors
 
 ## Later

index.js

@@ -1,4 +1,5 @@
 module.exports = {
-  WritableBulk: require('./lib/writable-bulk'),
-  ReadableSearch: require('./lib/readable-search')
+  WritableBulk   : require('./lib/writable-bulk'),
+  ReadableSearch : require('./lib/readable-search'),
+  PipableDocs    : require('./lib/transform-mget')
 };

lib/transform-mget.js

@@ -0,0 +1,50 @@
+/**
+ * Stream document IDs and transform them into Documents.
+ */
+'use strict';
+var Transform = require('stream').Transform;
+
+module.exports = PipableDocs;
+
+/**
+ * @param mgetExec an executable query functions that takes 2 arguments:
+ *       a list of document IDs and a callback.
+ * @param bulkSize number of bulk commands executed at once. 128 by default.
+ */
+function PipableDocs(mgetExec, corkSize) {
+  Transform.call(this, {objectMode:true});
+  this.mgetExec = mgetExec;
+  this.corkSize = corkSize;
+
+// current mget
+  this._docIds = [];
+}
+
+PipableDocs.prototype = Object.create(Transform.prototype, {constructor: {value: PipableDocs}});
+
+PipableDocs.prototype._transform = function(chunk, encoding, callback) {
+  this._docIds.push(chunk);
+  if (this._docIds.length < this.corkSize) {
+    return callback();
+  }
+  this._flush(callback);
+};
+
+PipableDocs.prototype._flush = function(callback) {
+  if (!this._docIds.length) {
+    return callback();
+  }
+  var docIds = this._docIds;
+  this._docIds = [];
+  var self = this;
+  this.mgetExec(docIds, function(e, resp) {
+    if (e) {
+      self.emit('error', e);
+      return callback();
+    }
+    for (var i = 0; i < resp.docs.length; i++) {
+      self.push(resp.docs[i]);
+    }
+    callback();
+  });
+};

package.json

@@ -1,7 +1,7 @@
 {
   "name": "elasticsearch-streams",
   "description": "Stream in and out of Elasticsearch",
-  "version": "0.0.2",
+  "version": "0.0.3",
   "repository": {
     "type": "git",
     "url": "https://github.com/hmalphettes/elasticsearch-streamer.git"

test/test-search.js

@@ -1,8 +1,10 @@
 'use strict';
 var expect = require('chai').expect;
 var ReadableSearch = require('..').ReadableSearch;
+var PipableDocs = require('..').PipableDocs;
 var Writable = require('stream').Writable;
-var client = new require('elasticsearch').Client();
+var Readable = require('stream').Readable;
+var client = new require('elasticsearch').Client({log: 'warning'});
 var random = require('random-document-stream');
 
 describe('When searching', function() {
@@ -25,7 +27,7 @@ describe('When searching', function() {
     });
   });
   it('Must find 42 records by searching them', function(done) {
-    checkRecords(rs, 42, done);
+    checkRecords(rs, null, 42, done);
   });
 });
 
@@ -60,18 +62,54 @@ describe('When scrolling', function() {
     });
   });
   it('Must find the 42 records by scroll', function(done) {
-    checkRecords(rs, 42, done);
+    checkRecords(rs, null, 42, done);
+  });
+});
+
+describe('When getting documents by id', function() {
+  var ids;
+  var ts;
+  before(function(done) {
+    ids = populateIndex(42, function(e) {
+      if (e) { return done(e); }
+      expect(ids.length).to.equal(42);
+      var mgetExec = function(docs, callback) {
+        client.mget({
+          index: 'myindex',
+          type: 'mytype',
+          body: {
+            docs: { ids: docs }
+          }
+        }, callback);
+      };
+      ts = new PipableDocs(mgetExec, 4);
+      done();
+    });
+  });
+  it('Must pipe', function(done) {
+    var i = -1;
+    var rs = new Readable({objectMode: true});
+    rs._read = function() {
+      i++;
+      if (i >= ids.length) {
+        return rs.push(null);
+      }
+      rs.push(ids[i]);
+    };
+    checkRecords(rs, ts, 42, done);
   });
 });
 
 function populateIndex(nb, done) {
+  var ids = [];
   client.indices.delete({index:'myindex'}, function() {
     var cmds = [];
     var generator = random(0);
     for (var i = 0; i < nb; i++) {
       var rec = generator.makeJunk();
       cmds.push({ index:  { _index: 'myindex', _type: 'mytype', _id: rec._id } });
       cmds.push(rec);
+      ids.push(rec._id);
     }
     client.bulk({
       body: cmds
@@ -80,28 +118,37 @@ function populateIndex(nb, done) {
       client.indices.refresh({index: 'myindex'}, done);
     });
   });
+  return ids;
 }
 
-function checkRecords(rs, nb, done) {
+function checkRecords(rs, ts, nb, done) {
   var hits = [];
   var err;
   var ws = new Writable({objectMode:true});
   ws._write = function(chunk, enc, next) {
     hits.push(chunk);
     expect(chunk._index).to.equal('myindex');
     expect(chunk._type).to.equal('mytype');
-    expect(chunk._score).to.equal(1);
     expect(chunk._id).to.exist;
     expect(chunk._source.name).to.exist;
     next();
   };
-  rs.on('error', function(e) {
+  var errClosure = function(e) {
     err = e;
-  });
-  rs.on('end', function() {
+  };
+  rs.on('error', errClosure);
+  ws.on('error', errClosure);
+  if (ts) {
+    ts.on('error', errClosure);
+  }
+  function onFinish() {
     if (err) { return done(err); }
     expect(hits.length).to.equal(nb);
     done();
-  });
-  rs.pipe(ws);
+  }
+  if (ts) {
+    rs.pipe(ts).pipe(ws).on('finish', onFinish);
+  } else {
+    rs.pipe(ws).on('finish', onFinish);
+  }
 }