From 8de72d6d3c4621d1d3cc7171bfc1da9d56ede11a Mon Sep 17 00:00:00 2001
From: ofekshenawa <ofek.shenawa@redis.com>
Date: Mon, 23 Sep 2024 01:32:35 +0300
Subject: [PATCH] Extract keys from string commands

---
 command.go         | 16 +++++++++++++++
 commands_test.go   | 51 ++++++++++++++++++++++++++++++++++++++++------
 string_commands.go |  4 ++--
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/command.go b/command.go
index 2e4d1a1e..a18589c9 100644
--- a/command.go
+++ b/command.go
@@ -234,6 +234,22 @@ func (cmd *baseCmd) readRawReply(rd *proto.Reader) (err error) {
 	return err
 }
 
+// getInterleavedArguments returns arguments at even indices starting from index 0.
+func (cmd *baseCmd) getInterleavedArguments() []string {
+	return cmd.getInterleavedArgumentsWithOffset(0)
+}
+
+// getInterleavedArgumentsWithOffset returns arguments at even indices starting from the specified offset.
+func (cmd *baseCmd) getInterleavedArgumentsWithOffset(offset int) []string {
+	var matchingArguments []string
+	for i := offset; i < len(cmd.args); i += 2 {
+		if arg, ok := cmd.args[i].(string); ok {
+			matchingArguments = append(matchingArguments, arg)
+		}
+	}
+	return matchingArguments
+}
+
 //------------------------------------------------------------------------------
 
 type Cmd struct {
diff --git a/commands_test.go b/commands_test.go
index dd78f1d1..cf722a9e 100644
--- a/commands_test.go
+++ b/commands_test.go
@@ -7279,7 +7279,7 @@ var _ = Describe("Commands", func() {
 	})
 })
 
-var _ = Describe("Get Keys from Commands", func() {
+var _ = Describe("Keys Extraction Tests", func() {
 	var client *redis.Client
 	var ctx = context.TODO()
 
@@ -7292,12 +7292,51 @@ var _ = Describe("Get Keys from Commands", func() {
 		Expect(client.Close()).NotTo(HaveOccurred())
 	})
 
-	It("should test string commands", func() {
-		setCmd := client.Set(ctx, "key1", "val1", 0)
-		Expect(setCmd.Keys()).To(Equal([]string{"key1"}))
+	// STRING COMMANDS
 
-		getCmd := client.MGet(ctx, "key1", "key2", "key3")
-		Expect(getCmd.Keys()).To(Equal([]string{"key1", "key2", "key3"}))
+	It("should test Append command", func() {
+		cmd := client.Append(ctx, "key1", "value")
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test Decr command", func() {
+		cmd := client.Decr(ctx, "key1")
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test Get command", func() {
+		cmd := client.Get(ctx, "key1")
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test MGet command", func() {
+		cmd := client.MGet(ctx, "key1", "key2", "key3")
+		Expect(cmd.Keys()).To(Equal([]string{"key1", "key2", "key3"}))
+	})
+
+	It("should test Set command", func() {
+		cmd := client.Set(ctx, "key1", "value", time.Second)
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test MSet command", func() {
+		cmd := client.MSet(ctx, "key1", "value1", "key2", "value2")
+		Expect(cmd.Keys()).To(Equal([]string{"key1", "key2"}))
+	})
+
+	It("should test IncrBy command", func() {
+		cmd := client.IncrBy(ctx, "key1", 10)
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test SetNX command", func() {
+		cmd := client.SetNX(ctx, "key1", "value", time.Second)
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
+	})
+
+	It("should test GetDel command", func() {
+		cmd := client.GetDel(ctx, "key1")
+		Expect(cmd.Keys()).To(Equal([]string{"key1"}))
 	})
 })
 
diff --git a/string_commands.go b/string_commands.go
index bf84418e..292dfb32 100644
--- a/string_commands.go
+++ b/string_commands.go
@@ -151,7 +151,7 @@ func (c cmdable) MSet(ctx context.Context, values ...interface{}) *StatusCmd {
 	args[0] = "mset"
 	args = appendArgs(args, values)
 	cmd := NewStatusCmd(ctx, args...)
-	// cmd.keys = append(cmd.keys, keys...)
+	cmd.keys = append(cmd.keys, cmd.getInterleavedArgumentsWithOffset(1)...)
 	_ = c(ctx, cmd)
 	return cmd
 }
@@ -166,7 +166,7 @@ func (c cmdable) MSetNX(ctx context.Context, values ...interface{}) *BoolCmd {
 	args[0] = "msetnx"
 	args = appendArgs(args, values)
 	cmd := NewBoolCmd(ctx, args...)
-	// cmd.keys = append(cmd.keys, keys...)
+	cmd.keys = append(cmd.keys, cmd.getInterleavedArgumentsWithOffset(1)...)
 	_ = c(ctx, cmd)
 	return cmd
 }