diff --git a/src/commands/blocking_commander.h b/src/commands/blocking_commander.h index 2749681984f..2571e456067 100644 --- a/src/commands/blocking_commander.h +++ b/src/commands/blocking_commander.h @@ -52,7 +52,7 @@ class BlockingCommander : public Commander, // to start the blocking process // usually put to the end of the Execute method Status StartBlocking(int64_t timeout, std::string *output) { - if (conn_->IsInExec()) { + if (conn_->IsInExec() || conn_->IsInScript()) { *output = NoopReply(conn_); return Status::OK(); // no blocking in multi-exec } diff --git a/src/server/redis_connection.cc b/src/server/redis_connection.cc index 3f8e0dd264f..a4926dfb7ac 100644 --- a/src/server/redis_connection.cc +++ b/src/server/redis_connection.cc @@ -496,6 +496,12 @@ void Connection::ExecuteCommands(std::deque *to_process_cmds) { continue; } + ScopeExit in_script_exit{[this] { in_script_ = false; }, false}; + if (attributes->category == CommandCategory::Script || attributes->category == CommandCategory::Function) { + in_script_ = true; + in_script_exit.Enable(); + } + SetLastCmd(cmd_name); { std::optional guard; diff --git a/src/server/redis_connection.h b/src/server/redis_connection.h index e7d277f6e8b..e8b44d94144 100644 --- a/src/server/redis_connection.h +++ b/src/server/redis_connection.h @@ -171,6 +171,7 @@ class Connection : public EvbufCallbackBase { // Multi exec void SetInExec() { in_exec_ = true; } bool IsInExec() const { return in_exec_; } + bool IsInScript() const { return in_script_; } bool IsMultiError() const { return multi_error_; } void ResetMultiExec(); std::deque *GetMultiExecCommands() { return &multi_cmds_; } @@ -210,6 +211,7 @@ class Connection : public EvbufCallbackBase { bool multi_error_ = false; std::atomic is_running_ = false; std::deque multi_cmds_; + bool in_script_ = false; bool importing_ = false; RESP protocol_version_ = RESP::v2; diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc index b2c0eebcf0a..6a93ad9f8ee 100644 --- a/src/storage/scripting.cc +++ b/src/storage/scripting.cc @@ -756,7 +756,7 @@ int RedisGenericCommand(lua_State *lua, int raise_error) { auto attributes = cmd->GetAttributes(); if (!attributes->CheckArity(argc)) { - PushError(lua, "Wrong number of args calling Redis command From Lua script"); + PushError(lua, "Wrong number of args while calling Redis command from Lua script"); return raise_error ? RaiseError(lua) : 1; } @@ -772,12 +772,6 @@ int RedisGenericCommand(lua_State *lua, int raise_error) { return raise_error ? RaiseError(lua) : 1; } - // TODO: fix blocking commands to make them work in scripting - if (cmd_flags & redis::kCmdBlocking) { - PushError(lua, "This Redis command is not allowed from scripts"); - return raise_error ? RaiseError(lua) : 1; - } - std::string cmd_name = attributes->name; auto *conn = script_run_ctx->conn; diff --git a/tests/gocase/unit/scripting/scripting_test.go b/tests/gocase/unit/scripting/scripting_test.go index cd04351adc5..c3753ecc158 100644 --- a/tests/gocase/unit/scripting/scripting_test.go +++ b/tests/gocase/unit/scripting/scripting_test.go @@ -198,10 +198,16 @@ return {type(foo),foo == false} }) t.Run("EVAL - Scripts can't run certain commands", func(t *testing.T) { - r := rdb.Eval(ctx, `return redis.pcall('blpop','x',0)`, []string{}) + r := rdb.Eval(ctx, `return redis.pcall('shutdown')`, []string{}) require.ErrorContains(t, r.Err(), "not allowed") }) + t.Run("EVAL - Scripts can run blocking commands and get immediate result", func(t *testing.T) { + r := rdb.Eval(ctx, `return redis.pcall('blpop', KEYS[1], 0)`, []string{"key_for_blpop_script"}) + require.Equal(t, r.Val(), nil) + require.ErrorContains(t, r.Err(), "nil") + }) + t.Run("EVAL - Scripts can run certain commands", func(t *testing.T) { r := rdb.Eval(ctx, `redis.pcall('randomkey'); return redis.pcall('set','x','ciao')`, []string{}) require.NoError(t, r.Err())