Support command cmd

apache · May 8, 2021 · bb60cea · bb60cea
1 parent fbc41f0
commit bb60cea
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 0 deletions.
diff --git a/src/redis_cmd.cc b/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->GetCommandsInfo(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),

diff --git a/src/server.cc b/src/server.cc
@@ -1161,3 +1161,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::GetCommandsInfo(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();
+}
+
diff --git a/src/server.h b/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 GetCommandsInfo(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();