feat: support cmd keys (#217)

issue: #30 Signed-off-by: HappyUncle <[email protected]>
OpenAtomFoundation · Mar 22, 2024 · 5e770eb · 5e770eb
1 parent 14f8135
commit 5e770eb
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 0 deletions.
diff --git a/src/base_cmd.h b/src/base_cmd.h
@@ -29,6 +29,7 @@ const std::string kCmdNamePExpire = "pexpire";
 const std::string kCmdNameExpireat = "expireat";
 const std::string kCmdNamePExpireat = "pexpireat";
 const std::string kCmdNamePersist = "persist";
+const std::string kCmdNameKeys = "keys";
 
 // string cmd
 const std::string kCmdNameSet = "set";

diff --git a/src/cmd_keys.cc b/src/cmd_keys.cc
@@ -146,4 +146,25 @@ void PersistCmd::DoCmd(PClient* client) {
   }
 }
 
+KeysCmd::KeysCmd(const std::string& name, int16_t arity)
+    : BaseCmd(name, arity, kCmdFlagsReadonly, kAclCategoryRead | kAclCategoryKeyspace) {}
+
+bool KeysCmd::DoInitial(PClient* client) {
+  client->SetKey(client->argv_[1]);
+  return true;
+}
+
+void KeysCmd::DoCmd(PClient* client) {
+  std::vector<std::string> keys;
+  auto s = PSTORE.GetBackend(client->GetCurrentDB())->Keys(storage::DataType::kAll, client->Key(), &keys);
+  if (s.ok()) {
+    client->AppendArrayLen(keys.size());
+    for (auto k : keys) {
+      client->AppendString(k);
+    }
+  } else {
+    client->SetRes(CmdRes::kErrOther, s.ToString());
+  }
+}
+
 }  // namespace pikiwidb
diff --git a/src/cmd_keys.h b/src/cmd_keys.h
@@ -76,4 +76,15 @@ class PersistCmd : public BaseCmd {
  private:
   void DoCmd(PClient* client) override;
 };
+
+class KeysCmd : public BaseCmd {
+ public:
+  KeysCmd(const std::string& name, int16_t arity);
+
+ protected:
+  bool DoInitial(PClient* client) override;
+
+ private:
+  void DoCmd(PClient* client) override;
+};
 }  // namespace pikiwidb
diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc
@@ -50,6 +50,8 @@ void CmdTableManager::InitCmdTable() {
   ADD_COMMAND(Expireat, 3);
   ADD_COMMAND(PExpireat, 3);
   ADD_COMMAND(Persist, 2);
+  ADD_COMMAND(Keys, 2);
+
   // kv
   ADD_COMMAND(Get, 2);
   ADD_COMMAND(Set, -3);

diff --git a/tests/key_test.go b/tests/key_test.go
@@ -211,5 +211,33 @@ var _ = Describe("Keyspace", Ordered, func() {
 		Expect(client.Set(ctx, DefaultKey, DefaultValue, 0).Val()).To(Equal(OK))
 		Expect(client.PExpireAt(ctx, DefaultKey, time.Now().Add(time.Second*1000)).Err()).NotTo(HaveOccurred())
 		Expect(client.Persist(ctx, DefaultKey).Err()).NotTo(HaveOccurred())
+
+		// del keys
+		Expect(client.PExpireAt(ctx, DefaultKey, time.Now().Add(time.Second*1)).Err()).NotTo(HaveOccurred())
+		time.Sleep(2 * time.Second)
+	})
+
+	It("keys", func() {
+		// empty
+		Expect(client.Keys(ctx, "*").Val()).To(Equal([]string{}))
+		Expect(client.Keys(ctx, "dummy").Val()).To(Equal([]string{}))
+		Expect(client.Keys(ctx, "dummy*").Val()).To(Equal([]string{}))
+
+		Expect(client.Set(ctx, "a1", "v1", 0).Val()).To(Equal(OK))
+		Expect(client.Set(ctx, "k1", "v1", 0).Val()).To(Equal(OK))
+		Expect(client.SAdd(ctx, "k2", "v2").Val()).To(Equal(int64(1)))
+		Expect(client.HSet(ctx, "k3", "k3", "v3").Val()).To(Equal(int64(1)))
+		Expect(client.LPush(ctx, "k4", "v4").Val()).To(Equal(int64(1)))
+		Expect(client.ZAdd(ctx, "k5", redis.Z{Score: 1, Member: "v5"}).Val()).To(Equal(int64(1)))
+
+		// all
+		Expect(client.Keys(ctx, "*").Val()).To(Equal([]string{"a1", "k1", "k3", "k4", "k5", "k2"}))
+
+		// pattern
+		Expect(client.Keys(ctx, "k*").Val()).To(Equal([]string{"k1", "k3", "k4", "k5", "k2"}))
+		Expect(client.Keys(ctx, "k1").Val()).To(Equal([]string{"k1"}))
+
+		// del keys
+		Expect(client.Del(ctx, "a1", "k1", "k2", "k3", "k4", "k5").Err()).NotTo(HaveOccurred())
 	})
 })