Skip to content

Commit

Permalink
support JSON.NUMINCRBY and JSON.NUMMULTBY command
Browse files Browse the repository at this point in the history
  • Loading branch information
skyitachi authored and Shiping Yao committed Nov 19, 2023
1 parent 358cc41 commit 66548eb
Show file tree
Hide file tree
Showing 6 changed files with 289 additions and 1 deletion.
38 changes: 37 additions & 1 deletion src/commands/cmd_json.cc
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,40 @@ class CommanderJsonArrIndex : public Commander {
ssize_t end_;
};

class CommandJsonNumIncrBy : public Commander {
public:
Status Execute(Server *svr, Connection *conn, std::string *output) override {
redis::Json json(svr->storage, conn->GetNamespace());

JsonValue result = JsonValue::FromString("[]").GetValue();
auto s = json.NumIncrBy(args_[1], args_[2], args_[3], &result);
if (!s.ok()) {
return {Status::RedisExecErr, s.ToString()};
}

*output = redis::BulkString(result.value.to_string());

return Status::OK();
}
};

class CommandJsonNumMultBy : public Commander {
public:
Status Execute(Server *svr, Connection *conn, std::string *output) override {
redis::Json json(svr->storage, conn->GetNamespace());

JsonValue result = JsonValue::FromString("[]").GetValue();
auto s = json.NumMultBy(args_[1], args_[2], args_[3], &result);
if (!s.ok()) {
return {Status::RedisExecErr, s.ToString()};
}

*output = redis::BulkString(result.value.to_string());

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 @@ -491,6 +525,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandJsonSet>("json.set", 4, "write", 1, 1
MakeCmdAttr<CommandJsonMerge>("json.merge", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonObjkeys>("json.objkeys", -2, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonArrPop>("json.arrpop", -2, "write", 1, 1, 1),
MakeCmdAttr<CommanderJsonArrIndex>("json.arrindex", -4, "read-only", 1, 1, 1), );
MakeCmdAttr<CommanderJsonArrIndex>("json.arrindex", -4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandJsonNumIncrBy>("json.numincrby", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandJsonNumMultBy>("json.nummultby", 4, "write", 1, 1, 1), );

} // namespace redis
43 changes: 43 additions & 0 deletions src/types/json.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,11 @@
constexpr ssize_t NOT_FOUND_INDEX = -1;
constexpr ssize_t NOT_ARRAY = -2;

enum class NumOpEnum : uint8_t {
Incr = 1,
Mul = 2,
};

struct JsonValue {
JsonValue() = default;
explicit JsonValue(jsoncons::basic_json<char> value) : value(std::move(value)) {}
Expand Down Expand Up @@ -475,6 +480,44 @@ struct JsonValue {
return Status::OK();
}

Status NumOp(std::string_view path, const JsonValue &number, NumOpEnum op, JsonValue *result) {
Status status = Status::OK();
try {
jsoncons::jsonpath::json_replace(value, path, [&](const std::string & /*path*/, jsoncons::json &origin) {
if (!status.IsOK()) {
return;
}
if (!origin.is_number()) {
result->value.push_back(jsoncons::json::null());
return;
}
if (number.value.is_double() || origin.is_double()) {
double v = 0;
if (op == NumOpEnum::Incr) {
v = origin.as_double() + number.value.as_double();
} else if (op == NumOpEnum::Mul) {
v = origin.as_double() * number.value.as_double();
}
if (std::isinf(v)) {
status = {Status::RedisExecErr, "result is an infinite number"};
return;
}
origin = v;
} else {
if (op == NumOpEnum::Incr) {
origin = origin.as_integer<int64_t>() + number.value.as_integer<int64_t>();
} else if (op == NumOpEnum::Mul) {
origin = origin.as_integer<int64_t>() * number.value.as_integer<int64_t>();
}
}
result->value.push_back(origin);
});
} catch (const jsoncons::jsonpath::jsonpath_error &e) {
return {Status::NotOK, e.what()};
}
return status;
}

JsonValue(const JsonValue &) = default;
JsonValue(JsonValue &&) = default;

Expand Down
34 changes: 34 additions & 0 deletions src/types/redis_json.cc
Original file line number Diff line number Diff line change
Expand Up @@ -390,4 +390,38 @@ rocksdb::Status Json::ArrTrim(const std::string &user_key, const std::string &pa
return write(ns_key, &metadata, json_val);
}

rocksdb::Status Json::NumIncrBy(const std::string &user_key, const std::string &path, const std::string &value,
JsonValue *result) {
return numop(NumOpEnum::Incr, user_key, path, value, result);
}

rocksdb::Status Json::NumMultBy(const std::string &user_key, const std::string &path, const std::string &value,
JsonValue *result) {
return numop(NumOpEnum::Mul, user_key, path, value, result);
}

rocksdb::Status Json::numop(NumOpEnum op, const std::string &user_key, const std::string &path,
const std::string &value, JsonValue *result) {
JsonValue number;
auto number_res = JsonValue::FromString(value);
if (!number_res || !number_res.GetValue().value.is_number()) {
return rocksdb::Status::InvalidArgument("should be a number");
}
number = std::move(number_res.GetValue());

auto ns_key = AppendNamespacePrefix(user_key);
JsonMetadata metadata;
JsonValue json_val;
auto s = read(ns_key, &metadata, &json_val);
if (!s.ok()) return s;

LockGuard guard(storage_->GetLockManager(), ns_key);

auto res = json_val.NumOp(path, number, op, result);
if (!res) {
return rocksdb::Status::InvalidArgument(res.Msg());
}
return write(ns_key, &metadata, json_val);
}

} // namespace redis
6 changes: 6 additions & 0 deletions src/types/redis_json.h
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@ class Json : public Database {
std::vector<std::optional<JsonValue>> *results);
rocksdb::Status ArrIndex(const std::string &user_key, const std::string &path, const std::string &needle,
ssize_t start, ssize_t end, std::vector<ssize_t> *result);
rocksdb::Status NumIncrBy(const std::string &user_key, const std::string &path, const std::string &value,
JsonValue *result);
rocksdb::Status NumMultBy(const std::string &user_key, const std::string &path, const std::string &value,
JsonValue *result);

rocksdb::Status ArrTrim(const std::string &user_key, const std::string &path, int64_t start, int64_t stop,
std::vector<std::optional<uint64_t>> &results);
Expand All @@ -61,6 +65,8 @@ class Json : public Database {
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 numop(NumOpEnum op, const std::string &user_key, const std::string &path, const std::string &value,
JsonValue *result);
};

} // namespace redis
117 changes: 117 additions & 0 deletions tests/cppunit/types/json_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -485,3 +485,120 @@ TEST_F(RedisJsonTest, ArrIndex) {
ASSERT_TRUE(json_->ArrIndex(key_, "$.arr", "3", 0, 2, &res).ok() && res.size() == 1);
ASSERT_EQ(res[0], -1);
}

TEST_F(RedisJsonTest, NumIncrBy) {
ASSERT_TRUE(json_->Set(key_, "$", R"({ "foo": 0, "bar": "baz" })").ok());
JsonValue res = JsonValue::FromString("[]").GetValue();
ASSERT_TRUE(json_->NumIncrBy(key_, "$.foo", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1]");
res.value.clear();
ASSERT_TRUE(json_->NumIncrBy(key_, "$.foo", "2", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[3]");
res.value.clear();
ASSERT_TRUE(json_->NumIncrBy(key_, "$.foo", "0.5", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[3.5]");
res.value.clear();

ASSERT_TRUE(json_->NumIncrBy(key_, "$.bar", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[null]");
res.value.clear();

ASSERT_TRUE(json_->NumIncrBy(key_, "$.fuzz", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[]");
res.value.clear();

ASSERT_TRUE(json_->Set(key_, "$", "0").ok());
ASSERT_TRUE(json_->NumIncrBy(key_, "$", "0", &res).ok());
res.value.clear();
ASSERT_TRUE(json_->NumIncrBy(key_, "$", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1]");
res.value.clear();
ASSERT_TRUE(json_->NumIncrBy(key_, "$", "1.5", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[2.5]");
res.value.clear();

// overflow case
ASSERT_TRUE(json_->Set(key_, "$", "1.6350000000001313e+308").ok());
ASSERT_TRUE(json_->NumIncrBy(key_, "$", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

ASSERT_TRUE(json_->NumIncrBy(key_, "$", "2", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

// nested big_num object
ASSERT_TRUE(json_->Set(key_, "$", R"({"l1":{"l2_a":1.6350000000001313e+308,"l2_b":2}})").ok());
ASSERT_TRUE(json_->NumIncrBy(key_, "$.l1.l2_a", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

ASSERT_TRUE(json_->NumIncrBy(key_, "$.l1.l2_a", "2", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

// nested big_num array
ASSERT_TRUE(json_->Set(key_, "$", R"({"l1":{"l2":[0,1.6350000000001313e+308]}})").ok());
ASSERT_FALSE(json_->NumIncrBy(key_, "$.l1.l2[1]", "1.6350000000001313e+308", &res).ok());
res.value.clear();

ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
ASSERT_EQ(json_val_.Dump().GetValue(), R"({"l1":{"l2":[0,1.6350000000001313e+308]}})");
}

TEST_F(RedisJsonTest, NumMultBy) {
ASSERT_TRUE(json_->Set(key_, "$", R"({ "foo": 1, "bar": "baz" })").ok());
JsonValue res = JsonValue::FromString("[]").GetValue();
ASSERT_TRUE(json_->NumMultBy(key_, "$.foo", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1]");
res.value.clear();
ASSERT_TRUE(json_->NumMultBy(key_, "$.foo", "2", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[2]");
res.value.clear();
ASSERT_TRUE(json_->NumMultBy(key_, "$.foo", "0.5", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.0]");
res.value.clear();

ASSERT_TRUE(json_->NumMultBy(key_, "$.bar", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[null]");
res.value.clear();

ASSERT_TRUE(json_->NumMultBy(key_, "$.fuzz", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[]");
res.value.clear();

// num object
ASSERT_TRUE(json_->Set(key_, "$", "1.0").ok());
ASSERT_TRUE(json_->NumMultBy(key_, "$", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.0]");
res.value.clear();
ASSERT_TRUE(json_->NumMultBy(key_, "$", "1.5", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.5]");
res.value.clear();

// overflow case
ASSERT_TRUE(json_->Set(key_, "$", "1.6350000000001313e+308").ok());
ASSERT_TRUE(json_->NumMultBy(key_, "$", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

ASSERT_FALSE(json_->NumMultBy(key_, "$", "2", &res).ok());
res.value.clear();

// nested big_num object
ASSERT_TRUE(json_->Set(key_, "$", R"({"l1":{"l2_a":1.6350000000001313e+308,"l2_b":2}})").ok());
ASSERT_TRUE(json_->NumMultBy(key_, "$.l1.l2_a", "1", &res).ok());
ASSERT_EQ(res.Print(0, true).GetValue(), "[1.6350000000001313e+308]");
res.value.clear();

ASSERT_FALSE(json_->NumMultBy(key_, "$.l1.l2_a", "2", &res).ok());
res.value.clear();

// nested big_num array
ASSERT_TRUE(json_->Set(key_, "$", R"({"l1":{"l2":[0,1.6350000000001313e+308]}})").ok());
ASSERT_FALSE(json_->NumMultBy(key_, "$.l1.l2[1]", "2", &res).ok());
res.value.clear();

ASSERT_TRUE(json_->Get(key_, {}, &json_val_).ok());
ASSERT_EQ(json_val_.Dump().GetValue(), R"({"l1":{"l2":[0,1.6350000000001313e+308]}})");
}
52 changes: 52 additions & 0 deletions tests/gocase/unit/type/json/json_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -462,4 +462,56 @@ func TestJson(t *testing.T) {
require.NoError(t, rdb.Do(ctx, "JSON.SET", "a1", "$", `{"arr":[[1],[2],[3]]}`).Err())
require.Equal(t, []interface{}{int64(0), int64(-1), int64(-1)}, rdb.Do(ctx, arrIndexCmd, "a1", "$.arr.*", `1`).Val())
})

t.Run("JSON.NUMOP basics", func(t *testing.T) {
require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{ "foo": 0, "bar": "baz" }`).Err())
require.Equal(t, `[1]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.foo", 1).Val())
require.Equal(t, `[1]`, rdb.Do(ctx, "JSON.GET", "a", "$.foo").Val())
require.Equal(t, `[3]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.foo", 2).Val())
require.Equal(t, `[3.5]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.foo", 0.5).Val())

// wrong type
require.Equal(t, `[null]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.bar", 1).Val())

require.Equal(t, `[]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.fuzz", 1).Val())

require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `0`).Err())
require.Equal(t, `[1]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$", 1).Val())
require.Equal(t, `[2.5]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$", 1.5).Val())

require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"foo":0,"bar":42}`).Err())
require.Equal(t, `[1]`, rdb.Do(ctx, "JSON.NUMINCRBY", "a", "$.foo", 1).Val())
require.Equal(t, `[84]`, rdb.Do(ctx, "JSON.NUMMULTBY", "a", "$.bar", 2).Val())

// overflow case
require.NoError(t, rdb.Do(ctx, "JSON.SET", "big_num", "$", "1.6350000000001313e+308").Err())
require.Equal(t, `[1.6350000000001313e+308]`,
rdb.Do(ctx, "JSON.NUMINCRBY", "big_num", "$", 1).Val())

require.Error(t, rdb.Do(ctx, "JSON.NUMMULTBY", "big_num", "$", 2).Err())

require.Equal(t, `1.6350000000001313e+308`, rdb.Do(ctx, "JSON.GET", "big_num").Val())

require.NoError(t, rdb.Do(ctx, "JSON.SET", "nested_obj_big_num", "$",
`{"l1":{"l2_a":1.6350000000001313e+308,"l2_b":2}}`).Err())

require.Equal(t, `[1.6350000000001313e+308]`,
rdb.Do(ctx, "JSON.NUMINCRBY", "nested_obj_big_num", "$.l1.l2_a", 1).Val())

require.Error(t, rdb.Do(ctx, "JSON.NUMMULTBY", "nested_obj_big_num", "$.l1.l2_a", 2).Err())

require.Equal(t, `{"l1":{"l2_a":1.6350000000001313e+308,"l2_b":2}}`,
rdb.Do(ctx, "JSON.GET", "nested_obj_big_num").Val())

require.NoError(t, rdb.Do(ctx, "JSON.SET", "nested_arr_big_num", "$",
`{"l1":{"l2":[0,1.6350000000001313e+308]}}`).Err())

require.Error(t, rdb.Do(ctx, "JSON.NUMINCRBY", "nested_arr_big_num", "$.l1.l2[1]",
`1.6350000000001313e+308`).Err())
require.Error(t, rdb.Do(ctx, "JSON.NUMMULTBY", "nested_arr_big_num", "$.l1.l2[1]", 2).Err())

require.Equal(t, `{"l1":{"l2":[0,1.6350000000001313e+308]}}`,
rdb.Do(ctx, "JSON.GET", "nested_arr_big_num").Val())
})

}

0 comments on commit 66548eb

Please sign in to comment.