fix: let set cmd support `SET key value [NX | XX] [EX seconds | PX mi…

…lliseconds]` OpenAtomFoundation#320 Signed-off-by: HappyUncle <[email protected]>
happy-v587 · May 26, 2024 · 0e45fc7 · 0e45fc7
1 parent 128d408
commit 0e45fc7
Show file tree

Hide file tree

Showing 3 changed files with 169 additions and 8 deletions.
diff --git a/src/cmd_kv.cc b/src/cmd_kv.cc
@@ -37,15 +37,72 @@ void GetCmd::DoCmd(PClient* client) {
 SetCmd::SetCmd(const std::string& name, int16_t arity)
     : BaseCmd(name, arity, kCmdFlagsWrite, kAclCategoryWrite | kAclCategoryString) {}
 
+// SET key value [NX | XX] [EX seconds | PX milliseconds]
 bool SetCmd::DoInitial(PClient* client) {
   client->SetKey(client->argv_[1]);
+
+  auto argv_ = client->argv_;
+  value_ = argv_[2];
+  condition_ = SetCmd::kNONE;
+  sec_ = 0;
+  size_t index = 3;
+
+  while (index != argv_.size()) {
+    std::string opt = argv_[index];
+    if (strcasecmp(opt.data(), "xx") == 0) {
+      condition_ = SetCmd::kXX;
+    } else if (strcasecmp(opt.data(), "nx") == 0) {
+      condition_ = SetCmd::kNX;
+    } else if ((strcasecmp(opt.data(), "ex") == 0) || (strcasecmp(opt.data(), "px") == 0)) {
+      condition_ = (condition_ == SetCmd::kNONE) ? SetCmd::kEXORPX : condition_;
+      index++;
+      if (index == argv_.size()) {
+        client->SetRes(CmdRes::kSyntaxErr);
+        return false;
+      }
+      if (pstd::String2int(argv_[index].data(), argv_[index].size(), &sec_) == 0) {
+        client->SetRes(CmdRes::kInvalidInt);
+        return false;
+      }
+
+      if (strcasecmp(opt.data(), "px") == 0) {
+        sec_ /= 1000;
+      }
+    } else {
+      client->SetRes(CmdRes::kSyntaxErr);
+      return false;
+    }
+    index++;
+  }
+
   return true;
 }
 
 void SetCmd::DoCmd(PClient* client) {
-  storage::Status s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->Set(client->Key(), client->argv_[2]);
-  if (s.ok()) {
-    client->SetRes(CmdRes::kOK);
+  int32_t res = 1;
+  storage::Status s;
+  auto key_ = client->Key();
+  switch (condition_) {
+    case SetCmd::kXX:
+      s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->Setxx(key_, value_, &res, sec_);
+      break;
+    case SetCmd::kNX:
+      s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->Setnx(key_, value_, &res, sec_);
+      break;
+    case SetCmd::kEXORPX:
+      s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->Setex(key_, value_, sec_);
+      break;
+    default:
+      s = PSTORE.GetBackend(client->GetCurrentDB())->GetStorage()->Set(key_, value_);
+      break;
+  }
+
+  if (s.ok() || s.IsNotFound()) {
+    if (res == 1) {
+      client->SetRes(CmdRes::kOK);
+    } else {
+      client->AppendStringLen(-1);
+    }
   } else {
     client->SetRes(CmdRes::kErrOther, s.ToString());
   }

diff --git a/src/cmd_kv.h b/src/cmd_kv.h
@@ -24,13 +24,19 @@ class GetCmd : public BaseCmd {
 
 class SetCmd : public BaseCmd {
  public:
+  enum SetCondition { kNONE, kNX, kXX, kEXORPX };
   SetCmd(const std::string &name, int16_t arity);
 
  protected:
   bool DoInitial(PClient *client) override;
 
  private:
   void DoCmd(PClient *client) override;
+
+  std::string value_;
+  std::string target_;
+  int64_t sec_ = 0;
+  SetCmd::SetCondition condition_{kNONE};
 };
 
 class BitOpCmd : public BaseCmd {

diff --git a/tests/key_test.go b/tests/key_test.go
@@ -65,6 +65,104 @@ var _ = Describe("Keyspace", Ordered, func() {
 		}
 	})
 
+	It("Set", func() {
+		{
+			// set px
+			res, err := client.Set(ctx, "a", "a", time.Millisecond*1001).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal("OK"))
+			time.Sleep(time.Millisecond * 2000)
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(0)))
+		}
+		{
+			// set ex
+			res, err := client.Set(ctx, "a", "a", time.Second*60).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal("OK"))
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(1)))
+		}
+		{
+			// set xx
+			res, err := client.SetXX(ctx, "a", "a", 0).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(1)))
+		}
+		{
+			// set ex xx
+			res, err := client.SetXX(ctx, "a", "a", time.Second*30).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(1)))
+		}
+		{
+			// set px xx
+			res, err := client.SetXX(ctx, "a", "a", time.Millisecond*1001).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+			time.Sleep(time.Millisecond * 2000)
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(0)))
+		}
+		{
+			// set ex nx
+			res, err := client.SetNX(ctx, "a", "a", time.Second).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+			time.Sleep(time.Second * 2)
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(0)))
+		}
+		{
+			// set px nx
+			res, err := client.SetNX(ctx, "a", "a", time.Millisecond*1001).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+			time.Sleep(time.Millisecond * 2000)
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(0)))
+		}
+		{
+			// set nx
+			res, err := client.SetNX(ctx, "a", "a", 0).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal(true))
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(1)))
+		}
+		{
+			// setex
+			res, err := client.SetEx(ctx, "a", "a", time.Second).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(res).To(Equal("OK"))
+			time.Sleep(time.Second * 2)
+
+			n, err := client.Exists(ctx, "a").Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(n).To(Equal(int64(0)))
+		}
+	})
+
 	//TODO(dingxiaoshuai) Add more test cases.
 	It("Exists", func() {
 		n, err := client.Exists(ctx, "key1").Result()
@@ -317,28 +415,28 @@ var _ = Describe("Keyspace", Ordered, func() {
 		Expect(client.Do(ctx, "pexpire", DefaultKey, "err").Err()).To(MatchError("ERR value is not an integer or out of range"))
 	})
 
-	It("should Rename", func ()  {
+	It("should Rename", func() {
 		client.Set(ctx, "mykey", "hello", 0)
 		client.Rename(ctx, "mykey", "mykey1")
 		client.Rename(ctx, "mykey1", "mykey2")
 		Expect(client.Get(ctx, "mykey2").Val()).To(Equal("hello"))
 
 		Expect(client.Exists(ctx, "mykey").Val()).To(Equal(int64(0)))
-	
+
 		client.Set(ctx, "mykey", "foo", 0)
 		Expect(client.Rename(ctx, "mykey", "mykey").Val()).To(Equal(OK))
-	
+
 		client.Del(ctx, "mykey", "mykey2")
 		client.Set(ctx, "mykey", "foo", 0)
 		client.Set(ctx, "mykey2", "bar", 0)
-		client.Expire(ctx, "mykey2", 100 * time.Second)
+		client.Expire(ctx, "mykey2", 100*time.Second)
 		Expect(client.TTL(ctx, "mykey").Val()).To(Equal(-1 * time.Nanosecond))
 		Expect(client.TTL(ctx, "mykey2").Val()).NotTo(Equal(-1 * time.Nanosecond))
 		client.Rename(ctx, "mykey", "mykey2")
 		Expect(client.TTL(ctx, "mykey2").Val()).To(Equal(-1 * time.Nanosecond))
 	})
 
-	It("should RenameNX", func ()  {
+	It("should RenameNX", func() {
 		client.Del(ctx, "mykey", "mykey1", "mykey2")
 		client.Set(ctx, "mykey", "hello", 0)
 		client.RenameNX(ctx, "mykey", "mykey1")