Implement json_set

core/function.rs

@@ -82,6 +82,7 @@ pub enum JsonFunc {
     JsonValid,
     JsonPatch,
     JsonRemove,
+    JsonSet,
 }
 
 #[cfg(feature = "json")]
@@ -103,6 +104,7 @@ impl Display for JsonFunc {
                 Self::JsonValid => "json_valid".to_string(),
                 Self::JsonPatch => "json_patch".to_string(),
                 Self::JsonRemove => "json_remove".to_string(),
+                Self::JsonSet => "json_set".to_string(),
             }
         )
     }
@@ -534,6 +536,8 @@ impl Func {
             "json_patch" => Ok(Self::Json(JsonFunc::JsonPatch)),
             #[cfg(feature = "json")]
             "json_remove" => Ok(Self::Json(JsonFunc::JsonRemove)),
+            #[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

@@ -148,6 +148,39 @@ 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)?,
+                    };
+
+                    create_and_mutate_json_by_path(&mut json_value, path, |val| match val {
+                        Target::Array(arr, index) => arr[index] = new_value.clone(),
+                        Target::Value(val) => *val = new_value.clone(),
+                    });
+                }
+
+                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> {
@@ -405,6 +438,38 @@ fn json_extract_single<'a>(
     Ok(Some(current_element))
 }
 
+fn json_path_from_owned_value(path: &OwnedValue, strict: bool) -> crate::Result<Option<JsonPath>> {
+    let json_path = if strict {
+        match path {
+            OwnedValue::Text(t) => json_path(t.value.as_str())?,
+            OwnedValue::Null => return Ok(None),
+            _ => crate::bail_constraint_error!("JSON path error near: {:?}", path.to_string()),
+        }
+    } else {
+        match path {
+            OwnedValue::Text(t) => {
+                if t.value.starts_with("$") {
+                    json_path(t.value.as_str())?
+                } else {
+                    JsonPath {
+                        elements: vec![PathElement::Root(), PathElement::Key(t.value.to_string())],
+                    }
+                }
+            }
+            OwnedValue::Null => return Ok(None),
+            OwnedValue::Integer(i) => JsonPath {
+                elements: vec![PathElement::Root(), PathElement::ArrayLocator(*i as i32)],
+            },
+            OwnedValue::Float(f) => JsonPath {
+                elements: vec![PathElement::Root(), PathElement::Key(f.to_string())],
+            },
+            _ => crate::bail_constraint_error!("JSON path error near: {:?}", path.to_string()),
+        }
+    };
+
+    Ok(Some(json_path))
+}
+
 enum Target<'a> {
     Array(&'a mut Vec<Val>, usize),
     Value(&'a mut Val),
@@ -463,6 +528,69 @@ 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 {
+                            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 {
+                        obj.push((key.clone(), Val::Object(vec![])));
+                        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) {

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

@@ -29,6 +29,7 @@ use crate::error::{LimboError, SQLITE_CONSTRAINT_PRIMARYKEY};
 use crate::ext::ExtValue;
 use crate::function::{AggFunc, ExtFunc, FuncCtx, MathFunc, MathFuncArity, ScalarFunc};
 use crate::info;
+use crate::json::json_set;
 use crate::pseudo::PseudoCursor;
 use crate::result::LimboResult;
 use crate::storage::sqlite3_ondisk::DatabaseHeader;
@@ -1787,6 +1788,18 @@ impl Program {
                                     &state.registers[*start_reg..*start_reg + arg_count],
                                 )?;
                             }
+                            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 => {