From fcb95bec6e20a05c020af6734e75ec9810a52d27 Mon Sep 17 00:00:00 2001 From: iko1 Date: Sat, 15 Oct 2022 18:36:47 +0200 Subject: [PATCH] feat(server): implement json.arrindex command (#104) (#376) Signed-off-by: iko1 --- src/server/json_family.cc | 124 +++++++++++++++++++++++++++++++++ src/server/json_family.h | 1 + src/server/json_family_test.cc | 35 ++++++++++ 3 files changed, 160 insertions(+) diff --git a/src/server/json_family.cc b/src/server/json_family.cc index 451452e..dabc104 100644 --- a/src/server/json_family.cc +++ b/src/server/json_family.cc @@ -29,6 +29,7 @@ using namespace jsoncons; using JsonExpression = jsonpath::jsonpath_expression; using OptBool = optional; +using OptLong = optional; using OptSizeT = optional; using OptString = optional; using JsonReplaceCb = function; @@ -771,8 +772,130 @@ OpResult> OpArrInsert(const OpArgs& op_args, string_view key, s return vec; } +// Returns a numeric vector representing each JSON value first index of the JSON scalar. +// An index value of -1 represents unfound in the array. +// JSON scalar has types of string, boolean, null, and number. +OpResult> OpArrIndex(const OpArgs& op_args, string_view key, + JsonExpression expression, const json& search_val, + int start_index, int end_index) { + OpResult result = GetJson(op_args, key); + if (!result) { + return result.status(); + } + + vector vec; + auto cb = [&](const string_view& path, const json& val) { + if (!val.is_array()) { + vec.emplace_back(nullopt); + return; + } + + if (val.empty()) { + vec.emplace_back(-1); + return; + } + + // Negative value or out-of-range index is handled by rounding the index to the array's start + // and end. example: for array size 9 and index -11 the index mapped to index 7. + if (start_index < 0) { + start_index %= val.size(); + start_index += val.size(); + } + + // See the comment above. + // A value index of 0 means searching until the end of the array. + if (end_index == 0) { + end_index = val.size(); + } else if (end_index < 0) { + end_index %= val.size(); + end_index += val.size(); + } + + if (start_index > end_index) { + vec.emplace_back(-1); + return; + } + + size_t pos = -1; + auto it = next(val.array_range().begin(), start_index); + while (it != val.array_range().end()) { + if (search_val == *it) { + pos = start_index; + break; + } + + ++it; + if (++start_index == end_index) { + break; + } + } + + vec.emplace_back(pos); + }; + + expression.evaluate(*result, cb); + return vec; +} + } // namespace +void JsonFamily::ArrIndex(CmdArgList args, ConnectionContext* cntx) { + string_view key = ArgS(args, 1); + string_view path = ArgS(args, 2); + + error_code ec; + JsonExpression expression = jsonpath::make_expression(path, ec); + + if (ec) { + VLOG(1) << "Invalid JSONPath syntax: " << ec.message(); + (*cntx)->SendError(kSyntaxErr); + return; + } + + optional search_value = ConstructJsonFromString(ArgS(args, 3)); + if (!search_value) { + (*cntx)->SendError(kSyntaxErr); + return; + } + + if (search_value->is_object() || search_value->is_array()) { + (*cntx)->SendError(kWrongTypeErr); + return; + } + + int start_index = 0; + if (args.size() >= 5) { + if (!absl::SimpleAtoi(ArgS(args, 4), &start_index)) { + VLOG(1) << "Failed to convert the start index to numeric" << ArgS(args, 4); + (*cntx)->SendError(kInvalidIntErr); + return; + } + } + + int end_index = 0; + if (args.size() >= 6) { + if (!absl::SimpleAtoi(ArgS(args, 5), &end_index)) { + VLOG(1) << "Failed to convert the stop index to numeric" << ArgS(args, 5); + (*cntx)->SendError(kInvalidIntErr); + return; + } + } + + auto cb = [&](Transaction* t, EngineShard* shard) { + return OpArrIndex(t->GetOpArgs(shard), key, move(expression), *search_value, start_index, + end_index); + }; + + Transaction* trans = cntx->transaction; + OpResult> result = trans->ScheduleSingleHopT(move(cb)); + + if (result) { + PrintOptVec(cntx, result); + } else { + (*cntx)->SendError(result.status()); + } +} + void JsonFamily::ArrInsert(CmdArgList args, ConnectionContext* cntx) { string_view key = ArgS(args, 1); string_view path = ArgS(args, 2); @@ -1218,6 +1341,7 @@ void JsonFamily::Register(CommandRegistry* registry) { *registry << CI{"JSON.ARRTRIM", CO::WRITE | CO::DENYOOM | CO::FAST, 5, 1, 1, 1}.HFUNC(ArrTrim); *registry << CI{"JSON.ARRINSERT", CO::WRITE | CO::DENYOOM | CO::FAST, -4, 1, 1, 1}.HFUNC( ArrInsert); + *registry << CI{"JSON.ARRINDEX", CO::READONLY | CO::FAST, -4, 1, 1, 1}.HFUNC(ArrIndex); } } // namespace dfly diff --git a/src/server/json_family.h b/src/server/json_family.h index 69ed932..c0d187c 100644 --- a/src/server/json_family.h +++ b/src/server/json_family.h @@ -35,6 +35,7 @@ class JsonFamily { static void ArrPop(CmdArgList args, ConnectionContext* cntx); static void ArrTrim(CmdArgList args, ConnectionContext* cntx); static void ArrInsert(CmdArgList args, ConnectionContext* cntx); + static void ArrIndex(CmdArgList args, ConnectionContext* cntx); }; } // namespace dfly diff --git a/src/server/json_family_test.cc b/src/server/json_family_test.cc index 2939193..6ea92f4 100644 --- a/src/server/json_family_test.cc +++ b/src/server/json_family_test.cc @@ -776,4 +776,39 @@ TEST_F(JsonFamilyTest, ArrInsert) { EXPECT_EQ(resp, R"([["b","c","a"],["a","c","b","a"],["a","c","a","b","b"]])"); } +TEST_F(JsonFamilyTest, ArrIndex) { + string json = R"( + [[], ["a"], ["a", "b"], ["a", "b", "c"]] + )"; + + auto resp = Run({"set", "json", json}); + ASSERT_THAT(resp, "OK"); + + resp = Run({"JSON.ARRINDEX", "json", "$[*]", R"("b")"}); + ASSERT_EQ(RespExpr::ARRAY, resp.type); + EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(-1), IntArg(-1), IntArg(1), IntArg(1))); + + json = R"( + {"a":["a","b","c","d"], "nested": {"a": ["c","d"]}} + )"; + + resp = Run({"set", "json", json}); + ASSERT_THAT(resp, "OK"); + + resp = Run({"JSON.ARRINDEX", "json", "$..a", R"("b")"}); + ASSERT_EQ(RespExpr::ARRAY, resp.type); + EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(1), IntArg(-1))); + + json = R"( + {"a":["a","b","c","d"], "nested": {"a": false}} + )"; + + resp = Run({"set", "json", json}); + ASSERT_THAT(resp, "OK"); + + resp = Run({"JSON.ARRINDEX", "json", "$..a", R"("b")"}); + ASSERT_EQ(RespExpr::ARRAY, resp.type); + EXPECT_THAT(resp.GetVec(), ElementsAre(IntArg(1), ArgType(RespExpr::NIL))); +} + } // namespace dfly