datastore: runQuery(stream mode) - use limit(n) logically - fixes #296

googleapis · Nov 10, 2014 · b091d03 · b091d03
1 parent 6022e19
commit b091d03
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 15 deletions.
diff --git a/lib/datastore/request.js b/lib/datastore/request.js
@@ -336,6 +336,7 @@ DatastoreRequest.prototype.delete = function(keys, callback) {
 DatastoreRequest.prototype.runQuery = function(q, callback) {
   var that = this;
   var stream;
+  var resultsToSend = q.limitVal;
 
   var req = {
     read_options: {},
@@ -353,7 +354,6 @@ DatastoreRequest.prototype.runQuery = function(q, callback) {
     stream.once('reading', runQuery);
     return stream;
   } else {
-    callback = callback || util.noop;
     runQuery();
   }
 
@@ -376,21 +376,29 @@ DatastoreRequest.prototype.runQuery = function(q, callback) {
         cursor = resp.batch.end_cursor.toBase64();
       }
 
-      if (stream) {
-        if (cursor && entities.length > 0) {
-          entities.forEach(function (entity) {
-            stream.push(entity);
-          });
+      if (!stream) {
+        callback(null, entities, cursor);
+        return;
+      }
 
-          req.query = entity.queryToQueryProto(q.start(cursor).offset(0));
+      if (!cursor || entities.length === 0) {
+        stream.end();
+        return;
+      }
 
-          runQuery();
-        } else {
-          stream.end();
-        }
-      } else {
-        callback(null, entities, cursor);
+      var result;
+      while ((result = entities.shift()) && resultsToSend !== 0) {
+        stream.push(result);
+        resultsToSend--;
       }
+
+      if (resultsToSend === 0) {
+        stream.end();
+        return;
+      }
+
+      req.query = entity.queryToQueryProto(q.start(cursor).offset(0));
+      runQuery();
     });
   }
 };

diff --git a/regression/datastore.js b/regression/datastore.js
@@ -249,8 +249,7 @@ describe('datastore', function() {
     });
 
     it('should run a query as a stream', function(done) {
-      var q = ds.createQuery('Character').hasAncestor(ancestor)
-          .limit(5);
+      var q = ds.createQuery('Character').hasAncestor(ancestor);
 
       var resultsReturned = 0;
 
@@ -263,6 +262,21 @@ describe('datastore', function() {
         });
     });
 
+    it('should not go over a limit with a stream', function(done) {
+      var limit = 3;
+      var q = ds.createQuery('Character').hasAncestor(ancestor).limit(limit);
+
+      var resultsReturned = 0;
+
+      ds.runQuery(q)
+        .on('error', done)
+        .on('data', function() { resultsReturned++; })
+        .on('end', function() {
+          assert.equal(resultsReturned, limit);
+          done();
+        });
+    });
+
     it('should filter queries with simple indexes', function(done) {
       var q = ds.createQuery('Character').hasAncestor(ancestor)
           .filter('appearances >=', 20);

diff --git a/test/datastore/request.js b/test/datastore/request.js
@@ -387,6 +387,25 @@ describe('Request', function() {
           });
       });
 
+      it('should only emit the limited number of results', function(done) {
+        var limit = 2;
+
+        query.limitVal = limit;
+
+        request.makeReq_ = function(method, req, callback) {
+          callback(null, mockResponse.withResultsAndEndCursor);
+        };
+
+        var resultsReturned = 0;
+
+        request.runQuery(query)
+          .on('data', function() { resultsReturned++; })
+          .on('end', function() {
+            assert.equal(resultsReturned, limit);
+            done();
+          });
+      });
+
       it('should emit an error', function(done) {
         var error = new Error('Error.');