Merge pull request #64 from b41sh/feat-object-insert

feat: support `object_insert` function

src/error.rs

@@ -86,6 +86,8 @@ pub enum Error {
     InvalidKeyPath,
 
     InvalidJsonType,
+    InvalidObject,
+    ObjectDuplicateKey,
 
     Syntax(ParseErrorCode, usize),
 }

src/functions.rs

@@ -2840,6 +2840,90 @@ fn array_overlap_jsonb(value1: &[u8], value2: &[u8]) -> Result<bool, Error> {
     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<u8>,
+) -> 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<u8>,
+) -> 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<u8>) -> Result<(), Error> {

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![