From 20e4eac6392f9347767183b7a97aaca5d3efd804 Mon Sep 17 00:00:00 2001
From: ptbxzrt <ptbxzrt@aliyun.com>
Date: Sat, 3 Jun 2023 20:13:15 +0800
Subject: [PATCH] =?UTF-8?q?fix=20issue=201517:=20scan=20=E5=91=BD=E4=BB=A4?=
 =?UTF-8?q?=E4=B8=8D=E6=94=AF=E6=8C=81=20type=20=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 include/pika_kv.h   |  2 ++
 src/pika_kv.cc      | 20 ++++++++++++++++++--
 tests/unit/scan.tcl | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/include/pika_kv.h b/include/pika_kv.h
index a87303c55a..44232180e0 100644
--- a/include/pika_kv.h
+++ b/include/pika_kv.h
@@ -600,10 +600,12 @@ class ScanCmd : public Cmd {
   int64_t cursor_ = 0;
   std::string pattern_ = "*";
   int64_t count_ = 10;
+  storage::DataType type_ = storage::kAll;
   void DoInitial() override;
   void Clear() override {
     pattern_ = "*";
     count_ = 10;
+    type_ = storage::kAll;
   }
 };
 
diff --git a/src/pika_kv.cc b/src/pika_kv.cc
index d0a6987589..d37460007a 100644
--- a/src/pika_kv.cc
+++ b/src/pika_kv.cc
@@ -1090,7 +1090,8 @@ void ScanCmd::DoInitial() {
 
   while (index < argc) {
     std::string opt = argv_[index];
-    if ((strcasecmp(opt.data(), "match") == 0) || (strcasecmp(opt.data(), "count") == 0)) {
+    if ((strcasecmp(opt.data(), "match") == 0) || (strcasecmp(opt.data(), "count") == 0) ||
+        (strcasecmp(opt.data(), "type") == 0)) {
       index++;
       if (index >= argc) {
         res_.SetRes(CmdRes::kSyntaxErr);
@@ -1098,6 +1099,21 @@ void ScanCmd::DoInitial() {
       }
       if (strcasecmp(opt.data(), "match") == 0) {
         pattern_ = argv_[index];
+      } else if (strcasecmp(opt.data(), "type") == 0) {
+        std::string str_type = argv_[index];
+        if (strcasecmp(str_type.data(), "string") == 0) {
+          type_ = storage::DataType::kStrings;
+        } else if (strcasecmp(str_type.data(), "zset") == 0) {
+          type_ = storage::DataType::kZSets;
+        } else if (strcasecmp(str_type.data(), "set") == 0) {
+          type_ = storage::DataType::kSets;
+        } else if (strcasecmp(str_type.data(), "list") == 0) {
+          type_ = storage::DataType::kLists;
+        } else if (strcasecmp(str_type.data(), "hash") == 0) {
+          type_ = storage::DataType::kHashes;
+        } else {
+          res_.SetRes(CmdRes::kSyntaxErr);
+        }
       } else if ((pstd::string2int(argv_[index].data(), argv_[index].size(), &count_) == 0) || count_ <= 0) {
         res_.SetRes(CmdRes::kInvalidInt);
         return;
@@ -1123,7 +1139,7 @@ void ScanCmd::Do(std::shared_ptr<Slot> slot) {
     keys.clear();
     batch_count = left < PIKA_SCAN_STEP_LENGTH ? left : PIKA_SCAN_STEP_LENGTH;
     left = left > PIKA_SCAN_STEP_LENGTH ? left - PIKA_SCAN_STEP_LENGTH : 0;
-    cursor_ret = slot->db()->Scan(storage::DataType::kAll, cursor_ret, pattern_, batch_count, &keys);
+    cursor_ret = slot->db()->Scan(type_, cursor_ret, pattern_, batch_count, &keys);
     for (const auto& key : keys) {
       RedisAppendLen(raw, key.size(), "$");
       RedisAppendContent(raw, key);
diff --git a/tests/unit/scan.tcl b/tests/unit/scan.tcl
index 1d84f128da..5d0c7ccc95 100644
--- a/tests/unit/scan.tcl
+++ b/tests/unit/scan.tcl
@@ -53,6 +53,51 @@ start_server {tags {"scan"}} {
         assert_equal 100 [llength $keys]
     }
 
+    test "SCAN TYPE" {
+        r flushdb
+        # populate only creates strings
+        r debug populate 1000
+
+        # Check non-strings are excluded
+        set cur 0
+        set keys {}
+        while 1 {
+            set res [r scan $cur type "list"]
+            set cur [lindex $res 0]
+            set k [lindex $res 1]
+            lappend keys {*}$k
+            if {$cur == 0} break
+        }
+
+        assert_equal 0 [llength $keys]
+
+        # Check strings are included
+        set cur 0
+        set keys {}
+        while 1 {
+            set res [r scan $cur type "string"]
+            set cur [lindex $res 0]
+            set k [lindex $res 1]
+            lappend keys {*}$k
+            if {$cur == 0} break
+        }
+
+        assert_equal 1000 [llength $keys]
+
+        # Check all three args work together
+        set cur 0
+        set keys {}
+        while 1 {
+            set res [r scan $cur type "string" match "key:*" count 10]
+            set cur [lindex $res 0]
+            set k [lindex $res 1]
+            lappend keys {*}$k
+            if {$cur == 0} break
+        }
+
+        assert_equal 1000 [llength $keys]
+    }
+
     foreach enc {intset hashtable} {
         test "SSCAN with encoding $enc" {
             # Create the Set