core (stream-router): allow setting maxApiCalls (#1332)

* core (stream-router): allow setting maxApiCalls

* move logic to iterator

* iterator -> limiter

* re-arrange stream-events use

* tests.

* lint

* re-sort tests so overrides dont mess with docs

* clean up maxApiCalls property

* add system test for bigquery

* docs

* explain interceptors quirk

lib/bigquery/dataset.js

@@ -279,6 +279,9 @@ Dataset.prototype.delete = function(options, callback) {
  * @resource [Tables: list API Documentation]{@link https://cloud.google.com/bigquery/docs/reference/v2/tables/list}
  *
  * @param {object=} query - Configuration object.
+ * @param {boolean} query.autoPaginate - Have pagination handled automatically.
+ *     Default: true.
+ * @param {number} query.maxApiCalls - Maximum number of API calls to make.
  * @param {number} query.maxResults - Maximum number of results to return.
  * @param {string} query.pageToken - Token returned from a previous call, to
  *     request the next page of results.

lib/bigquery/index.js

@@ -178,6 +178,7 @@ BigQuery.prototype.dataset = function(id) {
  * @param {boolean} query.all - List all datasets, including hidden ones.
  * @param {boolean} query.autoPaginate - Have pagination handled automatically.
  *     Default: true.
+ * @param {number} query.maxApiCalls - Maximum number of API calls to make.
  * @param {number} query.maxResults - Maximum number of results to return.
  * @param {string} query.pageToken - Token returned from a previous call, to
  *     request the next page of results.
@@ -279,6 +280,7 @@ BigQuery.prototype.getDatasets = function(query, callback) {
  *     project.
  * @param {boolean} options.autoPaginate - Have pagination handled
  *     automatically. Default: true.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number=} options.maxResults - Maximum number of results to return.
  * @param {string=} options.pageToken - Token returned from a previous call, to
  *     request the next page of results.
@@ -407,6 +409,7 @@ BigQuery.prototype.job = function(id) {
  * @param {string|object} options - A string SQL query or configuration object.
  * @param {boolean} options.autoPaginate - Have pagination handled
  *     automatically. Default: true.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of results to read.
  * @param {string} options.query - A query string, following the BigQuery query
  *     syntax, of the query to execute.

lib/bigquery/job.js

@@ -211,6 +211,7 @@ Job.prototype.cancel = function(callback) {
  * @param {object=} options - Configuration object.
  * @param {boolean} options.autoPaginate - Have pagination handled
  *     automatically. Default: true.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of results to read.
  * @param {string} options.pageToken - Page token, returned by a previous call,
  *     to request the next page of results. Note: This is automatically added to

lib/bigquery/table.js

@@ -607,6 +607,7 @@ Table.prototype.export = function(destination, options, callback) {
  * @param {object=} options - The configuration object.
  * @param {boolean} options.autoPaginate - Have pagination handled
  *     automatically. Default: true.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of results to return.
  * @param {function} callback - The callback function.
  * @param {?error} callback.err - An error returned while making this request

lib/common/stream-router.js

@@ -22,10 +22,15 @@
 
 var arrify = require('arrify');
 var concat = require('concat-stream');
-var split = require('split-array-stream');
+var extend = require('extend');
 var is = require('is');
-var streamEvents = require('stream-events');
-var through = require('through2');
+var split = require('split-array-stream');
+
+/**
+ * @type {module:common/util}
+ * @private
+ */
+var util = require('./util.js');
 
 /*! Developer Documentation
  *
@@ -72,9 +77,10 @@ streamRouter.extend = function(Class, methodNames) {
  */
 streamRouter.parseArguments_ = function(args) {
   var query;
-  var callback;
-  var maxResults = -1;
   var autoPaginate = true;
+  var maxApiCalls = -1;
+  var maxResults = -1;
+  var callback;
 
   var firstArgument = args[0];
   var lastArgument = args[args.length - 1];
@@ -90,6 +96,8 @@ streamRouter.parseArguments_ = function(args) {
   }
 
   if (is.object(query)) {
+    query = extend(true, {}, query);
+
     // Check if the user only asked for a certain amount of results.
     if (is.number(query.maxResults)) {
       // `maxResults` is used API-wide.
@@ -99,6 +107,11 @@ streamRouter.parseArguments_ = function(args) {
       maxResults = query.pageSize;
     }
 
+    if (is.number(query.maxApiCalls)) {
+      maxApiCalls = query.maxApiCalls;
+      delete query.maxApiCalls;
+    }
+
     if (callback &&
         (maxResults !== -1 || // The user specified a limit.
         query.autoPaginate === false)) {
@@ -108,9 +121,10 @@ streamRouter.parseArguments_ = function(args) {
 
   return {
     query: query || {},
-    callback: callback,
+    autoPaginate: autoPaginate,
+    maxApiCalls: maxApiCalls,
     maxResults: maxResults,
-    autoPaginate: autoPaginate
+    callback: callback
   };
 };
 
@@ -126,6 +140,7 @@ streamRouter.parseArguments_ = function(args) {
  *     string in some places.
  * @param {function=} parsedArguments.callback - Callback function.
  * @param {boolean} parsedArguments.autoPaginate - Auto-pagination enabled.
+ * @param {boolean} parsedArguments.maxApiCalls - Maximum API calls to make.
  * @param {number} parsedArguments.maxResults - Maximum results to return.
  * @param {function} originalMethod - The cached method that accepts a callback
  *     and returns `nextQuery` to receive more results.
@@ -163,6 +178,7 @@ streamRouter.router_ = function(parsedArguments, originalMethod) {
  *     string in some places.
  * @param {function=} parsedArguments.callback - Callback function.
  * @param {boolean} parsedArguments.autoPaginate - Auto-pagination enabled.
+ * @param {boolean} parsedArguments.maxApiCalls - Maximum API calls to make.
  * @param {number} parsedArguments.maxResults - Maximum results to return.
  * @param {function} originalMethod - The cached method that accepts a callback
  *     and returns `nextQuery` to receive more results.
@@ -172,12 +188,20 @@ streamRouter.runAsStream_ = function(parsedArguments, originalMethod) {
   var query = parsedArguments.query;
   var resultsToSend = parsedArguments.maxResults;
 
-  var stream = streamEvents(through.obj());
+  var limiter = util.createLimiter(makeRequest, {
+    maxApiCalls: parsedArguments.maxApiCalls
+  });
+
+  var stream = limiter.stream;
 
   stream.once('reading', function() {
-    originalMethod(query, onResultSet);
+    makeRequest(query);
   });
 
+  function makeRequest(query) {
+    originalMethod(query, onResultSet);
+  }
+
   function onResultSet(err, results, nextQuery) {
     if (err) {
       stream.destroy(err);
@@ -196,15 +220,15 @@ streamRouter.runAsStream_ = function(parsedArguments, originalMethod) {
       }
 
       if (nextQuery && resultsToSend !== 0) {
-        originalMethod(nextQuery, onResultSet);
+        limiter.makeRequest(nextQuery);
         return;
       }
 
       stream.push(null);
     });
   }
 
-  return stream;
+  return limiter.stream;
 };
 
 module.exports = streamRouter;

lib/common/util.js

@@ -33,6 +33,7 @@ var request = require('request').defaults({
   }
 });
 var retryRequest = require('retry-request');
+var streamEvents = require('stream-events');
 var through = require('through2');
 var uniq = require('array-uniq');
 
@@ -545,3 +546,43 @@ function normalizeArguments(globalContext, localConfig, options) {
 }
 
 util.normalizeArguments = normalizeArguments;
+
+/**
+ * Limit requests according to a `maxApiCalls` limit.
+ *
+ * @param {function} makeRequestFn - The function that will be called.
+ * @param {object=} options - Configuration object.
+ * @param {number} options.maxApiCalls - The maximum number of API calls to
+ *     make.
+ */
+function createLimiter(makeRequestFn, options) {
+  var stream = streamEvents(through.obj());
+
+  var requestsMade = 0;
+  var requestsToMake = -1;
+
+  options = options || {};
+
+  if (is.number(options.maxApiCalls)) {
+    requestsToMake = options.maxApiCalls;
+  }
+
+  return {
+    makeRequest: function() {
+      requestsMade++;
+
+      if (requestsToMake >= 0 && requestsMade > requestsToMake) {
+        stream.push(null);
+        return;
+      }
+
+      makeRequestFn.apply(null, arguments);
+
+      return stream;
+    },
+
+    stream: stream
+  };
+}
+
+util.createLimiter = createLimiter;

lib/compute/index.js

@@ -600,6 +600,7 @@ Compute.prototype.firewall = function(name) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of addresses to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -716,6 +717,7 @@ Compute.prototype.getAddresses = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of addresses to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -834,6 +836,7 @@ Compute.prototype.getAutoscalers = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of disks to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -947,6 +950,7 @@ Compute.prototype.getDisks = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of instance groups to
  *     return.
  * @param {string} options.pageToken - A previously-returned page token
@@ -1062,6 +1066,7 @@ Compute.prototype.getInstanceGroups = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of firewalls to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1170,6 +1175,7 @@ Compute.prototype.getFirewalls = function(options, callback) {
  *       can be a regular expression.
  * @param {boolean} options.https - List only HTTPs health checks. Default:
  *     `false`.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of networks to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1280,6 +1286,7 @@ Compute.prototype.getHealthChecks = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of networks to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1385,6 +1392,7 @@ Compute.prototype.getNetworks = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of operations to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1490,6 +1498,7 @@ Compute.prototype.getOperations = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of instances to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1592,6 +1601,7 @@ Compute.prototype.getRegions = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of rules to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1696,6 +1706,7 @@ Compute.prototype.getRules = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of snapshots to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1801,6 +1812,7 @@ Compute.prototype.getServices = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of snapshots to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -1906,6 +1918,7 @@ Compute.prototype.getSnapshots = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of instances to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -2019,6 +2032,7 @@ Compute.prototype.getVMs = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of instances to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.

lib/compute/instance-group.js

@@ -273,6 +273,8 @@ InstanceGroup.prototype.delete = function(callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
+ * @param {number} options.maxResults - Maximum number of VMs to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
  * @param {boolean} options.running - Only return instances which are running.

lib/compute/network.js

@@ -278,6 +278,7 @@ Network.prototype.firewall = function(name) {
  * @param {object=} options - Firewall search options.
  * @param {boolean} options.autoPaginate - Have pagination handled
  *     automatically. Default: true.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of firewalls to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.

lib/compute/region.js

@@ -282,6 +282,7 @@ Region.prototype.createRule = function(name, config, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of addresses to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -387,6 +388,7 @@ Region.prototype.getAddresses = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of operations to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.
@@ -491,6 +493,7 @@ Region.prototype.getOperations = function(options, callback) {
  *       (not equal)
  *     - **`filterString`**: the string to filter to. For string fields, this
  *       can be a regular expression.
+ * @param {number} options.maxApiCalls - Maximum number of API calls to make.
  * @param {number} options.maxResults - Maximum number of rules to return.
  * @param {string} options.pageToken - A previously-returned page token
  *     representing part of the larger set of results to view.