From 90e3267f968699c96b0c8441ee7a7406d14c3c8d Mon Sep 17 00:00:00 2001 From: sewenew Date: Wed, 2 Oct 2024 15:41:57 +0800 Subject: [PATCH] support Redis function (#598) --- src/sw/redis++/async_redis.h | 88 ++++++++++++++++++++++++ src/sw/redis++/async_redis_cluster.h | 48 +++++++++++++ src/sw/redis++/cmd_formatter.h | 50 ++++++++++++++ src/sw/redis++/command.h | 56 +++++++++++++++ src/sw/redis++/redis.cpp | 12 ++++ src/sw/redis++/redis.h | 56 +++++++++++++++ src/sw/redis++/redis.hpp | 82 ++++++++++++++++++++++ src/sw/redis++/redis_cluster.h | 52 ++++++++++++++ src/sw/redis++/redis_cluster.hpp | 84 ++++++++++++++++++++++ test/src/sw/redis++/script_cmds_test.h | 2 + test/src/sw/redis++/script_cmds_test.hpp | 58 ++++++++++++++++ 11 files changed, 588 insertions(+) diff --git a/src/sw/redis++/async_redis.h b/src/sw/redis++/async_redis.h index 425fdc34..8621f492 100644 --- a/src/sw/redis++/async_redis.h +++ b/src/sw/redis++/async_redis.h @@ -1613,6 +1613,94 @@ class AsyncRedis { return evalsha(script, keys.begin(), keys.end(), args.begin(), args.end(), std::forward(cb)); } + template + Future fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + return _command(fmt::fcall, + func, keys_first, keys_last, args_first, args_last); + } + + template + auto fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Callback &&cb) + -> typename std::enable_if::type, Future &&>::value, void>::type { + _callback_fmt_command(std::forward(cb), fmt::fcall, + func, keys_first, keys_last, args_first, args_last); + } + + template + Future fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall(func, + keys.begin(), keys.end(), + args.begin(), args.end()); + } + + template + auto fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Callback &&cb) + -> typename std::enable_if::type, Future &&>::value, void>::type { + return fcall(func, keys.begin(), keys.end(), args.begin(), args.end(), std::forward(cb)); + } + + template + Future fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + return _command(fmt::fcall_ro, + func, keys_first, keys_last, args_first, args_last); + } + + template + auto fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Callback &&cb) + -> typename std::enable_if::type, Future &&>::value, void>::type { + _callback_fmt_command(std::forward(cb), fmt::fcall_ro, + func, keys_first, keys_last, args_first, args_last); + } + + template + Future fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall_ro(func, + keys.begin(), keys.end(), + args.begin(), args.end()); + } + + template + auto fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Callback &&cb) + -> typename std::enable_if::type, Future &&>::value, void>::type { + return fcall_ro(func, keys.begin(), keys.end(), args.begin(), args.end(), std::forward(cb)); + } + + Future function_load(const StringView &code, bool replace = false) { + return _command(fmt::function_load, code, replace); + } + + Future function_delete(const StringView &lib_name) { + return _command(fmt::function_delete, lib_name); + } + // PUBSUB commands. Future publish(const StringView &channel, const StringView &message) { diff --git a/src/sw/redis++/async_redis_cluster.h b/src/sw/redis++/async_redis_cluster.h index 1e772ce9..8c2a41fd 100644 --- a/src/sw/redis++/async_redis_cluster.h +++ b/src/sw/redis++/async_redis_cluster.h @@ -1075,6 +1075,54 @@ class AsyncRedisCluster { args.begin(), args.end()); } + template + Future fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + return _generic_command(fmt::fcall, *keys_first, func, + keys_first, keys_last, + args_first, args_last); + } + + template + Future fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall(func, + keys.begin(), keys.end(), + args.begin(), args.end()); + } + + template + Future fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + return _generic_command(fmt::fcall_ro, *keys_first, func, + keys_first, keys_last, + args_first, args_last); + } + + template + Future fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall_ro(func, + keys.begin(), keys.end(), + args.begin(), args.end()); + } + // PUBSUB commands. Future publish(const StringView &channel, const StringView &message) { diff --git a/src/sw/redis++/cmd_formatter.h b/src/sw/redis++/cmd_formatter.h index abf1517f..2467c866 100644 --- a/src/sw/redis++/cmd_formatter.h +++ b/src/sw/redis++/cmd_formatter.h @@ -791,6 +791,56 @@ FormattedCommand evalsha(const StringView &script, return format_cmd(args); } +template +FormattedCommand fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + CmdArgs args; + auto keys_num = std::distance(keys_first, keys_last); + + args << "FCALL" << func << keys_num + << std::make_pair(keys_first, keys_last) + << std::make_pair(args_first, args_last); + + return format_cmd(args); +} + +template +FormattedCommand fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + CmdArgs args; + auto keys_num = std::distance(keys_first, keys_last); + + args << "FCALL_RO" << func << keys_num + << std::make_pair(keys_first, keys_last) + << std::make_pair(args_first, args_last); + + return format_cmd(args); +} + +inline FormattedCommand function_load(const StringView &code, bool replace) { + CmdArgs cmd_args; + + cmd_args << "FUNCTION" << "LOAD"; + + if (replace) { + cmd_args << "REPLACE"; + } + + cmd_args << code; + + return format_cmd(cmd_args); +} + +inline FormattedCommand function_delete(const StringView &lib_name) { + return format_cmd("FUNCTION DELETE %b", lib_name.data(), lib_name.size()); +} + // PUBSUB commands. inline FormattedCommand psubscribe(const StringView &pattern) { diff --git a/src/sw/redis++/command.h b/src/sw/redis++/command.h index c10a536d..5ee80e64 100644 --- a/src/sw/redis++/command.h +++ b/src/sw/redis++/command.h @@ -1568,6 +1568,62 @@ inline void script_load(Connection &connection, const StringView &script) { connection.send("SCRIPT LOAD %b", script.data(), script.size()); } +template +inline void fcall(Connection &connection, + const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + CmdArgs cmd_args; + + auto keys_num = std::distance(keys_first, keys_last); + + cmd_args << "FCALL" << func << keys_num + << std::make_pair(keys_first, keys_last) + << std::make_pair(args_first, args_last); + + connection.send(cmd_args); +} + +template +inline void fcall_ro(Connection &connection, + const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + CmdArgs cmd_args; + + auto keys_num = std::distance(keys_first, keys_last); + + cmd_args << "FCALL_RO" << func << keys_num + << std::make_pair(keys_first, keys_last) + << std::make_pair(args_first, args_last); + + connection.send(cmd_args); +} + +inline void function_load(Connection &connection, + const StringView &code, + bool replace) { + CmdArgs cmd_args; + + cmd_args << "FUNCTION" << "LOAD"; + + if (replace) { + cmd_args << "REPLACE"; + } + + cmd_args << code; + + connection.send(cmd_args); +} + +inline void function_delete(Connection &connection, const StringView &lib_name) { + connection.send("FUNCTION DELETE %b", lib_name.data(), lib_name.size()); +} + // PUBSUB commands. inline void psubscribe(Connection &connection, const StringView &pattern) { diff --git a/src/sw/redis++/redis.cpp b/src/sw/redis++/redis.cpp index 3c0ed4f4..761103f2 100644 --- a/src/sw/redis++/redis.cpp +++ b/src/sw/redis++/redis.cpp @@ -855,6 +855,18 @@ std::string Redis::script_load(const StringView &script) { return reply::parse(*reply); } +std::string Redis::function_load(const StringView &code, bool replace) { + auto reply = command(cmd::function_load, code, replace); + + return reply::parse(*reply); +} + +void Redis::function_delete(const StringView &lib_name) { + auto reply = command(cmd::function_delete, lib_name); + + reply::parse(*reply); +} + // PUBSUB commands. long long Redis::publish(const StringView &channel, const StringView &message) { diff --git a/src/sw/redis++/redis.h b/src/sw/redis++/redis.h index 816445a5..15070db1 100644 --- a/src/sw/redis++/redis.h +++ b/src/sw/redis++/redis.h @@ -3247,6 +3247,62 @@ class Redis { std::string script_load(const StringView &script); + template + Result fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last); + + template + Result fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args); + + template + void fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output); + + template + void fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output); + + template + Result fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last); + + template + Result fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args); + + template + void fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output); + + template + void fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output); + + std::string function_load(const StringView &code, bool replace = false); + + void function_delete(const StringView &lib_name); + // PUBSUB commands. long long publish(const StringView &channel, const StringView &message); diff --git a/src/sw/redis++/redis.hpp b/src/sw/redis++/redis.hpp index 3c97224c..5d40f503 100644 --- a/src/sw/redis++/redis.hpp +++ b/src/sw/redis++/redis.hpp @@ -1022,6 +1022,88 @@ void Redis::script_exists(Input first, Input last, Output output) { reply::to_array(*reply, output); } +template +Result Redis::fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + auto reply = command(cmd::fcall, func, keys_first, keys_last, args_first, args_last); + + return reply::parse(*reply); +} + +template +Result Redis::fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall(func, keys.begin(), keys.end(), args.begin(), args.end()); +} + +template +void Redis::fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output) { + auto reply = command(cmd::fcall, + func, + keys_first, keys_last, + args_first, args_last); + + reply::to_array(*reply, output); +} + +template +void Redis::fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output) { + fcall(func, keys.begin(), keys.end(), args.begin(), args.end(), output); +} + +template +Result Redis::fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + auto reply = command(cmd::fcall_ro, func, keys_first, keys_last, args_first, args_last); + + return reply::parse(*reply); +} + +template +Result Redis::fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall_ro(func, keys.begin(), keys.end(), args.begin(), args.end()); +} + +template +void Redis::fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output) { + auto reply = command(cmd::fcall_ro, + func, + keys_first, keys_last, + args_first, args_last); + + reply::to_array(*reply, output); +} + +template +void Redis::fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output) { + fcall_ro(func, keys.begin(), keys.end(), args.begin(), args.end(), output); +} + // Transaction commands. template diff --git a/src/sw/redis++/redis_cluster.h b/src/sw/redis++/redis_cluster.h index 89758088..cef76656 100644 --- a/src/sw/redis++/redis_cluster.h +++ b/src/sw/redis++/redis_cluster.h @@ -1065,6 +1065,58 @@ class RedisCluster { std::initializer_list args, Output output); + template + Result fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last); + + template + Result fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args); + + template + void fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output); + + template + void fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output); + + template + Result fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last); + + template + Result fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args); + + template + void fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output); + + template + void fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args, + Output output); + // PUBSUB commands. long long publish(const StringView &channel, const StringView &message); diff --git a/src/sw/redis++/redis_cluster.hpp b/src/sw/redis++/redis_cluster.hpp index 30688620..c0e8188b 100644 --- a/src/sw/redis++/redis_cluster.hpp +++ b/src/sw/redis++/redis_cluster.hpp @@ -1002,6 +1002,90 @@ void RedisCluster::evalsha(const StringView &script, evalsha(script, keys.begin(), keys.end(), args.begin(), args.end(), output); } +template +Result RedisCluster::fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + auto reply = _command(cmd::fcall, *keys_first, func, keys_first, keys_last, args_first, args_last); + + return reply::parse(*reply); +} + +template +Result RedisCluster::fcall(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall(func, keys.begin(), keys.end(), args.begin(), args.end()); +} + +template +void RedisCluster::fcall(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + auto reply = _command(cmd::fcall, + *keys_first, + func, + keys_first, keys_last, + args_first, args_last); + + reply::to_array(*reply, output); +} + +template +Result RedisCluster::fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + auto reply = _command(cmd::fcall_ro, *keys_first, func, keys_first, keys_last, args_first, args_last); + + return reply::parse(*reply); +} + +template +Result RedisCluster::fcall_ro(const StringView &func, + std::initializer_list keys, + std::initializer_list args) { + return fcall_ro(func, keys.begin(), keys.end(), args.begin(), args.end()); +} + +template +void RedisCluster::fcall_ro(const StringView &func, + Keys keys_first, + Keys keys_last, + Args args_first, + Args args_last, + Output output) { + if (keys_first == keys_last) { + throw Error("DO NOT support function without key"); + } + + auto reply = _command(cmd::fcall_ro, + *keys_first, + func, + keys_first, keys_last, + args_first, args_last); + + reply::to_array(*reply, output); +} + // Stream commands. template diff --git a/test/src/sw/redis++/script_cmds_test.h b/test/src/sw/redis++/script_cmds_test.h index f7adb11b..edd69b34 100644 --- a/test/src/sw/redis++/script_cmds_test.h +++ b/test/src/sw/redis++/script_cmds_test.h @@ -35,6 +35,8 @@ class ScriptCmdTest { private: void _run(Redis &instance); + void _run_function_test(Redis &instance); + RedisInstance &_redis; }; diff --git a/test/src/sw/redis++/script_cmds_test.hpp b/test/src/sw/redis++/script_cmds_test.hpp index 86a0e7df..43abb9b2 100644 --- a/test/src/sw/redis++/script_cmds_test.hpp +++ b/test/src/sw/redis++/script_cmds_test.hpp @@ -32,6 +32,10 @@ void ScriptCmdTest::run() { cluster_specializing_test(*this, &ScriptCmdTest::_run, _redis); + + cluster_specializing_test(*this, + &ScriptCmdTest::_run_function_test, + _redis); } template @@ -117,6 +121,60 @@ void ScriptCmdTest::_run(Redis &instance) { REDIS_ASSERT(!instance.script_exists("not exist"), "failed to test script exists"); } +template +void ScriptCmdTest::_run_function_test(Redis &instance) { + auto key1 = test_key("k1"); + auto key2 = test_key("k2"); + + KeyDeleter deleter(instance, {key1, key2}); + + try { + instance.function_delete("swredistestlib"); + } catch (const Error &) { + } + + std::string code = "#!lua name=swredistestlib\n" + "redis.register_function('my_func', function(keys, args) " + "redis.call('set', keys[1], 1);" + "redis.call('set', keys[2], 2);" + "local first = redis.call('get', keys[1]);" + "local second = redis.call('get', keys[2]);" + "return first + second\n" + "end)"; + + auto lib_name = instance.function_load(code); + REDIS_ASSERT(lib_name == "swredistestlib", "failed to test function_load"); + + std::initializer_list keys = {key1, key2}; + std::initializer_list empty_list = {}; + + auto num = instance.fcall("my_func", keys, empty_list); + REDIS_ASSERT(num == 3, "failed to test fcall"); + + num = instance.fcall("my_func", keys.begin(), keys.end(), + empty_list.begin(), empty_list.end()); + REDIS_ASSERT(num == 3, "failed to test fcall"); + + auto is_readonly = false; + try { + instance.fcall_ro("my_func", {}, {}); + } catch (const Error &) { + is_readonly = true; + } + REDIS_ASSERT(is_readonly, "failed to test fcall_ro"); + + code = "#!lua name=swredistestlib\n" + "local function readonly_func(keys, args) return 'hello' end\n" + "redis.register_function{function_name='my_func', callback=readonly_func, flags={ 'no-writes' }}"; + lib_name = instance.function_load(code, true); + REDIS_ASSERT(lib_name == "swredistestlib", "failed to test function_load"); + + auto res = instance.fcall_ro("my_func", {}, {}); + REDIS_ASSERT(res == "hello", "failed to test fcall_ro"); + + instance.function_delete("swredistestlib"); +} + } }