From 50a2c25688005022c0bb39f37a11240fa1eb85c7 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Sat, 22 Oct 2022 22:19:12 +0100
Subject: [PATCH] Ensure cancelled context does not prevent FETCH of records

Closes #1395
---
 lib/src/dbs/iterator.rs | 22 +++++-----
 lib/tests/limit.rs      | 95 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+), 11 deletions(-)
 create mode 100644 lib/tests/limit.rs

diff --git a/lib/src/dbs/iterator.rs b/lib/src/dbs/iterator.rs
index f9c62d4f..bf96a850 100644
--- a/lib/src/dbs/iterator.rs
+++ b/lib/src/dbs/iterator.rs
@@ -78,30 +78,30 @@ impl Iterator {
 		// Log the statement
 		trace!(target: LOG, "Iterating: {}", stm);
 		// Enable context override
-		let mut ctx = Context::new(ctx);
-		self.run = ctx.add_cancel();
+		let mut run = Context::new(ctx);
+		self.run = run.add_cancel();
 		// Process the query LIMIT clause
-		self.setup_limit(&ctx, opt, txn, stm).await?;
+		self.setup_limit(&run, opt, txn, stm).await?;
 		// Process the query START clause
-		self.setup_start(&ctx, opt, txn, stm).await?;
+		self.setup_start(&run, opt, txn, stm).await?;
 		// Process prepared values
-		self.iterate(&ctx, opt, txn, stm).await?;
+		self.iterate(&run, opt, txn, stm).await?;
 		// Return any document errors
 		if let Some(e) = self.error.take() {
 			return Err(e);
 		}
 		// Process any SPLIT clause
-		self.output_split(&ctx, opt, txn, stm).await?;
+		self.output_split(ctx, opt, txn, stm).await?;
 		// Process any GROUP clause
-		self.output_group(&ctx, opt, txn, stm).await?;
+		self.output_group(ctx, opt, txn, stm).await?;
 		// Process any ORDER clause
-		self.output_order(&ctx, opt, txn, stm).await?;
+		self.output_order(ctx, opt, txn, stm).await?;
 		// Process any START clause
-		self.output_start(&ctx, opt, txn, stm).await?;
+		self.output_start(ctx, opt, txn, stm).await?;
 		// Process any LIMIT clause
-		self.output_limit(&ctx, opt, txn, stm).await?;
+		self.output_limit(ctx, opt, txn, stm).await?;
 		// Process any FETCH clause
-		self.output_fetch(&ctx, opt, txn, stm).await?;
+		self.output_fetch(ctx, opt, txn, stm).await?;
 		// Output the results
 		Ok(mem::take(&mut self.results).into())
 	}
diff --git a/lib/tests/limit.rs b/lib/tests/limit.rs
new file mode 100644
index 00000000..d5be8c31
--- /dev/null
+++ b/lib/tests/limit.rs
@@ -0,0 +1,95 @@
+mod parse;
+use parse::Parse;
+use surrealdb::sql::Value;
+use surrealdb::Datastore;
+use surrealdb::Error;
+use surrealdb::Session;
+
+#[tokio::test]
+async fn select_limit_fetch() -> Result<(), Error> {
+	let sql = "
+		CREATE tag:rs SET name = 'Rust';
+		CREATE tag:go SET name = 'Golang';
+		CREATE tag:js SET name = 'JavaScript';
+		CREATE person:tobie SET tags = [tag:rs, tag:go, tag:js];
+		CREATE person:jaime SET tags = [tag:js];
+		SELECT * FROM person LIMIT 1 FETCH tags;
+	";
+	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);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: tag:rs,
+				name: 'Rust'
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: tag:go,
+				name: 'Golang'
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: tag:js,
+				name: 'JavaScript'
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:tobie,
+				tags: [tag:rs, tag:go, tag:js]
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:jaime,
+				tags: [tag:js]
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:jaime,
+				tags: [
+					{
+						id: tag:js,
+						name: 'JavaScript'
+					}
+				]
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	Ok(())
+}