Merge pull request #698 from mumumu/add_replace_option_to_restore_command

implemented REPLACE modifier of restore command

2 files changed, 15 insertions, 2 deletions

diff --git a/redis/client.py b/redis/client.py
index 7426ead..765eb87 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -1117,12 +1117,15 @@ class StrictRedis(object):
         "Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist"
         return self.execute_command('RENAMENX', src, dst)
 
-    def restore(self, name, ttl, value):
+    def restore(self, name, ttl, value, replace=False):
         """
         Create a key using the provided serialized value, previously obtained
         using DUMP.
         """
-        return self.execute_command('RESTORE', name, ttl, value)
+        params = [name, ttl, value]
+        if replace:
+            params.append('REPLACE')
+        return self.execute_command('RESTORE', *params)
 
     def set(self, name, value, ex=None, px=None, nx=False, xx=False):
         """
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 06b19b7..246ef65 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -285,6 +285,16 @@ class TestRedisCommands(object):
         r.restore('a', 0, dumped)
         assert r['a'] == b('foo')
 
+    @skip_if_server_version_lt('3.0.0')
+    def test_dump_and_restore_and_replace(self, r):
+        r['a'] = 'bar'
+        dumped = r.dump('a')
+        with pytest.raises(redis.ResponseError):
+            r.restore('a', 0, dumped)
+
+        r.restore('a', 0, dumped, replace=True)
+        assert r['a'] == b('bar')
+
     def test_exists(self, r):
         assert not r.exists('a')
         r['a'] = 'foo'