IDLETIME and FREQ support for RESTORE (#1580)

Author: chayim

redis/commands.py

@@ -1034,7 +1034,8 @@ def renamenx(self, src, dst):
         "Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist"
         return self.execute_command('RENAMENX', src, dst)
 
-    def restore(self, name, ttl, value, replace=False, absttl=False):
+    def restore(self, name, ttl, value, replace=False, absttl=False,
+                idletime=None, frequency=None):
         """
         Create a key using the provided serialized value, previously obtained
         using DUMP.
@@ -1045,12 +1046,32 @@ def restore(self, name, ttl, value, replace=False, absttl=False):
         ``absttl`` if True, specified ``ttl`` should represent an absolute Unix
         timestamp in milliseconds in which the key will expire. (Redis 5.0 or
         greater).
+
+        ``idletime`` Used for eviction, this is the number of seconds the
+        key must be idle, prior to execution.
+
+        ``frequency`` Used for eviction, this is the frequency counter of
+        the object stored at the key, prior to execution.
         """
         params = [name, ttl, value]
         if replace:
             params.append('REPLACE')
         if absttl:
             params.append('ABSTTL')
+        if idletime is not None:
+            params.append('IDLETIME')
+            try:
+                params.append(int(idletime))
+            except ValueError:
+                raise DataError("idletimemust be an integer")
+
+        if frequency is not None:
+            params.append('FREQ')
+            try:
+                params.append(int(frequency))
+            except ValueError:
+                raise DataError("frequency must be an integer")
+
         return self.execute_command('RESTORE', *params)
 
     def set(self, name, value,
@@ -3084,6 +3105,7 @@ def _geosearchgeneric(self, command, *args, **kwargs):
     def module_load(self, path, *args):
         """
         Loads the module from ``path``.
+        Passes all ``*args`` to the module, during loading.
         Raises ``ModuleError`` if a module is not found at ``path``.
         """
         pieces = list(args)

tests/test_commands.py

@@ -3471,6 +3471,48 @@ def test_module(self, r):
             r.module_load('/some/fake/path', 'arg1', 'arg2', 'arg3', 'arg4')
             assert "Error loading the extension." in str(excinfo.value)
 
+    @skip_if_server_version_lt('2.6.0')
+    def test_restore(self, r):
+
+        # standard restore
+        key = 'foo'
+        r.set(key, 'bar')
+        dumpdata = r.dump(key)
+        r.delete(key)
+        assert r.restore(key, 0, dumpdata)
+        assert r.get(key) == b'bar'
+
+        # overwrite restore
+        with pytest.raises(redis.exceptions.ResponseError):
+            assert r.restore(key, 0, dumpdata)
+        r.set(key, 'a new value!')
+        assert r.restore(key, 0, dumpdata, replace=True)
+        assert r.get(key) == b'bar'
+
+        # ttl check
+        key2 = 'another'
+        r.set(key2, 'blee!')
+        dumpdata = r.dump(key2)
+        r.delete(key2)
+        assert r.restore(key2, 0, dumpdata)
+        assert r.ttl(key2) == -1
+
+        # idletime
+        key = 'yayakey'
+        r.set(key, 'blee!')
+        dumpdata = r.dump(key)
+        r.delete(key)
+        assert r.restore(key, 0, dumpdata, idletime=5)
+        assert r.get(key) == b'blee!'
+
+        # frequency
+        key = 'yayakey'
+        r.set(key, 'blee!')
+        dumpdata = r.dump(key)
+        r.delete(key)
+        assert r.restore(key, 0, dumpdata, frequency=5)
+        assert r.get(key) == b'blee!'
+
 
 class TestBinarySave: