diff --git a/core/function.rs b/core/function.rs index 7385b237f..5357b68fb 100644 --- a/core/function.rs +++ b/core/function.rs @@ -26,6 +26,8 @@ pub enum JsonFunc { Json, JsonArray, JsonExtract, + JsonArrowExtract, + JsonArrowShiftExtract, JsonArrayLength, } @@ -40,6 +42,8 @@ impl Display for JsonFunc { Self::JsonArray => "json_array".to_string(), Self::JsonExtract => "json_extract".to_string(), Self::JsonArrayLength => "json_array_length".to_string(), + Self::JsonArrowExtract => "->".to_string(), + Self::JsonArrowShiftExtract => "->>".to_string(), } ) } diff --git a/core/json/mod.rs b/core/json/mod.rs index eda97bf2b..d359ccb60 100644 --- a/core/json/mod.rs +++ b/core/json/mod.rs @@ -143,6 +143,53 @@ pub fn json_array_length( } } +/// 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 { + if let OwnedValue::Null = value { + return Ok(OwnedValue::Null); + } + + let json = get_json_value(value)?; + + match path { + OwnedValue::Null => Ok(OwnedValue::Null), + OwnedValue::Text(p) => { + let extracted = json_extract_single(&json, p.value.as_str())?; + + let json = crate::json::to_string(&extracted).unwrap(); + Ok(OwnedValue::Text(LimboText::json(Rc::new(json)))) + } + _ => crate::bail_constraint_error!("JSON path error near: {:?}", path.to_string()), + } +} + +/// Implements the ->> operator. Always returns a SQL representation of the JSON subcomponent. +/// https://sqlite.org/json1.html#the_and_operators +pub fn json_arrow_shift_extract( + value: &OwnedValue, + path: &OwnedValue, +) -> crate::Result { + if let OwnedValue::Null = value { + return Ok(OwnedValue::Null); + } + + let json = get_json_value(value)?; + + match path { + OwnedValue::Null => Ok(OwnedValue::Null), + OwnedValue::Text(p) => { + let extracted = json_extract_single(&json, p.value.as_str())?; + + convert_json_to_db_type(&extracted) + } + _ => crate::bail_constraint_error!("JSON path error near: {:?}", path.to_string()), + } +} + +/// Extracts a JSON value from a JSON object or array. +/// If there's only a single path, the return value might be either a TEXT or a database type. +/// https://sqlite.org/json1.html#the_json_extract_function pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result { if let OwnedValue::Null = value { return Ok(OwnedValue::Null); @@ -153,12 +200,22 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result 1 { - result.push('['); + if paths.len() == 1 { + match &paths[0] { + OwnedValue::Null => return Ok(OwnedValue::Null), + OwnedValue::Text(p) => { + let extracted = json_extract_single(&json, p.value.as_str())?; + + return convert_json_to_db_type(&extracted); + } + _ => crate::bail_constraint_error!("JSON path error near: {:?}", paths[0].to_string()), + } } + // multiple paths - we should return an array + let mut result = "[".to_string(); + for path in paths { match path { OwnedValue::Text(p) => { @@ -186,6 +243,34 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result the SQL datatype of the result is NULL for a JSON null, +/// > INTEGER or REAL for a JSON numeric value, +/// > an INTEGER zero for a JSON false value, +/// > an INTEGER one for a JSON true value, +/// > the dequoted text for a JSON string value, +/// > and a text representation for JSON object and array values. +/// https://sqlite.org/json1.html#the_json_extract_function +fn convert_json_to_db_type(extracted: &Val) -> crate::Result { + match extracted { + Val::Null => Ok(OwnedValue::Null), + Val::Float(f) => Ok(OwnedValue::Float(*f)), + Val::Integer(i) => Ok(OwnedValue::Integer(*i)), + Val::Bool(b) => { + if *b { + Ok(OwnedValue::Integer(1)) + } else { + Ok(OwnedValue::Integer(0)) + } + } + Val::String(s) => Ok(OwnedValue::Text(LimboText::json(Rc::new(s.clone())))), + _ => { + let json = crate::json::to_string(&extracted).unwrap(); + Ok(OwnedValue::Text(LimboText::json(Rc::new(json)))) + } + } +} + fn json_extract_single(json: &Val, path: &str) -> crate::Result { let json_path = json_path(path)?; diff --git a/core/translate/expr.rs b/core/translate/expr.rs index 4c4ee72b2..0cbc7f73b 100644 --- a/core/translate/expr.rs +++ b/core/translate/expr.rs @@ -20,6 +20,50 @@ pub struct ConditionMetadata { pub jump_target_when_false: BranchOffset, } +macro_rules! expect_arguments_exact { + ( + $args:expr, + $expected_arguments:expr, + $func:ident + ) => {{ + let args = if let Some(args) = $args { + if args.len() != $expected_arguments { + crate::bail_parse_error!( + "{} function with not exactly 2 arguments", + $func.to_string() + ); + } + args + } else { + crate::bail_parse_error!("{} function with no arguments", $func.to_string()); + }; + + args + }}; +} + +macro_rules! expect_arguments_max { + ( + $args:expr, + $expected_arguments:expr, + $func:ident + ) => {{ + let args = if let Some(args) = $args { + if args.len() > $expected_arguments { + crate::bail_parse_error!( + "{} function with not exactly 2 arguments", + $func.to_string() + ); + } + args + } else { + crate::bail_parse_error!("{} function with no arguments", $func.to_string()); + }; + + args + }}; +} + pub fn translate_condition_expr( program: &mut ProgramBuilder, referenced_tables: &[TableReference], @@ -590,6 +634,24 @@ pub fn translate_expr( dest: target_register, }); } + #[cfg(feature = "json")] + op @ (ast::Operator::ArrowRight | ast::Operator::ArrowRightShift) => { + let json_func = match op { + ast::Operator::ArrowRight => JsonFunc::JsonArrowExtract, + ast::Operator::ArrowRightShift => JsonFunc::JsonArrowShiftExtract, + _ => unreachable!(), + }; + + program.emit_insn(Insn::Function { + constant_mask: 0, + start_reg: e1_reg, + dest: target_register, + func: FuncCtx { + func: Func::Json(json_func), + arg_count: 2, + }, + }) + } other_unimplemented => todo!("{:?}", other_unimplemented), } Ok(target_register) @@ -733,100 +795,41 @@ pub fn translate_expr( #[cfg(feature = "json")] Func::Json(j) => match j { JsonFunc::Json => { - let args = if let Some(args) = args { - if args.len() != 1 { - crate::bail_parse_error!( - "{} function with not exactly 1 argument", - j.to_string() - ); - } - args - } else { - crate::bail_parse_error!( - "{} function with no arguments", - j.to_string() - ); - }; - let regs = program.alloc_register(); - translate_expr(program, referenced_tables, &args[0], regs, resolver)?; - program.emit_insn(Insn::Function { - constant_mask: 0, - start_reg: regs, - dest: target_register, - func: func_ctx, - }); - Ok(target_register) - } - JsonFunc::JsonArray => { - let start_reg = translate_variable_sized_function_parameter_list( + let args = expect_arguments_exact!(args, 1, j); + + translate_function( program, args, referenced_tables, resolver, - )?; - - program.emit_insn(Insn::Function { - constant_mask: 0, - start_reg, - dest: target_register, - func: func_ctx, - }); - Ok(target_register) + target_register, + func_ctx, + ) } - JsonFunc::JsonExtract => { - let start_reg = translate_variable_sized_function_parameter_list( - program, - args, - referenced_tables, - resolver, - )?; - - program.emit_insn(Insn::Function { - constant_mask: 0, - start_reg, - dest: target_register, - func: func_ctx, - }); - Ok(target_register) + JsonFunc::JsonArray | JsonFunc::JsonExtract => 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" + ) } JsonFunc::JsonArrayLength => { - let args = if let Some(args) = args { - if args.len() > 2 { - crate::bail_parse_error!( - "{} function with wrong number of arguments", - j.to_string() - ) - } - args - } else { - crate::bail_parse_error!( - "{} function with no arguments", - j.to_string() - ); - }; + let args = expect_arguments_max!(args, 2, j); - let json_reg = program.alloc_register(); - let path_reg = program.alloc_register(); - - translate_expr(program, referenced_tables, &args[0], json_reg, resolver)?; - - if args.len() == 2 { - translate_expr( - program, - referenced_tables, - &args[1], - path_reg, - resolver, - )?; - } - - program.emit_insn(Insn::Function { - constant_mask: 0, - start_reg: json_reg, - dest: target_register, - func: func_ctx, - }); - Ok(target_register) + translate_function( + program, + args, + referenced_tables, + resolver, + target_register, + func_ctx, + ) } }, Func::Scalar(srf) => { @@ -850,22 +853,14 @@ pub fn translate_expr( }); Ok(target_register) } - ScalarFunc::Char => { - let start_reg = translate_variable_sized_function_parameter_list( - program, - args, - referenced_tables, - resolver, - )?; - - program.emit_insn(Insn::Function { - constant_mask: 0, - start_reg, - dest: target_register, - func: func_ctx, - }); - Ok(target_register) - } + ScalarFunc::Char => translate_function( + program, + args.as_deref().unwrap_or_default(), + referenced_tables, + resolver, + target_register, + func_ctx, + ), ScalarFunc::Coalesce => { let args = if let Some(args) = args { if args.len() < 2 { @@ -1902,18 +1897,19 @@ pub fn translate_expr( } } -// Returns the starting register for the function. -// TODO: Use this function for all functions with variable number of parameters in `translate_expr` -fn translate_variable_sized_function_parameter_list( +/// Emits a whole insn for a function call. +/// Assumes the number of parameters is valid for the given function. +/// Returns the target register for the function. +fn translate_function( program: &mut ProgramBuilder, - args: &Option>, + args: &[ast::Expr], referenced_tables: Option<&[TableReference]>, resolver: &Resolver, + target_register: usize, + func_ctx: FuncCtx, ) -> Result { - let args = args.as_deref().unwrap_or_default(); - - let reg = program.alloc_registers(args.len()); - let mut current_reg = reg; + let start_reg = program.alloc_registers(args.len()); + let mut current_reg = start_reg; for arg in args.iter() { translate_expr(program, referenced_tables, arg, current_reg, resolver)?; @@ -1921,7 +1917,14 @@ fn translate_variable_sized_function_parameter_list( current_reg += 1; } - Ok(reg) + program.emit_insn(Insn::Function { + constant_mask: 0, + start_reg, + dest: target_register, + func: func_ctx, + }); + + Ok(target_register) } fn wrap_eval_jump_expr( diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index c9151e934..d05318dec 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -28,6 +28,7 @@ use crate::error::{LimboError, SQLITE_CONSTRAINT_PRIMARYKEY}; #[cfg(feature = "uuid")] use crate::ext::{exec_ts_from_uuid7, exec_uuid, exec_uuidblob, exec_uuidstr, ExtFunc, UuidFunc}; use crate::function::{AggFunc, FuncCtx, MathFunc, MathFuncArity, ScalarFunc}; +use crate::json::{json_arrow_extract, json_arrow_shift_extract}; use crate::pseudo::PseudoCursor; use crate::result::LimboResult; use crate::schema::Table; @@ -1381,6 +1382,24 @@ impl Program { } } #[cfg(feature = "json")] + crate::function::Func::Json( + func @ (JsonFunc::JsonArrowExtract | JsonFunc::JsonArrowShiftExtract), + ) => { + assert_eq!(arg_count, 2); + let json = &state.registers[*start_reg]; + let path = &state.registers[*start_reg + 1]; + let func = match func { + JsonFunc::JsonArrowExtract => json_arrow_extract, + JsonFunc::JsonArrowShiftExtract => json_arrow_shift_extract, + _ => unreachable!(), + }; + let json_str = func(json, path); + match json_str { + Ok(json) => state.registers[*dest] = json, + Err(e) => return Err(e), + } + } + #[cfg(feature = "json")] crate::function::Func::Json(JsonFunc::JsonArrayLength) => { let json_value = &state.registers[*start_reg]; let path_value = if arg_count > 1 { diff --git a/testing/json.test b/testing/json.test index 5340f7049..28756de03 100755 --- a/testing/json.test +++ b/testing/json.test @@ -89,6 +89,18 @@ do_execsql_test json_extract_null { SELECT json_extract(null, '$') } {{}} +do_execsql_test json_extract_json_null_type { + SELECT typeof(json_extract('null', '$')) +} {{null}} + +do_execsql_test json_arrow_json_null_type { + SELECT typeof('null' -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_json_null_type { + SELECT typeof('null' ->> '$') +} {{null}} + do_execsql_test json_extract_empty { SELECT json_extract() } {{}} @@ -113,10 +125,38 @@ do_execsql_test json_extract_number { SELECT json_extract(1, '$') } {{1}} +do_execsql_test json_extract_number_type { + SELECT typeof(json_extract(1, '$')) +} {{integer}} + +do_execsql_test json_arrow_number { + SELECT 1 -> '$' +} {{1}} + +do_execsql_test json_arrow_number_type { + SELECT typeof(1 -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_number { + SELECT 1 -> '$' +} {{1}} + +do_execsql_test json_arrow_shift_number_type { + SELECT typeof(1 ->> '$') +} {{integer}} + do_execsql_test json_extract_object_1 { SELECT json_extract('{"a": [1,2,3]}', '$.a') } {{[1,2,3]}} +do_execsql_test json_arrow_object { + SELECT '{"a": [1,2,3]}' -> '$.a' +} {{[1,2,3]}} + +do_execsql_test json_arrow_shift_object { + SELECT '{"a": [1,2,3]}' ->> '$.a' +} {{[1,2,3]}} + do_execsql_test json_extract_object_2 { SELECT json_extract('{"a": [1,2,3]}', '$.a', '$.a[0]', '$.a[1]', '$.a[3]') } {{[[1,2,3],1,2,null]}} @@ -140,11 +180,147 @@ do_execsql_test json_extract_null_path { SELECT json_extract(1, null) } {{}} +do_execsql_test json_arrow_null_path { + SELECT 1 -> null +} {{}} + +do_execsql_test json_arrow_shift_null_path { + SELECT 1 ->> null +} {{}} + +do_execsql_test json_extract_float { + SELECT typeof(json_extract(1.0, '$')) +} {{real}} + +do_execsql_test json_arrow_float { + SELECT typeof(1.0 -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_float { + SELECT typeof(1.0 ->> '$') +} {{real}} + +do_execsql_test json_extract_true { + SELECT json_extract('true', '$') +} {{1}} + +do_execsql_test json_extract_true_type { + SELECT typeof(json_extract('true', '$')) +} {{integer}} + +do_execsql_test json_arrow_true { + SELECT 'true' -> '$' +} {{true}} + +do_execsql_test json_arrow_true_type { + SELECT typeof('true' -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_true { + SELECT 'true' ->> '$' +} {{1}} + +do_execsql_test json_arrow_shift_true_type { + SELECT typeof('true' ->> '$') +} {{integer}} + +do_execsql_test json_extract_false { + SELECT json_extract('false', '$') +} {{0}} + +do_execsql_test json_extract_false_type { + SELECT typeof(json_extract('false', '$')) +} {{integer}} + +do_execsql_test json_arrow_false { + SELECT 'false' -> '$' +} {{false}} + +do_execsql_test json_arrow_false_type { + SELECT typeof('false' -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_false { + SELECT 'false' ->> '$' +} {{0}} + +do_execsql_test json_arrow_shift_false_type { + SELECT typeof('false' ->> '$') +} {{integer}} + +do_execsql_test json_extract_string { + SELECT json_extract('"string"', '$') +} {{string}} + +do_execsql_test json_extract_string_type { + SELECT typeof(json_extract('"string"', '$')) +} {{text}} + +do_execsql_test json_arrow_string { + SELECT '"string"' -> '$' +} {{"string"}} + +do_execsql_test json_arrow_string_type { + SELECT typeof('"string"' -> '$') +} {{text}} + +do_execsql_test json_arrow_shift_string { + SELECT '"string"' ->> '$' +} {{string}} + +do_execsql_test json_arrow_shift_string_type { + SELECT typeof('"string"' ->> '$') +} {{text}} + +do_execsql_test json_arrow_implicit_root_path { + SELECT '{"a":1}' -> 'a'; +} {{1}} + +do_execsql_test json_arrow_shift_implicit_root_path { + SELECT '{"a":1}' ->> 'a'; +} {{1}} + +do_execsql_test json_arrow_implicit_root_path_undefined_key { + SELECT '{"a":1}' -> 'x'; +} {{}} + +do_execsql_test json_arrow_shift_implicit_root_path_undefined_key { + SELECT '{"a":1}' ->> 'x'; +} {{}} + +do_execsql_test json_arrow_implicit_root_path_array { + SELECT '[1,2,3]' -> 1; +} {{2}} + +do_execsql_test json_arrow_shift_implicit_root_path_array { + SELECT '[1,2,3]' ->> 1; +} {{2}} + +do_execsql_test json_arrow_implicit_root_path_array_negative_idx { + SELECT '[1,2,3]' -> -1; +} {{3}} + +do_execsql_test json_arrow_shift_implicit_root_path_array_negative_idx { + SELECT '[1,2,3]' ->> -1; +} {{3}} + # TODO: fix me - this passes on SQLite and needs to be fixed in Limbo. do_execsql_test json_extract_multiple_null_paths { SELECT json_extract(1, null, null, null) } {{}} +do_execsql_test json_extract_array { + SELECT json_extract('[1,2,3]', '$') +} {{[1,2,3]}} + +do_execsql_test json_arrow_array { + SELECT '[1,2,3]' -> '$' +} {{[1,2,3]}} + +do_execsql_test json_arrow_shift_array { + SELECT '[1,2,3]' ->> '$' +} {{[1,2,3]}} + # TODO: fix me - this passes on SQLite and needs to be fixed in Limbo. #do_execsql_test json_extract_quote { # SELECT json_extract('{"\"":1 }', '$.\"')