Skip to content

Add support for Redis 7 functions #1998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Feb 22, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions redis/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
126 changes: 126 additions & 0 deletions redis/commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -5474,6 +5599,7 @@ class CoreCommands(
ModuleCommands,
PubSubCommands,
ScriptCommands,
FunctionCommands,
):
"""
A class containing all of the implemented redis commands. This class is
Expand Down
11 changes: 9 additions & 2 deletions redis/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
22 changes: 11 additions & 11 deletions tests/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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

Expand Down
94 changes: 94 additions & 0 deletions tests/test_function.py
Original file line number Diff line number Diff line change
@@ -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")
4 changes: 2 additions & 2 deletions tests/test_scripting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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")

Expand Down