aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-02-19 23:15:28 +01:00
Code Issues Packages Projects Releases Wiki Activity

add jsonb_extract function

Browse Source
This commit is contained in:
Ihor Andrianov
2025-03-16 15:14:29 +02:00
parent a3a9376347
commit 23d7d82b6c
5 changed files with 96 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
core/function.rs
View File

@@ -77,6 +77,7 @@ pub enum JsonFunc {
JsonArrowExtract,
JsonArrowShiftExtract,
JsonExtract,
JsonbExtract,
JsonObject,
JsonType,
JsonErrorPosition,
@@ -99,6 +100,7 @@ impl Display for JsonFunc {
Self::Jsonb => "jsonb".to_string(),
Self::JsonArray => "json_array".to_string(),
Self::JsonExtract => "json_extract".to_string(),
Self::JsonbExtract => "jsonb_extract".to_string(),
Self::JsonArrayLength => "json_array_length".to_string(),
Self::JsonArrowExtract => "->".to_string(),
Self::JsonArrowShiftExtract => "->>".to_string(),
@@ -559,6 +561,8 @@ impl Func {
#[cfg(feature = "json")]
"json_extract" => Ok(Func::Json(JsonFunc::JsonExtract)),
#[cfg(feature = "json")]
"jsonb_extract" => Ok(Func::Json(JsonFunc::JsonbExtract)),
#[cfg(feature = "json")]
"json_object" => Ok(Func::Json(JsonFunc::JsonObject)),
#[cfg(feature = "json")]
"json_type" => Ok(Func::Json(JsonFunc::JsonType)),

2
core/json/jsonb.rs
View File

@@ -1563,7 +1563,7 @@ impl Jsonb {
}
}
bail_parse_error!("Did not find anything")
bail_parse_error!("Not found")
}
fn skip_element(&self, mut pos: usize) -> Result<usize> {

92
core/json/mod.rs
View File

@@ -259,7 +259,7 @@ pub fn json_arrow_shift_extract(
let json = convert_dbtype_to_jsonb(value)?;
let extracted = json.get_by_path(&path);
if let Ok((json, element_type)) = extracted {
Ok(json_string_to_db_type(json.to_string()?, element_type))
Ok(json_string_to_db_type(json, element_type, false)?)
} else {
Ok(OwnedValue::Null)
}
@@ -278,17 +278,44 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<O
if paths.is_empty() {
return Ok(OwnedValue::Null);
} else if paths.len() == 1 {
}
let (json, element_type) = jsonb_extract_internal(value, paths)?;
let result = json_string_to_db_type(json, element_type, false)?;
Ok(result)
}
pub fn jsonb_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<OwnedValue> {
if let OwnedValue::Null = value {
return Ok(OwnedValue::Null);
}
if paths.is_empty() {
return Ok(OwnedValue::Null);
}
let (json, element_type) = jsonb_extract_internal(value, paths)?;
let result = json_string_to_db_type(json, element_type, true)?;
Ok(result)
}
fn jsonb_extract_internal(
value: &OwnedValue,
paths: &[OwnedValue],
) -> crate::Result<(Jsonb, ElementType)> {
let null = Jsonb::from_raw_data(JsonbHeader::make_null().into_bytes().as_bytes());
if paths.len() == 1 {
if let Some(path) = json_path_from_owned_value(&paths[0], true)? {
let json = convert_dbtype_to_jsonb(value)?;
let (expected_value, value_type) = json.get_by_path(&path)?;
return Ok(json_string_to_db_type(
expected_value.to_string()?,
value_type,
));
if let Ok((json, value_type)) = json.get_by_path(&path) {
return Ok((json, value_type));
} else {
return Ok((null, ElementType::NULL));
}
} else {
return Ok(OwnedValue::Null);
return Ok((null, ElementType::NULL));
}
}
@@ -307,36 +334,41 @@ pub fn json_extract(value: &OwnedValue, paths: &[OwnedValue]) -> crate::Result<O
result.append_to_array_unsafe(JsonbHeader::make_null().into_bytes().as_bytes());
}
} else {
return Ok(OwnedValue::Null);
return Ok((null, ElementType::NULL));
}
}
result.finalize_array_unsafe()?;
Ok(json_string_to_db_type(
result.to_string()?,
ElementType::ARRAY,
))
Ok((result, ElementType::ARRAY))
}
fn json_string_to_db_type(mut json: String, element_type: ElementType) -> OwnedValue {
fn json_string_to_db_type(
json: Jsonb,
element_type: ElementType,
raw_flag: bool,
) -> crate::Result<OwnedValue> {
let mut json_string = json.to_string()?;
if raw_flag && matches!(element_type, ElementType::ARRAY | ElementType::OBJECT) {
return Ok(OwnedValue::Blob(Rc::new(json.data())));
}
match element_type {
ElementType::ARRAY | ElementType::OBJECT => OwnedValue::Text(Text::json(json)),
ElementType::ARRAY | ElementType::OBJECT => Ok(OwnedValue::Text(Text::json(json_string))),
ElementType::TEXT | ElementType::TEXT5 | ElementType::TEXTJ | ElementType::TEXTRAW => {
json.remove(json.len() - 1);
json.remove(0);
OwnedValue::Text(Text {
value: Rc::new(json.into_bytes()),
json_string.remove(json_string.len() - 1);
json_string.remove(0);
Ok(OwnedValue::Text(Text {
value: Rc::new(json_string.into_bytes()),
subtype: TextSubtype::Text,
})
}))
}
ElementType::FLOAT5 | ElementType::FLOAT => {
OwnedValue::Float(json.parse().expect("Should be valid f64"))
}
ElementType::INT | ElementType::INT5 => {
OwnedValue::Integer(json.parse().expect("Should be valid i64"))
}
ElementType::TRUE => OwnedValue::Integer(1),
ElementType::FALSE => OwnedValue::Integer(0),
ElementType::NULL => OwnedValue::Null,
ElementType::FLOAT5 | ElementType::FLOAT => Ok(OwnedValue::Float(
json_string.parse().expect("Should be valid f64"),
)),
ElementType::INT | ElementType::INT5 => Ok(OwnedValue::Integer(
json_string.parse().expect("Should be valid i64"),
)),
ElementType::TRUE => Ok(OwnedValue::Integer(1)),
ElementType::FALSE => Ok(OwnedValue::Integer(0)),
ElementType::NULL => Ok(OwnedValue::Null),
_ => unreachable!(),
}
}

21
core/translate/expr.rs
View File

@@ -894,16 +894,17 @@ pub fn translate_expr(
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::JsonArray
| JsonFunc::JsonExtract
| JsonFunc::JsonSet
| JsonFunc::JsonbExtract => 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"

18
core/vdbe/mod.rs
View File

@@ -53,6 +53,7 @@ use crate::{
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_quote, json::json_remove, json::json_set, json::json_type, json::jsonb,
json::jsonb_extract,
};
use crate::{info, CheckpointStatus};
use crate::{
@@ -2172,6 +2173,23 @@ impl Program {
Err(e) => return Err(e),
}
}
JsonFunc::JsonbExtract => {
let result = match arg_count {
0 => jsonb_extract(&OwnedValue::Null, &[]),
_ => {
let val = &state.registers[*start_reg];
let reg_values = &state.registers
[*start_reg + 1..*start_reg + arg_count];
jsonb_extract(val, reg_values)
}
};
match result {
Ok(json) => state.registers[*dest] = json,
Err(e) => return Err(e),
}
}
JsonFunc::JsonArrowExtract | JsonFunc::JsonArrowShiftExtract => {
assert_eq!(arg_count, 2);
let json = &state.registers[*start_reg];