Add SQL array::insert() function (#1504)

This commit is contained in:
Aman Sharma 2022-12-18 02:33:13 +05:30 committed by GitHub
parent 0c2fe82f32
commit e1e617cff7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 76 additions and 0 deletions

View file

@ -142,6 +142,29 @@ impl<A: FromArg, B: FromArg> FromArgs for (A, Option<B>) {
}
}
// Some functions take 2 or 3 arguments, so the third argument is optional.
impl<A: FromArg, B: FromArg, C: FromArg> FromArgs for (A, B, Option<C>) {
fn from_args(name: &str, args: Vec<Value>) -> Result<Self, Error> {
let err = || Error::InvalidArguments {
name: name.to_owned(),
message: String::from("Expected 2 or 3 arguments."),
};
let mut args = args.into_iter();
let a = A::from_arg(args.next().ok_or_else(err)?)?;
let b = B::from_arg(args.next().ok_or_else(err)?)?;
let c = match args.next() {
Some(c) => Some(C::from_arg(c)?),
None => None,
};
if args.next().is_some() {
// Too many.
return Err(err());
}
Ok((a, b, c))
}
}
// Some functions take 0, 1, or 2 arguments, so both arguments are optional.
// It is safe to assume that, if the first argument is None, the second argument will also be None.
impl<A: FromArg, B: FromArg> FromArgs for (Option<A>, Option<B>) {

View file

@ -51,6 +51,30 @@ pub fn flatten((arg,): (Value,)) -> Result<Value, Error> {
})
}
pub fn insert((array, value, index): (Value, Value, Option<Value>)) -> Result<Value, Error> {
match (array, index) {
(Value::Array(mut v), Some(Value::Number(i))) => {
let mut i = i.as_int();
// Negative index means start from the back
if i < 0 {
i += v.len() as i64;
}
// Invalid index so return array unaltered
if i > v.len() as i64 || i < 0 {
return Ok(v.into());
}
// Insert the value into the array
v.insert(i as usize, value);
Ok(v.into())
}
(Value::Array(mut v), None) => {
v.push(value);
Ok(v.into())
}
(_, _) => Ok(Value::None),
}
}
pub fn intersect(arrays: (Value, Value)) -> Result<Value, Error> {
Ok(match arrays {
(Value::Array(v), Value::Array(w)) => v.intersect(w).into(),

View file

@ -65,6 +65,7 @@ pub fn synchronous(ctx: &Context<'_>, name: &str, args: Vec<Value>) -> Result<Va
"array::difference" => array::difference,
"array::distinct" => array::distinct,
"array::flatten" => array::flatten,
"array::insert" => array::insert,
"array::intersect" => array::intersect,
"array::len" => array::len,
"array::sort" => array::sort,

View file

@ -238,6 +238,7 @@ fn function_array(i: &str) -> IResult<&str, &str> {
tag("array::difference"),
tag("array::flatten"),
tag("array::distinct"),
tag("array::insert"),
tag("array::intersect"),
tag("array::len"),
tag("array::sort::asc"),

View file

@ -172,6 +172,33 @@ async fn function_array_flatten() -> Result<(), Error> {
Ok(())
}
#[tokio::test]
async fn function_array_insert() -> Result<(), Error> {
let sql = r#"
RETURN array::insert([], 1);
RETURN array::insert([3], 1, 1);
RETURN array::insert([1,2,3,4], 5, -1);
"#;
let dbs = Datastore::new("memory").await?;
let ses = Session::for_kv().with_ns("test").with_db("test");
let res = &mut dbs.execute(&sql, &ses, None, false).await?;
assert_eq!(res.len(), 3);
//
let tmp = res.remove(0).result?;
let val = Value::parse("[1]");
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
let val = Value::parse("[3,1]");
assert_eq!(tmp, val);
//
let tmp = res.remove(0).result?;
let val = Value::parse("[1,2,3,5,4]");
assert_eq!(tmp, val);
//
Ok(())
}
#[tokio::test]
async fn function_array_intersect() -> Result<(), Error> {
let sql = r#"