add jsonb_set

core/function.rs

@@ -93,6 +93,7 @@ pub enum JsonFunc {
     JsonbInsert,
     JsonPretty,
     JsonSet,
+    JsonbSet,
     JsonQuote,
 }
 
@@ -126,6 +127,7 @@ impl Display for JsonFunc {
                 Self::JsonbInsert => "jsonb_insert".to_string(),
                 Self::JsonPretty => "json_pretty".to_string(),
                 Self::JsonSet => "json_set".to_string(),
+                Self::JsonbSet => "jsonb_set".to_string(),
                 Self::JsonQuote => "json_quote".to_string(),
             }
         )
@@ -599,6 +601,7 @@ impl Func {
             "jsonb_replace" => Ok(Self::Json(JsonFunc::JsonReplace)),
             "json_pretty" => Ok(Self::Json(JsonFunc::JsonPretty)),
             "json_set" => Ok(Self::Json(JsonFunc::JsonSet)),
+            "jsonb_set" => Ok(Self::Json(JsonFunc::JsonbSet)),
             "json_quote" => Ok(Self::Json(JsonFunc::JsonQuote)),
             "unixepoch" => Ok(Self::Scalar(ScalarFunc::UnixEpoch)),
             "julianday" => Ok(Self::Scalar(ScalarFunc::JulianDay)),

core/json/mod.rs

@@ -227,6 +227,30 @@ pub fn json_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<
     json_string_to_db_type(json, el_type, OutputVariant::String)
 }
 
+pub fn jsonb_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<OwnedValue> {
+    if args.is_empty() {
+        return Ok(OwnedValue::Null);
+    }
+
+    let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
+    let mut json = json_cache.get_or_insert_with(&args[0].get_owned_value(), make_jsonb_fn)?;
+    let other = args[1..].chunks_exact(2);
+
+    for chunk in other {
+        let path = json_path_from_owned_value(&chunk[0].get_owned_value(), true)?;
+
+        let value = convert_dbtype_to_jsonb(&chunk[1].get_owned_value(), Conv::NotStrict)?;
+        let mut op = SetOperation::new(value);
+        if let Some(path) = path {
+            let _ = json.operate_on_path(&path, &mut op);
+        }
+    }
+
+    let el_type = json.is_valid()?;
+
+    json_string_to_db_type(json, el_type, OutputVariant::Binary)
+}
+
 /// Implements the -> operator. Always returns a proper JSON value.
 /// https://sqlite.org/json1.html#the_and_operators
 pub fn json_arrow_extract(

core/translate/expr.rs

@@ -897,6 +897,7 @@ pub fn translate_expr(
                     | JsonFunc::JsonbArray
                     | JsonFunc::JsonExtract
                     | JsonFunc::JsonSet
+                    | JsonFunc::JsonbSet
                     | JsonFunc::JsonbExtract
                     | JsonFunc::JsonReplace
                     | JsonFunc::JsonbReplace

core/vdbe/mod.rs

@@ -59,7 +59,8 @@ use crate::{
     json::json_from_raw_bytes_agg, json::json_insert, json::json_object, json::json_patch,
     json::json_quote, json::json_remove, json::json_replace, json::json_set, json::json_type,
     json::jsonb, json::jsonb_array, json::jsonb_extract, json::jsonb_insert, json::jsonb_object,
-    json::jsonb_patch, json::jsonb_remove, json::jsonb_replace, json::JsonCacheCell,
+    json::jsonb_patch, json::jsonb_remove, json::jsonb_replace, json::jsonb_set,
+    json::JsonCacheCell,
 };
 use crate::{
     resolve_ext_path, Connection, MvCursor, MvStore, Result, TransactionState, DATABASE_VERSION,
@@ -2658,6 +2659,22 @@ impl Program {
                                     Err(e) => return Err(e),
                                 }
                             }
+                            JsonFunc::JsonbSet => {
+                                if arg_count % 2 == 0 {
+                                    bail_constraint_error!(
+                                        "json_set() needs an odd number of arguments"
+                                    )
+                                }
+                                let reg_values =
+                                    &state.registers[*start_reg..*start_reg + arg_count];
+
+                                let json_result = jsonb_set(reg_values, &state.json_cache);
+
+                                match json_result {
+                                    Ok(json) => state.registers[*dest] = Register::OwnedValue(json),
+                                    Err(e) => return Err(e),
+                                }
+                            }
                             JsonFunc::JsonQuote => {
                                 let json_value = &state.registers[*start_reg];