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

12 files changed, 95 insertions, 44 deletions

diff --git a/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml b/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
index c8c64ceda8f..8d6f68bff9a 100644
--- a/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
+++ b/buildscripts/resmokeconfig/suites/core_map_reduce_agg.yaml
@@ -12,6 +12,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/mr_use_this_object.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/accumulator_js_test.cpp b/src/mongo/db/pipeline/accumulator_js_test.cpp
index 8efb460cf27..ff6d24c49ec 100644
--- a/src/mongo/db/pipeline/accumulator_js_test.cpp
+++ b/src/mongo/db/pipeline/accumulator_js_test.cpp
@@ -65,7 +65,7 @@ private:
 void MapReduceFixture::setUp() {
     setTestCommandsEnabled(true);
     ServiceContextMongoDTest::setUp();
-    ScriptEngine::setup();
+    ScriptEngine::setup(false);
 }
 
 void MapReduceFixture::tearDown() {
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index c0b44e88da6..e82d7bef9c5 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -242,7 +242,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/expression_javascript_test.cpp b/src/mongo/db/pipeline/expression_javascript_test.cpp
index 96ff7295ddc..e921d15fa65 100644
--- a/src/mongo/db/pipeline/expression_javascript_test.cpp
+++ b/src/mongo/db/pipeline/expression_javascript_test.cpp
@@ -73,7 +73,7 @@ private:
 void MapReduceFixture::setUp() {
     setTestCommandsEnabled(true);
     ServiceContextMongoDTest::setUp();
-    ScriptEngine::setup();
+    ScriptEngine::setup(false);
 }
 
 void MapReduceFixture::tearDown() {
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..0573084c548 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"
@@ -47,9 +48,10 @@ class JsExecution {
 public:
     /**
      * Create or get a pointer to a JsExecution instance, capable of invoking Javascript functions
-     * and reading the return value.
+     * and reading the return value. This will load all stored procedures from database unless
+     * 'disableLoadStored' is set on the global ScriptEngine.
      */
-    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..f677475246e 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -68,7 +68,8 @@ static std::unique_ptr<ScriptEngine> globalScriptEngine;
 
 }  // namespace
 
-ScriptEngine::ScriptEngine() : _scopeInitCallback() {}
+ScriptEngine::ScriptEngine(bool disableLoadStored)
+    : _disableLoadStored(disableLoadStored), _scopeInitCallback() {}
 
 ScriptEngine::~ScriptEngine() {}
 
@@ -209,6 +210,8 @@ void Scope::validateObjectIdString(const string& str) {
 }
 
 void Scope::loadStored(OperationContext* opCtx, bool ignoreNotConnected) {
+    if (!getGlobalScriptEngine()->_disableLoadStored)
+        return;
     if (_localDBName.size() == 0) {
         if (ignoreNotConnected)
             return;
@@ -421,7 +424,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..72ec298af9a 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;
@@ -209,7 +210,7 @@ class ScriptEngine : public KillOpListenerInterface {
     ScriptEngine& operator=(const ScriptEngine&) = delete;
 
 public:
-    ScriptEngine();
+    ScriptEngine(bool disableLoadStored);
     virtual ~ScriptEngine();
 
     virtual Scope* newScope() {
@@ -233,7 +234,12 @@ public:
     virtual int getJSHeapLimitMB() const = 0;
     virtual void setJSHeapLimitMB(int limit) = 0;
 
-    static void setup();
+    /**
+     * Calls the constructor for the Global ScriptEngine. 'disableLoadStored' causes future calls to
+     * the function Scope::loadStored(), which would otherwise load stored procedures, to be
+     * ignored.
+     */
+    static void setup(bool disableLoadStored = true);
     static void dropScopeCache();
 
     /** gets a scope from the pool or a new one if pool is empty
@@ -263,6 +269,7 @@ public:
     virtual void interruptAll() {}
 
     static std::string getInterpreterVersionString();
+    const bool _disableLoadStored;
 
 protected:
     virtual Scope* createScope() = 0;
diff --git a/src/mongo/scripting/engine_none.cpp b/src/mongo/scripting/engine_none.cpp
index 0262fbf24fe..a0726d76015 100644
--- a/src/mongo/scripting/engine_none.cpp
+++ b/src/mongo/scripting/engine_none.cpp
@@ -30,7 +30,7 @@
 #include "engine.h"
 
 namespace mongo {
-void ScriptEngine::setup() {
+void ScriptEngine::setup(bool disableLoadStored) {
     // noop
 }
 
diff --git a/src/mongo/scripting/mozjs/engine.cpp b/src/mongo/scripting/mozjs/engine.cpp
index c09c97a3f0b..4a8427d7dd9 100644
--- a/src/mongo/scripting/mozjs/engine.cpp
+++ b/src/mongo/scripting/mozjs/engine.cpp
@@ -47,11 +47,11 @@ void DisableExtraThreads();
 
 namespace mongo {
 
-void ScriptEngine::setup() {
+void ScriptEngine::setup(bool disableLoadStored) {
     if (getGlobalScriptEngine())
         return;
 
-    setGlobalScriptEngine(new mozjs::MozJSScriptEngine());
+    setGlobalScriptEngine(new mozjs::MozJSScriptEngine(disableLoadStored));
 
     if (hasGlobalServiceContext()) {
         getGlobalServiceContext()->registerKillOpListener(getGlobalScriptEngine());
@@ -64,7 +64,7 @@ std::string ScriptEngine::getInterpreterVersionString() {
 
 namespace mozjs {
 
-MozJSScriptEngine::MozJSScriptEngine() {
+MozJSScriptEngine::MozJSScriptEngine(bool disableLoadStored) : ScriptEngine(disableLoadStored) {
     uassert(ErrorCodes::JSInterpreterFailure, "Failed to JS_Init()", JS_Init());
     js::DisableExtraThreads();
 }
diff --git a/src/mongo/scripting/mozjs/engine.h b/src/mongo/scripting/mozjs/engine.h
index d49103a7996..688d476282f 100644
--- a/src/mongo/scripting/mozjs/engine.h
+++ b/src/mongo/scripting/mozjs/engine.h
@@ -48,7 +48,7 @@ class MozJSImplScope;
  */
 class MozJSScriptEngine final : public mongo::ScriptEngine {
 public:
-    MozJSScriptEngine();
+    MozJSScriptEngine(bool disableLoadStored);
     ~MozJSScriptEngine() override;
 
     mongo::Scope* createScope() override;