Merge pull request #65 from b41sh/feat-object-delete

feat: support `object_delete` and `object_pick` function

src/functions.rs

@@ -2840,7 +2840,7 @@ 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.
+/// Insert a new value into a JSONB object value by the new key and new value.
 pub fn object_insert(
     value: &[u8],
     new_key: &str,
@@ -2924,6 +2924,68 @@ fn object_insert_jsonb(
     Ok(())
 }
 
+/// Delete keys and values from a JSONB object value by keys.
+pub fn object_delete(value: &[u8], keys: &BTreeSet<&str>, 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);
+        return object_delete_jsonb(&val_buf, keys, buf);
+    }
+    object_delete_jsonb(value, keys, buf)
+}
+
+fn object_delete_jsonb(
+    value: &[u8],
+    keys: &BTreeSet<&str>,
+    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 builder = ObjectBuilder::new();
+    for (key, jentry, item) in iterate_object_entries(value, header) {
+        if keys.contains(key) {
+            continue;
+        }
+        builder.push_raw(key, jentry, item);
+    }
+    builder.build_into(buf);
+
+    Ok(())
+}
+
+/// Pick keys and values from a JSONB object value by keys.
+pub fn object_pick(value: &[u8], keys: &BTreeSet<&str>, 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);
+        return object_pick_jsonb(&val_buf, keys, buf);
+    }
+    object_pick_jsonb(value, keys, buf)
+}
+
+fn object_pick_jsonb(value: &[u8], keys: &BTreeSet<&str>, 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 builder = ObjectBuilder::new();
+    for (key, jentry, item) in iterate_object_entries(value, header) {
+        if !keys.contains(key) {
+            continue;
+        }
+        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

@@ -15,16 +15,17 @@
 use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 
 use jsonb::{
     array_distinct, array_except, array_insert, array_intersection, array_length, array_overlap,
     array_values, as_bool, as_null, as_number, as_str, build_array, build_object, compare, concat,
     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_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,
+    get_by_path, get_by_path_array, is_array, is_object, keypath::parse_key_paths, object_delete,
+    object_each, object_insert, object_keys, object_pick, 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,
 };
 
 use jsonb::jsonpath::parse_json_path;
@@ -1818,6 +1819,120 @@ fn test_object_insert() {
     }
 }
 
+#[test]
+fn test_object_pick() {
+    let sources = vec![
+        (
+            r#"{"b":11,"d":22,"m":[1,2]}"#,
+            vec!["a", "b", "c"],
+            Some(r#"{"b":11}"#),
+        ),
+        (
+            r#"{"b":11,"d":22,"m":[1,2]}"#,
+            vec!["a", "x", "y"],
+            Some(r#"{}"#),
+        ),
+        (
+            r#"{"k1":"v1","k2":{"x":"y"}}"#,
+            vec!["k1"],
+            Some(r#"{"k1":"v1"}"#),
+        ),
+        (r#"1"#, vec!["a", "b"], None),
+    ];
+    for (val, keys, result) in sources {
+        let keys = BTreeSet::from_iter(keys);
+        {
+            let val = val.as_bytes();
+            let mut buf = Vec::new();
+            let ret = object_pick(val, &keys, &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 mut buf = Vec::new();
+            let ret = object_pick(&val, &keys, &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_object_delete() {
+    let sources = vec![
+        (
+            r#"{"b":11,"d":22,"m":[1,2]}"#,
+            vec!["a", "b", "c"],
+            Some(r#"{"d":22,"m":[1,2]}"#),
+        ),
+        (
+            r#"{"b":11,"d":22,"m":[1,2]}"#,
+            vec!["a", "x", "y"],
+            Some(r#"{"b":11,"d":22,"m":[1,2]}"#),
+        ),
+        (
+            r#"{"k1":"v1","k2":{"x":"y"}}"#,
+            vec!["k1"],
+            Some(r#"{"k2":{"x":"y"}}"#),
+        ),
+        (r#"1"#, vec!["a", "b"], None),
+    ];
+    for (val, keys, result) in sources {
+        let keys = BTreeSet::from_iter(keys);
+        {
+            let val = val.as_bytes();
+            let mut buf = Vec::new();
+            let ret = object_delete(val, &keys, &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 mut buf = Vec::new();
+            let ret = object_delete(&val, &keys, &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![