Add cluster support for functions (#2016)

dvora-h · web-flow · commit c4e408880c0f · 2022-03-06T21:25:34.000+02:00
* cluster support for functions

* fix test_list_on_cluster mark

* fix mark

* cluster unstable url

* fix

* fix cluster url

* skip tests

* linters

* linters
diff --git a/redis/cluster.py b/redis/cluster.py
@@ -292,12 +292,21 @@ class RedisCluster(RedisClusterCommands):
             [
                 "FLUSHALL",
                 "FLUSHDB",
+                "FUNCTION DELETE",
+                "FUNCTION FLUSH",
+                "FUNCTION LIST",
+                "FUNCTION LOAD",
+                "FUNCTION RESTORE",
                 "SCRIPT EXISTS",
                 "SCRIPT FLUSH",
                 "SCRIPT LOAD",
             ],
             PRIMARIES,
         ),
+        list_keys_to_dict(
+            ["FUNCTION DUMP"],
+            RANDOM,
+        ),
         list_keys_to_dict(
             [
                 "CLUSTER COUNTKEYSINSLOT",
@@ -916,6 +925,10 @@ def determine_slot(self, *args):
         else:
             keys = self._get_command_keys(*args)
             if keys is None or len(keys) == 0:
+                # FCALL can call a function with 0 keys, that means the function
+                #  can be run on any node so we can just return a random slot
+                if command in ("FCALL", "FCALL_RO"):
+                    return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
                 raise RedisClusterException(
                     "No way to dispatch this command to Redis Cluster. "
                     "Missing key.\nYou can execute the command by specifying "
diff --git a/redis/commands/cluster.py b/redis/commands/cluster.py
@@ -4,6 +4,7 @@
 from .core import (
     ACLCommands,
     DataAccessCommands,
+    FunctionCommands,
     ManagementCommands,
     PubSubCommands,
     ScriptCommands,
@@ -213,6 +214,7 @@ class RedisClusterCommands(
     PubSubCommands,
     ClusterDataAccessCommands,
     ScriptCommands,
+    FunctionCommands,
     RedisModuleCommands,
 ):
     """
diff --git a/tests/test_function.py b/tests/test_function.py
@@ -2,6 +2,8 @@
 
 from redis.exceptions import ResponseError
 
+from .conftest import skip_if_server_version_lt
+
 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) \
@@ -10,42 +12,42 @@
                 return redis.call('GET', keys[1]) end)"
 
 
-@pytest.mark.onlynoncluster
-# @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
+@skip_if_server_version_lt("7.0.0")
 class TestFunction:
     @pytest.fixture(autouse=True)
-    def reset_functions(self, unstable_r):
-        unstable_r.function_flush()
+    def reset_functions(self, r):
+        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)
+    def test_function_load(self, r):
+        assert r.function_load("Lua", "mylib", function)
+        assert r.function_load("Lua", "mylib", function, replace=True)
         with pytest.raises(ResponseError):
-            unstable_r.function_load("Lua", "mylib", function)
+            r.function_load("Lua", "mylib", function)
         with pytest.raises(ResponseError):
-            unstable_r.function_load("Lua", "mylib2", function)
+            r.function_load("Lua", "mylib2", function)
 
-    def test_function_delete(self, unstable_r):
-        unstable_r.function_load("Lua", "mylib", set_function)
+    def test_function_delete(self, r):
+        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")
+            r.function_load("Lua", "mylib", set_function)
+        assert r.fcall("set", 1, "foo", "bar") == "OK"
+        assert r.function_delete("mylib")
         with pytest.raises(ResponseError):
-            unstable_r.fcall("set", 1, "foo", "bar")
-        assert unstable_r.function_load("Lua", "mylib", set_function)
+            r.fcall("set", 1, "foo", "bar")
+        assert 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()
+    def test_function_flush(self, r):
+        r.function_load("Lua", "mylib", function)
+        assert r.fcall("myfunc", 0, "hello") == "hello"
+        assert r.function_flush()
         with pytest.raises(ResponseError):
-            unstable_r.fcall("myfunc", 0, "hello")
+            r.fcall("myfunc", 0, "hello")
         with pytest.raises(ResponseError):
-            unstable_r.function_flush("ABC")
+            r.function_flush("ABC")
 
-    def test_function_list(self, unstable_r):
-        unstable_r.function_load("Lua", "mylib", function)
+    @pytest.mark.onlynoncluster
+    def test_function_list(self, r):
+        r.function_load("Lua", "mylib", function)
         res = [
             [
                 "library_name",
@@ -58,37 +60,61 @@ def test_function_list(self, unstable_r):
                 [["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
+        assert r.function_list() == res
+        assert r.function_list(library="*lib") == res
+        assert r.function_list(withcode=True)[0][9] == function
+
+    @pytest.mark.onlycluster
+    def test_function_list_on_cluster(self, r):
+        r.function_load("Lua", "mylib", function)
+        function_list = [
+            [
+                "library_name",
+                "mylib",
+                "engine",
+                "LUA",
+                "description",
+                None,
+                "functions",
+                [["name", "myfunc", "description", None]],
+            ],
+        ]
+        primaries = r.get_primaries()
+        res = {}
+        for node in primaries:
+            res[node.name] = function_list
+        assert r.function_list() == res
+        assert r.function_list(library="*lib") == res
+        node = primaries[0].name
+        assert r.function_list(withcode=True)[node][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"
+    def test_fcall(self, r):
+        r.function_load("Lua", "mylib", set_function)
+        r.function_load("Lua", "mylib2", get_function)
+        assert r.fcall("set", 1, "foo", "bar") == "OK"
+        assert r.fcall("get", 1, "foo") == "bar"
         with pytest.raises(ResponseError):
-            unstable_r.fcall("myfunc", 0, "hello")
+            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)
+    def test_fcall_ro(self, r):
+        r.function_load("Lua", "mylib", function)
+        assert r.fcall_ro("myfunc", 0, "hello") == "hello"
+        r.function_load("Lua", "mylib2", set_function)
         with pytest.raises(ResponseError):
-            unstable_r.fcall_ro("set", 1, "foo", "bar")
+            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")
+    def test_function_dump_restore(self, r):
+        r.function_load("Lua", "mylib", set_function)
+        payload = r.function_dump()
+        assert r.fcall("set", 1, "foo", "bar") == "OK"
+        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")
+            r.fcall("set", 1, "foo", "bar")
+        assert r.function_restore(payload)
+        assert r.fcall("set", 1, "foo", "bar") == "OK"
+        r.function_load("Lua", "mylib2", get_function)
+        assert r.fcall("get", 1, "foo") == "bar"
+        r.function_delete("mylib")
+        assert r.function_restore(payload, "FLUSH")
         with pytest.raises(ResponseError):
-            unstable_r.fcall("get", 1, "foo")
+            r.fcall("get", 1, "foo")
diff --git a/tox.ini b/tox.ini
@@ -283,10 +283,11 @@ extras =
     ocsp: cryptography, pyopenssl, requests
 setenv =
     CLUSTER_URL = "redis://localhost:16379/0"
+    UNSTABLE_CLUSTER_URL = "redis://localhost:6372/0"
 commands =
     standalone: pytest --cov=./ --cov-report=xml:coverage_redis.xml -W always -m 'not onlycluster' {posargs}
     standalone-uvloop: pytest --cov=./ --cov-report=xml:coverage_redis.xml -W always -m 'not onlycluster' --uvloop {posargs}
-    cluster: pytest --cov=./ --cov-report=xml:coverage_cluster.xml -W always -m 'not onlynoncluster and not redismod' --redis-url={env:CLUSTER_URL:} --redismod-url={env:CLUSTER_URL:} {posargs}
+    cluster: pytest --cov=./ --cov-report=xml:coverage_cluster.xml -W always -m 'not onlynoncluster and not redismod' --redis-url={env:CLUSTER_URL:} --redis-unstable-url={env:UNSTABLE_CLUSTER_URL:} {posargs}
     cluster-uvloop: pytest --cov=./ --cov-report=xml:coverage_redis.xml -W always -m 'not onlycluster' --uvloop {posargs}
 
 [testenv:redis5]
