feat: add command sscan (#263)

* add command SSCAN * fix format error --------- Co-authored-by: shenmengju <[email protected]>
OpenAtomFoundation · Apr 10, 2024 · 78e2744 · 78e2744
1 parent b829e45
commit 78e2744
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 0 deletions.
diff --git a/src/base_cmd.h b/src/base_cmd.h
@@ -103,6 +103,7 @@ const std::string kCmdNameSPop = "spop";
 const std::string kCmdNameSMembers = "smembers";
 const std::string kCmdNameSDiff = "sdiff";
 const std::string kCmdNameSDiffstore = "sdiffstore";
+const std::string kCmdNameSScan = "sscan";
 
 // list cmd
 const std::string kCmdNameLPush = "lpush";

diff --git a/src/cmd_set.cc b/src/cmd_set.cc
@@ -316,4 +316,64 @@ void SDiffstoreCmd::DoCmd(PClient* client) {
   }
   client->AppendInteger(reply_num);
 }
+
+SScanCmd::SScanCmd(const std::string& name, int16_t arity)
+    : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategorySet) {}
+
+bool SScanCmd::DoInitial(PClient* client) {
+  if (auto size = client->argv_.size(); size != 3 && size != 5 && size != 7) {
+    client->SetRes(CmdRes::kSyntaxErr);
+    return false;
+  }
+  client->SetKey(client->argv_[1]);
+  return true;
+}
+
+void SScanCmd::DoCmd(PClient* client) {
+  const auto& argv = client->argv_;
+  // parse arguments
+  int64_t cursor = 0;
+  int64_t count = 10;
+  std::string pattern{"*"};
+  if (pstd::String2int(argv[2], &cursor) == 0) {
+    client->SetRes(CmdRes::kInvalidCursor, kCmdNameSScan);
+    return;
+  }
+  for (size_t i = 3; i < argv.size(); i += 2) {
+    if (auto lower = pstd::StringToLower(argv[i]); kMatchSymbol == lower) {
+      pattern = argv[i + 1];
+    } else if (kCountSymbol == lower) {
+      if (pstd::String2int(argv[i + 1], &count) == 0) {
+        client->SetRes(CmdRes::kInvalidInt, kCmdNameSScan);
+        return;
+      }
+      if (count < 0) {
+        client->SetRes(CmdRes::kSyntaxErr, kCmdNameSScan);
+        return;
+      }
+    } else {
+      client->SetRes(CmdRes::kSyntaxErr, kCmdNameSScan);
+      return;
+    }
+  }
+
+  // execute command
+  std::vector<std::string> members;
+  int64_t next_cursor{};
+  auto status = PSTORE.GetBackend(client->GetCurrentDB())
+                    ->GetStorage()
+                    ->SScan(client->Key(), cursor, pattern, count, &members, &next_cursor);
+  if (!status.ok() && !status.IsNotFound()) {
+    client->SetRes(CmdRes::kErrOther, status.ToString());
+    return;
+  }
+
+  // reply to client
+  client->AppendArrayLen(2);
+  client->AppendString(std::to_string(next_cursor));
+  client->AppendArrayLenUint64(members.size());
+  for (const auto& member : members) {
+    client->AppendString(member);
+  }
+}
 }  // namespace pikiwidb
diff --git a/src/cmd_set.h b/src/cmd_set.h
@@ -165,4 +165,17 @@ class SDiffstoreCmd : public BaseCmd {
   void DoCmd(PClient *client) override;
 };
 
+class SScanCmd : public BaseCmd {
+ public:
+  SScanCmd(const std::string &name, int16_t arity);
+
+ protected:
+  bool DoInitial(PClient *client) override;
+
+ private:
+  void DoCmd(PClient *client) override;
+
+  static constexpr const char *kMatchSymbol = "match";
+  static constexpr const char *kCountSymbol = "count";
+};
 }  // namespace pikiwidb
diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc
@@ -110,6 +110,7 @@ void CmdTableManager::InitCmdTable() {
   ADD_COMMAND(SMembers, 2);
   ADD_COMMAND(SDiff, -2);
   ADD_COMMAND(SDiffstore, -3);
+  ADD_COMMAND(SScan, -3);
 
   // list
   ADD_COMMAND(LPush, -3);

diff --git a/tests/set_test.go b/tests/set_test.go
@@ -373,4 +373,47 @@ var _ = Describe("Set", Ordered, func() {
 		Expect(sMembers.Val()).To(ConsistOf([]string{"a", "b"}))
 	})
 
+	It("should SScan", func() {
+		// add elements first
+		sAdd := client.SAdd(ctx, "setSScan1", "user1")
+		Expect(sAdd.Err()).NotTo(HaveOccurred())
+		Expect(sAdd.Val()).To(Equal(int64(1)))
+
+		sAdd = client.SAdd(ctx, "setSScan1", "user2")
+		Expect(sAdd.Err()).NotTo(HaveOccurred())
+		Expect(sAdd.Val()).To(Equal(int64(1)))
+
+		sAdd = client.SAdd(ctx, "setSScan1", "user3")
+		Expect(sAdd.Err()).NotTo(HaveOccurred())
+		Expect(sAdd.Val()).To(Equal(int64(1)))
+
+		set := []string{"Hello", "World", "World"}
+		sAdd = client.SAdd(ctx, "setSScan1", set)
+		Expect(sAdd.Err()).NotTo(HaveOccurred())
+		Expect(sAdd.Val()).To(Equal(int64(2)))
+
+		// func (c Client) SScan(ctx context.Context, key string, cursor uint64, match string, count int64) *ScanCmd
+		sScan:=client.SScan(ctx,"setSScan1",0,"*",5)
+		Expect(sScan.Err()).NotTo(HaveOccurred())
+		Expect(sScan.Val()).To(ConsistOf([]string{"user1", "user2","user3","Hello","World"}))
+
+		sScan=client.SScan(ctx,"setSScan1",0,"user*",5)
+		Expect(sScan.Err()).NotTo(HaveOccurred())
+		Expect(sScan.Val()).To(ConsistOf([]string{"user1", "user2","user3"}))
+
+		sScan=client.SScan(ctx,"setSScan1",0,"He*",5)
+		Expect(sScan.Err()).NotTo(HaveOccurred())
+		Expect(sScan.Val()).To(ConsistOf([]string{"Hello"}))
+
+		// sScan=client.SScan(ctx,"setSScan1",0,"*",-1)
+		// Expect(sScan.Err()).To(HaveOccurred())
+
+
+		//del
+		del := client.Del(ctx, "setSScan1")
+		Expect(del.Err()).NotTo(HaveOccurred())
+	})
+
+
+
 })