Skip to content

Commit

Permalink
Add the support of JSON.MGET command (#1930)
Browse files Browse the repository at this point in the history
  • Loading branch information
@skyitachi
skyitachi authored Dec 14, 2023
1 parent 55c55e7 commit 99703d6
Show file tree
Hide file tree
Showing 4 changed files with 151 additions and 1 deletion.
31 changes: 30 additions & 1 deletion src/commands/cmd_json.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string> user_keys;
for (size_t i = 1; i + 1 < args_.size(); i++) {
user_keys.emplace_back(args_[i]);
}

std::vector<JsonValue> json_values;
auto statuses = json.MGet(user_keys, path, json_values);

std::vector<std::string> 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<CommandJsonSet>("json.set", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonGet>("json.get", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonInfo>("json.info", 2, "read-only", 1, 1, 1),
Expand All @@ -587,6 +615,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", 4, "write", 1, 1
MakeCmdAttr<CommandJsonNumMultBy>("json.nummultby", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonObjLen>("json.objlen", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonStrAppend>("json.strappend", -3, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonStrLen>("json.strlen", -2, "read-only", 1, 1, 1), );
MakeCmdAttr<CommandJsonStrLen>("json.strlen", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonMGet>("json.mget", -3, "read-only", 1, 1, 1), );

} // namespace redis
83 changes: 83 additions & 0 deletions src/types/redis_json.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -512,4 +512,87 @@ rocksdb::Status Json::ObjLen(const std::string &user_key, const std::string &pat
return rocksdb::Status::OK();
}

std::vector<rocksdb::Status> Json::MGet(const std::vector<std::string> &user_keys, const std::string &path,
std::vector<JsonValue> &results) {
std::vector<Slice> ns_keys;
std::vector<std::string> 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<JsonValue> 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<rocksdb::Status> Json::readMulti(const std::vector<Slice> &ns_keys, std::vector<JsonValue> &values) {
std::vector<std::string> raw_values;
std::vector<JsonMetadata> 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<rocksdb::Status> Json::getRawMetaData(const std::vector<Slice> &ns_keys,
std::vector<JsonMetadata> &metadatas,
std::vector<std::string> *raw_values) {
rocksdb::ReadOptions read_options = storage_->DefaultMultiGetOptions();
LatestSnapShot ss(storage_);
read_options.snapshot = ss.GetSnapShot();
raw_values->resize(ns_keys.size());
std::vector<rocksdb::Status> statuses(ns_keys.size());
std::vector<rocksdb::PinnableSlice> 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
6 changes: 6 additions & 0 deletions src/types/redis_json.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,13 +64,19 @@ class Json : public Database {
rocksdb::Status StrLen(const std::string &user_key, const std::string &path, Optionals<uint64_t> *results);
rocksdb::Status ObjLen(const std::string &user_key, const std::string &path, Optionals<uint64_t> *results);

std::vector<rocksdb::Status> MGet(const std::vector<std::string> &user_keys, const std::string &path,
std::vector<JsonValue> &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);
rocksdb::Status create(const std::string &ns_key, JsonMetadata &metadata, const std::string &value);
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<rocksdb::Status> readMulti(const std::vector<Slice> &ns_keys, std::vector<JsonValue> &values);
std::vector<rocksdb::Status> getRawMetaData(const std::vector<Slice> &ns_keys, std::vector<JsonMetadata> &metadatas,
std::vector<std::string> *raw_values);
};

} // namespace redis
32 changes: 32 additions & 0 deletions tests/gocase/unit/type/json/json_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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])

})
}

0 comments on commit 99703d6

Please sign in to comment.