SERVER-15096 added modes to mongoshim. also added JS test and a number of test-related options.

3 files changed, 169 insertions, 9 deletions

diff --git a/src/mongo/tools/mongoshim_options.cpp b/src/mongo/tools/mongoshim_options.cpp
index e3e14300721..2e38b26ee2d 100644
--- a/src/mongo/tools/mongoshim_options.cpp
+++ b/src/mongo/tools/mongoshim_options.cpp
@@ -31,6 +31,7 @@
 #include "mongo/tools/mongoshim_options.h"
 
 #include "mongo/base/status.h"
+#include "mongo/db/json.h"
 #include "mongo/util/options_parser/startup_options.h"
 #include "mongo/util/text.h"
 
@@ -64,32 +65,44 @@ namespace mongo {
             return ret;
         }
 
-        // TODO(benety): Refactor some of these related options (load, upsert, ...) into modes.
+        const string modeDisplayFormat("find, insert, upsert, remove, repair or applyOps");
+        options->addOptionChaining("mode", "mode", moe::String,
+                "Runs shim in one of several modes: " + modeDisplayFormat)
+                                  .format("find|insert|upsert|remove|repair|applyOps",
+                                          modeDisplayFormat)
+                                  .incompatibleWith("load")
+                                  .incompatibleWith("remove")
+                                  .incompatibleWith("repair")
+                                  .incompatibleWith("applyOps");
+
         options->addOptionChaining("load", "load", moe::Switch,
-                                   "load data" );
+                                   "load data" )
+                                  .hidden();
 
         options->addOptionChaining("remove", "remove", moe::Switch,
                 "removes documents from collection matching query "
                 "(or all documents if query is not provided)" )
                                   .incompatibleWith("load")
-                                  .incompatibleWith("drop");
+                                  .incompatibleWith("drop")
+                                  .hidden();
 
         options->addOptionChaining("applyOps", "applyOps", moe::Switch,
                                    "apply oplog entries" )
                                   .incompatibleWith("load")
-                                  .incompatibleWith("drop");
+                                  .incompatibleWith("drop")
+                                  .hidden();
 
         options->addOptionChaining("drop", "drop", moe::Switch,
                                    "drop collection before import" );
 
         options->addOptionChaining("upsert", "upsert", moe::Switch,
                                    "upsert instead of insert" )
-                                  .requires("load");
+                                  .requires("load")
+                                  .hidden();
 
         options->addOptionChaining("upsertFields", "upsertFields", moe::String,
                 "comma-separated fields for the query part of the upsert. "
-                "Ensure these fields are indexed.")
-                                  .requires("upsert");
+                "Ensure these fields are indexed.");
 
         options->addOptionChaining("query", "query,q", moe::String,
                 "query filter, as a JSON string, e.g., '{x:{$gt:1}}'");
@@ -98,7 +111,25 @@ namespace mongo {
                 "try to recover a crashed collection")
                                   .incompatibleWith("applyOps")
                                   .incompatibleWith("load")
-                                  .incompatibleWith("remove");
+                                  .incompatibleWith("remove")
+                                  .hidden();
+
+        // Used for testing.
+        options->addOptionChaining("in", "in", moe::String,
+                "input file; if not specified, stdin is used")
+                                  .hidden();
+
+        // Used for testing.
+        options->addOptionChaining("inputDocuments", "inputDocuments", moe::String,
+                "input documents. If specified, stdin will be ignored. "
+                "Format: --inputDocuments=\"{in: [doc1, doc2, doc3, ...]}\". ")
+                                  .incompatibleWith("in")
+                                  .hidden();
+        // Used for testing.
+        options->addOptionChaining("out", "out", moe::String,
+                "output file; if not specified, stdout is used")
+                                  .incompatibleWith("in")
+                                  .hidden();
 
         options->addOptionChaining("slaveOk", "slaveOk,k", moe::Bool,
                 "use secondaries for export if available, default true")
@@ -150,6 +181,15 @@ namespace mongo {
                           "MongoShim requires a --dbpath value to proceed");
         }
 
+        // Ensure that collection is specified.
+        // Tool::getNs() validates --collection but error
+        // is not propagated to calling process because Tool::main()
+        // always exits cleanly when --dbpath is enabled.
+        if (toolGlobalParams.coll.size() == 0) {
+            return Status(ErrorCodes::BadValue,
+                          "MongoShim requires a --collection value to proceed");
+        }
+
         ret = storeFieldOptions(params, args);
         if (!ret.isOK()) {
             return ret;
@@ -168,6 +208,30 @@ namespace mongo {
         mongoShimGlobalParams.drop = params.count("drop") > 0;
         mongoShimGlobalParams.upsert = params.count("upsert") > 0;
 
+        // Process shim mode. Using --mode is preferred to --load, --applyOps, --remove, ...
+        if (params.count("mode") > 0) {
+            const string mode = getParam("mode");
+            if (mode == "find") {
+                // No change to mongoShimGlobalParams.
+            }
+            else if (mode == "insert") {
+                mongoShimGlobalParams.load = true;
+            }
+            else if (mode == "remove") {
+                mongoShimGlobalParams.remove = true;
+            }
+            else if (mode == "repair") {
+                mongoShimGlobalParams.repair = true;
+            }
+            else if (mode == "upsert") {
+                mongoShimGlobalParams.load = true;
+                mongoShimGlobalParams.upsert = true;
+            }
+            else if (mode == "applyOps") {
+                mongoShimGlobalParams.applyOps = true;
+            }
+        }
+
         if (mongoShimGlobalParams.upsert) {
             string uf = getParam("upsertFields");
             if (uf.empty()) {
@@ -178,6 +242,17 @@ namespace mongo {
             }
         }
 
+        // Used for testing. In normal operation, results will be written to stdout.
+        mongoShimGlobalParams.outputFile = getParam("out");
+        mongoShimGlobalParams.outputFileSpecified = params.count("out") > 0;
+
+        // Used for testing. In normal operation, documents will be read from stdin.
+        mongoShimGlobalParams.inputFile = getParam("in");
+        mongoShimGlobalParams.inputFileSpecified = params.count("in") > 0;
+
+        // Used for testing. In normal operation, documents will be read from stdin.
+        mongoShimGlobalParams.inputDocuments = mongo::fromjson(getParam("inputDocuments", "{}"));
+
         mongoShimGlobalParams.query = getParam("query", "");
         mongoShimGlobalParams.snapShotQuery = false;
 
diff --git a/src/mongo/tools/mongoshim_options.h b/src/mongo/tools/mongoshim_options.h
index b1cd61d89bc..5878022e35e 100644
--- a/src/mongo/tools/mongoshim_options.h
+++ b/src/mongo/tools/mongoshim_options.h
@@ -35,6 +35,7 @@
 #include <vector>
 
 #include "mongo/base/status.h"
+#include "mongo/bson/bsonobj.h"
 #include "mongo/tools/tool_options.h"
 
 namespace mongo {
@@ -59,6 +60,22 @@ namespace mongo {
         std::vector<std::string> upsertFields;
 
         std::string query;
+
+        // If --in is specified, reads documents from 'outputFile' instead of stdin.
+        // Used primarily for testing.
+        std::string inputFile;
+        bool inputFileSpecified;
+
+        // If --inputDocuments is specified, documents will be parsed from option value
+        // instead of being read from stdin.
+        // Used primarily for testing.
+        BSONObj inputDocuments;
+
+        // If --out is specified, writes output to 'outputFile' instead of stdout.
+        // Used primarily for testing.
+        std::string outputFile;
+        bool outputFileSpecified;
+
         bool slaveOk;
         bool snapShotQuery;
         unsigned int skip;
diff --git a/src/mongo/tools/shim.cpp b/src/mongo/tools/shim.cpp
index 4868818f7d4..994ffe27c41 100644
--- a/src/mongo/tools/shim.cpp
+++ b/src/mongo/tools/shim.cpp
@@ -28,8 +28,16 @@
 *    then also delete it in the license file.
 */
 
+#include "mongo/platform/basic.h"
+
 #include <boost/scoped_ptr.hpp>
 
+#include <boost/filesystem/convenience.hpp>
+#include <boost/filesystem/operations.hpp>
+#include <fstream>
+#include <iostream>
+#include <memory>
+
 #include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/client/dbclientcursor.h"
 #include "mongo/db/catalog/collection.h"
@@ -45,6 +53,9 @@
 #include "mongo/util/assert_util.h"
 #include "mongo/util/options_parser/option_section.h"
 
+using std::auto_ptr;
+using std::ios_base;
+using std::ofstream;
 using std::string;
 using std::vector;
 
@@ -126,7 +137,30 @@ public:
             if ( mongoShimGlobalParams.drop ) {
                 conn().dropCollection( _ns );
             }
-            processFile( "-" );
+            // --inputDocuments and --in are used primarily for testing.
+            if (!mongoShimGlobalParams.inputDocuments.isEmpty()) {
+                BSONElement firstElement = mongoShimGlobalParams.inputDocuments.firstElement();
+                if (firstElement.type() != Array) {
+                    toolError() << "first element of --inputDocuments has to be an array: "
+                                << firstElement;
+                    return -1;
+                }
+                BSONObjIterator i(firstElement.Obj());
+                while ( i.more() ) {
+                   BSONElement e = i.next();
+                   if (!e.isABSONObj()) {
+                       toolError() << "skipping non-object in input documents: " << e;
+                       continue;
+                   }
+                   gotObject(e.Obj());
+                }
+            }
+            else if (mongoShimGlobalParams.inputFileSpecified) {
+                processFile(mongoShimGlobalParams.inputFile);
+            }
+            else {
+                processFile("-");
+            }
         }
         else if (mongoShimGlobalParams.remove) {
             // Removes all documents matching query
@@ -136,10 +170,27 @@ public:
         else if (mongoShimGlobalParams.repair) {
             // Repairs collection before writing documents to output.
             ostream *out = &cout;
+            auto_ptr<ofstream> fileStream = _createOutputFile();
+            if (fileStream.get()) {
+                if (!fileStream->good()) {
+                    toolError() << "couldn't open [" << mongoShimGlobalParams.outputFile << "]";
+                    return -1;
+                }
+                out = fileStream.get();
+            }
             _repair(*out);
         }
         else {
+            // Write results to stdout unless output file is specified using --out option.
             ostream *out = &cout;
+            auto_ptr<ofstream> fileStream = _createOutputFile();
+            if (fileStream.get()) {
+                if (!fileStream->good()) {
+                    toolError() << "couldn't open [" << mongoShimGlobalParams.outputFile << "]";
+                    return -1;
+                }
+                out = fileStream.get();
+            }
 
             Query q(mongoShimGlobalParams.query);
             if (mongoShimGlobalParams.sort != "") {
@@ -231,6 +282,23 @@ private:
         cx.commit();
     }
 
+    /**
+     * Returns a valid filestream if output file is specified and is not "-".
+     */
+    auto_ptr<ofstream> _createOutputFile() {
+        auto_ptr<ofstream> fileStream;
+        if (mongoShimGlobalParams.outputFileSpecified && mongoShimGlobalParams.outputFile != "-") {
+            size_t idx = mongoShimGlobalParams.outputFile.rfind("/");
+            if (idx != string::npos) {
+                string dir = mongoShimGlobalParams.outputFile.substr(0 , idx + 1);
+                boost::filesystem::create_directories(dir);
+            }
+            fileStream.reset(new ofstream(mongoShimGlobalParams.outputFile.c_str(),
+                                          ios_base::out | ios_base::binary));
+        }
+        return fileStream;
+    }
+
     string _ns;
 };