From 10462cacfe769e7e046ae8722ecf8aa9af8b40c3 Mon Sep 17 00:00:00 2001 From: baishen Date: Thu, 17 Oct 2024 10:53:24 +0800 Subject: [PATCH] feat: support `object_insert` function --- src/error.rs | 2 ++ src/functions.rs | 84 +++++++++++++++++++++++++++++++++++++++++++ tests/it/functions.rs | 84 +++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 168 insertions(+), 2 deletions(-) diff --git a/src/error.rs b/src/error.rs index a1880e4..0e76a1b 100644 --- a/src/error.rs +++ b/src/error.rs @@ -86,6 +86,8 @@ pub enum Error { InvalidKeyPath, InvalidJsonType, + InvalidObject, + ObjectDuplicateKey, Syntax(ParseErrorCode, usize), } diff --git a/src/functions.rs b/src/functions.rs index 208d267..f429621 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -2840,6 +2840,90 @@ fn array_overlap_jsonb(value1: &[u8], value2: &[u8]) -> Result { Ok(false) } +/// Insert a new value into a JSONB array value by the specified position. +pub fn object_insert( + value: &[u8], + new_key: &str, + new_value: &[u8], + update_flag: bool, + buf: &mut Vec, +) -> Result<(), Error> { + if !is_jsonb(value) { + let value = parse_value(value)?; + let mut val_buf = Vec::new(); + value.write_to_vec(&mut val_buf); + if !is_jsonb(new_value) { + let new_value = parse_value(new_value)?; + let mut new_val_buf = Vec::new(); + new_value.write_to_vec(&mut new_val_buf); + return object_insert_jsonb(&val_buf, new_key, &new_val_buf, update_flag, buf); + } + return object_insert_jsonb(&val_buf, new_key, new_value, update_flag, buf); + } + object_insert_jsonb(value, new_key, new_value, update_flag, buf) +} + +fn object_insert_jsonb( + value: &[u8], + new_key: &str, + new_value: &[u8], + update_flag: bool, + buf: &mut Vec, +) -> Result<(), Error> { + let header = read_u32(value, 0)?; + if header & CONTAINER_HEADER_TYPE_MASK != OBJECT_CONTAINER_TAG { + return Err(Error::InvalidObject); + } + + let mut idx = 0; + let mut duplicate_key = false; + for (i, obj_key) in iteate_object_keys(value, header).enumerate() { + if new_key.eq(obj_key) { + if !update_flag { + return Err(Error::ObjectDuplicateKey); + } + idx = i; + duplicate_key = true; + break; + } else if new_key > obj_key { + idx = i + 1; + } else { + break; + } + } + + let mut builder = ObjectBuilder::new(); + let mut obj_iter = iterate_object_entries(value, header); + for _ in 0..idx { + if let Some((key, jentry, item)) = obj_iter.next() { + builder.push_raw(key, jentry, item); + } + } + // insert new key and value + let new_header = read_u32(new_value, 0)?; + match new_header & CONTAINER_HEADER_TYPE_MASK { + ARRAY_CONTAINER_TAG | OBJECT_CONTAINER_TAG => { + let new_jentry = JEntry::make_container_jentry(new_value.len()); + builder.push_raw(new_key, new_jentry, new_value); + } + _ => { + let encoded = read_u32(new_value, 4)?; + let new_jentry = JEntry::decode_jentry(encoded); + builder.push_raw(new_key, new_jentry, &new_value[8..]); + } + } + // if the key is duplicated, ignore the original key and value. + if duplicate_key { + let _ = obj_iter.next(); + } + for (key, jentry, item) in obj_iter { + builder.push_raw(key, jentry, item); + } + builder.build_into(buf); + + Ok(()) +} + /// Deletes all object fields that have null values from the given JSON value, recursively. /// Null values that are not object fields are untouched. pub fn strip_nulls(value: &[u8], buf: &mut Vec) -> Result<(), Error> { diff --git a/tests/it/functions.rs b/tests/it/functions.rs index ddead8a..bf089d9 100644 --- a/tests/it/functions.rs +++ b/tests/it/functions.rs @@ -22,8 +22,8 @@ use jsonb::{ contains, convert_to_comparable, delete_by_index, delete_by_keypath, delete_by_name, exists_all_keys, exists_any_keys, from_slice, get_by_index, get_by_keypath, get_by_name, get_by_path, get_by_path_array, is_array, is_object, keypath::parse_key_paths, object_each, - object_keys, parse_value, path_exists, path_match, strip_nulls, to_bool, to_f64, to_i64, - to_pretty_string, to_serde_json, to_serde_json_object, to_str, to_string, to_u64, + object_insert, object_keys, parse_value, path_exists, path_match, strip_nulls, to_bool, to_f64, + to_i64, to_pretty_string, to_serde_json, to_serde_json_object, to_str, to_string, to_u64, traverse_check_string, type_of, Error, Number, Object, Value, }; @@ -1738,6 +1738,86 @@ fn test_array_overlap() { } } +#[test] +fn test_object_insert() { + let sources = vec![ + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + "a", + r#""hello""#, + false, + Some(r#"{"a":"hello","b":11,"d":22,"m":[1,2]}"#), + ), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + "e", + r#"{"k":"v"}"#, + false, + Some(r#"{"b":11,"d":22,"e":{"k":"v"},"m":[1,2]}"#), + ), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + "z", + r#"["z1","z2"]"#, + false, + Some(r#"{"b":11,"d":22,"m":[1,2],"z":["z1","z2"]}"#), + ), + (r#"{"b":11,"d":22,"m":[1,2]}"#, "d", r#"100"#, false, None), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + "d", + r#"100"#, + true, + Some(r#"{"b":11,"d":100,"m":[1,2]}"#), + ), + (r#"{"b":11,"d":22,"m":[1,2]}"#, "m", r#"true"#, false, None), + ( + r#"{"b":11,"d":22,"m":[1,2]}"#, + "m", + r#"true"#, + true, + Some(r#"{"b":11,"d":22,"m":true}"#), + ), + (r#"1"#, "xx", r#"{"k":"v"}"#, true, None), + ]; + for (val, new_key, new_val, update_flag, result) in sources { + { + let val = val.as_bytes(); + let new_val = new_val.as_bytes(); + let mut buf = Vec::new(); + let ret = object_insert(val, new_key, new_val, update_flag, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + { + let val = parse_value(val.as_bytes()).unwrap().to_vec(); + let new_val = parse_value(new_val.as_bytes()).unwrap().to_vec(); + let mut buf = Vec::new(); + let ret = object_insert(&val, new_key, &new_val, update_flag, &mut buf); + match result { + Some(result) => { + assert!(ret.is_ok()); + let actual = from_slice(&buf).unwrap(); + let expected = parse_value(result.as_bytes()).unwrap(); + assert_eq!(actual, expected); + } + None => { + assert!(ret.is_err()); + } + } + } + } +} + #[test] fn test_to_serde_json() { let sources = vec![