diff --git a/src/server/json_family.cc b/src/server/json_family.cc index 85ba048fa..0736a268d 100644 --- a/src/server/json_family.cc +++ b/src/server/json_family.cc @@ -40,6 +40,7 @@ namespace dfly { using namespace std; using namespace jsoncons; using facade::kSyntaxErrType; +using facade::WrongNumArgsError; using JsonExpression = jsonpath::jsonpath_expression; using OptBool = optional; @@ -1355,6 +1356,24 @@ void JsonFamily::Set(CmdArgList args, ConnectionContext* cntx) { } } +// JSON.MSET key path value [key path value ...] +void JsonFamily::MSet(CmdArgList args, ConnectionContext* cntx) { + DCHECK_GE(args.size(), 3u); + if (args.size() % 3 != 0) { + return cntx->SendError(facade::WrongNumArgsError("json.mset")); + } + + auto cb = [&](Transaction* t, EngineShard* shard) { + ArgSlice args = t->GetShardArgs(shard->shard_id()); + LOG(INFO) << shard->shard_id() << " " << args; + return OpStatus::OK; + }; + + Transaction* trans = cntx->transaction; + trans->ScheduleSingleHop(cb); + cntx->SendOk(); +} + void JsonFamily::Resp(CmdArgList args, ConnectionContext* cntx) { string_view key = ArgS(args, 0); string_view path = DefaultJsonPath; @@ -2026,9 +2045,12 @@ void JsonFamily::Register(CommandRegistry* registry) { ArrAppend); *registry << CI{"JSON.ARRINDEX", CO::READONLY | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(ArrIndex); // TODO: Support negative first_key index to revive the debug sub-command - *registry << CI{"JSON.DEBUG", CO::READONLY | CO::FAST, -3, 2, 2, acl::JSON}.HFUNC(Debug); - *registry << CI{"JSON.RESP", CO::READONLY | CO::FAST, -2, 1, 1, acl::JSON}.HFUNC(Resp); - *registry << CI{"JSON.SET", CO::WRITE | CO::DENYOOM | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(Set); + *registry << CI{"JSON.DEBUG", CO::READONLY | CO::FAST, -3, 2, 2, acl::JSON}.HFUNC(Debug) + << CI{"JSON.RESP", CO::READONLY | CO::FAST, -2, 1, 1, acl::JSON}.HFUNC(Resp) + << CI{"JSON.SET", CO::WRITE | CO::DENYOOM | CO::FAST, -4, 1, 1, acl::JSON}.HFUNC(Set) + << CI{"JSON.MSET", CO::WRITE | CO::DENYOOM | CO::FAST | CO::INTERLEAVED_KEYS, -4, 1, -1, + acl::JSON} + .HFUNC(MSet); } } // namespace dfly diff --git a/src/server/json_family.h b/src/server/json_family.h index c95d81f98..0c2848fe0 100644 --- a/src/server/json_family.h +++ b/src/server/json_family.h @@ -41,6 +41,7 @@ class JsonFamily { static void Debug(CmdArgList args, ConnectionContext* cntx); static void Resp(CmdArgList args, ConnectionContext* cntx); static void Set(CmdArgList args, ConnectionContext* cntx); + static void MSet(CmdArgList args, ConnectionContext* cntx); }; } // namespace dfly diff --git a/src/server/json_family_test.cc b/src/server/json_family_test.cc index c0560debc..470e2e02e 100644 --- a/src/server/json_family_test.cc +++ b/src/server/json_family_test.cc @@ -1081,4 +1081,18 @@ TEST_F(JsonFamilyTest, Set) { EXPECT_EQ(resp, R"([{"a":2,"b":8,"c":[1,2,3]}])"); } +TEST_F(JsonFamilyTest, MSet) { + string json = R"( + {"a":{"a":1, "b":2, "c":3}} + )"; + + auto resp = Run({"JSON.MSET", "j1", "$"}); + EXPECT_THAT(resp, ErrArg("wrong number")); + resp = Run({"JSON.MSET", "j1", "$", json, "j3", "$"}); + EXPECT_THAT(resp, ErrArg("wrong number")); + + resp = Run({"JSON.MSET", "j1", "$", json, "j3", "$", json}); + EXPECT_EQ(resp, "OK"); +} + } // namespace dfly diff --git a/src/server/transaction.cc b/src/server/transaction.cc index 5d8227ee7..a850ba136 100644 --- a/src/server/transaction.cc +++ b/src/server/transaction.cc @@ -208,8 +208,8 @@ void Transaction::BuildShardIndex(const KeyIndex& key_index, std::vectorsize()); // shard_data isn't sparse, so we must allocate for all :( - DCHECK(key_index.step == 1 || key_index.step == 2); - DCHECK(key_index.step != 2 || (full_args_.size() % 2) == 0); + DCHECK_EQ(full_args_.size() % key_index.step, 0u); // Safe, because flow below is not preemptive. auto& shard_index = tmp_space.GetShardIndex(shard_data_.size()); @@ -1578,7 +1577,15 @@ OpResult DetermineKeys(const CommandId* cid, CmdArgList args) { } else { key_index.end = last > 0 ? last : (int(args.size()) + last + 1); } - key_index.step = cid->opt_mask() & CO::INTERLEAVED_KEYS ? 2 : 1; + if (cid->opt_mask() & CO::INTERLEAVED_KEYS) { + if (cid->name() == "JSON.MSET") { + key_index.step = 3; + } else { + key_index.step = 2; + } + } else { + key_index.step = 1; + } if (cid->opt_mask() & CO::STORE_LAST_KEY) { string_view name{cid->name()};