diff --git a/lib/datastore/request.js b/lib/datastore/request.js index 7907189b035..30155f60e71 100644 --- 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 index 149c91801f1..8dfe75a141f 100644 --- 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 index 85a9a0e4035..6a990b1c8dc 100644 --- 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.');