diff --git a/core/function.rs b/core/function.rs index 333266eea..6a309ffc5 100644 --- a/core/function.rs +++ b/core/function.rs @@ -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)), diff --git a/core/json/jsonb.rs b/core/json/jsonb.rs index a497d691e..fa8218c4b 100644 --- a/core/json/jsonb.rs +++ b/core/json/jsonb.rs @@ -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 { diff --git a/core/json/mod.rs b/core/json/mod.rs index 3690b2a5e..692dddb7a 100644 --- a/core/json/mod.rs +++ b/core/json/mod.rs @@ -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 crate::Result { + 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 OwnedValue { +fn json_string_to_db_type( + json: Jsonb, + element_type: ElementType, + raw_flag: bool, +) -> crate::Result { + 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!(), } } diff --git a/core/translate/expr.rs b/core/translate/expr.rs index 24e7418b3..ab9496c0e 100644 --- a/core/translate/expr.rs +++ b/core/translate/expr.rs @@ -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" diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 72eb20645..3f5729d0e 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -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];