diff --git a/src/commands/cmd_json.cc b/src/commands/cmd_json.cc index 50dacba2535..4f316b652e9 100644 --- a/src/commands/cmd_json.cc +++ b/src/commands/cmd_json.cc @@ -346,6 +346,47 @@ class CommandJsonArrPop : public Commander { int64_t index_ = -1; }; +class CommandJsonArrTrim : public Commander { + public: + Status Parse(const std::vector &args) override { + path_ = args_[2]; + start_ = GET_OR_RET(ParseInt(args_[3], 10)); + stop_ = GET_OR_RET(ParseInt(args_[4], 10)); + + return Status::OK(); + } + + Status Execute(Server *srv, Connection *conn, std::string *output) override { + redis::Json json(srv->storage, conn->GetNamespace()); + + std::vector> results; + + auto s = json.ArrTrim(args_[1], path_, start_, stop_, results); + + if (s.IsNotFound()) { + *output = redis::NilString(); + return Status::OK(); + } + if (!s.ok()) return {Status::RedisExecErr, s.ToString()}; + + *output = redis::MultiLen(results.size()); + for (const auto &len : results) { + if (len.has_value()) { + *output += redis::Integer(len.value()); + } else { + *output += redis::NilString(); + } + } + + return Status::OK(); + } + + private: + std::string path_; + int64_t start_ = 0; + int64_t stop_ = 0; +}; + class CommanderJsonArrIndex : public Commander { public: Status Parse(const std::vector &args) override { @@ -399,6 +440,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("json.set", 4, "write", 1, 1 MakeCmdAttr("json.info", 2, "read-only", 1, 1, 1), MakeCmdAttr("json.type", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.arrappend", -4, "write", 1, 1, 1), + MakeCmdAttr("json.arrtrim", 5, "write", 1, 1, 1), MakeCmdAttr("json.clear", -2, "write", 1, 1, 1), MakeCmdAttr("json.toggle", -2, "write", 1, 1, 1), MakeCmdAttr("json.arrlen", -2, "read-only", 1, 1, 1), diff --git a/src/types/json.h b/src/types/json.h index 6050e1cc011..62d49bfc913 100644 --- a/src/types/json.h +++ b/src/types/json.h @@ -20,6 +20,7 @@ #pragma once +#include #include #include #include @@ -413,6 +414,38 @@ struct JsonValue { return popped_values; } + Status ArrTrim(std::string_view path, int64_t start, int64_t stop, std::vector> &results) { + try { + jsoncons::jsonpath::json_replace( + value, path, [&results, start, stop](const std::string & /*path*/, jsoncons::json &val) { + if (val.is_array()) { + auto len = static_cast(val.size()); + auto begin_index = start < 0 ? std::max(len + start, static_cast(0)) : start; + auto end_index = std::min(stop < 0 ? std::max(len + stop, static_cast(0)) : stop, len - 1); + + if (begin_index >= len || begin_index > end_index) { + val = jsoncons::json::array(); + results.emplace_back(0); + return; + } + + auto n_val = jsoncons::json::array(); + auto begin_iter = val.array_range().begin(); + + n_val.insert(n_val.end(), begin_iter + begin_index, begin_iter + end_index + 1); + val = n_val; + results.emplace_back(static_cast(n_val.size())); + } else { + results.emplace_back(std::nullopt); + } + }); + } catch (const jsoncons::jsonpath::jsonpath_error &e) { + return {Status::NotOK, e.what()}; + } + + return Status::OK(); + } + JsonValue(const JsonValue &) = default; JsonValue(JsonValue &&) = default; diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc index e16fbbdd9c6..de410ce6967 100644 --- a/src/types/redis_json.cc +++ b/src/types/redis_json.cc @@ -339,4 +339,23 @@ rocksdb::Status Json::ObjKeys(const std::string &user_key, const std::string &pa return rocksdb::Status::OK(); } +rocksdb::Status Json::ArrTrim(const std::string &user_key, const std::string &path, int64_t start, int64_t stop, + std::vector> &results) { + auto ns_key = AppendNamespacePrefix(user_key); + + LockGuard guard(storage_->GetLockManager(), ns_key); + + JsonMetadata metadata; + JsonValue json_val; + auto s = read(ns_key, &metadata, &json_val); + if (!s.ok()) return s; + + auto len_res = json_val.ArrTrim(path, start, stop, results); + if (!len_res) return rocksdb::Status::InvalidArgument(len_res.Msg()); + bool is_write = + std::any_of(results.begin(), results.end(), [](const std::optional &val) { return val.has_value(); }); + if (!is_write) return rocksdb::Status::OK(); + return write(ns_key, &metadata, json_val); +} + } // namespace redis diff --git a/src/types/redis_json.h b/src/types/redis_json.h index a34ff19d67a..1a181e599c9 100644 --- a/src/types/redis_json.h +++ b/src/types/redis_json.h @@ -52,6 +52,9 @@ class Json : public Database { rocksdb::Status ArrIndex(const std::string &user_key, const std::string &path, const std::string &needle, ssize_t start, ssize_t end, std::vector *result); + rocksdb::Status ArrTrim(const std::string &user_key, const std::string &path, int64_t start, int64_t stop, + std::vector> &results); + private: rocksdb::Status write(Slice ns_key, JsonMetadata *metadata, const JsonValue &json_val); rocksdb::Status read(const Slice &ns_key, JsonMetadata *metadata, JsonValue *value); diff --git a/tests/gocase/unit/type/json/json_test.go b/tests/gocase/unit/type/json/json_test.go index 0208066bca0..ba5806b58e1 100644 --- a/tests/gocase/unit/type/json/json_test.go +++ b/tests/gocase/unit/type/json/json_test.go @@ -264,6 +264,90 @@ func TestJson(t *testing.T) { require.ErrorContains(t, rdb.Do(ctx, "JSON.ARRPOP", "a", "$", "0", "1").Err(), "wrong number of arguments") }) + t.Run("JSON.ARRTRIM basics", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "a").Err()) + // key no exists + require.EqualError(t, rdb.Do(ctx, "JSON.ARRTRIM", "not_exists", "$", 0, 0).Err(), redis.Nil.Error()) + // key not json + require.NoError(t, rdb.Do(ctx, "SET", "no_json", "1").Err()) + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "no_json", "$", 0, 0).Err()) + // json path no exists + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"a1":{}}`).Err()) + require.EqualValues(t, []interface{}{}, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.not_exists", 0, 0).Val()) + // json path not array + require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 0, 0).Val()) + // json path has one array + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"a1":[1,2,3,4,5,6,7,8,9]}`).Err()) + require.EqualValues(t, []interface{}{int64(5)}, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 2, 6).Val()) + // json path has many array + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"a":{"a1":[1,2,3,4,5,6]},"b":{"a1":["a",{},"b"]},"c":{"a1":[7,8,9,10,11]}}`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(3), int64(2), int64(3)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$..a1", 1, 3).Val()) + require.EqualValues(t, "[{\"a\":{\"a1\":[2,3,4]},\"b\":{\"a1\":[{},\"b\"]},\"c\":{\"a1\":[8,9,10]}}]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // json path has many array and one is not array + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"a":{"a1":[1,2,3,4,5,6]},"b":{"a1":{"b":1,"c":1}},"c":{"a1":[7,8,9,10]}}`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(3), interface{}(nil), int64(3)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$..a1", 1, 3).Val()) + require.EqualValues(t, "[{\"a\":{\"a1\":[2,3,4]},\"b\":{\"a1\":{\"b\":1,\"c\":1}},\"c\":{\"a1\":[8,9,10]}}]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start not a integer + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", "no", 1).Err()) + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 1.1, 1).Err()) + // stop not a integer + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 1, 1.1).Err()) + // args size != 5 + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 0).Err()) + require.Error(t, rdb.Do(ctx, "JSON.ARRTRIM", "a", "$.a1", 0, 2, 3).Err()) + }) + + t.Run("JSON.ARRTRIM special and args", func(t *testing.T) { + // start < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(4)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -5, 8).Val()) + require.EqualValues(t, "[[6,7,8,9]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start + len < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(6)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -20, 5).Val()) + require.EqualValues(t, "[[1,2,3,4,5,6]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start > len + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(0)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 15, 25).Val()) + require.EqualValues(t, "[[]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start = 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(9)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 0, 8).Val()) + require.EqualValues(t, "[[1,2,3,4,5,6,7,8,9]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // stop = 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(1)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -12, 0).Val()) + require.EqualValues(t, "[[1]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // stop < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(9)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 0, -2).Val()) + require.EqualValues(t, "[[1,2,3,4,5,6,7,8,9]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // len + stop < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(1)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 0, -20).Val()) + require.EqualValues(t, "[[1]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // stop > len + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(10)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 0, 20).Val()) + require.EqualValues(t, "[[1,2,3,4,5,6,7,8,9,10]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start > stop + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(0)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", 8, 5).Val()) + require.EqualValues(t, "[[]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start < 0 and stop < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(4)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -8, -5).Val()) + require.EqualValues(t, "[[3,4,5,6]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start < 0 , stop < 0 and start > end + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(0)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -5, -8).Val()) + require.EqualValues(t, "[[]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + // start + len < 0 , stop + len < 0 + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `[1,2,3,4,5,6,7,8,9,10]`).Err()) + require.EqualValues(t, []interface{}([]interface{}{int64(1)}), rdb.Do(ctx, "JSON.ARRTRIM", "a", "$", -30, -20).Val()) + require.EqualValues(t, "[[1]]", rdb.Do(ctx, "JSON.GET", "a", "$").Val()) + }) + t.Run("JSON.TOGGLE basics", func(t *testing.T) { require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `true`).Err()) require.EqualValues(t, []interface{}{int64(0)}, rdb.Do(ctx, "JSON.TOGGLE", "a", "$").Val())