Implement RedisClient::Cluster::Command#extract_all_keys

lib/redis_client/cluster/command.rb

@@ -10,10 +10,13 @@ class Cluster
     class Command
       EMPTY_STRING = ''
       EMPTY_HASH = {}.freeze
+      EMPTY_ARRAY = [].freeze
 
       Detail = Struct.new(
         'RedisCommand',
         :first_key_position,
+        :last_key_position,
+        :key_step,
         :write?,
         :readonly?,
         keyword_init: true
@@ -49,6 +52,8 @@ def parse_command_reply(rows)
 
             acc[row[0].downcase] = ::RedisClient::Cluster::Command::Detail.new(
               first_key_position: row[3],
+              last_key_position: row[4],
+              key_step: row[5],
               write?: row[2].include?('write'),
               readonly?: row[2].include?('readonly')
             )
@@ -67,6 +72,17 @@ def extract_first_key(command)
         (command[i].is_a?(Array) ? command[i].flatten.first : command[i]).to_s
       end
 
+      def extract_all_keys(command)
+        keys_start = determine_first_key_position(command)
+        keys_end = determine_last_key_position(command, keys_start)
+        keys_step = determine_key_step(command)
+        return EMPTY_ARRAY if [keys_start, keys_end, keys_step].any?(&:zero?)
+
+        keys_end = [keys_end, command.size - 1].min
+        # use .. inclusive range because keys_end is a valid index.
+        (keys_start..keys_end).step(keys_step).map { |i| command[i] }
+      end
+
       def should_send_to_primary?(command)
         name = ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
         @commands[name]&.write?
@@ -98,6 +114,41 @@ def determine_first_key_position(command) # rubocop:disable Metrics/CyclomaticCo
         end
       end
 
+      # IMPORTANT: this determines the last key position INCLUSIVE of the last key -
+      # i.e. command[determine_last_key_position(command)] is a key.
+      # This is in line with what Redis returns from COMMANDS.
+      def determine_last_key_position(command, keys_start) # rubocop:disable Metrics/AbcSize
+        case name = ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
+        when 'eval', 'evalsha', 'zinterstore', 'zunionstore'
+          # EVALSHA sha1 numkeys [key [key ...]] [arg [arg ...]]
+          # ZINTERSTORE destination numkeys key [key ...] [WEIGHTS weight [weight ...]] [AGGREGATE <SUM | MIN | MAX>]
+          command[2].to_i + 2
+        when 'object', 'memory'
+          # OBJECT [ENCODING | FREQ | IDLETIME | REFCOUNT] key
+          # MEMORY USAGE key [SAMPLES count]
+          keys_start
+        when 'migrate'
+          # MIGRATE host port <key | ""> destination-db timeout [COPY] [REPLACE] [AUTH password | AUTH2 username password] [KEYS key [key ...]]
+          command[3].empty? ? (command.length - 1) : 3
+        when 'xread', 'xreadgroup'
+          # XREAD [COUNT count] [BLOCK milliseconds] STREAMS key [key ...] id [id ...]
+          keys_start + ((command.length - keys_start) / 2) - 1
+        else
+          # If there is a fixed, non-variable number of keys, don't iterate past that.
+          if @commands[name].last_key_position >= 0
+            @commands[name].last_key_position
+          else
+            command.length + @commands[name].last_key_position
+          end
+        end
+      end
+
+      def determine_key_step(command)
+        name = ::RedisClient::Cluster::NormalizedCmdName.instance.get_by_command(command)
+        # Some commands like EVALSHA have zero as the step in COMMANDS somehow.
+        @commands[name].key_step == 0 ? 1 : @commands[name].key_step
+      end
+
       def determine_optional_key_position(command, option_name) # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
         idx = command&.flatten&.map(&:to_s)&.map(&:downcase)&.index(option_name&.downcase)
         idx.nil? ? 0 : idx + 1

test/redis_client/cluster/test_command.rb

@@ -49,20 +49,20 @@ def test_parse_command_reply
         [
           {
             rows: [
-              ['get', 2, Set['readonly', 'fast'], 1, 1, 1, Set['@read', '@string', '@fast'], Set[], Set[], Set[]],
-              ['set', -3, Set['write', 'denyoom', 'movablekeys'], 1, 1, 1, Set['@write', '@string', '@slow'], Set[], Set[], Set[]]
+              ['get', 2, Set['readonly', 'fast'], 1, -1, 1, Set['@read', '@string', '@fast'], Set[], Set[], Set[]],
+              ['set', -3, Set['write', 'denyoom', 'movablekeys'], 1, -1, 2, Set['@write', '@string', '@slow'], Set[], Set[], Set[]]
             ],
             want: {
-              'get' => { first_key_position: 1, write?: false, readonly?: true },
-              'set' => { first_key_position: 1, write?: true, readonly?: false }
+              'get' => { first_key_position: 1, last_key_position: -1, key_step: 1, write?: false, readonly?: true },
+              'set' => { first_key_position: 1, last_key_position: -1, key_step: 2, write?: true, readonly?: false }
             }
           },
           {
             rows: [
-              ['GET', 2, Set['readonly', 'fast'], 1, 1, 1, Set['@read', '@string', '@fast'], Set[], Set[], Set[]]
+              ['GET', 2, Set['readonly', 'fast'], 1, -1, 1, Set['@read', '@string', '@fast'], Set[], Set[], Set[]]
             ],
             want: {
-              'get' => { first_key_position: 1, write?: false, readonly?: true }
+              'get' => { first_key_position: 1, last_key_position: -1, key_step: 1, write?: false, readonly?: true }
             }
           },
           { rows: [[]], want: {} },
@@ -190,6 +190,37 @@ def test_determine_optional_key_position
           assert_equal(c[:want], got, msg)
         end
       end
+
+      def test_extract_all_keys
+        cmd = ::RedisClient::Cluster::Command.load(@raw_clients)
+        [
+          { command: ['EVAL', 'return ARGV[1]', '0', 'hello'], want: [] },
+          { command: ['EVAL', 'return ARGV[1]', '3', 'key1', 'key2', 'key3', 'arg1', 'arg2'], want: %w[key1 key2 key3] },
+          { command: [['EVAL'], '"return ARGV[1]"', 0, 'hello'], want: [] },
+          { command: %w[EVALSHA sha1 2 foo bar baz zap], want: %w[foo bar] },
+          { command: %w[MIGRATE host port key 0 5 COPY], want: %w[key] },
+          { command: ['MIGRATE', 'host', 'port', '', '0', '5', 'COPY', 'KEYS', 'key1'], want: %w[key1] },
+          { command: ['MIGRATE', 'host', 'port', '', '0', '5', 'COPY', 'KEYS', 'key1', 'key2'], want: %w[key1 key2] },
+          { command: %w[ZINTERSTORE out 2 zset1 zset2 WEIGHTS 2 3], want: %w[zset1 zset2] },
+          { command: %w[ZUNIONSTORE out 2 zset1 zset2 WEIGHTS 2 3], want: %w[zset1 zset2] },
+          { command: %w[OBJECT HELP], want: [] },
+          { command: %w[MEMORY HELP], want: [] },
+          { command: %w[MEMORY USAGE key], want: %w[key] },
+          { command: %w[XREAD COUNT 2 STREAMS mystream writers 0-0 0-0], want: %w[mystream writers] },
+          { command: %w[XREADGROUP GROUP group consumer STREAMS key id], want: %w[key] },
+          { command: %w[SET foo 1], want: %w[foo] },
+          { command: %w[set foo 1], want: %w[foo] },
+          { command: [['SET'], 'foo', 1], want: %w[foo] },
+          { command: %w[GET foo], want: %w[foo] },
+          { command: %w[MGET foo bar baz], want: %w[foo bar baz] },
+          { command: %w[MSET foo val bar val baz val], want: %w[foo bar baz] },
+          { command: %w[BLPOP foo bar 0], want: %w[foo bar] }
+        ].each_with_index do |c, idx|
+          msg = "Case: #{idx}"
+          got = cmd.send(:extract_all_keys, c[:command])
+          assert_equal(c[:want], got, msg)
+        end
+      end
     end
   end
 end