diff --git a/src/base_cmd.h b/src/base_cmd.h index 169a6f3ed..4893b0025 100644 --- a/src/base_cmd.h +++ b/src/base_cmd.h @@ -129,6 +129,10 @@ const std::string kCmdNameZScore = "zscore"; const std::string kCmdNameZRange = "zrange"; const std::string kCmdNameZRangebylex = "zrangebylex"; const std::string kCmdNameZRevrangebylex = "zrevrangebylex"; +const std::string kCmdNameZRank = "zrank"; +const std::string kCmdNameZRevrank = "zrevrank"; +const std::string kCmdNameZRem = "zrem"; +const std::string kCmdNameZIncrby = "zincrby"; enum CmdFlags { kCmdFlagsWrite = (1 << 0), // May modify the dataset diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc index e0dadbf49..2fc79d2be 100644 --- a/src/cmd_table_manager.cc +++ b/src/cmd_table_manager.cc @@ -136,6 +136,10 @@ void CmdTableManager::InitCmdTable() { ADD_COMMAND(ZRange, -4); ADD_COMMAND(ZRangebylex, -3); ADD_COMMAND(ZRevrangebylex, -3); + ADD_COMMAND(ZRank, 3); + ADD_COMMAND(ZRevrank, 3); + ADD_COMMAND(ZRem, -3); + ADD_COMMAND(ZIncrby, 4); } std::pair CmdTableManager::GetCommand(const std::string& cmdName, PClient* client) { diff --git a/src/cmd_zset.cc b/src/cmd_zset.cc index 3b2cac4de..8fc8d02fb 100644 --- a/src/cmd_zset.cc +++ b/src/cmd_zset.cc @@ -254,24 +254,6 @@ void ZRangebyscoreCmd::DoCmd(PClient* client) { } } -ZCardCmd::ZCardCmd(const std::string& name, int16_t arity) - : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} - -bool ZCardCmd::DoInitial(PClient* client) { - client->SetKey(client->argv_[1]); - return true; -} - -void ZCardCmd::DoCmd(PClient* client) { - int32_t reply_Num = 0; - storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZCard(client->Key(), &reply_Num); - if (!s.ok()) { - client->SetRes(CmdRes::kSyntaxErr, "ZCard cmd error"); - return; - } - client->AppendInteger(reply_Num); -} - ZRemrangebyrankCmd::ZRemrangebyrankCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategoryString) {} @@ -303,6 +285,109 @@ void ZRemrangebyrankCmd::DoCmd(PClient* client) { } } +ZRevrangebyscoreCmd::ZRevrangebyscoreCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} + +bool ZRevrangebyscoreCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZRevrangebyscoreCmd::DoCmd(PClient* client) { + double min_score = 0; + double max_score = 0; + bool right_close = true; + bool left_close = true; + bool with_scores = false; + int64_t offset = 0, count = -1; + int32_t ret = DoScoreStrRange(client->argv_[3], client->argv_[2], &left_close, &right_close, &min_score, &max_score); + if (ret == -1) { + client->SetRes(CmdRes::kErrOther, "min or max is not a float"); + return; + } + size_t argc = client->argv_.size(); + if (argc >= 5) { + size_t index = 4; + while (index < argc) { + if (strcasecmp(client->argv_[index].data(), "withscores") == 0) { + with_scores = true; + } else if (strcasecmp(client->argv_[index].data(), "limit") == 0) { + if (index + 3 > argc) { + client->SetRes(CmdRes::kSyntaxErr); + return; + } + index++; + if (pstd::String2int(client->argv_[index].data(), client->argv_[index].size(), &offset) == 0) { + client->SetRes(CmdRes::kInvalidInt); + return; + } + index++; + if (pstd::String2int(client->argv_[index].data(), client->argv_[index].size(), &count) == 0) { + client->SetRes(CmdRes::kInvalidInt); + return; + } + } else { + client->SetRes(CmdRes::kSyntaxErr); + return; + } + index++; + } + } + + if (min_score == storage::ZSET_SCORE_MAX || max_score == storage::ZSET_SCORE_MIN) { + client->AppendContent("*0"); + return; + } + std::vector score_members; + storage::Status s = PSTORE.GetBackend(client->GetCurrentDB()) + ->GetStorage() + ->ZRevrangebyscore(client->Key(), min_score, max_score, left_close, right_close, count, + offset, &score_members); + if (!s.ok() && !s.IsNotFound()) { + client->SetRes(CmdRes::kErrOther, s.ToString()); + return; + } + FitLimit(count, offset, static_cast(score_members.size())); + size_t start = offset; + size_t end = offset + count; + if (with_scores) { + char buf[32]; + int64_t len = 0; + client->AppendArrayLen(count * 2); + for (; start < end; start++) { + client->AppendStringLenUint64(score_members[start].member.size()); + client->AppendContent(score_members[start].member); + len = pstd::D2string(buf, sizeof(buf), score_members[start].score); + client->AppendStringLen(len); + client->AppendContent(buf); + } + } else { + client->AppendArrayLen(count); + for (; start < end; start++) { + client->AppendStringLenUint64(score_members[start].member.size()); + client->AppendContent(score_members[start].member); + } + } +} + +ZCardCmd::ZCardCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} + +bool ZCardCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZCardCmd::DoCmd(PClient* client) { + int32_t reply_Num = 0; + storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZCard(client->Key(), &reply_Num); + if (!s.ok()) { + client->SetRes(CmdRes::kSyntaxErr, "ZCard cmd error"); + return; + } + client->AppendInteger(reply_Num); +} + ZRangeCmd::ZRangeCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} @@ -466,91 +551,6 @@ void ZScoreCmd::DoCmd(PClient* client) { } } -ZRevrangebyscoreCmd::ZRevrangebyscoreCmd(const std::string& name, int16_t arity) - : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} - -bool ZRevrangebyscoreCmd::DoInitial(PClient* client) { - client->SetKey(client->argv_[1]); - return true; -} - -void ZRevrangebyscoreCmd::DoCmd(PClient* client) { - double min_score = 0; - double max_score = 0; - bool right_close = true; - bool left_close = true; - bool with_scores = false; - int64_t offset = 0, count = -1; - int32_t ret = DoScoreStrRange(client->argv_[3], client->argv_[2], &left_close, &right_close, &min_score, &max_score); - if (ret == -1) { - client->SetRes(CmdRes::kErrOther, "min or max is not a float"); - return; - } - size_t argc = client->argv_.size(); - if (argc >= 5) { - size_t index = 4; - while (index < argc) { - if (strcasecmp(client->argv_[index].data(), "withscores") == 0) { - with_scores = true; - } else if (strcasecmp(client->argv_[index].data(), "limit") == 0) { - if (index + 3 > argc) { - client->SetRes(CmdRes::kSyntaxErr); - return; - } - index++; - if (pstd::String2int(client->argv_[index].data(), client->argv_[index].size(), &offset) == 0) { - client->SetRes(CmdRes::kInvalidInt); - return; - } - index++; - if (pstd::String2int(client->argv_[index].data(), client->argv_[index].size(), &count) == 0) { - client->SetRes(CmdRes::kInvalidInt); - return; - } - } else { - client->SetRes(CmdRes::kSyntaxErr); - return; - } - index++; - } - } - - if (min_score == storage::ZSET_SCORE_MAX || max_score == storage::ZSET_SCORE_MIN) { - client->AppendContent("*0"); - return; - } - std::vector score_members; - storage::Status s = PSTORE.GetBackend(client->GetCurrentDB()) - ->GetStorage() - ->ZRevrangebyscore(client->Key(), min_score, max_score, left_close, right_close, count, - offset, &score_members); - if (!s.ok() && !s.IsNotFound()) { - client->SetRes(CmdRes::kErrOther, s.ToString()); - return; - } - FitLimit(count, offset, static_cast(score_members.size())); - size_t start = offset; - size_t end = offset + count; - if (with_scores) { - char buf[32]; - int64_t len = 0; - client->AppendArrayLen(count * 2); - for (; start < end; start++) { - client->AppendStringLenUint64(score_members[start].member.size()); - client->AppendContent(score_members[start].member); - len = pstd::D2string(buf, sizeof(buf), score_members[start].score); - client->AppendStringLen(len); - client->AppendContent(buf); - } - } else { - client->AppendArrayLen(count); - for (; start < end; start++) { - client->AppendStringLenUint64(score_members[start].member.size()); - client->AppendContent(score_members[start].member); - } - } -} - ZRangebylexCmd::ZRangebylexCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} @@ -668,6 +668,97 @@ void ZRevrangebylexCmd::DoCmd(PClient* client) { } } +ZRankCmd::ZRankCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} + +bool ZRankCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZRankCmd::DoCmd(PClient* client) { + int32_t rank = 0; + storage::Status s = + PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZRank(client->Key(), client->argv_[2], &rank); + if (s.ok()) { + client->AppendInteger(rank); + } else if (s.IsNotFound()) { + client->AppendContent("$-1"); + } else { + client->SetRes(CmdRes::kErrOther, s.ToString()); + } +} + +ZRevrankCmd::ZRevrankCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} + +bool ZRevrankCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZRevrankCmd::DoCmd(PClient* client) { + int32_t revrank = 0; + storage::Status s = + PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZRevrank(client->Key(), client->argv_[2], &revrank); + if (s.ok()) { + client->AppendInteger(revrank); + } else if (s.IsNotFound()) { + client->AppendContent("$-1"); + } else { + client->SetRes(CmdRes::kErrOther, s.ToString()); + } +} + +ZRemCmd::ZRemCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} + +bool ZRemCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZRemCmd::DoCmd(PClient* client) { + auto iter = client->argv_.begin() + 2; + std::vector members(iter, client->argv_.end()); + int32_t deleted = 0; + storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZRem(client->Key(), members, &deleted); + if (s.ok() || s.IsNotFound()) { + client->AppendInteger(deleted); + } else { + client->SetRes(CmdRes::kErrOther, s.ToString()); + } +} + +ZIncrbyCmd::ZIncrbyCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} + +bool ZIncrbyCmd::DoInitial(PClient* client) { + client->SetKey(client->argv_[1]); + return true; +} + +void ZIncrbyCmd::DoCmd(PClient* client) { + double by = .0f; + double score = .0f; + if (pstd::String2d(client->argv_[2].data(), client->argv_[2].size(), &by) == 0) { + client->SetRes(CmdRes::kInvalidFloat); + return; + } + + std::string member = client->argv_[3]; + storage::Status s = + PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->ZIncrby(client->Key(), member, by, &score); + if (s.ok()) { + char buf[32]; + int64_t len = pstd::D2string(buf, sizeof(buf), score); + client->AppendStringLen(len); + client->AppendContent(buf); + } else { + client->SetRes(CmdRes::kErrOther, s.ToString()); + } +} + ZRemrangebyscoreCmd::ZRemrangebyscoreCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} @@ -698,4 +789,4 @@ void ZRemrangebyscoreCmd::DoCmd(PClient* client) { } } -} // namespace pikiwidb \ No newline at end of file +} // namespace pikiwidb diff --git a/src/cmd_zset.h b/src/cmd_zset.h index c4ed6ba55..13049eaa9 100644 --- a/src/cmd_zset.h +++ b/src/cmd_zset.h @@ -122,6 +122,50 @@ class ZRevrangebylexCmd : public BaseCmd { void DoCmd(PClient *client) override; }; +class ZRankCmd : public BaseCmd { + public: + ZRankCmd(const std::string &name, int16_t arity); + + protected: + bool DoInitial(PClient *client) override; + + private: + void DoCmd(PClient *client) override; +}; + +class ZRevrankCmd : public BaseCmd { + public: + ZRevrankCmd(const std::string &name, int16_t arity); + + protected: + bool DoInitial(PClient *client) override; + + private: + void DoCmd(PClient *client) override; +}; + +class ZRemCmd : public BaseCmd { + public: + ZRemCmd(const std::string &name, int16_t arity); + + protected: + bool DoInitial(PClient *client) override; + + private: + void DoCmd(PClient *client) override; +}; + +class ZIncrbyCmd : public BaseCmd { + public: + ZIncrbyCmd(const std::string &name, int16_t arity); + + protected: + bool DoInitial(PClient *client) override; + + private: + void DoCmd(PClient *client) override; +}; + class ZRemrangebyscoreCmd : public BaseCmd { public: ZRemrangebyscoreCmd(const std::string &name, int16_t arity); @@ -133,4 +177,4 @@ class ZRemrangebyscoreCmd : public BaseCmd { void DoCmd(PClient *client) override; }; -} // namespace pikiwidb \ No newline at end of file +} // namespace pikiwidb diff --git a/tests/zset_test.go b/tests/zset_test.go index e60c1ca7c..eab3e7838 100644 --- a/tests/zset_test.go +++ b/tests/zset_test.go @@ -387,6 +387,74 @@ var _ = Describe("Zset", Ordered, func() { Expect(zRangeByLex.Val()).To(Equal([]string{})) }) + It("should ZRank", func() { + err := client.ZAdd(ctx, "zrank", redis.Z{ + Score: 1, + Member: "one", + }).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zrank", redis.Z{ + Score: 2, + Member: "two", + }).Err() + Expect(err).NotTo(HaveOccurred()) + + rank, err := client.ZRank(ctx, "zrank", "two").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(rank).To(Equal(int64(1))) + }) + + It("should ZRevrank", func() { + err := client.ZAdd(ctx, "zrevrank", redis.Z{ + Score: 1, + Member: "one", + }).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zrevrank", redis.Z{ + Score: 2, + Member: "two", + }).Err() + Expect(err).NotTo(HaveOccurred()) + + revrank, err := client.ZRevRank(ctx, "zrevrank", "one").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(revrank).To(Equal(int64(1))) + }) + + It("should ZRem", func() { + err := client.ZAdd(ctx, "zrem", redis.Z{ + Score: 1, + Member: "one", + }).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zrem", redis.Z{ + Score: 2, + Member: "two", + }).Err() + Expect(err).NotTo(HaveOccurred()) + + rem, err := client.ZRem(ctx, "zrem", "one").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(rem).To(Equal(int64(1))) + }) + + It("should ZIncrby", func() { + err := client.ZAdd(ctx, "zincrby", redis.Z{ + Score: 1, + Member: "one", + }).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zincrby", redis.Z{ + Score: 2, + Member: "two", + }).Err() + Expect(err).NotTo(HaveOccurred()) + + rem, err := client.ZIncrBy(ctx, "zincrby", 5, "one").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(rem).To(Equal(float64(6))) + }) + It("should ZRemRangeByScore", func() { err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err() Expect(err).NotTo(HaveOccurred())