From 4c62cfa75e151b0934f83518d6feddb42170845b Mon Sep 17 00:00:00 2001 From: panlei-coder Date: Tue, 19 Mar 2024 00:40:19 +0800 Subject: [PATCH 1/4] feat: support zrank, zrevrank, zrem, zincrby --- src/base_cmd.h | 4 ++ src/cmd_table_manager.cc | 4 ++ src/cmd_zset.cc | 89 ++++++++++++++++++++++++++++++++++++++++ src/cmd_zset.h | 44 ++++++++++++++++++++ 4 files changed, 141 insertions(+) diff --git a/src/base_cmd.h b/src/base_cmd.h index f3c98ca3f..3b4a5a2f9 100644 --- a/src/base_cmd.h +++ b/src/base_cmd.h @@ -120,6 +120,10 @@ const std::string kCmdNameZRevrange = "zrevrange"; const std::string kCmdNameZRangebyscore = "zrangebyscore"; const std::string kCmdNameZRevRangeByScore = "zrevrangebyscore"; const std::string kCmdNameZCard = "zcard"; +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 286b1d076..521676af0 100644 --- a/src/cmd_table_manager.cc +++ b/src/cmd_table_manager.cc @@ -127,6 +127,10 @@ void CmdTableManager::InitCmdTable() { ADD_COMMAND(ZRangebyscore, -4); ADD_COMMAND(ZRevRangeByScore, -4); ADD_COMMAND(ZCard, 2); + 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 04247c800..5561c43e8 100644 --- a/src/cmd_zset.cc +++ b/src/cmd_zset.cc @@ -319,4 +319,93 @@ void ZRevRangeByScoreCmd::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())->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())->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())->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())->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()); + } +} + } // namespace pikiwidb \ No newline at end of file diff --git a/src/cmd_zset.h b/src/cmd_zset.h index ddef7b956..bd564ddf3 100644 --- a/src/cmd_zset.h +++ b/src/cmd_zset.h @@ -67,4 +67,48 @@ class ZCardCmd : 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; +}; + } // namespace pikiwidb \ No newline at end of file From 27383e40f6684d3d7b194a060e06c48d98b222d2 Mon Sep 17 00:00:00 2001 From: panlei-coder Date: Tue, 19 Mar 2024 13:25:26 +0800 Subject: [PATCH 2/4] fix: add test for zrank, zrevrank, zrem, zincrby --- tests/zset_test.go | 68 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/tests/zset_test.go b/tests/zset_test.go index 5cdf309f6..740281ce8 100644 --- a/tests/zset_test.go +++ b/tests/zset_test.go @@ -215,4 +215,72 @@ var _ = Describe("Zset", Ordered, func() { Expect(err).NotTo(HaveOccurred()) Expect(card).To(Equal(int64(2))) }) + + 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))) + }) }) From 47505d77f078222382e42350dacde8465a6a3466 Mon Sep 17 00:00:00 2001 From: panlei-coder Date: Tue, 26 Mar 2024 12:18:05 +0800 Subject: [PATCH 3/4] fix: code format --- src/cmd_zset.cc | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/src/cmd_zset.cc b/src/cmd_zset.cc index 5561c43e8..535352759 100644 --- a/src/cmd_zset.cc +++ b/src/cmd_zset.cc @@ -319,15 +319,15 @@ void ZRevRangeByScoreCmd::DoCmd(PClient* client) { } } -ZRankCmd::ZRankCmd(const std::string &name, int16_t arity) - : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet){} +ZRankCmd::ZRankCmd(const std::string& name, int16_t arity) + : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} -bool ZRankCmd::DoInitial(PClient *client) { +bool ZRankCmd::DoInitial(PClient* client) { client->SetKey(client->argv_[1]); return true; } -void ZRankCmd::DoCmd(PClient *client) { +void ZRankCmd::DoCmd(PClient* client) { int32_t rank = 0; storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->ZRank(client->Key(), client->argv_[2], &rank); if (s.ok()) { @@ -339,15 +339,15 @@ void ZRankCmd::DoCmd(PClient *client) { } } -ZRevrankCmd::ZRevrankCmd(const std::string &name, int16_t arity) +ZRevrankCmd::ZRevrankCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySortedSet) {} -bool ZRevrankCmd::DoInitial(PClient *client) { +bool ZRevrankCmd::DoInitial(PClient* client) { client->SetKey(client->argv_[1]); return true; } -void ZRevrankCmd::DoCmd(PClient *client) { +void ZRevrankCmd::DoCmd(PClient* client) { int32_t revrank = 0; storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->ZRevrank(client->Key(), client->argv_[2], &revrank); if (s.ok()) { @@ -359,16 +359,15 @@ void ZRevrankCmd::DoCmd(PClient *client) { } } - -ZRemCmd::ZRemCmd(const std::string &name, int16_t arity) +ZRemCmd::ZRemCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} -bool ZRemCmd::DoInitial(PClient *client) { +bool ZRemCmd::DoInitial(PClient* client) { client->SetKey(client->argv_[1]); return true; } -void ZRemCmd::DoCmd(PClient *client) { +void ZRemCmd::DoCmd(PClient* client) { auto iter = client->argv_.begin() + 2; std::vector members(iter, client->argv_.end()); int32_t deleted = 0; @@ -380,15 +379,15 @@ void ZRemCmd::DoCmd(PClient *client) { } } -ZIncrbyCmd::ZIncrbyCmd(const std::string &name, int16_t arity) +ZIncrbyCmd::ZIncrbyCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} -bool ZIncrbyCmd::DoInitial(PClient *client) { +bool ZIncrbyCmd::DoInitial(PClient* client) { client->SetKey(client->argv_[1]); return true; } -void ZIncrbyCmd::DoCmd(PClient *client) { +void ZIncrbyCmd::DoCmd(PClient* client) { double by = .0f; double score = .0f; if (pstd::String2d(client->argv_[2].data(), client->argv_[2].size(), &by) == 0) { From 733eaa83c337c3ba30d4df93cf610b13c58a895d Mon Sep 17 00:00:00 2001 From: panlei-coder Date: Fri, 29 Mar 2024 21:59:21 +0800 Subject: [PATCH 4/4] fix: fix format --- src/cmd_zset.cc | 2 +- tests/zset_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cmd_zset.cc b/src/cmd_zset.cc index 932c53044..8fc8d02fb 100644 --- a/src/cmd_zset.cc +++ b/src/cmd_zset.cc @@ -760,7 +760,7 @@ void ZIncrbyCmd::DoCmd(PClient* client) { } ZRemrangebyscoreCmd::ZRemrangebyscoreCmd(const std::string& name, int16_t arity) - : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} + : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategorySortedSet) {} bool ZRemrangebyscoreCmd::DoInitial(PClient* client) { client->SetKey(client->argv_[1]); diff --git a/tests/zset_test.go b/tests/zset_test.go index a3f96f23a..eab3e7838 100644 --- a/tests/zset_test.go +++ b/tests/zset_test.go @@ -453,7 +453,7 @@ var _ = Describe("Zset", Ordered, func() { 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()