This is an automated email from the ASF dual-hosted git repository.

hulk pushed a commit to branch unstable
in repository https://gitbox.apache.org/repos/asf/kvrocks.git


The following commit(s) were added to refs/heads/unstable by this push:
     new 2cab8f71 Support for the JSON.ARRINSERT command  (#1867)
2cab8f71 is described below

commit 2cab8f71458f2961be92a41bd0b9bb7d9c366843
Author: 纪华裕 <[email protected]>
AuthorDate: Fri Nov 17 22:23:09 2023 +0800

    Support for the JSON.ARRINSERT command  (#1867)
---
 src/commands/cmd_json.cc                 | 44 ++++++++++++++++++++++++++++++
 src/types/json.h                         | 29 ++++++++++++++++++++
 src/types/redis_json.cc                  | 32 ++++++++++++++++++++++
 src/types/redis_json.h                   |  2 ++
 tests/gocase/unit/type/json/json_test.go | 46 ++++++++++++++++++++++++++++++++
 5 files changed, 153 insertions(+)

diff --git a/src/commands/cmd_json.cc b/src/commands/cmd_json.cc
index 4f316b65..edc3f549 100644
--- a/src/commands/cmd_json.cc
+++ b/src/commands/cmd_json.cc
@@ -22,7 +22,9 @@
 
 #include "commander.h"
 #include "commands/command_parser.h"
+#include "commands/error_constants.h"
 #include "error_constants.h"
+#include "parse_util.h"
 #include "server/redis_reply.h"
 #include "server/server.h"
 #include "storage/redis_metadata.h"
@@ -142,6 +144,47 @@ class CommandJsonArrAppend : public Commander {
   }
 };
 
+class CommandJsonArrInsert : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_result = ParseInt<int>(args[3], 10);
+    if (!parse_result) {
+      return {Status::RedisParseErr, errValueNotInteger};
+    }
+
+    index_ = *parse_result;
+    return Commander::Parse(args);
+  }
+
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::Json json(srv->storage, conn->GetNamespace());
+
+    std::vector<std::optional<uint64_t>> result_count;
+    auto parse_result = ParseInt<int>(args_[3], 10);
+
+    auto s = json.ArrInsert(args_[1], args_[2], index_, {args_.begin() + 4, 
args_.end()}, &result_count);
+    if (s.IsNotFound()) {
+      *output = redis::NilString();
+      return Status::OK();
+    }
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(result_count.size());
+    for (auto c : result_count) {
+      if (c.has_value()) {
+        *output += redis::Integer(c.value());
+      } else {
+        *output += redis::NilString();
+      }
+    }
+
+    return Status::OK();
+  }
+
+ private:
+  int index_;
+};
+
 class CommandJsonType : public Commander {
  public:
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
@@ -440,6 +483,7 @@ 
REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", 4, "write", 1, 1
                         MakeCmdAttr<CommandJsonInfo>("json.info", 2, 
"read-only", 1, 1, 1),
                         MakeCmdAttr<CommandJsonType>("json.type", -2, 
"read-only", 1, 1, 1),
                         MakeCmdAttr<CommandJsonArrAppend>("json.arrappend", 
-4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandJsonArrInsert>("json.arrinsert", 
-5, "write", 1, 1, 1),
                         MakeCmdAttr<CommandJsonArrTrim>("json.arrtrim", 5, 
"write", 1, 1, 1),
                         MakeCmdAttr<CommandJsonClear>("json.clear", -2, 
"write", 1, 1, 1),
                         MakeCmdAttr<CommandJsonToggle>("json.toggle", -2, 
"write", 1, 1, 1),
diff --git a/src/types/json.h b/src/types/json.h
index 62d49bfc..9df5f592 100644
--- a/src/types/json.h
+++ b/src/types/json.h
@@ -179,6 +179,35 @@ struct JsonValue {
     return result_count;
   }
 
+  StatusOr<std::vector<std::optional<uint64_t>>> ArrInsert(std::string_view 
path, const int64_t &index,
+                                                           const 
std::vector<jsoncons::json> &insert_values) {
+    std::vector<std::optional<uint64_t>> result_count;
+
+    try {
+      jsoncons::jsonpath::json_replace(
+          value, path, [&insert_values, &result_count, index](const 
std::string & /*path*/, jsoncons::json &val) {
+            if (val.is_array()) {
+              auto len = static_cast<int64_t>(val.size());
+              // When index > 0, we need index < len
+              // when index < 0, we need index >= -len.
+              if (index >= len || index < -len) {
+                result_count.emplace_back(std::nullopt);
+                return;
+              }
+              auto base_iter = index >= 0 ? val.array_range().begin() : 
val.array_range().end();
+              val.insert(base_iter + index, insert_values.begin(), 
insert_values.end());
+              result_count.emplace_back(val.size());
+            } else {
+              result_count.emplace_back(std::nullopt);
+            }
+          });
+    } catch (const jsoncons::jsonpath::jsonpath_error &e) {
+      return {Status::NotOK, e.what()};
+    }
+
+    return result_count;
+  }
+
   static std::pair<ssize_t, ssize_t> NormalizeArrIndices(ssize_t start, 
ssize_t end, ssize_t len) {
     if (start < 0) {
       start = std::max<ssize_t>(0, len + start);
diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc
index de410ce6..c31fac0d 100644
--- a/src/types/redis_json.cc
+++ b/src/types/redis_json.cc
@@ -286,6 +286,38 @@ rocksdb::Status Json::ArrLen(const std::string &user_key, 
const std::string &pat
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status Json::ArrInsert(const std::string &user_key, const std::string 
&path, const int64_t &index,
+                                const std::vector<std::string> &values,
+                                std::vector<std::optional<uint64_t>> 
*result_count) {
+  auto ns_key = AppendNamespacePrefix(user_key);
+
+  std::vector<jsoncons::json> insert_values;
+  insert_values.reserve(values.size());
+  for (auto &v : values) {
+    auto value_res = JsonValue::FromString(v, 
storage_->GetConfig()->json_max_nesting_depth);
+    if (!value_res) return rocksdb::Status::InvalidArgument(value_res.Msg());
+    auto value = *std::move(value_res);
+    insert_values.emplace_back(std::move(value.value));
+  }
+
+  LockGuard guard(storage_->GetLockManager(), ns_key);
+
+  JsonMetadata metadata;
+  JsonValue value;
+  auto s = read(ns_key, &metadata, &value);
+  if (!s.ok()) return s;
+
+  auto insert_res = value.ArrInsert(path, index, insert_values);
+  if (!insert_res) return rocksdb::Status::InvalidArgument(insert_res.Msg());
+  *result_count = *insert_res;
+
+  bool is_write =
+      std::any_of(result_count->begin(), result_count->end(), 
[](std::optional<uint64_t> c) { return c.has_value(); });
+  if (!is_write) return rocksdb::Status::OK();
+
+  return write(ns_key, &metadata, value);
+}
+
 rocksdb::Status Json::Toggle(const std::string &user_key, const std::string 
&path,
                              std::vector<std::optional<bool>> &result) {
   auto ns_key = AppendNamespacePrefix(user_key);
diff --git a/src/types/redis_json.h b/src/types/redis_json.h
index 1a181e59..29af781d 100644
--- a/src/types/redis_json.h
+++ b/src/types/redis_json.h
@@ -43,6 +43,8 @@ class Json : public Database {
   rocksdb::Status Clear(const std::string &user_key, const std::string &path, 
size_t *result);
   rocksdb::Status ArrLen(const std::string &user_key, const std::string &path,
                          std::vector<std::optional<uint64_t>> &arr_lens);
+  rocksdb::Status ArrInsert(const std::string &user_key, const std::string 
&path, const int64_t &index,
+                            const std::vector<std::string> &values, 
std::vector<std::optional<uint64_t>> *result_count);
   rocksdb::Status Toggle(const std::string &user_key, const std::string &path,
                          std::vector<std::optional<bool>> &result);
   rocksdb::Status ObjKeys(const std::string &user_key, const std::string &path,
diff --git a/tests/gocase/unit/type/json/json_test.go 
b/tests/gocase/unit/type/json/json_test.go
index ba5806b5..c226c327 100644
--- a/tests/gocase/unit/type/json/json_test.go
+++ b/tests/gocase/unit/type/json/json_test.go
@@ -215,6 +215,52 @@ func TestJson(t *testing.T) {
                require.EqualValues(t, []uint64{}, lens)
        })
 
+       t.Run("JSON.ARRINSERT basics", func(t *testing.T) {
+               arrInsertCmd := "JSON.ARRINSERT"
+               require.NoError(t, rdb.Del(ctx, "a").Err())
+               // key no exists
+               require.EqualError(t, rdb.Do(ctx, arrInsertCmd, "not_exists", 
"$", 0, 1).Err(), redis.Nil.Error())
+               // key not json
+               require.NoError(t, rdb.Do(ctx, "SET", "no_json", "1").Err())
+               require.Error(t, rdb.Do(ctx, arrInsertCmd, "no_json", "$", 0, 
1).Err())
+               // json path no exists
+               require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", 
`{"a":{}}`).Err())
+               require.EqualValues(t, []interface{}{}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.not_exists", 0, 1).Val())
+               // json path not array
+               require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.a", 0, 1).Val())
+               // index not a integer
+               require.Error(t, rdb.Do(ctx, arrInsertCmd, "a", "$", "no", 
1).Err())
+               require.Error(t, rdb.Do(ctx, arrInsertCmd, "a", "$", 1.1, 
1).Err())
+               // args size < 4
+               require.Error(t, rdb.Do(ctx, arrInsertCmd, "a", "$", 0).Err())
+               // json path has one array
+               require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", 
`{"a":[1,2,3], 
"b":{"a":[4,5,6,7],"c":2},"c":[1,2,3,4],"e":[6,7,8],"f":{"a":[10,11,12,13,14], 
"g":2}}`).Err())
+               require.EqualValues(t, []interface{}{int64(4)}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.e", 1, 90).Val())
+               require.Equal(t, 
"[{\"a\":[1,2,3],\"b\":{\"a\":[4,5,6,7],\"c\":2},\"c\":[1,2,3,4],\"e\":[6,90,7,8],\"f\":{\"a\":[10,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // insert many value
+               require.EqualValues(t, []interface{}{int64(8)}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.e", 2, 80, 81, 82, 83).Val())
+               require.Equal(t, 
"[{\"a\":[1,2,3],\"b\":{\"a\":[4,5,6,7],\"c\":2},\"c\":[1,2,3,4],\"e\":[6,90,80,81,82,83,7,8],\"f\":{\"a\":[10,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // json path has many array
+               require.EqualValues(t, []interface{}{int64(6), int64(5), 
int64(4)}, rdb.Do(ctx, arrInsertCmd, "a", "$..a", 1, 91).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[1,2,3,4],\"e\":[6,90,80,81,82,83,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // json path has many array and one is not array
+               require.EqualValues(t, []interface{}{int64(5), nil}, 
rdb.Do(ctx, arrInsertCmd, "a", "$..c", 0, 92).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[92,1,2,3,4],\"e\":[6,90,80,81,82,83,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // index = 0
+               require.EqualValues(t, []interface{}{int64(9)}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.e", 0, 93).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[92,1,2,3,4],\"e\":[93,6,90,80,81,82,83,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // index < 0
+               require.EqualValues(t, []interface{}{int64(10)}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.e", -2, 94).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[92,1,2,3,4],\"e\":[93,6,90,80,81,82,83,94,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // index >= len
+               require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, 
arrInsertCmd, "a", "$.e", 15, 95).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[92,1,2,3,4],\"e\":[93,6,90,80,81,82,83,94,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+               // index + len < 0
+               require.EqualValues(t, []interface{}{nil}, rdb.Do(ctx, 
arrInsertCmd, "a", "$", -15, 96).Val())
+               require.Equal(t, 
"[{\"a\":[1,91,2,3],\"b\":{\"a\":[4,91,5,6,7],\"c\":2},\"c\":[92,1,2,3,4],\"e\":[93,6,90,80,81,82,83,94,7,8],\"f\":{\"a\":[10,91,11,12,13,14],\"g\":2}}]",
 rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+
+       })
+
        t.Run("JSON.OBJKEYS basics", func(t *testing.T) {
                require.NoError(t, rdb.Del(ctx, "a").Err())
                // key no exists

Reply via email to