Fixes for paged iterations. (#126)

- Fix #123: introduces maxResults parameter
- Fix #124: fixed the behaviors of the streaming
  it seems better to be back of using through2.

Author: jmuk

packages/packages/google-gax/lib/gax.js

@@ -51,11 +51,15 @@
  * @property {RetryOptions=} retry - determines whether and how to retry
  *   on transient errors. When set to null, the call will not retry.
  * @property {boolean=} autoPaginate - If set to false and the call is
- *   configured for page streaming, page streaming is not performed, instead
+ *   configured for paged iteration, page unrolling is not performed, instead
  *   the callback will be called with the response object.
  * @property {Object=} pageToken - If set and the call is configured for
- *   page streaming, page streaming is not performed and requested with this
+ *   paged iteration, paged iteration is not performed and requested with this
  *   pageToken.
+ * @property {number} maxResults - If set and the call is configured for
+ *   paged iteration, the call will stop when the number of response elements
+ *   reaches to the specified size. By default, it will unroll the page to
+ *   the end of the list.
  * @property {boolean=} isBundling - If set to false and the call is configured
  *   for bundling, bundling is not performed.
  * @property {BackoffSettings=} longrunning - BackoffSettings used for polling.
@@ -164,6 +168,7 @@ function CallSettings(settings) {
   this.autoPaginate =
       ('autoPaginate' in settings) ? settings.autoPaginate : true;
   this.pageToken = settings.pageToken;
+  this.maxResults = settings.maxResults;
   this.otherArgs = settings.otherArgs || {};
   this.bundleOptions = settings.bundleOptions;
   this.isBundling = ('isBundling' in settings) ? settings.isBundling : true;
@@ -189,6 +194,7 @@ CallSettings.prototype.merge = function merge(options) {
   var retry = this.retry;
   var autoPaginate = this.autoPaginate;
   var pageToken = this.pageToken;
+  var maxResults = this.maxResults;
   var otherArgs = this.otherArgs;
   var isBundling = this.isBundling;
   var longrunning = this.longrunning;
@@ -209,6 +215,10 @@ CallSettings.prototype.merge = function merge(options) {
     pageToken = options.pageToken;
   }
 
+  if ('maxResults' in options) {
+    maxResults = options.maxResults;
+  }
+
   if ('otherArgs' in options) {
     otherArgs = {};
     for (var key in this.otherArgs) {
@@ -243,6 +253,7 @@ CallSettings.prototype.merge = function merge(options) {
     longrunning: longrunning,
     autoPaginate: autoPaginate,
     pageToken: pageToken,
+    maxResults: maxResults,
     otherArgs: otherArgs,
     isBundling: isBundling,
     promise: promise

packages/packages/google-gax/lib/paged_iteration.js

@@ -34,7 +34,8 @@ var extend = require('extend');
 var nextTick = require('process-nextick-args');
 var util = require('util');
 var NormalApiCaller = require('./api_callable').NormalApiCaller;
-var ReadableStream = require('readable-stream');
+var through2 = require('through2');
+var ended = require('is-stream-ended');
 
 /**
  * Creates an API caller that returns a stream to performs page-streaming.
@@ -97,49 +98,31 @@ PagedIteration.prototype.call = function(
     return;
   }
 
+  var maxResults = settings.maxResults || -1;
   var allResources = [];
   function pushResources(err, resources, next) {
     if (err) {
       canceller.callback(err);
       return;
     }
 
-    allResources.push.apply(allResources, resources);
+    for (var i = 0; i < resources.length; ++i) {
+      allResources.push(resources[i]);
+      if (allResources.length === maxResults) {
+        next = null;
+        break;
+      }
+    }
     if (!next) {
       canceller.callback(null, allResources);
       return;
     }
-    nextTick(apiCall, argument, pushResources);
+    nextTick(apiCall, next, pushResources);
   }
 
   nextTick(apiCall, argument, pushResources);
 };
 
-/**
- * An implementation of readalbe stream which fits for the usage of paged iteration.
- * @private
- * @constructor
- */
-function PagedStream() {
-  ReadableStream.call(this, {objectMode: true});
-}
-
-util.inherits(PagedStream, ReadableStream);
-
-PagedStream.prototype._read = function(n) {
-};
-
-/**
- */
-PagedStream.prototype.end = function() {
-  // pushing a null will cause ending the stream.
-  this.push(null);
-
-  // onEof callback of ReadableStream does not update 'readable' field immediately,
-  // thus settings here explicitly.
-  this.readable = false;
-};
-
 /**
  * Describes the structure of a page-streaming call.
  *
@@ -174,24 +157,33 @@ exports.PageDescriptor = PageDescriptor;
  */
 PageDescriptor.prototype.createStream = function(
     apiCall, request, options) {
-  var stream = new PagedStream();
-  options = extend({}, options);
-  options.autoPaginate = false;
-  function callback(response) {
-    var resources = response[0];
+  var stream = through2.obj();
+  options = extend({}, options, {autoPaginate: false});
+  var maxResults = ('maxResults' in options) ? options.maxResults : -1;
+  var pushCount = 0;
+  var started = false;
+  function callback(err, resources, next) {
+    if (err) {
+      stream.emit('error', err);
+      return;
+    }
     for (var i = 0; i < resources.length; ++i) {
-      if (!stream.readable) {
+      if (ended(stream)) {
         return;
       }
       if (resources[i] === null) {
         continue;
       }
       stream.push(resources[i]);
+      pushCount++;
+      if (pushCount === maxResults) {
+        stream.end();
+      }
     }
-    if (!stream.readable) {
+    if (ended(stream)) {
       return;
     }
-    if (!response[1]) {
+    if (!next) {
       stream.end();
       return;
     }
@@ -200,17 +192,19 @@ PageDescriptor.prototype.createStream = function(
     if ('pageToken' in options) {
       delete options.pageToken;
     }
-    return apiCall(response[1], options)
-        .then(callback)
-        .catch(function(err) {
-          stream.emit('error', err);
-        });
+    if (stream.isPaused()) {
+      request = next;
+      started = false;
+    } else {
+      nextTick(apiCall, next, options, callback);
+    }
   }
-  apiCall(request, options)
-      .then(callback)
-      .catch(function(err) {
-        stream.emit('error', err);
-      });
+  stream.on('resume', function() {
+    if (!started) {
+      started = true;
+      apiCall(request, options, callback);
+    }
+  });
   return stream;
 };
 

packages/packages/google-gax/package.json

@@ -11,9 +11,11 @@
     "google-auto-auth": "^0.5.2",
     "google-proto-files": "^0.9.1",
     "grpc": "^1.1",
+    "is-stream-ended": "^0.1.0",
     "lodash": "^4.17.2",
     "process-nextick-args": "^1.0.7",
-    "readable-stream": "^2.2.2"
+    "readable-stream": "^2.2.2",
+    "through2": "^2.0.3"
   },
   "devDependencies": {
     "chai": "*",
@@ -24,8 +26,9 @@
     "jsdoc": "~3.4.0",
     "mocha": "~2.2.5",
     "pegjs": "~0.9.0",
+    "pumpify": "^1.3.5",
     "sinon": "~1.17.3",
-    "through2": "~2.0.1"
+    "stream-events": "^1.0.1"
   },
   "scripts": {
     "codecov": "istanbul cover _mocha -- --reporter spec --slow 500 --timeout 5000 && codecov",

packages/packages/google-gax/test/paged_iteration.js

@@ -35,7 +35,11 @@
 var util = require('./utils');
 var PageDescriptor = require('../lib/paged_iteration').PageDescriptor;
 var expect = require('chai').expect;
+var process = require('process');
+var pumpify = require('pumpify');
 var sinon = require('sinon');
+var streamEvents = require('stream-events');
+var through2 = require('through2');
 
 describe('paged iteration', function() {
   var pageSize = 3;
@@ -155,6 +159,22 @@ describe('paged iteration', function() {
     }).catch(done);
   });
 
+  it('caps the results by maxResults', function() {
+    var spy = sinon.spy(func);
+    var apiCall = util.createApiCall(spy, createOptions);
+    return apiCall({}, {maxResults: pageSize * 2 + 2}).then(function(response) {
+      expect(response).to.be.an('array');
+      expect(response[0]).to.be.an('array');
+      expect(response[0].length).to.eq(pageSize * 2 + 2);
+      var expected = 0;
+      for (var i = 0; i < response[0].length; ++i) {
+        expect(response[0][i]).to.eq(expected);
+        expected++;
+      }
+      expect(spy.callCount).to.eq(3);
+    });
+  });
+
   describe('stream conversion', function() {
     var spy;
     var apiCall;
@@ -171,7 +191,10 @@ describe('paged iteration', function() {
       }).on('end', function() {
         onEnd();
         done();
-      }).on('error', done);
+      }).on('error', function(err) {
+        console.error(err);
+        done(err);
+      });
     }
 
     it('returns a stream', function(done) {
@@ -200,5 +223,59 @@ describe('paged iteration', function() {
         expect(spy.callCount).to.eq(pagesToStream);
       }, done, pageSize);
     });
+
+    it('caps the elements by maxResults', function(done) {
+      var onData = sinon.spy();
+      var stream = descriptor.createStream(
+        apiCall, {}, {maxResults: pageSize * 2 + 2});
+      stream.on('data', onData);
+      streamChecker(stream, function() {
+        expect(spy.callCount).to.eq(3);
+        expect(onData.callCount).to.eq(pageSize * 2 + 2);
+      }, done, 0);
+    });
+
+    it('does not call API eagerly', function(done) {
+      var stream = descriptor.createStream(apiCall, {}, null);
+      setTimeout(function() {
+        expect(spy.callCount).to.eq(0);
+        streamChecker(stream, function() {
+          expect(spy.callCount).to.eq(pagesToStream + 1);
+        }, done, 0);
+      }, 50);
+    });
+
+    it('does not start calls when it is already started', function(done) {
+      var stream = descriptor.createStream(apiCall, {}, null);
+      stream.on('end', function() {
+        expect(spy.callCount).to.eq(pagesToStream + 1);
+        done();
+      });
+      stream.resume();
+      process.nextTick(function() {
+        stream.pause();
+        stream.resume();
+      });
+    });
+
+    it('cooperates with google-cloud-node usage', function(done) {
+      var stream;
+      var output = streamEvents(pumpify.obj());
+      output.once('reading', function() {
+        stream = descriptor.createStream(apiCall, {}, null);
+        output.setPipeline(stream, through2.obj());
+      });
+      var count = 0;
+      output.on('data', function() {
+        count++;
+        if (count === pageSize + 1) {
+          output.end();
+        }
+      }).on('end', function() {
+        expect(count).to.eq(pageSize + 1);
+        expect(spy.callCount).to.eq(2);
+        done();
+      }).on('error', done);
+    });
   });
 });