aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-02-12 03:34:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

Implement -> and ->> operators for json

Browse Source
This commit is contained in:
Kacper Madej
2025-01-09 11:23:48 +07:00
parent 0ef4def900
commit dd533414ef
5 changed files with 401 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
core/function.rs
View File

@@ -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(),
}
)
}

91
core/json/mod.rs
View File

@@ -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<OwnedValue> {
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<OwnedValue> {
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<OwnedValue> {
if let OwnedValue::Null = value {
return Ok(OwnedValue::Null);
@@ -153,12 +200,22 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<O
}
let json = get_json_value(value)?;
let mut result = "".to_string();
if paths.len() > 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<O
Ok(OwnedValue::Text(LimboText::json(Rc::new(result))))
}
/// Returns a value with type defined by SQLite documentation:
/// > 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<OwnedValue> {
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<Val> {
let json_path = json_path(path)?;

225
core/translate/expr.rs
View File

@@ -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<Vec<ast::Expr>>,
args: &[ast::Expr],
referenced_tables: Option<&[TableReference]>,
resolver: &Resolver,
target_register: usize,
func_ctx: FuncCtx,
) -> Result<usize> {
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(

19
core/vdbe/mod.rs
View File

@@ -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 {

176
testing/json.test
View File

@@ -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 }', '$.\"')