Merge pull request #1 from tilfin/fix/throttle

Update JSONTransform

Author: tilfin

lib/index.js

@@ -1,10 +1,12 @@
 'use strict';
 
 const KLReadStream = require('./read_stream');
-
+const JSONTransform = require('./json_transform');
 
 exports.reader = function(event, opts) {
   return new KLReadStream(event, opts);
 }
 
-exports.parseJSON = require('./json_transform');
+exports.parseJSON = function(opts) {
+  return new JSONTransform(opts);
+}

lib/json_transform.js

@@ -1,40 +1,69 @@
 'use strict';
 
-const through2 = require('through2');
-
-
-module.exports = function(options) {
-  const opts = options || {};
-  const countBy = opts.countBy || 1;
-  const expanding = opts.expandArray || false;
-
-  let items = [];
-
-  const stream = through2({ objectMode: true },
-    function(chunk, enc, cb) {
-      try {
-        const data = JSON.parse(chunk);
-        if (expanding && (data instanceof Array)) {
-          items = items.concat(data);
-        } else {
-          items.push(data);
-        }
-      } catch (err) {
-        cb(err, null);
-        return;
-      }
-
-      while (items.length >= countBy) {
-        const enqItems = items.splice(0, countBy);
-        this.push(countBy > 1 ? enqItems : enqItems[0]);
-      }
-      cb();
-    }, function(cb) {
-      if (items.length) {
-        this.push(items);
-      }
-      cb();
-    });
-
-  return stream;
+const Transform = require('stream').Transform;
+const util = require('util');
+
+
+/**
+ * JSONTransform
+ *
+ * @param {Object} opts - Transform options
+ * @param {Number} opts.countBy - number of block items
+ * @param {Boolean} opts.expandArray - Flatten an array into items
+ */
+function JSONTransform(opts) {
+  if (!(this instanceof JSONTransform))
+    return new JSONTransform(opts);
+
+  const opts_ = opts || {};
+  this._countBy = opts_.countBy || 1;
+  this._expanding = opts_.expandArray || false;
+  delete opts_.countBy;
+  delete opts_.expandArray;
+
+  this._items = [];
+
+  opts_.objectMode = true;
+  Transform.call(this, opts_);
 }
+util.inherits(JSONTransform, Transform);
+
+JSONTransform.prototype._transform = function(chunk, encoding, cb) {
+  try {
+    const data = JSON.parse(chunk);
+    if (this._expanding && (data instanceof Array)) {
+      this._items = this._items.concat(data);
+    } else {
+      this._items.push(data);
+    }
+  } catch (err) {
+    cb(err, null);
+    return;
+  }
+
+  while (this._items.length >= this._countBy) {
+    if (!this._enqueue(this._shiftItems())) break;
+  }
+  cb();
+}
+
+JSONTransform.prototype._flush = function(cb) {
+  const items = this._items;
+
+  while (items.length >= this._countBy) {
+    this._enqueue(this._shiftItems());
+  }
+
+  if (items.length) this._enqueue(items);
+  cb();
+}
+
+JSONTransform.prototype._enqueue = function(items) {
+  return this.push(this._countBy > 1 ? items : items[0]);
+}
+
+JSONTransform.prototype._shiftItems = function() {
+  return this._items.splice(0, this._countBy);
+}
+
+module.exports = JSONTransform;

package.json

@@ -18,9 +18,6 @@
     "type": "git",
     "url": "https://github.com/tilfin/kinesis-stream-lambda"
   },
-  "dependencies": {
-    "through2": "^2.0.1"
-  },
   "devDependencies": {
     "aws-kinesis-agg": "^2.2.2",
     "chai": "^3.5.0",

test/fixtures/data/data_json.txt

@@ -0,0 +1 @@
+[{"id":1},{"id":2},{"id":3},{"id":4},{"id":5},{"id":6},{"id":7},{"id":8},{"id":9}]

test/json_stream.test.js

@@ -1,6 +1,7 @@
 'use strict';
 
 const fs = require('fs');
+const es = require('event-stream');
 const chai = require('chai');
 const assert = chai.assert;
 
@@ -25,6 +26,28 @@ describe('KinesisLambda.parseJSON', () => {
     });
   });
 
+  context('passed 9 items JSON with count by 2 and expanding array', () => {
+    it('flush 5 items', (done) => {
+      const readStream = fs.createReadStream(__dirname + '/fixtures/data/data_json.txt');
+      const jsonStream = JSONTransform({ expandArray: true, countBy: 2 })
+        .on('error', function(err) {
+          assert.ifError(err);
+          done();
+        });
+
+      const writeStream = es.writeArray(function (err, array) {
+        assert.deepEqual(array[0], [{ id: 1 }, { id: 2 }]);
+        assert.deepEqual(array[1], [{ id: 3 }, { id: 4 }]);
+        assert.deepEqual(array[2], [{ id: 5 }, { id: 6 }]);
+        assert.deepEqual(array[3], [{ id: 7 }, { id: 8 }]);
+        assert.deepEqual(array[4], [{ id: 9 }]);
+        done();
+      });
+
+      readStream.pipe(jsonStream).pipe(writeStream);
+    });
+  });
+
   context('passed invalid JSON', () => {
     it('raises error event', (done) => {
       const readStream = fs.createReadStream(__dirname + '/fixtures/data/invalid_json.txt');