datastore: fix runQuery end logic

lib/datastore/request.js

@@ -442,17 +442,19 @@ DatastoreRequest.prototype.runQuery = function(query, callback) {
     };
   }
 
+  var originalLimitVal = query.limitVal;
+  var entities = [];
+
   function onResponse(err, resp) {
     if (err) {
       callback(err, null, null, resp);
       return;
     }
 
-    var entities = [];
     var nextQuery = null;
 
     if (resp.batch.entityResults) {
-      entities = entity.formatArray(resp.batch.entityResults);
+      entities = entities.concat(entity.formatArray(resp.batch.entityResults));
     }
 
     var notFinished = resp.batch.moreResults === 'NOT_FINISHED';
@@ -469,11 +471,22 @@ DatastoreRequest.prototype.runQuery = function(query, callback) {
     if (notFinished) {
       // Run the query again to make sure all of the requested entities are
       // returned.
+      var limit = reqOpts.query.limit.value;
+      if (limit > -1) {
+        // Update the limit on the nextQuery to return only the amount of
+        // results originally asked for.
+        nextQuery.limit(limit - resp.batch.entityResults.length);
+      }
       reqOpts.query = entity.queryToQueryProto(nextQuery);
       self.request_(protoOpts, reqOpts, onResponse);
       return;
     }
 
+    if (nextQuery && originalLimitVal > -1) {
+      // Restore the original limit value for the query.
+      nextQuery.limit(originalLimitVal);
+    }
+
     callback(null, entities, nextQuery, resp);
   }
 

test/datastore/request.js

@@ -565,44 +565,62 @@ describe('Request', function() {
 
         request.runQuery({}, function(err, entities) {
           assert.ifError(err);
-          assert.strictEqual(entities, entityResults);
+          assert.deepEqual(entities, entityResults);
           done();
         });
       });
 
       it('should re-run query if not finished', function(done) {
-        entityOverrides.formatArray = util.noop;
-
+        var entityResults = {
+          1: ['a'],
+          2: ['b', 'c']
+        };
         var nextQuery;
-        var queryProto = {};
         var query = {
+          limitVal: 1,
           offsetVal: 8
         };
+        var queryProto = {
+          limit: {
+            value: query.limitVal
+          }
+        };
 
         var timesRequestCalled = 0;
         var startCalled = false;
         var offsetCalled = false;
 
+        entityOverrides.formatArray = function(array) {
+          assert.strictEqual(array, entityResults[timesRequestCalled]);
+          return entityResults[timesRequestCalled];
+        };
+
         request.request_ = function(protoOpts, reqOpts, callback) {
           timesRequestCalled++;
 
+          var resp = extend(true, {}, apiResponse);
+          resp.batch.entityResults = entityResults[timesRequestCalled];
+
           if (timesRequestCalled === 1) {
             assert.strictEqual(protoOpts.service, 'Datastore');
             assert.strictEqual(protoOpts.method, 'runQuery');
 
-            var resp = extend(true, {}, apiResponse);
             resp.batch.moreResults = 'NOT_FINISHED';
 
             callback(null, resp);
           } else {
             assert.strictEqual(startCalled, true);
             assert.strictEqual(offsetCalled, true);
             assert.strictEqual(reqOpts.query, queryProto);
-            done();
+
+            resp.batch.moreResults = 'MORE_RESULTS_AFTER_LIMIT';
+
+            callback(null, resp);
           }
         };
 
         FakeQuery.prototype.start = function(endCursor_) {
+          nextQuery = this;
           assert.strictEqual(endCursor_, endCursor);
           startCalled = true;
           return this;
@@ -615,14 +633,35 @@ describe('Request', function() {
           return this;
         };
 
+        FakeQuery.prototype.limit = function(limit_) {
+          if (timesRequestCalled === 1) {
+            assert.strictEqual(
+              limit_,
+              entityResults[1].length - query.limitVal
+            );
+          } else {
+            // Should restore the original limit.
+            assert.strictEqual(limit_, query.limitVal);
+          }
+          return this;
+        };
+
         entityOverrides.queryToQueryProto = function(query_) {
           if (timesRequestCalled > 1) {
             assert.strictEqual(query_, nextQuery);
           }
           return queryProto;
         };
 
-        request.runQuery(query, assert.ifError);
+        request.runQuery(query, function(err, entities, nextQuery_) {
+          assert.ifError(err);
+          assert.deepEqual(
+            entities,
+            [].slice.call(entityResults[1]).concat(entityResults[2])
+          );
+          assert.strictEqual(nextQuery_, nextQuery);
+          done();
+        });
       });
 
       it('should return nextQuery', function(done) {