Add support of COPY, REPLACE and KEYS options for the migrate command

Author: supercaracal

lib/redis.rb

@@ -499,22 +499,27 @@ def restore(key, ttl, serialized_value, options = {})
 
   # Transfer a key from the connected instance to another instance.
   #
-  # @param [String] key
+  # @param [String, Array<String>] key
   # @param [Hash] options
   #   - `:host => String`: host of instance to migrate to
   #   - `:port => Integer`: port of instance to migrate to
   #   - `:db => Integer`: database to migrate to (default: same as source)
   #   - `:timeout => Integer`: timeout (default: same as connection timeout)
+  #   - `:copy => Boolean`: Do not remove the key from the local instance.
+  #   - `:replace => Boolean`: Replace existing key on the remote instance.
   # @return [String] `"OK"`
   def migrate(key, options)
-    host = options[:host] || raise(RuntimeError, ":host not specified")
-    port = options[:port] || raise(RuntimeError, ":port not specified")
-    db = (options[:db] || @client.db).to_i
-    timeout = (options[:timeout] || @client.timeout).to_i
+    args = [:migrate]
+    args << (options[:host] || raise(':host not specified'))
+    args << (options[:port] || raise(':port not specified'))
+    args << (key.is_a?(String) ? key : '')
+    args << (options[:db] || @client.db).to_i
+    args << (options[:timeout] || @client.timeout).to_i
+    args << 'COPY' if options[:copy]
+    args << 'REPLACE' if options[:replace]
+    args += ['KEYS', *key] if key.is_a?(Array)
 
-    synchronize do |client|
-      client.call([:migrate, host, port, key, db, timeout])
-    end
+    synchronize { |client| client.call(args) }
   end
 
   # Delete one or more keys.

test/commands_on_value_types_test.rb

@@ -202,6 +202,21 @@ def test_migrate
       actual = redis.migrate("foo", options.merge(:timeout => default_timeout + 1))
       expected = ["127.0.0.1", "1234", "foo", default_db.to_s, (default_timeout + 1).to_s]
       assert_equal expected, actual
+
+      # Test copy override
+      actual = redis.migrate('foo', options.merge(copy: true))
+      expected = ['127.0.0.1', '1234', 'foo', default_db.to_s, default_timeout.to_s, 'COPY']
+      assert_equal expected, actual
+
+      # Test replace override
+      actual = redis.migrate('foo', options.merge(replace: true))
+      expected = ['127.0.0.1', '1234', 'foo', default_db.to_s, default_timeout.to_s, 'REPLACE']
+      assert_equal expected, actual
+
+      # Test multiple keys
+      actual = redis.migrate(%w[foo bar baz], options)
+      expected = ['127.0.0.1', '1234', '', default_db.to_s, default_timeout.to_s, 'KEYS', 'foo', 'bar', 'baz']
+      assert_equal expected, actual
     end
   end
 end