bigtable/table/createReadStream: allow early ending

Author: stephenplusplus

packages/bigtable/package.json

@@ -59,6 +59,7 @@
     "extend": "^3.0.0",
     "google-proto-files": "^0.11.0",
     "is": "^3.0.1",
+    "is-stream-ended": "^0.1.2",
     "lodash.flatten": "^4.2.0",
     "node-int64": "^0.4.0",
     "prop-assign": "^1.0.0",

packages/bigtable/src/table.js

@@ -26,6 +26,7 @@ var commonGrpc = require('@google-cloud/common-grpc');
 var concat = require('concat-stream');
 var flatten = require('lodash.flatten');
 var is = require('is');
+var isStreamEnded = require('is-stream-ended');
 var propAssign = require('prop-assign');
 var pumpify = require('pumpify');
 var through = require('through2');
@@ -494,24 +495,27 @@ Table.prototype.createReadStream = function(options) {
     reqOpts.rowsLimit = options.limit;
   }
 
-  return pumpify.obj([
+  var stream = pumpify.obj([
     this.requestStream(grpcOpts, reqOpts),
+
     through.obj(function(data, enc, next) {
-      var throughStream = this;
-      var rows = Row.formatChunks_(data.chunks, {
-        decode: options.decode
-      });
+      var transformStream = this;
 
-      rows.forEach(function(rowData) {
-        var row = self.row(rowData.key);
+      Row.formatChunks_(data.chunks, { decode: options.decode })
+        .forEach(function(rowData) {
+          var row = self.row(rowData.key);
+          row.data = rowData.data;
 
-        row.data = rowData.data;
-        throughStream.push(row);
-      });
+          if (!isStreamEnded(stream)) {
+            transformStream.push(row);
+          }
+        });
 
       next();
     })
   ]);
+
+  return stream;
 };
 
 /**

packages/bigtable/system-test/bigtable.js

@@ -379,37 +379,36 @@ describe('Bigtable', function() {
   });
 
   describe('rows', function() {
+    var ROWS = [{
+      key: 'gwashington',
+      data: {
+        follows: {
+          jadams: 1
+        }
+      }
+    }, {
+      key: 'tjefferson',
+      data: {
+        follows: {
+          gwashington: 1,
+          jadams: 1
+        }
+      }
+    }, {
+      key: 'jadams',
+      data: {
+        follows: {
+          gwashington: 1,
+          tjefferson: 1
+        }
+      }
+    }];
 
-    describe('inserting data', function() {
-
-      it('should insert rows', function(done) {
-        var rows = [{
-          key: 'gwashington',
-          data: {
-            follows: {
-              jadams: 1
-            }
-          }
-        }, {
-          key: 'tjefferson',
-          data: {
-            follows: {
-              gwashington: 1,
-              jadams: 1
-            }
-          }
-        }, {
-          key: 'jadams',
-          data: {
-            follows: {
-              gwashington: 1,
-              tjefferson: 1
-            }
-          }
-        }];
+    before(function(done) {
+      TABLE.insert(ROWS, done);
+    });
 
-        TABLE.insert(rows, done);
-      });
+    describe('inserting data', function() {
 
       it('should create an individual row', function(done) {
         var row = TABLE.row('alincoln');
@@ -510,7 +509,7 @@ describe('Bigtable', function() {
       it('should get rows', function(done) {
         TABLE.getRows(function(err, rows) {
           assert.ifError(err);
-          assert.strictEqual(rows.length, 4);
+          assert(rows.length >= ROWS.length);
           assert(rows[0] instanceof Row);
           done();
         });
@@ -524,9 +523,10 @@ describe('Bigtable', function() {
           .on('data', function(row) {
             assert(row instanceof Row);
             rows.push(row);
+            this.end();
           })
           .on('end', function() {
-            assert.strictEqual(rows.length, 4);
+            assert.strictEqual(rows.length, 1);
             done();
           });
       });

packages/bigtable/test/table.js

@@ -636,6 +636,21 @@ describe('Bigtable/Table', function() {
             done();
           });
       });
+
+      it('should allow a stream to end early', function(done) {
+        var rows = [];
+
+        table.createReadStream()
+          .on('error', done)
+          .on('data', function(row) {
+            rows.push(row);
+            this.end();
+          })
+          .on('end', function() {
+            assert.strictEqual(rows.length, 1);
+            done();
+          });
+      });
     });
 
     describe('error', function() {