Refs #5125: Fix LimitDocumentSource coalesce & extras

lib/pipeline/documentSources/CursorDocumentSource.js

@@ -1,6 +1,7 @@
 "use strict";
 
-var DocumentSource = require('./DocumentSource');
+var DocumentSource = require('./DocumentSource'),
+	LimitDocumentSource = require('./LimitDocumentSource');
 
 // Mimicking max memory size from mongo/db/query/new_find.cpp
 // Need to actually decide some size for this?
@@ -98,11 +99,13 @@ proto.getNext = function getNext(callback) {
 };
 
 proto.coalesce = function coalesce(nextSource) {
-	if (!this._limit) {
+	if (this._limit) {
+		return this._limit.coalesce(nextSource);
+	} else if (nextSource instanceof LimitDocumentSource) {
 		this._limit = nextSource;
 		return this._limit;
 	} else {
-		return this._limit.coalesce(nextSource);
+		return false;
 	}
 };
 
@@ -203,7 +206,7 @@ proto.serialize = function serialize(explain) {
 
 // LimitDocumentSource has the setLimit function which trickles down to any documentsource
 proto.getLimit = function getLimit() {
-	return this._limit;
+	return this._limit ? this._limit.getLimit() : -1;
 };
 
 //----------------private--------------

lib/pipeline/documentSources/LimitDocumentSource.js

@@ -86,3 +86,21 @@ klass.createFromJson = function createFromJson(jsonElement, ctx) {
 
 	return nextLimit;
 };
+
+proto.getLimit = function getLimit(newLimit) {
+	return this.limit;
+};
+
+proto.setLimit = function setLimit(newLimit) {
+	this.limit = newLimit;
+};
+
+proto.getDependencies = function(deps) {
+	return DocumentSource.GetDepsReturn.SEE_NEXT;
+};
+
+proto.serialize = function(explain) {
+	var out = {};
+	out[this.getSourceName()] = this.limit;
+	return out;
+};

test/lib/pipeline/documentSources/CursorDocumentSource.js

@@ -2,6 +2,8 @@
 var assert = require("assert"),
 	async = require("async"),
 	CursorDocumentSource = require("../../../../lib/pipeline/documentSources/CursorDocumentSource"),
+	LimitDocumentSource = require("../../../../lib/pipeline/documentSources/LimitDocumentSource"),
+	SkipDocumentSource = require("../../../../lib/pipeline/documentSources/SkipDocumentSource"),
 	Cursor = require("../../../../lib/Cursor");
 
 
@@ -25,6 +27,28 @@ module.exports = {
 			}
 		},
 
+		"#coalesce": {
+			"should be able to coalesce a limit into itself": function (){
+				var cwc = new CursorDocumentSource.CursorWithContext();
+				cwc._cursor = new Cursor( [] );
+
+				var lds = new LimitDocumentSource();
+				lds.limit = 1;
+
+				var cds = new CursorDocumentSource(cwc);
+				assert.equal(cds.coalesce(lds) instanceof LimitDocumentSource, true);
+			},
+			"should leave non-limit alone": function () {
+				var cwc = new CursorDocumentSource.CursorWithContext();
+				cwc._cursor = new Cursor( [] );
+
+				var sds = new SkipDocumentSource(),
+					cds = new CursorDocumentSource(cwc);
+
+				assert.equal(cds.coalesce(sds), false);
+			}
+		},
+
 		"#getNext": {
 			"should return the current cursor value sync": function(){
 				var cwc = new CursorDocumentSource.CursorWithContext();