Support COMMAND command (#251)

* Support command cmd
* Add testcase for COMMAND command
* FIX runtest return rm command's status instead of tcl test script's status

src/redis_cmd.cc

@@ -3653,6 +3653,44 @@ class CommandDebug : public Commander {
   uint64_t microsecond_ = 0;
 };
 
+class CommandCommand : public Commander {
+ public:
+  Status Execute(Server *svr, Connection *conn, std::string *output) override {
+    if (args_.size() == 1) {
+      svr->GetAllCommandsInfo(output);
+    } else {
+      std::string sub_command = Util::ToLower(args_[1]);
+      if ((sub_command == "count" && args_.size() != 2) ||
+          (sub_command == "getkeys" && args_.size() < 3) ||
+          (sub_command == "info" && args_.size() < 3)) {
+        *output = Redis::Error(errWrongNumOfArguments);
+        return Status::OK();
+      }
+      if (sub_command == "count") {
+        *output = Redis::Integer(GetCommandNum());
+      } else if (sub_command == "info") {
+        svr->GetCommandsInfo(output, std::vector<std::string>(args_.begin() + 2, args_.end()));
+      } else if (sub_command == "getkeys") {
+        std::vector<int> keys_indexes;
+        auto s = svr->GetKeysFromCommand(args_[2], args_.size() - 2, &keys_indexes);
+        if (!s.IsOK()) return s;
+        if (keys_indexes.size() == 0) {
+          *output = Redis::Error("Invalid arguments specified for command");
+          return Status::OK();
+        }
+        std::vector<std::string> keys;
+        for (const auto &key_index : keys_indexes) {
+          keys.emplace_back(args_[key_index + 2]);
+        }
+        *output = Redis::MultiBulkString(keys);
+      } else {
+        *output = Redis::Error("Command subcommand must be one of COUNT, GETKEYS, INFO");
+      }
+    }
+    return Status::OK();
+  }
+};
+
 class CommandScanBase : public Commander {
  public:
   Status ParseMatchAndCountParam(const std::string &type, std::string value) {
@@ -4053,6 +4091,7 @@ CommandAttributes redisCommandTable[] = {
     ADD_CMD("scan", -2, false, 0, 0, 0, CommandScan),
     ADD_CMD("randomkey", 1, false, 0, 0, 0, CommandRandomKey),
     ADD_CMD("debug", -2, false, 0, 0, 0, CommandDebug),
+    ADD_CMD("command", -1, false, 0, 0, 0, CommandCommand),
 
     ADD_CMD("ttl", 2, false, 1, 1, 1, CommandTTL),
     ADD_CMD("pttl", 2, false, 1, 1, 1, CommandPTTL),

src/server.cc

@@ -1163,3 +1163,62 @@ void Server::populateCommands() {
 bool Server::IsCommandExists(const std::string &name) {
   return commands_.find(name) != commands_.end();
 }
+
+std::string Server::GetCommandInfo(const Redis::CommandAttributes* command_attributes) {
+  std::string command, command_flags;
+  command.append(Redis::MultiLen(6));
+  command.append(Redis::BulkString(command_attributes->name));
+  command.append(Redis::Integer(command_attributes->arity));
+  command_flags.append(Redis::MultiLen(1));
+  command_flags.append(Redis::BulkString(command_attributes->is_write ? "write" : "readonly"));
+  command.append(command_flags);
+  command.append(Redis::Integer(command_attributes->first_key));
+  command.append(Redis::Integer(command_attributes->last_key));
+  command.append(Redis::Integer(command_attributes->key_step));
+  return command;
+}
+
+void Server::GetAllCommandsInfo(std::string *info) {
+  info->append(Redis::MultiLen(commands_.size()));
+  for (const auto &iter : commands_) {
+    auto command_attribute = iter.second;
+    auto command_info = GetCommandInfo(command_attribute);
+    info->append(command_info);
+  }
+}
+
+void Server::GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names) {
+  info->append(Redis::MultiLen(cmd_names.size()));
+  for (const auto &cmd_name : cmd_names) {
+    auto cmd_iter = commands_.find(Util::ToLower(cmd_name));
+    if (cmd_iter == commands_.end()) {
+      info->append(Redis::NilString());
+    } else {
+      auto command_attribute = cmd_iter->second;
+      auto command_info = GetCommandInfo(command_attribute);
+      info->append(command_info);
+    }
+  }
+}
+
+Status Server::GetKeysFromCommand(const std::string &cmd_name, int argc, std::vector<int> *keys_indexes) {
+  auto cmd_iter = commands_.find(Util::ToLower(cmd_name));
+  if (cmd_iter == commands_.end()) {
+    return Status(Status::RedisUnknownCmd, "Invalid command specified");
+  }
+  auto command_attribute = cmd_iter->second;
+  if (command_attribute->first_key == 0) {
+    return Status(Status::NotOK, "The command has no key arguments");
+  }
+  if ((command_attribute->arity > 0 && command_attribute->arity != argc) || argc < -command_attribute->arity) {
+    return Status(Status::NotOK, "Invalid number of arguments specified for command");
+  }
+  auto last = command_attribute->last_key;
+  if (last < 0) last = argc + last;
+
+  for (int j = command_attribute->first_key; j <= last; j += command_attribute->key_step) {
+    keys_indexes->emplace_back(j);
+  }
+  return Status::OK();
+}
+

src/server.h

@@ -53,6 +53,10 @@ class Server {
   void populateCommands();
   bool IsCommandExists(const std::string &name);
   Status LookupAndCreateCommand(const std::string &cmd_name, std::unique_ptr<Redis::Commander> *cmd);
+  void GetAllCommandsInfo(std::string *info);
+  void GetCommandsInfo(std::string *info, const std::vector<std::string> &cmd_names);
+  std::string GetCommandInfo(const Redis::CommandAttributes* command_attributes);
+  Status GetKeysFromCommand(const std::string &name, int argc, std::vector<int> *keys_indexes);
 
   Status AddMaster(std::string host, uint32_t port);
   Status RemoveMaster();

tests/tcl/runtest

@@ -11,6 +11,8 @@ then
     echo "You need tcl 8.5 or newer in order to run the Redis test"
     exit 1
 fi
-cp ../../src/kvrocks redis-server
-$TCLSH tests/test_helper.tcl "${@}"
+
+cp ../../src/kvrocks redis-server && $TCLSH tests/test_helper.tcl "${@}"
+status=$?
 rm ./redis-server
+exit $status

tests/tcl/tests/test_helper.tcl

@@ -31,6 +31,7 @@ set ::all_tests {
     unit/introspection
     unit/limits
     unit/geo
+    unit/command
     integration/replication
 }
 

tests/tcl/tests/unit/command.tcl

@@ -0,0 +1,27 @@
+start_server {tags {"command"}} {
+    test {kvrocks has 155 commands currently} {
+        r command count
+    } {155}
+
+    test {acquire GET command info by COMMAND INFO} {
+        set e [lindex [r command info get] 0]
+        assert_equal [llength $e] 6
+        assert_equal [lindex $e 0] get
+        assert_equal [lindex $e 1] 2
+        assert_equal [lindex $e 2] {readonly}
+        assert_equal [lindex $e 3] 1
+        assert_equal [lindex $e 4] 1
+        assert_equal [lindex $e 5] 1
+    }
+
+    test {COMMAND - command entry length check} {
+        set e [lindex [r command] 0]
+        assert_equal [llength $e] 6
+    }
+
+    test {get keys of commands by COMMAND GETKEYS} {
+        assert_equal {test} [r command getkeys get test]
+        assert_equal {test test2} [r command getkeys mget test test2]
+        assert_equal {test} [r command getkeys zadd test 1 m1]
+    }
+}