From e5d89ff7c764f1230a1b167b2ba99d865ccaab21 Mon Sep 17 00:00:00 2001 From: skyitachi Date: Fri, 10 Nov 2023 18:58:06 +0800 Subject: [PATCH] support JSON.MGET command --- src/commands/cmd_json.cc | 31 ++++++++- src/types/redis_json.cc | 83 ++++++++++++++++++++++++ src/types/redis_json.h | 6 ++ tests/gocase/unit/type/json/json_test.go | 32 +++++++++ 4 files changed, 151 insertions(+), 1 deletion(-) diff --git a/src/commands/cmd_json.cc b/src/commands/cmd_json.cc index 41763bdbc50..737581f81e7 100644 --- a/src/commands/cmd_json.cc +++ b/src/commands/cmd_json.cc @@ -566,6 +566,34 @@ class CommandJsonStrLen : public Commander { } }; +class CommandJsonMGet : public Commander { + public: + Status Execute(Server *svr, Connection *conn, std::string *output) override { + redis::Json json(svr->storage, conn->GetNamespace()); + + std::string path = args_[args_.size() - 1]; + std::vector user_keys; + for (size_t i = 1; i + 1 < args_.size(); i++) { + user_keys.emplace_back(args_[i]); + } + + std::vector json_values; + auto statuses = json.MGet(user_keys, path, json_values); + + std::vector values; + values.resize(user_keys.size()); + + for (size_t i = 0; i < statuses.size(); i++) { + if (statuses[i].ok()) { + values[i] = json_values[i].value.to_string(); + } + } + + *output = MultiBulkString(values, statuses); + return Status::OK(); + } +}; + REDIS_REGISTER_COMMANDS(MakeCmdAttr("json.set", 4, "write", 1, 1, 1), MakeCmdAttr("json.get", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.info", 2, "read-only", 1, 1, 1), @@ -587,6 +615,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("json.set", 4, "write", 1, 1 MakeCmdAttr("json.nummultby", 4, "write", 1, 1, 1), MakeCmdAttr("json.objlen", -2, "read-only", 1, 1, 1), MakeCmdAttr("json.strappend", -3, "write", 1, 1, 1), - MakeCmdAttr("json.strlen", -2, "read-only", 1, 1, 1), ); + MakeCmdAttr("json.strlen", -2, "read-only", 1, 1, 1), + MakeCmdAttr("json.mget", -3, "read-only", 1, 1, 1), ); } // namespace redis diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc index 8c0f28bfaf2..b1755e3a997 100644 --- a/src/types/redis_json.cc +++ b/src/types/redis_json.cc @@ -512,4 +512,87 @@ rocksdb::Status Json::ObjLen(const std::string &user_key, const std::string &pat return rocksdb::Status::OK(); } +std::vector Json::MGet(const std::vector &user_keys, const std::string &path, + std::vector &results) { + std::vector ns_keys; + std::vector ns_keys_string; + ns_keys.resize(user_keys.size()); + ns_keys_string.resize(user_keys.size()); + + for (size_t i = 0; i < user_keys.size(); i++) { + ns_keys_string[i] = AppendNamespacePrefix(user_keys[i]); + ns_keys[i] = Slice(ns_keys_string[i]); + } + + std::vector json_vals; + json_vals.resize(ns_keys.size()); + auto statuses = readMulti(ns_keys, json_vals); + + results.resize(ns_keys.size()); + for (size_t i = 0; i < ns_keys.size(); i++) { + if (!statuses[i].ok()) { + continue; + } + auto res = json_vals[i].Get(path); + + if (!res) { + statuses[i] = rocksdb::Status::Corruption(res.Msg()); + } else { + results[i] = *std::move(res); + } + } + return statuses; +} + +std::vector Json::readMulti(const std::vector &ns_keys, std::vector &values) { + std::vector raw_values; + std::vector meta_data; + raw_values.resize(ns_keys.size()); + meta_data.resize(ns_keys.size()); + + auto statuses = getRawMetaData(ns_keys, meta_data, &raw_values); + for (size_t i = 0; i < ns_keys.size(); i++) { + if (!statuses[i].ok()) continue; + if (meta_data[i].format == JsonStorageFormat::JSON) { + auto res = JsonValue::FromString(raw_values[i]); + if (!res) { + statuses[i] = rocksdb::Status::Corruption(res.Msg()); + continue; + } + values[i] = *std::move(res); + } else if (meta_data[i].format == JsonStorageFormat::CBOR) { + auto res = JsonValue::FromCBOR(raw_values[i]); + if (!res) { + statuses[i] = rocksdb::Status::Corruption(res.Msg()); + continue; + } + values[i] = *std::move(res); + } else { + statuses[i] = rocksdb::Status::NotSupported("JSON storage format not supported"); + } + } + return statuses; +} + +std::vector Json::getRawMetaData(const std::vector &ns_keys, + std::vector &metadatas, + std::vector *raw_values) { + rocksdb::ReadOptions read_options = storage_->DefaultMultiGetOptions(); + LatestSnapShot ss(storage_); + read_options.snapshot = ss.GetSnapShot(); + raw_values->resize(ns_keys.size()); + std::vector statuses(ns_keys.size()); + std::vector pin_values(ns_keys.size()); + storage_->MultiGet(read_options, metadata_cf_handle_, ns_keys.size(), ns_keys.data(), pin_values.data(), + statuses.data()); + for (size_t i = 0; i < ns_keys.size(); i++) { + if (!statuses[i].ok()) continue; + Slice slice(pin_values[i].data(), pin_values[i].size()); + statuses[i] = ParseMetadata(kRedisJson, &slice, &metadatas[i]); + if (!statuses[i].ok()) continue; + (*raw_values)[i].assign(slice.data(), slice.size()); + } + return statuses; +} + } // namespace redis diff --git a/src/types/redis_json.h b/src/types/redis_json.h index e0e5ec31a12..a39937863da 100644 --- a/src/types/redis_json.h +++ b/src/types/redis_json.h @@ -64,6 +64,9 @@ class Json : public Database { rocksdb::Status StrLen(const std::string &user_key, const std::string &path, Optionals *results); rocksdb::Status ObjLen(const std::string &user_key, const std::string &path, Optionals *results); + std::vector MGet(const std::vector &user_keys, const std::string &path, + 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); @@ -71,6 +74,9 @@ class Json : public Database { rocksdb::Status del(const Slice &ns_key); rocksdb::Status numop(JsonValue::NumOpEnum op, const std::string &user_key, const std::string &path, const std::string &value, JsonValue *result); + std::vector readMulti(const std::vector &ns_keys, std::vector &values); + std::vector getRawMetaData(const std::vector &ns_keys, std::vector &metadatas, + std::vector *raw_values); }; } // namespace redis diff --git a/tests/gocase/unit/type/json/json_test.go b/tests/gocase/unit/type/json/json_test.go index 37aa168ccc2..a849d46d4d5 100644 --- a/tests/gocase/unit/type/json/json_test.go +++ b/tests/gocase/unit/type/json/json_test.go @@ -581,4 +581,36 @@ func TestJson(t *testing.T) { err = rdb.Do(ctx, "JSON.OBJLEN", "no-such-json-key", "$").Err() require.EqualError(t, err, redis.Nil.Error()) }) + + t.Run("JSON.MGET basics", func(t *testing.T) { + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a0", "$", `{"a":1, "b": 2, "nested": {"a": 3}, "c": null}`).Err()) + require.NoError(t, rdb.Do(ctx, "JSON.SET", "a1", "$", `{"a":4, "b": 5, "nested": {"a": 6}, "c": null}`).Err()) + require.NoError(t, rdb.Do(ctx, "SET", "a2", `{"a": 100}`).Err()) + + vals, err := rdb.Do(ctx, "JSON.MGET", "a0", "a1", "$.a").Slice() + require.NoError(t, err) + require.Equal(t, 2, len(vals)) + require.EqualValues(t, "[1]", vals[0]) + require.EqualValues(t, "[4]", vals[1]) + + vals, err = rdb.Do(ctx, "JSON.MGET", "a0", "a1", "a2", "$.a").Slice() + require.NoError(t, err) + require.Equal(t, 3, len(vals)) + require.EqualValues(t, "[1]", vals[0]) + require.EqualValues(t, "[4]", vals[1]) + require.EqualValues(t, nil, vals[2]) + + vals, err = rdb.Do(ctx, "JSON.MGET", "a0", "a1", "$.c").Slice() + require.NoError(t, err) + require.Equal(t, 2, len(vals)) + require.EqualValues(t, "[null]", vals[0]) + require.EqualValues(t, "[null]", vals[1]) + + vals, err = rdb.Do(ctx, "JSON.MGET", "a0", "a1", "$.nonexists").Slice() + require.NoError(t, err) + require.Equal(t, 2, len(vals)) + require.EqualValues(t, "[]", vals[0]) + require.EqualValues(t, "[]", vals[1]) + + }) }