SERVER-34148 TypedCommand

1 files changed, 300 insertions, 2 deletions

diff --git a/src/mongo/db/commands_test.cpp b/src/mongo/db/commands_test.cpp
index cd74dff528e..9a7f36cd60a 100644
--- a/src/mongo/db/commands_test.cpp
+++ b/src/mongo/db/commands_test.cpp
@@ -26,11 +26,13 @@
  * then also delete it in the license file.
  */
 
-#include "mongo/db/commands.h"
+#include "mongo/platform/basic.h"
+
 #include "mongo/db/catalog/collection_mock.h"
+#include "mongo/db/commands.h"
+#include "mongo/db/commands_test_example_gen.h"
 #include "mongo/db/dbmessage.h"
 #include "mongo/db/service_context_noop.h"
-#include "mongo/platform/basic.h"
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -140,5 +142,301 @@ TEST_F(ParseNsOrUUID, ParseValidUUID) {
     ASSERT_EQUALS(uuid, *parsedNsOrUUID.uuid());
 }
 
+/**
+ * TypedCommand test
+ */
+class ExampleIncrementCommand final : public TypedCommand<ExampleIncrementCommand> {
+private:
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return Command::AllowedOnSecondary::kNever;
+    }
+
+    std::string help() const override {
+        return "Return an incremented request.i. Example of a simple TypedCommand.";
+    }
+
+public:
+    using Request = commands_test_example::ExampleIncrement;
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        /**
+         * Reply with an incremented 'request.i'.
+         */
+        auto typedRun(OperationContext* opCtx) {
+            commands_test_example::ExampleIncrementReply r;
+            r.setIPlusOne(request().getI() + 1);
+            return r;
+        }
+
+    private:
+        bool supportsWriteConcern() const override {
+            return true;
+        }
+
+        AllowedOnSecondary secondaryAllowed(ServiceContext* context) const override {
+            return definition()->secondaryAllowed(context);
+        }
+
+        void doCheckAuthorization(OperationContext*) const override {}
+
+        /**
+         * The ns() for when Request's IDL specifies "namespace: concatenate_with_db".
+         */
+        NamespaceString ns() const override {
+            return request().getNamespace();
+        }
+    };
+};
+
+// Just like ExampleIncrementCommand, but using the MinimalInvocationBase.
+class ExampleMinimalCommand final : public TypedCommand<ExampleMinimalCommand> {
+private:
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return Command::AllowedOnSecondary::kNever;
+    }
+
+    std::string help() const override {
+        return "Return an incremented request.i. Example of a simple TypedCommand.";
+    }
+
+public:
+    using Request = commands_test_example::ExampleMinimal;
+
+    class Invocation final : public MinimalInvocationBase {
+    public:
+        using MinimalInvocationBase::MinimalInvocationBase;
+
+        /**
+         * Reply with an incremented 'request.i'.
+         */
+        void run(OperationContext* opCtx, CommandReplyBuilder* reply) override {
+            commands_test_example::ExampleIncrementReply r;
+            r.setIPlusOne(request().getI() + 1);
+            reply->fillFrom(r);
+        }
+
+    private:
+        bool supportsWriteConcern() const override {
+            return true;
+        }
+
+        AllowedOnSecondary secondaryAllowed(ServiceContext* context) const override {
+            return definition()->secondaryAllowed(context);
+        }
+
+        void explain(OperationContext* opCtx,
+                     ExplainOptions::Verbosity verbosity,
+                     BSONObjBuilder* result) override {}
+
+        void doCheckAuthorization(OperationContext*) const override {}
+
+        /**
+         * The ns() for when Request's IDL specifies "namespace: concatenate_with_db".
+         */
+        NamespaceString ns() const override {
+            return request().getNamespace();
+        }
+    };
+};
+
+// Just like ExampleIncrementCommand, but with a void typedRun.
+class ExampleVoidCommand final : public TypedCommand<ExampleVoidCommand> {
+private:
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return Command::AllowedOnSecondary::kNever;
+    }
+
+    std::string help() const override {
+        return "Accepts Request and returns void.";
+    }
+
+public:
+    using Request = commands_test_example::ExampleVoid;
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        /**
+         * Have some testable side-effect.
+         */
+        void typedRun(OperationContext*) {
+            static_cast<const ExampleVoidCommand*>(definition())->iCapture = request().getI() + 1;
+        }
+
+    private:
+        bool supportsWriteConcern() const override {
+            return true;
+        }
+
+        AllowedOnSecondary secondaryAllowed(ServiceContext* context) const override {
+            return definition()->secondaryAllowed(context);
+        }
+
+        void explain(OperationContext* opCtx,
+                     ExplainOptions::Verbosity verbosity,
+                     BSONObjBuilder* result) override {}
+
+        void doCheckAuthorization(OperationContext*) const override {}
+
+        /**
+         * The ns() for when Request's IDL specifies "namespace: concatenate_with_db".
+         */
+        NamespaceString ns() const override {
+            return request().getNamespace();
+        }
+    };
+
+    mutable std::int32_t iCapture = 0;
+};
+
+template <typename Fn>
+class MyCommand final : public TypedCommand<MyCommand<Fn>> {
+public:
+    class Invocation final : public TypedCommand<MyCommand>::InvocationBase {
+    public:
+        using Base = typename TypedCommand<MyCommand>::InvocationBase;
+        using Base::Base;
+
+        auto typedRun(OperationContext*) const {
+            return _command()->_fn();
+        }
+
+    private:
+        NamespaceString ns() const override {
+            return Base::request().getNamespace();
+        }
+        bool supportsWriteConcern() const override {
+            return false;
+        }
+        Command::AllowedOnSecondary secondaryAllowed(ServiceContext* context) const override {
+            return Base::definition()->secondaryAllowed(context);
+        }
+        void doCheckAuthorization(OperationContext* opCtx) const override {}
+
+        const MyCommand* _command() const {
+            return static_cast<const MyCommand*>(Base::definition());
+        }
+    };
+
+    using Request = commands_test_example::ExampleVoid;
+
+    MyCommand(StringData name, Fn fn) : TypedCommand<MyCommand<Fn>>(name), _fn{std::move(fn)} {}
+
+private:
+    Command::AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return Command::AllowedOnSecondary::kAlways;
+    }
+
+    std::string help() const override {
+        return "Accepts Request and returns void.";
+    }
+
+    Fn _fn;
+};
+
+template <typename Fn>
+using CmdT = MyCommand<typename std::decay<Fn>::type>;
+
+auto okFn = [] { return Status::OK(); };
+auto errFn = [] { return Status(ErrorCodes::UnknownError, "some error"); };
+auto throwFn = []() -> Status { uasserted(ErrorCodes::UnknownError, "some error"); };
+
+ExampleIncrementCommand exampleIncrementCommand;
+ExampleMinimalCommand exampleMinimalCommand;
+ExampleVoidCommand exampleVoidCommand;
+CmdT<decltype(okFn)> okStatusCommand("okStatus", okFn);
+CmdT<decltype(errFn)> errStatusCommand("notOkStatus", errFn);
+CmdT<decltype(throwFn)> throwStatusCommand("throwsStatus", throwFn);
+
+struct IncrementTestCommon {
+    template <typename T>
+    void run(T& command, std::function<void(int, const BSONObj&)> postAssert) {
+        const NamespaceString ns("testdb.coll");
+        auto client = getGlobalServiceContext()->makeClient("commands_test");
+        for (std::int32_t i : {123, 12345, 0, -456}) {
+            const OpMsgRequest request = [&] {
+                typename T::Request incr(ns);
+                incr.setI(i);
+                return incr.serialize(BSON("$db" << ns.db()));
+            }();
+            auto opCtx = client->makeOperationContext();
+            auto invocation = command.parse(opCtx.get(), request);
+
+            ASSERT_EQ(invocation->ns(), ns);
+
+            const BSONObj reply = [&] {
+                BufBuilder bb;
+                CommandReplyBuilder crb{BSONObjBuilder{bb}};
+                try {
+                    invocation->run(opCtx.get(), &crb);
+                    auto bob = crb.getBodyBuilder();
+                    CommandHelpers::extractOrAppendOk(bob);
+                } catch (const DBException& e) {
+                    auto bob = crb.getBodyBuilder();
+                    CommandHelpers::appendCommandStatus(bob, e.toStatus());
+                }
+                return BSONObj(bb.release());
+            }();
+
+            postAssert(i, reply);
+        }
+    }
+};
+
+TEST(TypedCommand, runTyped) {
+    IncrementTestCommon{}.run(exampleIncrementCommand, [](int i, const BSONObj& reply) {
+        ASSERT_EQ(reply["ok"].Double(), 1.0);
+        ASSERT_EQ(reply["iPlusOne"].Int(), i + 1);
+    });
+}
+
+TEST(TypedCommand, runMinimal) {
+    IncrementTestCommon{}.run(exampleMinimalCommand, [](int i, const BSONObj& reply) {
+        ASSERT_EQ(reply["ok"].Double(), 1.0);
+        ASSERT_EQ(reply["iPlusOne"].Int(), i + 1);
+    });
+}
+
+TEST(TypedCommand, runVoid) {
+    IncrementTestCommon{}.run(exampleVoidCommand, [](int i, const BSONObj& reply) {
+        ASSERT_EQ(reply["ok"].Double(), 1.0);
+        ASSERT_EQ(exampleVoidCommand.iCapture, i + 1);
+    });
+}
+
+TEST(TypedCommand, runOkStatus) {
+    IncrementTestCommon{}.run(
+        okStatusCommand, [](int i, const BSONObj& reply) { ASSERT_EQ(reply["ok"].Double(), 1.0); });
+}
+
+TEST(TypedCommand, runErrStatus) {
+    IncrementTestCommon{}.run(errStatusCommand, [](int i, const BSONObj& reply) {
+        Status status = errFn();
+        ASSERT_EQ(reply["ok"].Double(), 0.0);
+        ASSERT_EQ(reply["errmsg"].String(), status.reason());
+        ASSERT_EQ(reply["code"].Int(), status.code());
+        ASSERT_EQ(reply["codeName"].String(), ErrorCodes::errorString(status.code()));
+    });
+}
+
+TEST(TypedCommand, runThrowStatus) {
+    IncrementTestCommon{}.run(throwStatusCommand, [](int i, const BSONObj& reply) {
+        Status status = Status::OK();
+        try {
+            (void)throwFn();
+        } catch (const DBException& e) {
+            status = e.toStatus();
+        }
+        ASSERT_EQ(reply["ok"].Double(), 0.0);
+        ASSERT_EQ(reply["errmsg"].String(), status.reason());
+        ASSERT_EQ(reply["code"].Int(), status.code());
+        ASSERT_EQ(reply["codeName"].String(), ErrorCodes::errorString(status.code()));
+    });
+}
+
 }  // namespace
 }  // namespace mongo