From e1e617cff799cd58fd64f182400789531130f78e Mon Sep 17 00:00:00 2001
From: Aman Sharma <35651798+lunchspider@users.noreply.github.com>
Date: Sun, 18 Dec 2022 02:33:13 +0530
Subject: [PATCH] Add SQL array::insert() function (#1504)

---
 lib/src/fnc/args.rs     | 23 +++++++++++++++++++++++
 lib/src/fnc/array.rs    | 24 ++++++++++++++++++++++++
 lib/src/fnc/mod.rs      |  1 +
 lib/src/sql/function.rs |  1 +
 lib/tests/function.rs   | 27 +++++++++++++++++++++++++++
 5 files changed, 76 insertions(+)

diff --git a/lib/src/fnc/args.rs b/lib/src/fnc/args.rs
index 66118643..199b7991 100644
--- a/lib/src/fnc/args.rs
+++ b/lib/src/fnc/args.rs
@@ -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>) {
diff --git a/lib/src/fnc/array.rs b/lib/src/fnc/array.rs
index 7981b0d1..c42d7d5b 100644
--- a/lib/src/fnc/array.rs
+++ b/lib/src/fnc/array.rs
@@ -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(),
diff --git a/lib/src/fnc/mod.rs b/lib/src/fnc/mod.rs
index be0e2c70..174f9ef6 100644
--- a/lib/src/fnc/mod.rs
+++ b/lib/src/fnc/mod.rs
@@ -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,
diff --git a/lib/src/sql/function.rs b/lib/src/sql/function.rs
index a1e5ca9a..cf9cc212 100644
--- a/lib/src/sql/function.rs
+++ b/lib/src/sql/function.rs
@@ -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"),
diff --git a/lib/tests/function.rs b/lib/tests/function.rs
index 35f3bf28..71b4e11b 100644
--- a/lib/tests/function.rs
+++ b/lib/tests/function.rs
@@ -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#"