From a8e3ae7aa3074decfd1484eb15f0ba4b62198172 Mon Sep 17 00:00:00 2001 From: Marcus Nilsson Date: Tue, 4 Feb 2025 17:51:51 +0100 Subject: [PATCH] add support for json_set Test cases are included. Related to #127 --- COMPAT.md | 2 +- Cargo.lock | 4 +- core/function.rs | 4 + core/json/mod.rs | 372 ++++++++++++++++++++++++++++++++++++++++- core/translate/expr.rs | 18 +- core/vdbe/mod.rs | 14 +- testing/json.test | 48 ++++++ 7 files changed, 446 insertions(+), 16 deletions(-) diff --git a/COMPAT.md b/COMPAT.md index 7224075d..1c2b9b22 100644 --- a/COMPAT.md +++ b/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 | | diff --git a/Cargo.lock b/Cargo.lock index 63f0cc3e..655743a8 100644 --- a/Cargo.lock +++ b/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", diff --git a/core/function.rs b/core/function.rs index d9084cab..8fbe9cb7 100644 --- a/core/function.rs +++ b/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)), diff --git a/core/json/mod.rs b/core/json/mod.rs index 5816eb7f..5d08576b 100644 --- a/core/json/mod.rs +++ b/core/json/mod.rs @@ -155,6 +155,39 @@ pub fn json_array_length( } } +pub fn json_set(json: &OwnedValue, values: &[OwnedValue]) -> crate::Result { + 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::>()?; + + 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 { @@ -479,6 +512,92 @@ fn find_target<'a>(json: &'a mut Val, path: &JsonPath) -> Option> { Some(Target::Value(current)) } +fn create_and_mutate_json_by_path(json: &mut Val, path: JsonPath, closure: F) -> Option +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> { + 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 { match json { OwnedValue::Text(t) => match from_str::(&t.value) { @@ -1239,7 +1358,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 +1410,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 +1427,253 @@ 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][0]".to_string())), + OwnedValue::Integer(123), + ], + ); + + assert!(result.is_ok()); + + assert_eq!( + result.unwrap(), + OwnedValue::build_text(Rc::new(r#"{"object":[[123]]}"#.to_string())) + ); + } } diff --git a/core/translate/expr.rs b/core/translate/expr.rs index fcbebed6..c765aa6b 100644 --- a/core/translate/expr.rs +++ b/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" diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 00c0a8d3..f65cdb6b 100644 --- a/core/vdbe/mod.rs +++ b/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 => { diff --git a/testing/json.test b/testing/json.test index 1c9b9d13..c47d0e10 100755 --- a/testing/json.test +++ b/testing/json.test @@ -826,3 +826,51 @@ 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]]}}}