add support for json_set

Test cases are included.
Related to #127

COMPAT.md

@@ -374,7 +374,7 @@ Modifiers:
 | jsonb_remove(json,path,...)        |         |                                                                                                                                              |
 | json_replace(json,path,value,...)  |         |                                                                                                                                              |
 | jsonb_replace(json,path,value,...) |         |                                                                                                                                              |
-| json_set(json,path,value,...)      |         |                                                                                                                                              |
+| json_set(json,path,value,...)      | Yes     |                                                                                                                                              |
 | jsonb_set(json,path,value,...)     |         |                                                                                                                                              |
 | json_type(json)                    | Yes     |                                                                                                                                              |
 | json_type(json,path)               | Yes     |                                                                                                                                              |

Cargo.lock

@@ -1528,7 +1528,7 @@ dependencies = [
 
 [[package]]
 name = "limbo"
-version = "0.0.14"
+version = "0.0.13"
 dependencies = [
  "anyhow",
  "clap",
@@ -1681,7 +1681,7 @@ dependencies = [
 
 [[package]]
 name = "limbo_time"
-version = "0.0.13"
+version = "0.0.14"
 dependencies = [
  "chrono",
  "limbo_ext",

core/function.rs

@@ -83,6 +83,7 @@ pub enum JsonFunc {
     JsonPatch,
     JsonRemove,
     JsonPretty,
+    JsonSet,
 }
 
 #[cfg(feature = "json")]
@@ -105,6 +106,7 @@ impl Display for JsonFunc {
                 Self::JsonPatch => "json_patch".to_string(),
                 Self::JsonRemove => "json_remove".to_string(),
                 Self::JsonPretty => "json_pretty".to_string(),
+                Self::JsonSet => "json_set".to_string(),
             }
         )
     }
@@ -540,6 +542,8 @@ impl Func {
             "json_remove" => Ok(Self::Json(JsonFunc::JsonRemove)),
             #[cfg(feature = "json")]
             "json_pretty" => Ok(Self::Json(JsonFunc::JsonPretty)),
+            #[cfg(feature = "json")]
+            "json_set" => Ok(Self::Json(JsonFunc::JsonSet)),
             "unixepoch" => Ok(Self::Scalar(ScalarFunc::UnixEpoch)),
             "julianday" => Ok(Self::Scalar(ScalarFunc::JulianDay)),
             "hex" => Ok(Self::Scalar(ScalarFunc::Hex)),

core/json/mod.rs

@@ -155,6 +155,45 @@ pub fn json_array_length(
     }
 }
 
+pub fn json_set(json: &OwnedValue, values: &[OwnedValue]) -> crate::Result<OwnedValue> {
+    let mut json_value = get_json_value(json)?;
+
+    values
+        .chunks(2)
+        .map(|chunk| match chunk {
+            [path, value] => {
+                let path = json_path_from_owned_value(path, true)?;
+
+                if let Some(path) = path {
+                    let new_value = match value {
+                        OwnedValue::Text(LimboText {
+                            value,
+                            subtype: TextSubtype::Text,
+                        }) => Val::String(value.to_string()),
+                        _ => get_json_value(value)?,
+                    };
+
+                    let mut new_json_value = json_value.clone();
+
+                    match create_and_mutate_json_by_path(&mut new_json_value, path, |val| match val
+                    {
+                        Target::Array(arr, index) => arr[index] = new_value.clone(),
+                        Target::Value(val) => *val = new_value.clone(),
+                    }) {
+                        Some(_) => json_value = new_json_value,
+                        _ => {}
+                    }
+                }
+
+                Ok(())
+            }
+            _ => crate::bail_constraint_error!("json_set needs an odd number of arguments"),
+        })
+        .collect::<crate::Result<()>>()?;
+
+    convert_json_to_db_type(&json_value, true)
+}
+
 /// Implements the -> operator. Always returns a proper JSON value.
 /// https://sqlite.org/json1.html#the_and_operators
 pub fn json_arrow_extract(value: &OwnedValue, path: &OwnedValue) -> crate::Result<OwnedValue> {
@@ -479,6 +518,92 @@ fn find_target<'a>(json: &'a mut Val, path: &JsonPath) -> Option<Target<'a>> {
     Some(Target::Value(current))
 }
 
+fn create_and_mutate_json_by_path<F, R>(json: &mut Val, path: JsonPath, closure: F) -> Option<R>
+where
+    F: FnOnce(Target) -> R,
+{
+    find_or_create_target(json, &path).map(closure)
+}
+
+fn find_or_create_target<'a>(json: &'a mut Val, path: &JsonPath) -> Option<Target<'a>> {
+    let mut current = json;
+    for (i, key) in path.elements.iter().enumerate() {
+        let is_last = i == path.elements.len() - 1;
+        match key {
+            PathElement::Root() => continue,
+            PathElement::ArrayLocator(index) => match current {
+                Val::Array(arr) => {
+                    if let Some(index) = match index {
+                        i if *i < 0 => arr.len().checked_sub(i.unsigned_abs() as usize),
+                        i => Some(*i as usize),
+                    } {
+                        if is_last {
+                            if index == arr.len() {
+                                arr.push(Val::Null);
+                            }
+
+                            if index >= arr.len() {
+                                return None;
+                            }
+
+                            return Some(Target::Array(arr, index));
+                        } else {
+                            if index == arr.len() {
+                                arr.push(
+                                    if matches!(path.elements[i + 1], PathElement::ArrayLocator(_))
+                                    {
+                                        Val::Array(vec![])
+                                    } else {
+                                        Val::Object(vec![])
+                                    },
+                                );
+                            }
+
+                            if index >= arr.len() {
+                                return None;
+                            }
+
+                            current = &mut arr[index];
+                        }
+                    } else {
+                        return None;
+                    }
+                }
+                _ => {
+                    *current = Val::Array(vec![]);
+                }
+            },
+            PathElement::Key(key) => match current {
+                Val::Object(obj) => {
+                    if let Some(pos) = &obj
+                        .iter()
+                        .position(|(k, v)| k == key && !matches!(v, Val::Removed))
+                    {
+                        let val = &mut obj[*pos].1;
+                        current = val;
+                    } else {
+                        let element = if !is_last
+                            && matches!(path.elements[i + 1], PathElement::ArrayLocator(_))
+                        {
+                            Val::Array(vec![])
+                        } else {
+                            Val::Object(vec![])
+                        };
+
+                        obj.push((key.clone(), element));
+                        let index = obj.len() - 1;
+                        current = &mut obj[index].1;
+                    }
+                }
+                _ => {
+                    return None;
+                }
+            },
+        }
+    }
+    Some(Target::Value(current))
+}
+
 pub fn json_error_position(json: &OwnedValue) -> crate::Result<OwnedValue> {
     match json {
         OwnedValue::Text(t) => match from_str::<Val>(&t.value) {
@@ -1239,7 +1364,7 @@ mod tests {
 
         let result = result.unwrap();
         match &result.elements[..] {
-            [PathElement::Root(), PathElement::Key(field)] if *field == "field".to_string() => {}
+            [PathElement::Root(), PathElement::Key(field)] if *field == "field" => {}
             _ => panic!("Expected root and field"),
         }
     }
@@ -1291,14 +1416,14 @@ mod tests {
 
     #[test]
     fn test_json_path_from_owned_value_float_strict() {
-        let path = OwnedValue::Float(3.14);
+        let path = OwnedValue::Float(1.23);
 
         assert!(json_path_from_owned_value(&path, true).is_err());
     }
 
     #[test]
     fn test_json_path_from_owned_value_float_non_strict() {
-        let path = OwnedValue::Float(3.14);
+        let path = OwnedValue::Float(1.23);
 
         let result = json_path_from_owned_value(&path, false);
         assert!(result.is_ok());
@@ -1308,8 +1433,255 @@ mod tests {
 
         let result = result.unwrap();
         match &result.elements[..] {
-            [PathElement::Root(), PathElement::Key(field)] if *field == "3.14".to_string() => {}
+            [PathElement::Root(), PathElement::Key(field)] if *field == "1.23" => {}
             _ => panic!("Expected root and field"),
         }
     }
+
+    #[test]
+    fn test_json_set_field_empty_object() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.field".to_string())),
+                OwnedValue::build_text(Rc::new("value".to_string())),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"field":"value"}"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_replace_field() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new(r#"{"field":"old_value"}"#.to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.field".to_string())),
+                OwnedValue::build_text(Rc::new("new_value".to_string())),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"field":"new_value"}"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_set_deeply_nested_key() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.object.doesnt.exist".to_string())),
+                OwnedValue::build_text(Rc::new("value".to_string())),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(
+                r#"{"object":{"doesnt":{"exist":"value"}}}"#.to_string()
+            ))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_value_to_empty_array() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[]".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$[0]".to_string())),
+                OwnedValue::build_text(Rc::new("value".to_string())),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"["value"]"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_value_to_nonexistent_array() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.some_array[0]".to_string())),
+                OwnedValue::Integer(123),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"some_array":[123]}"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_value_to_array() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[123]".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$[1]".to_string())),
+                OwnedValue::Integer(456),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new("[123,456]".to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_value_to_array_out_of_bounds() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[123]".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$[200]".to_string())),
+                OwnedValue::Integer(456),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new("[123]".to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_replace_value_in_array() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[123]".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$[0]".to_string())),
+                OwnedValue::Integer(456),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new("[456]".to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_null_path() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[OwnedValue::Null, OwnedValue::Integer(456)],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new("{}".to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_multiple_keys() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[123]".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$[0]".to_string())),
+                OwnedValue::Integer(456),
+                OwnedValue::build_text(Rc::new("$[1]".to_string())),
+                OwnedValue::Integer(789),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new("[456,789]".to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_missing_value() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("[123]".to_string())),
+            &[OwnedValue::build_text(Rc::new("$[0]".to_string()))],
+        );
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_json_set_add_array_in_nested_object() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.object[0].field".to_string())),
+                OwnedValue::Integer(123),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"object":[{"field":123}]}"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_array_in_array_in_nested_object() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.object[0][0]".to_string())),
+                OwnedValue::Integer(123),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"object":[[123]]}"#.to_string()))
+        );
+    }
+
+    #[test]
+    fn test_json_set_add_array_in_array_in_nested_object_out_of_bounds() {
+        let result = json_set(
+            &OwnedValue::build_text(Rc::new("{}".to_string())),
+            &[
+                OwnedValue::build_text(Rc::new("$.object[123].another".to_string())),
+                OwnedValue::build_text(Rc::new("value".to_string())),
+                OwnedValue::build_text(Rc::new("$.field".to_string())),
+                OwnedValue::build_text(Rc::new("value".to_string())),
+            ],
+        );
+
+        assert!(result.is_ok());
+
+        assert_eq!(
+            result.unwrap(),
+            OwnedValue::build_text(Rc::new(r#"{"field":"value"}"#.to_string()))
+        );
+    }
 }

core/translate/expr.rs

@@ -918,14 +918,16 @@ pub fn translate_expr(
                             func_ctx,
                         )
                     }
-                    JsonFunc::JsonArray | JsonFunc::JsonExtract => translate_function(
-                        program,
-                        args.as_deref().unwrap_or_default(),
-                        referenced_tables,
-                        resolver,
-                        target_register,
-                        func_ctx,
-                    ),
+                    JsonFunc::JsonArray | JsonFunc::JsonExtract | JsonFunc::JsonSet => {
+                        translate_function(
+                            program,
+                            args.as_deref().unwrap_or_default(),
+                            referenced_tables,
+                            resolver,
+                            target_register,
+                            func_ctx,
+                        )
+                    }
                     JsonFunc::JsonArrowExtract | JsonFunc::JsonArrowShiftExtract => {
                         unreachable!(
                             "These two functions are only reachable via the -> and ->> operators"

core/vdbe/mod.rs

@@ -47,7 +47,7 @@ use crate::{
     function::JsonFunc, json::get_json, json::is_json_valid, json::json_array,
     json::json_array_length, json::json_arrow_extract, json::json_arrow_shift_extract,
     json::json_error_position, json::json_extract, json::json_object, json::json_patch,
-    json::json_remove, json::json_type,
+    json::json_remove, json::json_set, json::json_type,
 };
 use crate::{resolve_ext_path, Connection, Result, TransactionState, DATABASE_VERSION};
 use datetime::{
@@ -1837,6 +1837,18 @@ impl Program {
                                 let json_str = get_json(json_value, Some(indent))?;
                                 state.registers[*dest] = json_str;
                             }
+                            JsonFunc::JsonSet => {
+                                let reg_values =
+                                    &state.registers[*start_reg + 1..*start_reg + arg_count];
+
+                                let json_result =
+                                    json_set(&state.registers[*start_reg], reg_values);
+
+                                match json_result {
+                                    Ok(json) => state.registers[*dest] = json,
+                                    Err(e) => return Err(e),
+                                }
+                            }
                         },
                         crate::function::Func::Scalar(scalar_func) => match scalar_func {
                             ScalarFunc::Cast => {

testing/json.test

@@ -826,3 +826,55 @@ do_execsql_test json-remove-6 {
 do_execsql_test json-remove-7 {
     SELECT json_remove('{"a": 1, "b": [1,2], "c": {"d": 3}}', '$.a', '$.b[0]', '$.c.d');
 } {{{"b":[2],"c":{}}}}
+
+do_execsql_test json_set_field_empty_object {
+   SELECT json_set('{}', '$.field', 'value');
+} {{{"field":"value"}}}
+
+do_execsql_test json_set_replace_field {
+   SELECT json_set('{"field":"old_value"}', '$.field', 'new_value');
+} {{{"field":"new_value"}}}
+
+do_execsql_test json_set_set_deeply_nested_key {
+   SELECT json_set('{}', '$.object.doesnt.exist', 'value');
+} {{{"object":{"doesnt":{"exist":"value"}}}}}
+
+do_execsql_test json_set_add_value_to_empty_array {
+   SELECT json_set('[]', '$[0]', 'value');
+} {{["value"]}}
+
+do_execsql_test json_set_add_value_to_nonexistent_array {
+   SELECT json_set('{}', '$.some_array[0]', 123);
+} {{{"some_array":[123]}}}
+
+do_execsql_test json_set_add_value_to_array {
+   SELECT json_set('[123]', '$[1]', 456);
+} {{[123,456]}}
+
+do_execsql_test json_set_add_value_to_array_out_of_bounds {
+   SELECT json_set('[123]', '$[200]', 456);
+} {{[123]}}
+
+do_execsql_test json_set_replace_value_in_array {
+   SELECT json_set('[123]', '$[0]', 456);
+} {{[456]}}
+
+do_execsql_test json_set_null_path {
+   SELECT json_set('{}', NULL, 456);
+} {{{}}}
+
+do_execsql_test json_set_multiple_keys {
+   SELECT json_set('[123]', '$[0]', 456, '$[1]', 789);
+} {{[456,789]}}
+
+do_execsql_test json_set_add_array_in_nested_object {
+   SELECT json_set('{}', '$.object[0].field', 123);
+} {{{"object":[{"field":123}]}}}
+
+do_execsql_test json_set_add_array_in_array_in_nested_object {
+   SELECT json_set('{}', '$.object[0][0]', 123);
+} {{{"object":[[123]]}}}
+
+do_execsql_test json_set_add_array_in_array_in_nested_object_out_of_bounds {
+   SELECT json_set('{}', '$.object[123].another', 'value', '$.field', 'value');
+} {{{"field":"value"}}}