From 223d119b2c0943529100a859d20ab5a45783ced9 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@surrealdb.com>
Date: Mon, 20 Feb 2023 20:40:05 +0000
Subject: [PATCH] Allow SQL subqueries to be ended witrh `;` character

Related to #247
Related to #225
Related to #1319
---
 lib/src/sql/ending.rs | 36 ++++++++++++++++++++++++------------
 lib/tests/subquery.rs | 40 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 13 deletions(-)

diff --git a/lib/src/sql/ending.rs b/lib/src/sql/ending.rs
index 67d55288..f3de23e0 100644
--- a/lib/src/sql/ending.rs
+++ b/lib/src/sql/ending.rs
@@ -1,5 +1,5 @@
 use crate::sql::comment::comment;
-use crate::sql::comment::shouldbespace;
+use crate::sql::comment::{mightbespace, shouldbespace};
 use crate::sql::error::IResult;
 use crate::sql::operator::{assigner, operator};
 use nom::branch::alt;
@@ -67,15 +67,27 @@ pub fn duration(i: &str) -> IResult<&str, ()> {
 }
 
 pub fn subquery(i: &str) -> IResult<&str, ()> {
-	peek(alt((
-		map(preceded(shouldbespace, tag_no_case("THEN")), |_| ()),
-		map(preceded(shouldbespace, tag_no_case("ELSE")), |_| ()),
-		map(preceded(shouldbespace, tag_no_case("END")), |_| ()),
-		map(comment, |_| ()),
-		map(char(']'), |_| ()),
-		map(char('}'), |_| ()),
-		map(char(';'), |_| ()),
-		map(char(','), |_| ()),
-		map(eof, |_| ()),
-	)))(i)
+	alt((
+		|i| {
+			let (i, _) = mightbespace(i)?;
+			let (i, _) = char(';')(i)?;
+			let (i, _) = peek(alt((
+				preceded(shouldbespace, tag_no_case("THEN")),
+				preceded(shouldbespace, tag_no_case("ELSE")),
+				preceded(shouldbespace, tag_no_case("END")),
+			)))(i)?;
+			Ok((i, ()))
+		},
+		peek(alt((
+			map(preceded(shouldbespace, tag_no_case("THEN")), |_| ()),
+			map(preceded(shouldbespace, tag_no_case("ELSE")), |_| ()),
+			map(preceded(shouldbespace, tag_no_case("END")), |_| ()),
+			map(comment, |_| ()),
+			map(char(']'), |_| ()),
+			map(char('}'), |_| ()),
+			map(char(';'), |_| ()),
+			map(char(','), |_| ()),
+			map(eof, |_| ()),
+		))),
+	))(i)
 }
diff --git a/lib/tests/subquery.rs b/lib/tests/subquery.rs
index 73c96f31..0914895d 100644
--- a/lib/tests/subquery.rs
+++ b/lib/tests/subquery.rs
@@ -124,11 +124,21 @@ async fn subquery_ifelse() -> Result<(), Error> {
 		ELSE
 			UPDATE person:test SET sport = ['basketball'] RETURN sport
 		END;
+		-- Check if the record exists
+		LET $record = SELECT *, count() AS count FROM person:test;
+		-- Return the specified record
+		RETURN $record;
+		-- Update the record field if it exists
+		IF $record.count THEN
+			UPDATE person:test SET sport += 'football' RETURN sport;
+		ELSE
+			UPDATE person:test SET sport = ['basketball'] RETURN sport;
+		END;
 	";
 	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(), 6);
+	assert_eq!(res.len(), 9);
 	//
 	let tmp = res.remove(0).result?;
 	let val = Value::None;
@@ -175,5 +185,33 @@ async fn subquery_ifelse() -> Result<(), Error> {
 	);
 	assert_eq!(tmp, val);
 	//
+	let tmp = res.remove(0).result?;
+	let val = Value::None;
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"{
+			count: 1,
+			id: person:test,
+			sport: [
+				'basketball',
+				'football'
+			]
+		}",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"{
+			sport: [
+				'basketball',
+				'football'
+			]
+		}",
+	);
+	assert_eq!(tmp, val);
+	//
 	Ok(())
 }