Add the TYPE argument to scan/scan_each, targeting 6.0.0+

Author: tednology

lib/redis.rb

@@ -2638,12 +2638,13 @@ def evalsha(*args)
     _eval(:evalsha, args)
   end
 
-  def _scan(command, cursor, args, match: nil, count: nil, &block)
+  def _scan(command, cursor, args, match: nil, count: nil, type: nil, &block)
     # SSCAN/ZSCAN/HSCAN already prepend the key to +args+.
 
     args << cursor
     args << "MATCH" << match if match
     args << "COUNT" << count if count
+    args << "TYPE" << type if type
 
     synchronize do |client|
       client.call([command] + args, &block)
@@ -2658,11 +2659,15 @@ def _scan(command, cursor, args, match: nil, count: nil, &block)
   # @example Retrieve a batch of keys matching a pattern
   #   redis.scan(4, :match => "key:1?")
   #     # => ["92", ["key:13", "key:18"]]
+  # @example Retrieve a batch of keys of a certain type
+  #   redis.scan(92, :type => "zset")
+  #     # => ["173", ["sortedset:14", "sortedset:78"]]
   #
   # @param [String, Integer] cursor the cursor of the iteration
   # @param [Hash] options
   #   - `:match => String`: only return keys matching the pattern
   #   - `:count => Integer`: return count keys at most per iteration
+  #   - `:type => String`: return keys only of the given type
   #
   # @return [String, Array<String>] the next cursor and all found keys
   def scan(cursor, **options)
@@ -2678,10 +2683,15 @@ def scan(cursor, **options)
   #   redis.scan_each(:match => "key:1?") {|key| puts key}
   #     # => key:13
   #     # => key:18
+  # @example Execute block for each key of a type
+  #   redis.scan_each(:type => "hash") {|key| puts redis.type(key)}
+  #     # => "hash"
+  #     # => "hash"
   #
   # @param [Hash] options
   #   - `:match => String`: only return keys matching the pattern
   #   - `:count => Integer`: return count keys at most per iteration
+  #   - `:type => String`: return keys only of the given type
   #
   # @return [Enumerator] an enumerator for all found keys
   def scan_each(**options, &block)

test/scanning_test.rb

@@ -53,6 +53,25 @@ def test_scan_match
     end
   end
 
+  def test_scan_type
+    target_version "6.0.0" do
+      r.debug :populate, 1000
+      r.zadd("foo", [1, "s1", 2, "s2", 3, "s3"])
+      r.zadd("bar", [6, "s1", 5, "s2", 4, "s3"])
+      r.hset("baz", "k1", "v1")
+
+      cursor = 0
+      all_keys = []
+      loop do
+        cursor, keys = r.scan cursor, type: "zset"
+        all_keys += keys
+        break if cursor == "0"
+      end
+
+      assert_equal 2, all_keys.uniq.size
+    end
+  end
+
   def test_scan_each_enumerator
     target_version "2.7.105" do
       r.debug :populate, 1000
@@ -78,6 +97,20 @@ def test_scan_each_enumerator_match
     end
   end
 
+  def test_scan_each_enumerator_type
+    target_version "6.0.0" do
+      r.debug :populate, 1000
+      r.zadd("key:zset", [1, "s1", 2, "s2", 3, "s3"])
+      r.hset("key:hash:1", "k1", "v1")
+      r.hset("key:hash:2", "k2", "v2")
+
+      keys_from_scan = r.scan_each(type: "hash").to_a.uniq
+      all_keys = r.keys "key:hash:*"
+
+      assert all_keys.sort == keys_from_scan.sort
+    end
+  end
+
   def test_scan_each_block
     target_version "2.7.105" do
       r.debug :populate, 100