From 86f8d57d02604d1fa9a1619c84903f1c647c8c2e Mon Sep 17 00:00:00 2001 From: Binbin Date: Fri, 26 Jan 2024 15:19:04 +0800 Subject: [PATCH] Add store key support in key range For store key related commands, we cannot extract the store key, see #2051. This PR adds store_key to CommandKeyRange to represent its index. This PR handle these store key commands: - zunionstore - zinterstore - zdiffstore - georadius - georadiusbymember - geosearchstore before: ``` > command getkeys zinterstore dst 2 src1 src2 1) "src1" 2) "src2" > command getkeys GEORADIUS src 1 1 1 km store dst 1) "src" > command getkeys GEORADIUSBYMEMBER src member radius m store dst 1) "src" > command getkeys GEOSEARCHSTORE dst src frommember member byradius 10 m 1) "dst" ``` after: ``` > command getkeys zinterstore dst 2 src1 src2 1) "dst" 2) "src1" 3) "src2" > command getkeys GEORADIUS src 1 1 1 km store dst 1) "src" 2) "dst" > command getkeys GEORADIUSBYMEMBER src member radius m store dst 1) "src" 2) "dst" > command getkeys GEOSEARCHSTORE dst src frommember member byradius 10 m 1) "dst" 2) "src" ``` --- src/commands/cmd_geo.cc | 42 ++++++++++++- src/commands/cmd_list.cc | 2 + src/commands/cmd_zset.cc | 6 +- src/commands/commander.cc | 10 +++ src/commands/commander.h | 37 +++++++++-- tests/gocase/unit/command/command_test.go | 76 ++++++++++++++++++++--- 6 files changed, 152 insertions(+), 21 deletions(-) diff --git a/src/commands/cmd_geo.cc b/src/commands/cmd_geo.cc index 0f4d98ebf25..4c9f73c635c 100644 --- a/src/commands/cmd_geo.cc +++ b/src/commands/cmd_geo.cc @@ -346,7 +346,25 @@ class CommandGeoRadius : public CommandGeoBase { return redis::Array(list); } + static CommandKeyRange Range(const std::vector &args) { + int store_key = 0; + + // Check for the presence of the stored key in the command args. + for (size_t i = 6; i < args.size(); i++) { + // For the case when a user specifies both "store" and "storedist" options, + // the second key will override the first key. The behavior is kept the same + // as in ParseRadiusExtraOption method. + if ((util::ToLower(args[i]) == "store" || util::ToLower(args[i]) == "storedist") && i + 1 < args.size()) { + store_key = (int)i + 1; + i++; + } + } + + return {1, 1, 1, store_key}; + } + protected: + int option_start_index_ = 6; double radius_ = 0; bool with_coord_ = false; bool with_dist_ = false; @@ -649,6 +667,23 @@ class CommandGeoRadiusByMember : public CommandGeoRadius { return Status::OK(); } + + static CommandKeyRange Range(const std::vector &args) { + int store_key = 0; + + // Check for the presence of the stored key in the command args. + for (size_t i = 5; i < args.size(); i++) { + // For the case when a user specifies both "store" and "storedist" options, + // the second key will override the first key. The behavior is kept the same + // as in ParseRadiusExtraOption method. + if ((util::ToLower(args[i]) == "store" || util::ToLower(args[i]) == "storedist") && i + 1 < args.size()) { + store_key = (int)i + 1; + i++; + } + } + + return {1, 1, 1, store_key}; + } }; class CommandGeoRadiusReadonly : public CommandGeoRadius { @@ -665,11 +700,12 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("geoadd", -5, "write", 1, 1, MakeCmdAttr("geodist", -4, "read-only", 1, 1, 1), MakeCmdAttr("geohash", -3, "read-only", 1, 1, 1), MakeCmdAttr("geopos", -3, "read-only", 1, 1, 1), - MakeCmdAttr("georadius", -6, "write", 1, 1, 1), - MakeCmdAttr("georadiusbymember", -5, "write", 1, 1, 1), + MakeCmdAttr("georadius", -6, "write", CommandGeoRadius::Range), + MakeCmdAttr("georadiusbymember", -5, "write", + CommandGeoRadiusByMember::Range), MakeCmdAttr("georadius_ro", -6, "read-only", 1, 1, 1), MakeCmdAttr("georadiusbymember_ro", -5, "read-only", 1, 1, 1), MakeCmdAttr("geosearch", -7, "read-only", 1, 1, 1), - MakeCmdAttr("geosearchstore", -8, "write", 1, 1, 1)) + MakeCmdAttr("geosearchstore", -8, "write", 2, 2, 1, 1)) } // namespace redis diff --git a/src/commands/cmd_list.cc b/src/commands/cmd_list.cc index 726d2a70889..ce20a60e385 100644 --- a/src/commands/cmd_list.cc +++ b/src/commands/cmd_list.cc @@ -225,6 +225,7 @@ class CommandLMPop : public Commander { // This parsing would always succeed as this cmd has been parsed before. auto num_key = *ParseInt(args[1], 10); range.last_key = range.first_key + num_key - 1; + range.store_key = 0; return range; }; @@ -446,6 +447,7 @@ class CommandBLMPop : public BlockingCommander { // This parsing would always succeed as this cmd has been parsed before. auto num_key = *ParseInt(args[2], 10); range.last_key = range.first_key + num_key - 1; + range.store_key = 0; return range; }; diff --git a/src/commands/cmd_zset.cc b/src/commands/cmd_zset.cc index 12927fa6451..043f66d6b32 100644 --- a/src/commands/cmd_zset.cc +++ b/src/commands/cmd_zset.cc @@ -1224,7 +1224,7 @@ class CommandZUnionStore : public Commander { static CommandKeyRange Range(const std::vector &args) { int num_key = *ParseInt(args[2], 10); - return {3, 2 + num_key, 1}; + return {3, 2 + num_key, 1, 1}; } protected: @@ -1251,7 +1251,7 @@ class CommandZInterStore : public CommandZUnionStore { static CommandKeyRange Range(const std::vector &args) { int num_key = *ParseInt(args[2], 10); - return {3, 2 + num_key, 1}; + return {3, 2 + num_key, 1, 1}; } }; @@ -1505,7 +1505,7 @@ class CommandZDiffStore : public Commander { static CommandKeyRange Range(const std::vector &args) { int num_key = *ParseInt(args[2], 10); - return {3, 2 + num_key, 1}; + return {3, 2 + num_key, 1, 1}; } protected: diff --git a/src/commands/commander.cc b/src/commands/commander.cc index 6ac10581606..e8170985592 100644 --- a/src/commands/commander.cc +++ b/src/commands/commander.cc @@ -81,9 +81,19 @@ void DumpKeyRange(std::vector &keys_index, int argc, const CommandKeyRange auto last = range.last_key; if (last < 0) last = argc + last; + // The index of store key is before the first key + if (range.store_key > 0 && range.store_key < range.first_key) { + keys_index.emplace_back(range.store_key); + } + for (int i = range.first_key; i <= last; i += range.key_step) { keys_index.emplace_back(i); } + + // The index of store key is after the last key + if (range.store_key > 0 && range.store_key > last) { + keys_index.emplace_back(range.store_key); + } } Status CommandTable::GetKeysFromCommand(const CommandAttributes *attributes, const std::vector &cmd_tokens, diff --git a/src/commands/commander.h b/src/commands/commander.h index 3330337c746..2161daed73d 100644 --- a/src/commands/commander.h +++ b/src/commands/commander.h @@ -93,18 +93,22 @@ class CommanderWithParseMove : Commander { using CommanderFactory = std::function()>; struct CommandKeyRange { - // index of the first key in command tokens + // index of the first key (non-store key) in command tokens // 0 stands for no key, since the first index of command arguments is command name int first_key; - // index of the last key in command tokens + // index of the last key (non-store key) in command tokens // in normal one-key commands, first key and last key index are both 1 // -n stands for the n-th last index of the sequence, i.e. args.size() - n int last_key; - // step length of key position + // step length of key (non-store key) position // e.g. key step 2 means "key other key other ..." sequence int key_step; + + // index of the store key in command tokens + // 0 stands for no store key, since the first index of command arguments is command name + int store_key; }; using CommandKeyRangeGen = std::function &)>; @@ -197,7 +201,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript description, ParseCommandFlags(description, name), flag_gen, - {first_key, last_key, key_step}, + {first_key, last_key, key_step, 0}, {}, {}, []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; @@ -210,6 +214,27 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript return attr; } +template +auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, int first_key, int last_key, + int key_step, int store_key, const AdditionalFlagGen &flag_gen = {}) { + CommandAttributes attr{name, + arity, + description, + ParseCommandFlags(description, name), + flag_gen, + {first_key, last_key, key_step, store_key}, + {}, + {}, + []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; + + if ((first_key > 0 && key_step <= 0) || (first_key > 0 && last_key >= 0 && last_key < first_key) || (store_key < 0)) { + std::cout << fmt::format("Encountered invalid key range in command {}", name) << std::endl; + std::abort(); + } + + return attr; +} + template auto MakeCmdAttr(const std::string &name, int arity, const std::string &description, const CommandKeyRangeGen &gen, const AdditionalFlagGen &flag_gen = {}) { @@ -218,7 +243,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript description, ParseCommandFlags(description, name), flag_gen, - {-1, 0, 0}, + {-1, 0, 0, 0}, gen, {}, []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; @@ -234,7 +259,7 @@ auto MakeCmdAttr(const std::string &name, int arity, const std::string &descript description, ParseCommandFlags(description, name), flag_gen, - {-2, 0, 0}, + {-2, 0, 0, 0}, {}, vec_gen, []() -> std::unique_ptr { return std::unique_ptr(new T()); }}; diff --git a/tests/gocase/unit/command/command_test.go b/tests/gocase/unit/command/command_test.go index 89d87a58c5c..8da82510a1b 100644 --- a/tests/gocase/unit/command/command_test.go +++ b/tests/gocase/unit/command/command_test.go @@ -88,9 +88,10 @@ func TestCommand(t *testing.T) { r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZINTERSTORE", "dst", "2", "src1", "src2") vs, err := r.Slice() require.NoError(t, err) - require.Len(t, vs, 2) - require.Equal(t, "src1", vs[0]) - require.Equal(t, "src2", vs[1]) + require.Len(t, vs, 3) + require.Equal(t, "dst", vs[0]) + require.Equal(t, "src1", vs[1]) + require.Equal(t, "src2", vs[2]) }) t.Run("COMMAND GETKEYS ZINTERCARD", func(t *testing.T) { @@ -115,9 +116,10 @@ func TestCommand(t *testing.T) { r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZUNIONSTORE", "dst", "2", "src1", "src2") vs, err := r.Slice() require.NoError(t, err) - require.Len(t, vs, 2) - require.Equal(t, "src1", vs[0]) - require.Equal(t, "src2", vs[1]) + require.Len(t, vs, 3) + require.Equal(t, "dst", vs[0]) + require.Equal(t, "src1", vs[1]) + require.Equal(t, "src2", vs[2]) }) t.Run("COMMAND GETKEYS ZDIFF", func(t *testing.T) { @@ -133,9 +135,10 @@ func TestCommand(t *testing.T) { r := rdb.Do(ctx, "COMMAND", "GETKEYS", "ZDIFFSTORE", "dst", "2", "src1", "src2") vs, err := r.Slice() require.NoError(t, err) - require.Len(t, vs, 2) - require.Equal(t, "src1", vs[0]) - require.Equal(t, "src2", vs[1]) + require.Len(t, vs, 3) + require.Equal(t, "dst", vs[0]) + require.Equal(t, "src1", vs[1]) + require.Equal(t, "src2", vs[2]) }) t.Run("COMMAND GETKEYS ZMPOP", func(t *testing.T) { @@ -173,4 +176,59 @@ func TestCommand(t *testing.T) { require.Equal(t, "key1", vs[0]) require.Equal(t, "key2", vs[1]) }) + + t.Run("COMMAND GETKEYS GEORADIUS", func(t *testing.T) { + r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "store", "dst") + vs, err := r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst", vs[1]) + + r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "storedist", "dst") + vs, err = r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst", vs[1]) + + r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUS", "src", "1", "1", "1", "km", "store", "dst1", "storedist", "dst2") + vs, err = r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst2", vs[1]) + }) + + t.Run("COMMAND GETKEYS GEORADIUSBYMEMBER", func(t *testing.T) { + r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "store", "dst") + vs, err := r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst", vs[1]) + + r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "storedist", "dst") + vs, err = r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst", vs[1]) + + r = rdb.Do(ctx, "COMMAND", "GETKEYS", "GEORADIUSBYMEMBER", "src", "member", "radius", "m", "store", "dst1", "storedist", "dst2") + vs, err = r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "src", vs[0]) + require.Equal(t, "dst2", vs[1]) + }) + + t.Run("COMMAND GETKEYS GEOSEARCHSTORE", func(t *testing.T) { + r := rdb.Do(ctx, "COMMAND", "GETKEYS", "GEOSEARCHSTORE", "dst", "src", "frommember", "member", "byradius", "10", "m") + vs, err := r.Slice() + require.NoError(t, err) + require.Len(t, vs, 2) + require.Equal(t, "dst", vs[0]) + require.Equal(t, "src", vs[1]) + }) }