SERVER-35485 Add compressor capabilities for shell with URI input

7 files changed, 76 insertions, 42 deletions

diff --git a/src/mongo/db/server_options_server_helpers.cpp b/src/mongo/db/server_options_server_helpers.cpp
index fb6191b3410..231b0268689 100644
--- a/src/mongo/db/server_options_server_helpers.cpp
+++ b/src/mongo/db/server_options_server_helpers.cpp
@@ -726,9 +726,12 @@ Status storeServerOptions(const moe::Environment& params) {
                       "--transitionToAuth must be used with keyFile or x509 authentication");
     }
 
-    ret = storeMessageCompressionOptions(params);
-    if (!ret.isOK()) {
-        return ret;
+    if (params.count("net.compression.compressors")) {
+        const auto ret =
+            storeMessageCompressionOptions(params["net.compression.compressors"].as<string>());
+        if (!ret.isOK()) {
+            return ret;
+        }
     }
 
     return Status::OK();
diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index cb3cafbda83..e53e09d6ed5 100644
--- a/src/mongo/shell/dbshell.cpp
+++ b/src/mongo/shell/dbshell.cpp
@@ -815,6 +815,11 @@ int _main(int argc, char* argv[], char** envp) {
         mongo::shell_utils::_dbConnect = ss.str();
 
         if (cmdlineURI.size()) {
+            const auto compressionKey = parsedURI.getOptions().find("compressors");
+            if (compressionKey != end(parsedURI.getOptions()) &&
+                shellGlobalParams.networkMessageCompressors.empty()) {
+                shellGlobalParams.networkMessageCompressors = compressionKey->second;
+            }
             const auto mechanismKey = parsedURI.getOptions().find("authMechanism");
             if (mechanismKey != end(parsedURI.getOptions()) &&
                 shellGlobalParams.authenticationMechanism.empty()) {
diff --git a/src/mongo/shell/shell_options.cpp b/src/mongo/shell/shell_options.cpp
index 981b965d964..b92ad6f5637 100644
--- a/src/mongo/shell/shell_options.cpp
+++ b/src/mongo/shell/shell_options.cpp
@@ -343,6 +343,13 @@ Status storeMongoShellOptions(const moe::Environment& params,
         shellGlobalParams.gssapiHostName = params["gssapiHostName"].as<string>();
     }
 
+    if (params.count("net.compression.compressors")) {
+        auto compressors = params["net.compression.compressors"].as<string>();
+        if (compressors != "disabled") {
+            shellGlobalParams.networkMessageCompressors = std::move(compressors);
+        }
+    }
+
     if (params.count("shell")) {
         shellGlobalParams.runShell = true;
     }
@@ -458,37 +465,55 @@ Status storeMongoShellOptions(const moe::Environment& params,
 
         auto cs = cs_status.getValue();
         auto uriOptions = cs.getOptions();
-        StringBuilder sb;
-        sb << "ERROR: Cannot specify ";
-
-        if (!shellGlobalParams.username.empty() && !cs.getUser().empty() &&
-            shellGlobalParams.username != cs.getUser()) {
-            sb << "different usernames";
-        } else if (!shellGlobalParams.password.empty() && !cs.getPassword().empty() &&
-                   shellGlobalParams.password != cs.getPassword()) {
-            sb << "different passwords";
-        } else if (!shellGlobalParams.authenticationMechanism.empty() &&
-                   uriOptions.count("authMechanism") &&
-                   uriOptions["authMechanism"] != shellGlobalParams.authenticationMechanism) {
-            sb << "different authentication mechanisms";
-        } else if (!shellGlobalParams.authenticationDatabase.empty() &&
-                   uriOptions.count("authSource") &&
-                   uriOptions["authSource"] != shellGlobalParams.authenticationDatabase) {
-            sb << "different authentication databases ";
-        } else if (shellGlobalParams.gssapiServiceName != saslDefaultServiceName &&
-                   uriOptions.count("gssapiServiceName")) {
-            sb << "the GSSAPI service name";
-        } else {
-            return Status::OK();
-        }
 
-        sb << " in connection URI and as a command-line option";
-        return Status(ErrorCodes::InvalidOptions, sb.str());
+        auto handleURIOptions = [&] {
+            StringBuilder sb;
+            sb << "ERROR: Cannot specify ";
+
+            if (!shellGlobalParams.username.empty() && !cs.getUser().empty() &&
+                shellGlobalParams.username != cs.getUser()) {
+                sb << "different usernames";
+            } else if (!shellGlobalParams.password.empty() && !cs.getPassword().empty() &&
+                       shellGlobalParams.password != cs.getPassword()) {
+                sb << "different passwords";
+            } else if (!shellGlobalParams.authenticationMechanism.empty() &&
+                       uriOptions.count("authMechanism") &&
+                       uriOptions["authMechanism"] != shellGlobalParams.authenticationMechanism) {
+                sb << "different authentication mechanisms";
+            } else if (!shellGlobalParams.authenticationDatabase.empty() &&
+                       uriOptions.count("authSource") &&
+                       uriOptions["authSource"] != shellGlobalParams.authenticationDatabase) {
+                sb << "different authentication databases";
+            } else if (shellGlobalParams.gssapiServiceName != saslDefaultServiceName &&
+                       uriOptions.count("gssapiServiceName")) {
+                sb << "the GSSAPI service name";
+            } else if (!shellGlobalParams.networkMessageCompressors.empty() &&
+                       uriOptions.count("compressors") &&
+                       uriOptions["compressors"] != shellGlobalParams.networkMessageCompressors) {
+                sb << "different network message compressors";
+            } else {
+                return Status::OK();
+            }
+
+            sb << " in connection URI and as a command-line option";
+            return Status(ErrorCodes::InvalidOptions, sb.str());
+        };
+
+        auto uriStatus = handleURIOptions();
+        if (!uriStatus.isOK())
+            return uriStatus;
+
+        if (uriOptions.count("compressors"))
+            shellGlobalParams.networkMessageCompressors = uriOptions["compressors"];
     }
 
-    auto ret = storeMessageCompressionOptions(params);
-    if (!ret.isOK())
-        return ret;
+    if (!shellGlobalParams.networkMessageCompressors.empty()) {
+        const auto ret =
+            storeMessageCompressionOptions(shellGlobalParams.networkMessageCompressors);
+        if (!ret.isOK()) {
+            return ret;
+        }
+    }
 
     if (params.count("setShellParameter")) {
         auto ssp = params["setShellParameter"].as<std::map<std::string, std::string>>();
diff --git a/src/mongo/shell/shell_options.h b/src/mongo/shell/shell_options.h
index 601aa80e42f..d28bf44d55f 100644
--- a/src/mongo/shell/shell_options.h
+++ b/src/mongo/shell/shell_options.h
@@ -59,6 +59,7 @@ struct ShellGlobalParams {
     std::string authenticationDatabase;
     std::string gssapiServiceName;
     std::string gssapiHostName;
+    std::string networkMessageCompressors;
 
     bool runShell;
     bool nodb;
diff --git a/src/mongo/tools/mongobridge_options_init.cpp b/src/mongo/tools/mongobridge_options_init.cpp
index f023bb5a870..20371223b10 100644
--- a/src/mongo/tools/mongobridge_options_init.cpp
+++ b/src/mongo/tools/mongobridge_options_init.cpp
@@ -65,10 +65,13 @@ MONGO_STARTUP_OPTIONS_STORE(MongoBridgeOptions)(InitializerContext* context) {
         quickExit(EXIT_BADOPTIONS);
     }
 
-    ret = storeMessageCompressionOptions(moe::startupOptionsParsed);
-    if (!ret.isOK()) {
-        std::cerr << ret.toString() << std::endl;
-        quickExit(EXIT_BADOPTIONS);
+    if (moe::startupOptionsParsed.count("net.compression.compressors")) {
+        const auto ret = storeMessageCompressionOptions(
+            moe::startupOptionsParsed["net.compression.compressors"].as<std::string>());
+        if (!ret.isOK()) {
+            std::cerr << ret.toString() << std::endl;
+            quickExit(EXIT_BADOPTIONS);
+        }
     }
 
     return Status::OK();
diff --git a/src/mongo/transport/message_compressor_registry.cpp b/src/mongo/transport/message_compressor_registry.cpp
index 9bdb640e5a9..2a38b4a3f6f 100644
--- a/src/mongo/transport/message_compressor_registry.cpp
+++ b/src/mongo/transport/message_compressor_registry.cpp
@@ -129,13 +129,10 @@ Status addMessageCompressionOptions(moe::OptionSection* options, bool forShell)
     return Status::OK();
 }
 
-Status storeMessageCompressionOptions(const moe::Environment& params) {
+Status storeMessageCompressionOptions(const std::string& compressors) {
     std::vector<std::string> restrict;
-    if (params.count("net.compression.compressors")) {
-        auto compressorListStr = params["net.compression.compressors"].as<std::string>();
-        if (compressorListStr != kDisabledConfigValue) {
-            boost::algorithm::split(restrict, compressorListStr, boost::is_any_of(", "));
-        }
+    if (compressors != kDisabledConfigValue) {
+        boost::algorithm::split(restrict, compressors, boost::is_any_of(", "));
     }
 
     auto& compressorFactory = MessageCompressorRegistry::get();
diff --git a/src/mongo/transport/message_compressor_registry.h b/src/mongo/transport/message_compressor_registry.h
index 933c247e928..fb980bbab0a 100644
--- a/src/mongo/transport/message_compressor_registry.h
+++ b/src/mongo/transport/message_compressor_registry.h
@@ -120,6 +120,6 @@ private:
 };
 
 Status addMessageCompressionOptions(moe::OptionSection* options, bool forShell);
-Status storeMessageCompressionOptions(const moe::Environment& params);
+Status storeMessageCompressionOptions(const std::string& compressors);
 void appendMessageCompressionStats(BSONObjBuilder* b);
 }  // namespace mongo