Add support for Redis 7 functions (#1998)

* add function support

* linters

* test fcall

* decode reponses for unstable_r

* linters

* fix evalsho_ro test

* fix eval_ro test

* add response callbaks

* linters

Author: dvora-h

redis/client.py

@@ -733,6 +733,10 @@ class AbstractRedis:
         "CONFIG RESETSTAT": bool_ok,
         "CONFIG SET": bool_ok,
         "DEBUG OBJECT": parse_debug_object,
+        "FUNCTION DELETE": bool_ok,
+        "FUNCTION FLUSH": bool_ok,
+        "FUNCTION LOAD": bool_ok,
+        "FUNCTION RESTORE": bool_ok,
         "GEOHASH": lambda r: list(map(str_if_bytes, r)),
         "GEOPOS": lambda r: list(
             map(lambda ll: (float(ll[0]), float(ll[1])) if ll is not None else None, r)

redis/commands/core.py

@@ -5429,6 +5429,131 @@ def readonly(self, **kwargs) -> ResponseT:
         return self.execute_command("READONLY", **kwargs)
 
 
+class FunctionCommands:
+    """
+    Redis Function commands
+    """
+
+    def function_load(
+        self,
+        engine: str,
+        library: str,
+        code: str,
+        replace: Optional[bool] = False,
+        description: Optional[str] = None,
+    ) -> str:
+        """
+        Load a library to Redis.
+        :param engine: the name of the execution engine for the library
+        :param library: the unique name of the library
+        :param code: the source code
+        :param replace: changes the behavior to replace the library if a library called
+         ``library`` already exists
+        :param description: description to the library
+
+        For more information check https://redis.io/commands/function-load
+        """
+        pieces = [engine, library]
+        if replace:
+            pieces.append("REPLACE")
+        if description is not None:
+            pieces.append(description)
+        pieces.append(code)
+        return self.execute_command("FUNCTION LOAD", *pieces)
+
+    def function_delete(self, library: str) -> str:
+        """
+        Delete the library called ``library`` and all its functions.
+
+        For more information check https://redis.io/commands/function-delete
+        """
+        return self.execute_command("FUNCTION DELETE", library)
+
+    def function_flush(self, mode: str = "SYNC") -> str:
+        """
+        Deletes all the libraries.
+
+        For more information check https://redis.io/commands/function-flush
+        """
+        return self.execute_command("FUNCTION FLUSH", mode)
+
+    def function_list(
+        self, library: Optional[str] = "*", withcode: Optional[bool] = False
+    ) -> List:
+        """
+        Return information about the functions and libraries.
+        :param library: pecify a pattern for matching library names
+        :param withcode: cause the server to include the libraries source
+         implementation in the reply
+        """
+        args = ["LIBRARYNAME", library]
+        if withcode:
+            args.append("WITHCODE")
+        return self.execute_command("FUNCTION LIST", *args)
+
+    def _fcall(
+        self, command: str, function, numkeys: int, *keys_and_args: Optional[List]
+    ) -> str:
+        return self.execute_command(command, function, numkeys, *keys_and_args)
+
+    def fcall(self, function, numkeys: int, *keys_and_args: Optional[List]) -> str:
+        """
+        Invoke a function.
+
+        For more information check https://redis.io/commands/fcall
+        """
+        return self._fcall("FCALL", function, numkeys, *keys_and_args)
+
+    def fcall_ro(self, function, numkeys: int, *keys_and_args: Optional[List]) -> str:
+        """
+        This is a read-only variant of the FCALL command that cannot
+        execute commands that modify data.
+
+        For more information check https://redis.io/commands/fcal_ro
+        """
+        return self._fcall("FCALL_RO", function, numkeys, *keys_and_args)
+
+    def function_dump(self) -> str:
+        """
+        Return the serialized payload of loaded libraries.
+
+        For more information check https://redis.io/commands/function-dump
+        """
+        from redis.client import NEVER_DECODE
+
+        options = {}
+        options[NEVER_DECODE] = []
+
+        return self.execute_command("FUNCTION DUMP", **options)
+
+    def function_restore(self, payload: str, policy: Optional[str] = "APPEND") -> str:
+        """
+        Restore libraries from the serialized ``payload``.
+        You can use the optional policy argument to provide a policy
+        for handling existing libraries.
+
+        For more information check https://redis.io/commands/function-restore
+        """
+        return self.execute_command("FUNCTION RESTORE", payload, policy)
+
+    def function_kill(self) -> str:
+        """
+        Kill a function that is currently executing.
+
+        For more information check https://redis.io/commands/function-kill
+        """
+        return self.execute_command("FUNCTION KILL")
+
+    def function_stats(self) -> list:
+        """
+        Return information about the function that's currently running
+        and information about the available execution engines.
+
+        For more information check https://redis.io/commands/function-stats
+        """
+        return self.execute_command("FUNCTION STATS")
+
+
 AsyncClusterCommands = ClusterCommands
 
 
@@ -5474,6 +5599,7 @@ class CoreCommands(
     ModuleCommands,
     PubSubCommands,
     ScriptCommands,
+    FunctionCommands,
 ):
     """
     A class containing all of the implemented redis commands. This class is

redis/connection.py

@@ -463,10 +463,17 @@ def read_response(self, disable_decoding=False):
             self._next_response = False
             return response
 
-        response = self._reader.gets()
+        if disable_decoding:
+            response = self._reader.gets(False)
+        else:
+            response = self._reader.gets()
+
         while response is False:
             self.read_from_socket()
-            response = self._reader.gets()
+            if disable_decoding:
+                response = self._reader.gets(False)
+            else:
+                response = self._reader.gets()
         # if an older version of hiredis is installed, we need to attempt
         # to convert ResponseErrors to their appropriate types.
         if not HIREDIS_SUPPORTS_CALLABLE_ERRORS:

tests/conftest.py

@@ -434,7 +434,9 @@ def master_host(request):
 @pytest.fixture()
 def unstable_r(request):
     url = request.config.getoption("--redis-unstable-url")
-    with _get_client(redis.Redis, request, from_url=url) as client:
+    with _get_client(
+        redis.Redis, request, from_url=url, decode_responses=True
+    ) as client:
         yield client
 
 

tests/test_commands.py

@@ -612,7 +612,7 @@ def test_client_unpause(self, r):
     @pytest.mark.onlynoncluster
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_client_no_evict(self, unstable_r):
-        assert unstable_r.client_no_evict("ON") == b"OK"
+        assert unstable_r.client_no_evict("ON") == "OK"
         with pytest.raises(TypeError):
             unstable_r.client_no_evict()
 
@@ -985,9 +985,9 @@ def test_unlink_with_multiple_keys(self, r):
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_lcs(self, unstable_r):
         unstable_r.mset({"foo": "ohmytext", "bar": "mynewtext"})
-        assert unstable_r.lcs("foo", "bar") == b"mytext"
+        assert unstable_r.lcs("foo", "bar") == "mytext"
         assert unstable_r.lcs("foo", "bar", len=True) == 6
-        result = [b"matches", [[[4, 7], [5, 8]]], b"len", 6]
+        result = ["matches", [[[4, 7], [5, 8]]], "len", 6]
         assert unstable_r.lcs("foo", "bar", idx=True, minmatchlen=3) == result
         with pytest.raises(redis.ResponseError):
             assert unstable_r.lcs("foo", "bar", len=True, idx=True)
@@ -1522,24 +1522,24 @@ def test_brpoplpush_empty_string(self, r):
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_blmpop(self, unstable_r):
         unstable_r.rpush("a", "1", "2", "3", "4", "5")
-        res = [b"a", [b"1", b"2"]]
+        res = ["a", ["1", "2"]]
         assert unstable_r.blmpop(1, "2", "b", "a", direction="LEFT", count=2) == res
         with pytest.raises(TypeError):
             unstable_r.blmpop(1, "2", "b", "a", count=2)
         unstable_r.rpush("b", "6", "7", "8", "9")
-        assert unstable_r.blmpop(0, "2", "b", "a", direction="LEFT") == [b"b", [b"6"]]
+        assert unstable_r.blmpop(0, "2", "b", "a", direction="LEFT") == ["b", ["6"]]
         assert unstable_r.blmpop(1, "2", "foo", "bar", direction="RIGHT") is None
 
     @pytest.mark.onlynoncluster
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_lmpop(self, unstable_r):
         unstable_r.rpush("foo", "1", "2", "3", "4", "5")
-        result = [b"foo", [b"1", b"2"]]
+        result = ["foo", ["1", "2"]]
         assert unstable_r.lmpop("2", "bar", "foo", direction="LEFT", count=2) == result
         with pytest.raises(redis.ResponseError):
             unstable_r.lmpop("2", "bar", "foo", direction="up", count=2)
         unstable_r.rpush("bar", "a", "b", "c", "d")
-        assert unstable_r.lmpop("2", "bar", "foo", direction="LEFT") == [b"bar", [b"a"]]
+        assert unstable_r.lmpop("2", "bar", "foo", direction="LEFT") == ["bar", ["a"]]
 
     def test_lindex(self, r):
         r.rpush("a", "1", "2", "3")
@@ -2148,23 +2148,23 @@ def test_bzpopmin(self, r):
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_zmpop(self, unstable_r):
         unstable_r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
-        res = [b"a", [[b"a1", b"1"], [b"a2", b"2"]]]
+        res = ["a", [["a1", "1"], ["a2", "2"]]]
         assert unstable_r.zmpop("2", ["b", "a"], min=True, count=2) == res
         with pytest.raises(redis.DataError):
             unstable_r.zmpop("2", ["b", "a"], count=2)
         unstable_r.zadd("b", {"b1": 10, "ab": 9, "b3": 8})
-        assert unstable_r.zmpop("2", ["b", "a"], max=True) == [b"b", [[b"b1", b"10"]]]
+        assert unstable_r.zmpop("2", ["b", "a"], max=True) == ["b", [["b1", "10"]]]
 
     @pytest.mark.onlynoncluster
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_bzmpop(self, unstable_r):
         unstable_r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
-        res = [b"a", [[b"a1", b"1"], [b"a2", b"2"]]]
+        res = ["a", [["a1", "1"], ["a2", "2"]]]
         assert unstable_r.bzmpop(1, "2", ["b", "a"], min=True, count=2) == res
         with pytest.raises(redis.DataError):
             unstable_r.bzmpop(1, "2", ["b", "a"], count=2)
         unstable_r.zadd("b", {"b1": 10, "ab": 9, "b3": 8})
-        res = [b"b", [[b"b1", b"10"]]]
+        res = ["b", [["b1", "10"]]]
         assert unstable_r.bzmpop(0, "2", ["b", "a"], max=True) == res
         assert unstable_r.bzmpop(1, "2", ["foo", "bar"], max=True) is None
 

tests/test_function.py

@@ -0,0 +1,94 @@
+import pytest
+
+from redis.exceptions import ResponseError
+
+function = "redis.register_function('myfunc', function(keys, args) return args[1] end)"
+function2 = "redis.register_function('hello', function() return 'Hello World' end)"
+set_function = "redis.register_function('set', function(keys, args) \
+                return redis.call('SET', keys[1], args[1]) end)"
+get_function = "redis.register_function('get', function(keys, args) \
+                return redis.call('GET', keys[1]) end)"
+
+
+@pytest.mark.onlynoncluster
+# @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
+class TestFunction:
+    @pytest.fixture(autouse=True)
+    def reset_functions(self, unstable_r):
+        unstable_r.function_flush()
+
+    def test_function_load(self, unstable_r):
+        assert unstable_r.function_load("Lua", "mylib", function)
+        assert unstable_r.function_load("Lua", "mylib", function, replace=True)
+        with pytest.raises(ResponseError):
+            unstable_r.function_load("Lua", "mylib", function)
+        with pytest.raises(ResponseError):
+            unstable_r.function_load("Lua", "mylib2", function)
+
+    def test_function_delete(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", set_function)
+        with pytest.raises(ResponseError):
+            unstable_r.function_load("Lua", "mylib", set_function)
+        assert unstable_r.fcall("set", 1, "foo", "bar") == "OK"
+        assert unstable_r.function_delete("mylib")
+        with pytest.raises(ResponseError):
+            unstable_r.fcall("set", 1, "foo", "bar")
+        assert unstable_r.function_load("Lua", "mylib", set_function)
+
+    def test_function_flush(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", function)
+        assert unstable_r.fcall("myfunc", 0, "hello") == "hello"
+        assert unstable_r.function_flush()
+        with pytest.raises(ResponseError):
+            unstable_r.fcall("myfunc", 0, "hello")
+        with pytest.raises(ResponseError):
+            unstable_r.function_flush("ABC")
+
+    def test_function_list(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", function)
+        res = [
+            [
+                "library_name",
+                "mylib",
+                "engine",
+                "LUA",
+                "description",
+                None,
+                "functions",
+                [["name", "myfunc", "description", None]],
+            ],
+        ]
+        assert unstable_r.function_list() == res
+        assert unstable_r.function_list(library="*lib") == res
+        assert unstable_r.function_list(withcode=True)[0][9] == function
+
+    def test_fcall(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", set_function)
+        unstable_r.function_load("Lua", "mylib2", get_function)
+        assert unstable_r.fcall("set", 1, "foo", "bar") == "OK"
+        assert unstable_r.fcall("get", 1, "foo") == "bar"
+        with pytest.raises(ResponseError):
+            unstable_r.fcall("myfunc", 0, "hello")
+
+    def test_fcall_ro(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", function)
+        assert unstable_r.fcall_ro("myfunc", 0, "hello") == "hello"
+        unstable_r.function_load("Lua", "mylib2", set_function)
+        with pytest.raises(ResponseError):
+            unstable_r.fcall_ro("set", 1, "foo", "bar")
+
+    def test_function_dump_restore(self, unstable_r):
+        unstable_r.function_load("Lua", "mylib", set_function)
+        payload = unstable_r.function_dump()
+        assert unstable_r.fcall("set", 1, "foo", "bar") == "OK"
+        unstable_r.function_delete("mylib")
+        with pytest.raises(ResponseError):
+            unstable_r.fcall("set", 1, "foo", "bar")
+        assert unstable_r.function_restore(payload)
+        assert unstable_r.fcall("set", 1, "foo", "bar") == "OK"
+        unstable_r.function_load("Lua", "mylib2", get_function)
+        assert unstable_r.fcall("get", 1, "foo") == "bar"
+        unstable_r.function_delete("mylib")
+        assert unstable_r.function_restore(payload, "FLUSH")
+        with pytest.raises(ResponseError):
+            unstable_r.fcall("get", 1, "foo")

tests/test_scripting.py

@@ -35,7 +35,7 @@ def test_eval(self, r):
     # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
     def test_eval_ro(self, unstable_r):
         unstable_r.set("a", "b")
-        assert unstable_r.eval_ro("return redis.call('GET', KEYS[1])", 1, "a") == b"b"
+        assert unstable_r.eval_ro("return redis.call('GET', KEYS[1])", 1, "a") == "b"
         with pytest.raises(redis.ResponseError):
             unstable_r.eval_ro("return redis.call('DEL', KEYS[1])", 1, "a")
 
@@ -79,7 +79,7 @@ def test_evalsha_ro(self, unstable_r):
         unstable_r.set("a", "b")
         get_sha = unstable_r.script_load("return redis.call('GET', KEYS[1])")
         del_sha = unstable_r.script_load("return redis.call('DEL', KEYS[1])")
-        assert unstable_r.evalsha_ro(get_sha, 1, "a") == b"b"
+        assert unstable_r.evalsha_ro(get_sha, 1, "a") == "b"
         with pytest.raises(redis.ResponseError):
             unstable_r.evalsha_ro(del_sha, 1, "a")