Skip to content

Commit 47e05e4

Browse files
gkorlandgkorland
authored
add json.mset (RedisJSON#944)
* add json.mset
1 parent 0b898be commit 47e05e4

File tree

4 files changed

+134
-20
lines changed

4 files changed

+134
-20
lines changed

‎src/commands.rs

Lines changed: 85 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -172,27 +172,10 @@ pub fn json_set<M: Manager>(manager: M, ctx: &Context, args: Vec<RedisString>) -
172172
Ok(RedisValue::Null)
173173
}
174174
} else {
175-
let mut update_info = KeyValue::new(doc).find_paths(path.get_path(), &op)?;
175+
let update_info = KeyValue::new(doc).find_paths(path.get_path(), &op)?;
176176
if !update_info.is_empty() {
177-
let mut res = false;
178-
if update_info.len() == 1 {
179-
res = match update_info.pop().unwrap() {
180-
UpdateInfo::SUI(sui) => redis_key.set_value(sui.path, val)?,
181-
UpdateInfo::AUI(aui) => redis_key.dict_add(aui.path, &aui.key, val)?,
182-
}
183-
} else {
184-
for ui in update_info {
185-
res = match ui {
186-
UpdateInfo::SUI(sui) => {
187-
redis_key.set_value(sui.path, val.clone())?
188-
}
189-
UpdateInfo::AUI(aui) => {
190-
redis_key.dict_add(aui.path, &aui.key, val.clone())?
191-
}
192-
}
193-
}
194-
}
195-
if res {
177+
let updated = apply_updates::<M>(&mut redis_key, val, update_info);
178+
if updated {
196179
redis_key.apply_changes(ctx, "json.set")?;
197180
REDIS_OK
198181
} else {
@@ -218,6 +201,88 @@ pub fn json_set<M: Manager>(manager: M, ctx: &Context, args: Vec<RedisString>) -
218201
}
219202
}
220203

204+
///
205+
/// JSON.MSET <key> <path> <json> [[<key> <path> <json>]...]
206+
///
207+
pub fn json_mset<M: Manager>(manager: M, ctx: &Context, args: Vec<RedisString>) -> RedisResult {
208+
let mut args = args.into_iter().skip(1);
209+
210+
if args.len() < 3 {
211+
return Err(RedisError::WrongArity);
212+
}
213+
214+
// Collect all the actions from the args (redis_key, update_info, value)
215+
let mut actions = Vec::new();
216+
while let Ok(key) = args.next_arg() {
217+
let mut redis_key = manager.open_key_write(ctx, key)?;
218+
219+
// Verify the key is a JSON type
220+
let key_value = redis_key.get_value()?;
221+
222+
// Verify the path is valid and get all the update info
223+
let path = Path::new(args.next_str()?);
224+
let update_info = if path.get_path() == JSON_ROOT_PATH {
225+
None
226+
} else if let Some(value) = key_value {
227+
Some(KeyValue::new(value).find_paths(path.get_path(), &SetOptions::None)?)
228+
} else {
229+
return Err(RedisError::Str(
230+
"ERR new objects must be created at the root",
231+
));
232+
};
233+
234+
// Parse the input and validate it's valid JSON
235+
let value_str = args.next_str()?;
236+
let value = manager.from_str(value_str, Format::JSON, true)?;
237+
238+
actions.push((redis_key, update_info, value));
239+
}
240+
241+
actions
242+
.drain(..)
243+
.fold(REDIS_OK, |res, (mut redis_key, update_info, value)| {
244+
let updated = if let Some(update_info) = update_info {
245+
!update_info.is_empty() && apply_updates::<M>(&mut redis_key, value, update_info)
246+
} else {
247+
// In case it is a root path
248+
redis_key.set_value(Vec::new(), value)?
249+
};
250+
if updated {
251+
redis_key.apply_changes(ctx, "json.mset")?
252+
}
253+
res
254+
})
255+
}
256+
257+
fn apply_updates<M: Manager>(
258+
redis_key: &mut M::WriteHolder,
259+
value: M::O,
260+
mut update_info: Vec<UpdateInfo>,
261+
) -> bool {
262+
// If there is only one update info, we can avoid cloning the value
263+
if update_info.len() == 1 {
264+
match update_info.pop().unwrap() {
265+
UpdateInfo::SUI(sui) => redis_key.set_value(sui.path, value).unwrap_or(false),
266+
UpdateInfo::AUI(aui) => redis_key
267+
.dict_add(aui.path, &aui.key, value)
268+
.unwrap_or(false),
269+
}
270+
} else {
271+
let mut updated = false;
272+
for ui in update_info {
273+
updated = match ui {
274+
UpdateInfo::SUI(sui) => redis_key
275+
.set_value(sui.path, value.clone())
276+
.unwrap_or(false),
277+
UpdateInfo::AUI(aui) => redis_key
278+
.dict_add(aui.path, &aui.key, value.clone())
279+
.unwrap_or(false),
280+
} || updated
281+
}
282+
updated
283+
}
284+
}
285+
221286
fn find_paths<T: SelectValue, F: FnMut(&T) -> bool>(
222287
path: &str,
223288
doc: &T,

‎src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ macro_rules! redis_json_module_create {(
209209
["json.get", json_command!(json_get), "readonly", 1,1,1],
210210
["json.mget", json_command!(json_mget), "readonly", 1,1,1],
211211
["json.set", json_command!(json_set), "write deny-oom", 1,1,1],
212+
["json.mset", json_command!(json_mset), "write deny-oom", 1,0,3],
212213
["json.type", json_command!(json_type), "readonly", 1,1,1],
213214
["json.numincrby", json_command!(json_num_incrby), "write", 1,1,1],
214215
["json.toggle", json_command!(json_bool_toggle), "write deny-oom", 1,1,1],

‎tests/pytest/test.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1072,6 +1072,48 @@ def testMultiPathResults(env):
10721072
# make sure legacy json path returns single result
10731073
env.expect("JSON.GET", "k", '.*[0,2]').equal('1')
10741074

1075+
1076+
def testMSET(env):
1077+
env.expect("JSON.MSET", "a{s}", '$', '"a_val"').ok()
1078+
env.expect("JSON.GET", "a{s}", '$').equal('["a_val"]')
1079+
1080+
env.expect("JSON.MSET", "a{s}", '$', '{"aa":"a_val"}', "b{s}", '$', '{"bb":"b_val"}').ok()
1081+
env.expect("JSON.MGET", "a{s}", "b{s}", '$').equal(['[{"aa":"a_val"}]', '[{"bb":"b_val"}]'])
1082+
1083+
env.expect("JSON.MSET", "a{s}", '$.ab', '"a_val2"', "b{s}", '$..bb', '"b_val2"').ok()
1084+
env.expect("JSON.MGET", "a{s}", "b{s}", '$').equal(['[{"aa":"a_val","ab":"a_val2"}]', '[{"bb":"b_val2"}]'])
1085+
1086+
1087+
def testMSET_Partial(env):
1088+
# Make sure MSET doesn't stop on the first update that can't be updated
1089+
env.expect("JSON.SET", "a{s}", '$', '{"x": {"y":[10,20], "z":[30,40]}}').ok()
1090+
env.expect("JSON.SET", "b{s}", '$', '{"x": 60}').ok()
1091+
env.expect("JSON.MSET", "a{s}", '$.x', '{}', "a{s}", '$.x.z[1]', '50', 'b{s}', '$.x', '70').ok()
1092+
env.expect("JSON.GET", "a{s}", '$').equal('[{"x":{}}]')
1093+
env.expect("JSON.GET", "b{s}", '$').equal('[{"x":70}]')
1094+
1095+
# Update the same key twice with a failure in the middle
1096+
env.expect("JSON.SET", "a{s}", '$', '{"x": {"y":[10,20], "z":[30,40]}}').ok()
1097+
env.expect("JSON.MSET", "a{s}", '$.x', '{}', "a{s}", '$.x.z[1]', '50', "a{s}", '$.u', '70').ok()
1098+
env.expect("JSON.GET", "a{s}", '$').equal('[{"x":{},"u":70}]')
1099+
1100+
def testMSET_Error(env):
1101+
env.expect("JSON.SET", "a{s}", '$', '"a_val"').ok()
1102+
1103+
env.expect("JSON.MSET", "a{s}").raiseError()
1104+
# make sure value was not changed
1105+
env.expect("JSON.GET", "a{s}", '$').equal('["a_val"]')
1106+
1107+
env.expect("JSON.MSET", "a{s}", '$', '"a_val2"', "b{s}", '$').raiseError()
1108+
# make sure value was not changed
1109+
env.expect("JSON.GET", "a{s}", '$').equal('["a_val"]')
1110+
env.expect("JSON.GET", "b{s}", '$').equal(None)
1111+
1112+
env.expect("JSON.MSET", "a{s}", '$', '"a_val2"', "b{s}", '$.....a', '"b_val"').raiseError()
1113+
# make sure value was not changed
1114+
env.expect("JSON.GET", "a{s}", '$').equal('["a_val"]')
1115+
env.expect("JSON.GET", "b{s}", '$').equal(None)
1116+
10751117
def testIssue_597(env):
10761118
env.expect("JSON.SET", "test", ".", "[0]").ok()
10771119
env.assertEqual(env.execute_command("JSON.SET", "test", ".[0]", "[0]", "NX"), None)

‎tests/pytest/test_keyspace_notifications.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ def test_keyspace_set(env):
3030
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'json.set')
3131
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'test_key')
3232

33+
env.assertEqual('OK', r.execute_command('JSON.MSET', 'test_key', '$.foo', '"gogo"', 'test_key{test_key}', '$', '{"a":"fufu"}'))
34+
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'json.mset')
35+
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'test_key')
36+
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'json.mset')
37+
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'test_key{test_key}')
38+
3339
env.assertEqual([8], r.execute_command('JSON.STRAPPEND', 'test_key', '$.foo', '"toto"'))
3440
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'json.strappend')
3541
assert_msg(env, pubsub.get_message(timeout=1), 'pmessage', 'test_key')

0 commit comments

Comments
 (0)