diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index 56b8a1777f07..a8d7ceb7eff8 100644 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -2976,21 +2976,29 @@ void GeoSearchStoreGeneric(Transaction* tx, SinkReplyBuilder* builder, const Geo rb->SendLong(smvec.size()); } } -} // namespace -void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) { - // parse arguments - string_view key = ArgS(args, 0); - GeoShape shape = {}; +void GeoSearchGeneric(CmdArgList args, bool store, const CommandContext& cmd_cntx) { + string_view key; GeoSearchOpts geo_ops; - string_view member; + if (store) { + geo_ops.store = GeoStoreType::kStoreHash; + geo_ops.store_key = ArgS(args, 0); + key = ArgS(args, 1); + } else { + key = ArgS(args, 0); + } // FROMMEMBER or FROMLONLAT is set bool from_set = false; + string_view member; + // BYRADIUS or BYBOX is set bool by_set = false; + GeoShape shape = {}; + auto* builder = cmd_cntx.rb; - for (size_t i = 1; i < args.size(); ++i) { + size_t next_arg = (store) ? 2 : 1; + for (size_t i = next_arg; i < args.size(); ++i) { string cur_arg = absl::AsciiStrToUpper(ArgS(args, i)); if (cur_arg == "FROMMEMBER") { @@ -3085,13 +3093,15 @@ void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) { geo_ops.any = true; i++; } - } else if (cur_arg == "WITHCOORD") { + } else if (!store && cur_arg == "WITHCOORD") { geo_ops.withcoord = true; - } else if (cur_arg == "WITHDIST") { + } else if (!store && cur_arg == "WITHDIST") { geo_ops.withdist = true; - } else if (cur_arg == "WITHHASH") + } else if (!store && cur_arg == "WITHHASH") { geo_ops.withhash = true; - else { + } else if (store && cur_arg == "STOREDIST") { + geo_ops.store = GeoStoreType::kStoreDist; + } else { return builder->SendError(kSyntaxErr); } } @@ -3108,6 +3118,16 @@ void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) { GeoSearchStoreGeneric(cmd_cntx.tx, builder, shape, key, member, geo_ops); } +} // namespace + +void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) { + GeoSearchGeneric(args, false, cmd_cntx); +} + +void ZSetFamily::GeoSearchStore(CmdArgList args, const CommandContext& cmd_cntx) { + GeoSearchGeneric(args, true, cmd_cntx); +} + void ZSetFamily::GeoRadiusByMember(CmdArgList args, const CommandContext& cmd_cntx) { GeoShape shape = {}; GeoSearchOpts geo_ops; @@ -3243,6 +3263,7 @@ constexpr uint32_t kGeoHash = READ | GEO | SLOW; constexpr uint32_t kGeoPos = READ | GEO | SLOW; constexpr uint32_t kGeoDist = READ | GEO | SLOW; constexpr uint32_t kGeoSearch = READ | GEO | SLOW; +constexpr uint32_t kGeoSearchStore = WRITE | GEO | SLOW; constexpr uint32_t kGeoRadiusByMember = WRITE | GEO | SLOW; } // namespace acl @@ -3296,6 +3317,8 @@ void ZSetFamily::Register(CommandRegistry* registry) { << CI{"GEOPOS", CO::FAST | CO::READONLY, -2, 1, 1, acl::kGeoPos}.HFUNC(GeoPos) << CI{"GEODIST", CO::READONLY, -4, 1, 1, acl::kGeoDist}.HFUNC(GeoDist) << CI{"GEOSEARCH", CO::READONLY, -4, 1, 1, acl::kGeoSearch}.HFUNC(GeoSearch) + << CI{"GEOSEARCHSTORE", CO::WRITE | CO::DENYOOM, -5, 1, 2, acl::kGeoSearchStore}.HFUNC( + GeoSearchStore) << CI{"GEORADIUSBYMEMBER", CO::WRITE | CO::STORE_LAST_KEY, -4, 1, 1, acl::kGeoRadiusByMember} .HFUNC(GeoRadiusByMember); } diff --git a/src/server/zset_family.h b/src/server/zset_family.h index ec678597a791..4ddb9b048ec2 100644 --- a/src/server/zset_family.h +++ b/src/server/zset_family.h @@ -98,6 +98,7 @@ class ZSetFamily { static void GeoPos(CmdArgList args, const CommandContext& cmd_cntx); static void GeoDist(CmdArgList args, const CommandContext& cmd_cntx); static void GeoSearch(CmdArgList args, const CommandContext& cmd_cntx); + static void GeoSearchStore(CmdArgList args, const CommandContext& cmd_cntx); static void GeoRadiusByMember(CmdArgList args, const CommandContext& cmd_cntx); }; diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc index 140488741de7..9bbc5fcf9d4f 100644 --- a/src/server/zset_family_test.cc +++ b/src/server/zset_family_test.cc @@ -1132,6 +1132,29 @@ TEST_F(ZSetFamilyTest, GeoSearch) { RespArray(ElementsAre(DoubleArg(9.1427), DoubleArg(38.7369)))))))); } +TEST_F(ZSetFamilyTest, GeoSearchStore) { + EXPECT_EQ(10, CheckedInt({"GEOADD", "Europe", "13.4050", "52.5200", "Berlin", "3.7038", + "40.4168", "Madrid", "9.1427", "38.7369", "Lisbon", "2.3522", + "48.8566", "Paris", "16.3738", "48.2082", "Vienna", "4.8952", + "52.3702", "Amsterdam", "10.7522", "59.9139", "Oslo", "23.7275", + "37.9838", "Athens", "19.0402", "47.4979", "Budapest", "6.2603", + "53.3498", "Dublin"})); + + EXPECT_EQ(2, CheckedInt({"GEOSEARCHSTORE", "key1", "Europe", "FROMLONLAT", "13.4050", "52.5200", + "BYRADIUS", "500", "KM"})); + + auto resp = Run({"ZRANGE", "key1", "0", "-1", "WITHSCORES"}); + EXPECT_THAT(resp, + RespArray(ElementsAre("Berlin", "3673983950397063", "Dublin", "3678981558208417"))); + + EXPECT_EQ(2, CheckedInt({"GEOSEARCHSTORE", "key2", "Europe", "FROMLONLAT", "13.4050", "52.5200", + "BYRADIUS", "500", "KM", "STOREDIST"})); + + resp = Run({"ZRANGE", "key2", "0", "-1", "WITHSCORES"}); + EXPECT_THAT(resp, RespArray(ElementsAre("Berlin", DoubleArg(0.00017343178521311378), "Dublin", + DoubleArg(487.5619030644293)))); +} + TEST_F(ZSetFamilyTest, GeoRadiusByMember) { EXPECT_EQ(10, CheckedInt({"geoadd", "Europe", "13.4050", "52.5200", "Berlin", "3.7038", "40.4168", "Madrid", "9.1427", "38.7369", "Lisbon", "2.3522",