feat: zset cmd zrevrangebyscore (#188)

* feat: zset cmd zrevrangebyscore * format * format * format * fix
OpenAtomFoundation · Mar 11, 2024 · 0c81bc2 · 0c81bc2
1 parent 2c7aadd
commit 0c81bc2
Show file tree

Hide file tree

Showing 5 changed files with 133 additions and 12 deletions.
diff --git a/src/base_cmd.h b/src/base_cmd.h
@@ -109,6 +109,7 @@ const std::string kCmdNameLInsert = "linsert";
 const std::string kCmdNameZAdd = "zadd";
 const std::string kCmdNameZRevrange = "zrevrange";
 const std::string kCmdNameZRangebyscore = "zrangebyscore";
+const std::string kCmdNameZRevRangeByScore = "zrevrangebyscore";
 const std::string kCmdNameZCard = "zcard";
 
 enum CmdFlags {

diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc
@@ -116,6 +116,7 @@ void CmdTableManager::InitCmdTable() {
   ADD_COMMAND(ZAdd, -4);
   ADD_COMMAND(ZRevrange, -4);
   ADD_COMMAND(ZRangebyscore, -4);
+  ADD_COMMAND(ZRevRangeByScore, -4);
   ADD_COMMAND(ZCard, 2);
 }
 

diff --git a/src/cmd_zset.cc b/src/cmd_zset.cc
@@ -234,4 +234,89 @@ void ZCardCmd::DoCmd(PClient* client) {
   }
   client->AppendInteger(reply_Num);
 }
-}  // namespace pikiwidb
+
+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<storage::ScoreMember> score_members;
+  storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())
+                          ->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<int64_t>(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);
+    }
+  }
+}
+
+}  // namespace pikiwidb
diff --git a/src/cmd_zset.h b/src/cmd_zset.h
@@ -45,6 +45,17 @@ class ZRangebyscoreCmd : public BaseCmd {
   void DoCmd(PClient *client) override;
 };
 
+class ZRevRangeByScoreCmd : public BaseCmd {
+ public:
+  ZRevRangeByScoreCmd(const std::string &name, int16_t arity);
+
+ protected:
+  bool DoInitial(PClient *client) override;
+
+ private:
+  void DoCmd(PClient *client) override;
+};
+
 class ZCardCmd : public BaseCmd {
  public:
   ZCardCmd(const std::string &name, int16_t arity);
@@ -56,4 +67,4 @@ class ZCardCmd : public BaseCmd {
   void DoCmd(PClient *client) override;
 };
 
-}  // namespace pikiwidb
+}  // namespace pikiwidb
diff --git a/tests/zset_test.go b/tests/zset_test.go
@@ -175,21 +175,44 @@ var _ = Describe("Zset", Ordered, func() {
 		Expect(zRevRange.Val()).To(Equal([]string{"two", "one"}))
 	})
 
-        It("should ZCard", func() {
-		err :=client.ZAdd(ctx, "zsetZCard", redis.Z{
-			Score: 1,
+	It("should ZRevRangeByScore", func() {
+		err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
+		Expect(err).NotTo(HaveOccurred())
+		err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err()
+		Expect(err).NotTo(HaveOccurred())
+		err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err()
+		Expect(err).NotTo(HaveOccurred())
+
+		vals, err := client.ZRevRangeByScore(
+			ctx, "zset", &redis.ZRangeBy{Max: "+inf", Min: "-inf"}).Result()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(vals).To(Equal([]string{"three", "two", "one"}))
+
+		vals, err = client.ZRevRangeByScore(
+			ctx, "zset", &redis.ZRangeBy{Max: "2", Min: "(1"}).Result()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(vals).To(Equal([]string{"two"}))
+
+		vals, err = client.ZRevRangeByScore(
+			ctx, "zset", &redis.ZRangeBy{Max: "(2", Min: "(1"}).Result()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(vals).To(Equal([]string{}))
+	})
+
+	It("should ZCard", func() {
+		err := client.ZAdd(ctx, "zsetZCard", redis.Z{
+			Score:  1,
 			Member: "one",
 		}).Err()
 		Expect(err).NotTo(HaveOccurred())
-                err =client.ZAdd(ctx, "zsetZCard", redis.Z{
-                        Score: 2,
-                        Member: "two",
-                }).Err()
+		err = client.ZAdd(ctx, "zsetZCard", redis.Z{
+			Score:  2,
+			Member: "two",
+		}).Err()
 		Expect(err).NotTo(HaveOccurred())
 
-		card, err :=client.ZCard(ctx, "zsetZCard").Result()
+		card, err := client.ZCard(ctx, "zsetZCard").Result()
 		Expect(err).NotTo(HaveOccurred())
 		Expect(card).To(Equal(int64(2)))
-        })
-
+	})
 })