Add the support of the json.arrappend command (#1837)

Co-authored-by: mwish <maplewish117@gmail.com>
apache · Oct 20, 2023 · 411ebfe · 411ebfe
1 parent a618574
commit 411ebfe
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 1 deletion.
diff --git a/src/commands/cmd_json.cc b/src/commands/cmd_json.cc
@@ -52,7 +52,31 @@ class CommandJsonGet : public Commander {
   }
 };
 
+class CommandJsonArrAppend : public Commander {
+ public:
+  Status Execute(Server *svr, Connection *conn, std::string *output) override {
+    redis::Json json(svr->storage, conn->GetNamespace());
+
+    std::vector<uint64_t> result_count;
+
+    auto s = json.ArrAppend(args_[1], args_[2], {args_.begin() + 3, args_.end()}, &result_count);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(result_count.size());
+    for (uint64_t c : result_count) {
+      if (c != 0) {
+        *output += redis::Integer(c);
+      } else {
+        *output += redis::NilString();
+      }
+    }
+
+    return Status::OK();
+  }
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", -3, "write", 1, 1, 1),
-                        MakeCmdAttr<CommandJsonGet>("json.get", -2, "read-only", 1, 1, 1), );
+                        MakeCmdAttr<CommandJsonGet>("json.get", -2, "read-only", 1, 1, 1),
+                        MakeCmdAttr<CommandJsonArrAppend>("json.arrappend", -4, "write", 1, 1, 1), );
 
 }  // namespace redis
diff --git a/src/types/json.h b/src/types/json.h
@@ -73,6 +73,24 @@ struct JsonValue {
     }
   }
 
+  Status ArrAppend(std::string_view path, const std::vector<jsoncons::json> &append_values,
+                   std::vector<uint64_t> *result_count) {
+    try {
+      jsoncons::jsonpath::json_replace(
+          value, path, [&append_values, result_count](const std::string &path, jsoncons::json &val) {
+            if (val.is_array()) {
+              val.insert(val.array_range().end(), append_values.begin(), append_values.end());
+              result_count->emplace_back(val.size());
+            } else {
+              result_count->emplace_back(0);
+            }
+          });
+    } 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
@@ -117,4 +117,41 @@ rocksdb::Status Json::Get(const std::string &user_key, const std::vector<std::st
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status Json::ArrAppend(const std::string &user_key, const std::string &path,
+                                const std::vector<std::string> &values, std::vector<uint64_t> *result_count) {
+  auto ns_key = AppendNamespacePrefix(user_key);
+
+  std::vector<jsoncons::json> append_values;
+  append_values.reserve(values.size());
+  for (auto &v : values) {
+    auto value_res = JsonValue::FromString(v);
+    if (!value_res) return rocksdb::Status::InvalidArgument(value_res.Msg());
+    auto value = *std::move(value_res);
+    append_values.emplace_back(std::move(value.value));
+  }
+
+  LockGuard guard(storage_->GetLockManager(), ns_key);
+
+  std::string bytes;
+  JsonMetadata metadata;
+  Slice rest;
+  auto s = GetMetadata(kRedisJson, ns_key, &bytes, &metadata, &rest);
+  if (!s.ok()) return s;
+
+  if (metadata.format != JsonStorageFormat::JSON)
+    return rocksdb::Status::NotSupported("JSON storage format not supported");
+
+  auto value_res = JsonValue::FromString(rest.ToStringView());
+  if (!value_res) return rocksdb::Status::Corruption(value_res.Msg());
+  auto value = *std::move(value_res);
+
+  auto append_res = value.ArrAppend(path, append_values, result_count);
+  if (!append_res) return rocksdb::Status::InvalidArgument(append_res.Msg());
+
+  bool is_write = std::any_of(result_count->begin(), result_count->end(), [](uint64_t c) { return c > 0; });
+  if (!is_write) return rocksdb::Status::OK();
+
+  return write(ns_key, &metadata, value);
+}
+
 }  // namespace redis
diff --git a/src/types/redis_json.h b/src/types/redis_json.h
@@ -34,6 +34,8 @@ class Json : public Database {
 
   rocksdb::Status Set(const std::string &user_key, const std::string &path, const std::string &value);
   rocksdb::Status Get(const std::string &user_key, const std::vector<std::string> &paths, JsonValue *result);
+  rocksdb::Status ArrAppend(const std::string &user_key, const std::string &path,
+                            const std::vector<std::string> &values, std::vector<uint64_t> *result_count);
 
  private:
   rocksdb::Status write(Slice ns_key, JsonMetadata *metadata, const JsonValue &json_val);

diff --git a/tests/cppunit/types/json_test.cc b/tests/cppunit/types/json_test.cc
@@ -119,3 +119,44 @@ TEST_F(RedisJsonTest, Get) {
   ASSERT_TRUE(json_->Get(key_, {"$..x", "$..y", "$..z"}, &json_val_).ok());
   ASSERT_EQ(json_val_.Dump(), R"({"$..x":[{"y":1},4],"$..y":[[2,{"z":3}],1],"$..z":[{"a":{"x":4}},3]})");
 }
+
+TEST_F(RedisJsonTest, ArrAppend) {
+  std::vector<uint64_t> res;
+
+  ASSERT_FALSE(json_->ArrAppend(key_, "$", {"1"}, &res).ok());
+
+  ASSERT_TRUE(json_->Set(key_, "$", R"({"x":1,"y":[]})").ok());
+  ASSERT_TRUE(json_->ArrAppend(key_, "$.x", {"1"}, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_EQ(res[0], 0);
+  ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
+  ASSERT_EQ(json_val_.Dump(), R"({"x":1,"y":[]})");
+  res.clear();
+
+  ASSERT_TRUE(json_->Set(key_, "$", R"({"x":[1,2,{"z":3}],"y":[]})").ok());
+  ASSERT_TRUE(json_->ArrAppend(key_, "$.x", {"1"}, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_EQ(res[0], 4);
+  ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
+  ASSERT_EQ(json_val_.Dump(), R"({"x":[1,2,{"z":3},1],"y":[]})");
+  res.clear();
+
+  ASSERT_TRUE(json_->ArrAppend(key_, "$..y", {"1", "2", "3"}, &res).ok());
+  ASSERT_EQ(res.size(), 1);
+  ASSERT_EQ(res[0], 3);
+  ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
+  ASSERT_EQ(json_val_.Dump(), R"({"x":[1,2,{"z":3},1],"y":[1,2,3]})");
+  res.clear();
+
+  ASSERT_TRUE(json_->Set(key_, "$.x[2]", R"({"x":[1,2,{"z":3,"y":[]}],"y":[{"y":1}]})").ok());
+  ASSERT_TRUE(json_->ArrAppend(key_, "$..y", {"1", "2", "3"}, &res).ok());
+  ASSERT_EQ(res.size(), 4);
+  std::sort(res.begin(), res.end());
+  ASSERT_EQ(res[0], 0);
+  ASSERT_EQ(res[1], 3);
+  ASSERT_EQ(res[2], 4);
+  ASSERT_EQ(res[3], 6);
+  ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
+  ASSERT_EQ(json_val_.Dump(), R"({"x":[1,2,{"x":[1,2,{"y":[1,2,3],"z":3}],"y":[{"y":1},1,2,3]},1],"y":[1,2,3,1,2,3]})");
+  res.clear();
+}
diff --git a/tests/gocase/unit/type/json/json_test.go b/tests/gocase/unit/type/json/json_test.go
@@ -55,4 +55,33 @@ func TestJson(t *testing.T) {
 		require.Equal(t, rdb.Do(ctx, "JSON.GET", "a", "$..x").Val(), `[1,{"y":2}]`)
 		require.Equal(t, rdb.Do(ctx, "JSON.GET", "a", "$..x", "$..y").Val(), `{"$..x":[1,{"y":2}],"$..y":[{"x":{"y":2},"y":3},3,2]}`)
 	})
+
+	t.Run("JSON.ARRAPPEND basics", func(t *testing.T) {
+		require.NoError(t, rdb.Do(ctx, "SET", "a", `1`).Err())
+		require.Error(t, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$", `1`).Err())
+		require.NoError(t, rdb.Do(ctx, "DEL", "a").Err())
+
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", ` {"x":1, "y": {"x":1} } `).Err())
+		require.Equal(t, []interface{}{}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$..k", `1`).Val())
+		require.Error(t, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$").Err())
+		require.Error(t, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$", ` 1, 2, 3`).Err())
+		require.Error(t, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$", `1`, ` 1, 2, 3`).Err())
+		require.Equal(t, []interface{}{nil, nil}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$..x", `1`).Val())
+
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", ` {"x":1, "y": {"x":[]} } `).Err())
+		require.Equal(t, []interface{}{int64(1), nil}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$..x", `1`).Val())
+		require.Equal(t, `[{"x":1,"y":{"x":[1]}}]`, rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", ` {"x":[], "y":[]} `).Err())
+		require.Equal(t, []interface{}{int64(1)}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$.x", `1`).Val())
+		require.Equal(t, `{"x":[1],"y":[]}`, rdb.Do(ctx, "JSON.GET", "a").Val())
+		require.Equal(t, []interface{}{int64(4)}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$.x", `1`, `2`, `3`).Val())
+		require.Equal(t, []interface{}{int64(1)}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$.y", ` {"x":[], "y":[]} `).Val())
+		require.Equal(t, `[{"x":[1,1,2,3],"y":[{"x":[],"y":[]}]}]`, rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+
+		require.Equal(t, []interface{}{int64(2), int64(6)}, rdb.Do(ctx, "JSON.ARRAPPEND", "a", "$..x", `1`, `2`).Val())
+		require.Equal(t, `[[1,2]]`, rdb.Do(ctx, "JSON.GET", "a", "$.y[0].x").Val())
+		require.Equal(t, `[]`, rdb.Do(ctx, "JSON.GET", "a", "$.x.x").Val())
+		require.Equal(t, `[{"x":[1,1,2,3,1,2],"y":[{"x":[1,2],"y":[]}]}]`, rdb.Do(ctx, "JSON.GET", "a", "$").Val())
+	})
 }