SERVER-42748 Support using stored procedures (system.js) in map/reduce

7 files changed, 74 insertions, 33 deletions

diff --git a/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml b/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
index aa403c4b48d..7348dda6682 100644
--- a/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
+++ b/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
@@ -13,6 +13,7 @@ selector:
   - jstests/core/mr_outreduce.js
   - jstests/core/mr_outreduce2.js
   - jstests/core/mr_sort.js
+  - jstests/core/mr_stored.js
   - jstests/core/mr_tolerates_js_exception.js
   - jstests/core/mr_undef.js
   - jstests/core/mr4.js
diff --git a/jstests/core/mr_stored.js b/jstests/core/mr_stored.js
index d3ee89a5cf5..b1da36c1784 100644
--- a/jstests/core/mr_stored.js
+++ b/jstests/core/mr_stored.js
@@ -8,6 +8,9 @@
 //   requires_non_retryable_writes,
 //   uses_map_reduce_with_temp_collections,
 // ]
+/**
+ * Tests that map reduce works with stored javascript.
+ */
 (function() {
 "use strict";
 
@@ -28,7 +31,9 @@ let map = function(obj) {
     emit(obj.partner, {stats: [obj.visits]});
 };
 
-let reduce = function(k, v) {
+let notStoredMap = `function() {(${map.toString()})(this);}`;
+
+const reduce = function(k, v) {
     var stats = [];
     var total = 0;
     for (var i = 0; i < v.length; i++) {
@@ -40,52 +45,84 @@ let reduce = function(k, v) {
     return {stats: stats, total: total};
 };
 
-// Test that map reduce works with stored javascript
+const finalize = function(key, reducedValue) {
+    reducedValue.avg = reducedValue.total / reducedValue.stats.length;
+    return reducedValue;
+};
+
 assert.commandWorked(testDB.system.js.insert({_id: "mr_stored_map", value: map}));
 assert.commandWorked(testDB.system.js.insert({_id: "mr_stored_reduce", value: reduce}));
+assert.commandWorked(testDB.system.js.insert({_id: "mr_stored_finalize", value: finalize}));
 
 const out = testDB.mr_stored_out;
 
-assert.commandWorked(coll.mapReduce(
-    function() {
+function assertCorrect(results) {
+    assert.eq(2, Object.keySet(results).length);
+    assert.eq([9, 11, 30], results["1"].stats);
+    assert.eq([9, 41, 41], results["2"].stats);
+}
+
+// Stored Map.
+assert.commandWorked(testDB.runCommand({
+    mapReduce: coll.getName(),
+    map: function() {
         mr_stored_map(this);
     },
-    function(k, v) {
+    reduce: reduce,
+    finalize: finalize,
+    out: "mr_stored_out"
+}));
+
+assertCorrect(out.convertToSingleObject("value"));
+out.drop();
+
+// Stored Reduce.
+assert.commandWorked(testDB.runCommand({
+    mapReduce: coll.getName(),
+    map: notStoredMap,
+    reduce: function(k, v) {
         return mr_stored_reduce(k, v);
     },
-    {out: "mr_stored_out", scope: {xx: 1}}));
-
-let z = out.convertToSingleObject("value");
-assert.eq(2, Object.keySet(z).length);
-assert.eq([9, 11, 30], z["1"].stats);
-assert.eq([9, 41, 41], z["2"].stats);
+    finalize: finalize,
+    out: "mr_stored_out"
+}));
 
+assertCorrect(out.convertToSingleObject("value"));
 out.drop();
 
-map = function(obj) {
-    var x = "partner";
-    var y = "visits";
-    emit(obj[x], {stats: [obj[y]]});
-};
+// Stored Finalize.
+assert.commandWorked(testDB.runCommand({
+    mapReduce: coll.getName(),
+    map: notStoredMap,
+    reduce: reduce,
+    finalize: function(key, reducedValue) {
+        return mr_stored_finalize(key, reducedValue);
+    },
+    out: "mr_stored_out"
+}));
 
-assert.commandWorked(testDB.system.js.save({_id: "mr_stored_map", value: map}));
+assertCorrect(out.convertToSingleObject("value"));
+out.drop();
 
-assert.commandWorked(coll.mapReduce(
-    function() {
+// All Stored.
+assert.commandWorked(testDB.runCommand({
+    mapReduce: coll.getName(),
+    map: function() {
         mr_stored_map(this);
     },
-    function(k, v) {
+    reduce: function(k, v) {
         return mr_stored_reduce(k, v);
     },
-    {out: "mr_stored_out", scope: {xx: 1}}));
+    finalize: function(key, reducedValue) {
+        return mr_stored_finalize(key, reducedValue);
+    },
+    out: "mr_stored_out"
+}));
 
-z = out.convertToSingleObject("value");
-assert.eq(2, Object.keySet(z).length);
-assert.eq([9, 11, 30], z["1"].stats);
-assert.eq([9, 41, 41], z["2"].stats);
+assertCorrect(out.convertToSingleObject("value"));
+out.drop();
 
 assert.commandWorked(testDB.system.js.remove({_id: "mr_stored_map"}));
 assert.commandWorked(testDB.system.js.remove({_id: "mr_stored_reduce"}));
-
-out.drop();
+assert.commandWorked(testDB.system.js.remove({_id: "mr_stored_finalize"}));
 }());
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index 7e129f08d28..c48455935c0 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -229,7 +229,7 @@ public:
     auto getJsExecWithScope() const {
         RuntimeConstants runtimeConstants = getRuntimeConstants();
         const boost::optional<mongo::BSONObj>& scope = runtimeConstants.getJsScope();
-        return JsExecution::get(opCtx, scope.get_value_or(BSONObj()));
+        return JsExecution::get(opCtx, scope.get_value_or(BSONObj()), ns.db());
     }
 
     // The explain verbosity requested by the user, or boost::none if no explain was requested.
diff --git a/src/mongo/db/pipeline/javascript_execution.cpp b/src/mongo/db/pipeline/javascript_execution.cpp
index ac8e474fb9c..720fdad7f8e 100644
--- a/src/mongo/db/pipeline/javascript_execution.cpp
+++ b/src/mongo/db/pipeline/javascript_execution.cpp
@@ -37,10 +37,11 @@ namespace {
 const auto getExec = OperationContext::declareDecoration<std::unique_ptr<JsExecution>>();
 }  // namespace
 
-JsExecution* JsExecution::get(OperationContext* opCtx, const BSONObj& scope) {
+JsExecution* JsExecution::get(OperationContext* opCtx, const BSONObj& scope, StringData database) {
     auto& exec = getExec(opCtx);
     if (!exec) {
         exec = std::make_unique<JsExecution>(scope);
+        exec->getScope()->setLocalDB(database);
         exec->getScope()->loadStored(opCtx, true);
     }
     return exec.get();
diff --git a/src/mongo/db/pipeline/javascript_execution.h b/src/mongo/db/pipeline/javascript_execution.h
index b62a2d2df12..0a2d6f1219c 100644
--- a/src/mongo/db/pipeline/javascript_execution.h
+++ b/src/mongo/db/pipeline/javascript_execution.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/base/string_data.h"
 #include "mongo/db/client.h"
 #include "mongo/db/exec/document_value/value.h"
 #include "mongo/db/operation_context.h"
@@ -49,7 +50,7 @@ public:
      * Create or get a pointer to a JsExecution instance, capable of invoking Javascript functions
      * and reading the return value.
      */
-    static JsExecution* get(OperationContext* opCtx, const BSONObj& scope);
+    static JsExecution* get(OperationContext* opCtx, const BSONObj& scope, StringData database);
 
     /**
      * Construct with a thread-local scope and initialize with the given scope variables.
diff --git a/src/mongo/scripting/engine.cpp b/src/mongo/scripting/engine.cpp
index 149da7f5969..eba9c50c9be 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -421,7 +421,7 @@ public:
     void init(const BSONObj* data) {
         _real->init(data);
     }
-    void setLocalDB(const string& dbName) {
+    void setLocalDB(StringData dbName) {
         _real->setLocalDB(dbName);
     }
     void loadStored(OperationContext* opCtx, bool ignoreNotConnected = false) {
diff --git a/src/mongo/scripting/engine.h b/src/mongo/scripting/engine.h
index b50c75baed6..4f572cad026 100644
--- a/src/mongo/scripting/engine.h
+++ b/src/mongo/scripting/engine.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/base/string_data.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/service_context.h"
 #include "mongo/platform/atomic_word.h"
@@ -65,8 +66,8 @@ public:
     }
 
     virtual void externalSetup() = 0;
-    virtual void setLocalDB(const std::string& localDBName) {
-        _localDBName = localDBName;
+    virtual void setLocalDB(StringData localDBName) {
+        _localDBName = localDBName.toString();
     }
 
     virtual BSONObj getObject(const char* field) = 0;