fix issue 1517: scan 命令不支持 type 参数

OpenAtomFoundation · Jun 3, 2023 · 20e4eac · 20e4eac
1 parent 162c661
commit 20e4eac
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 2 deletions.
diff --git 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
@@ -1090,14 +1090,30 @@ 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);
         return;
       }
       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
@@ -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