From 072bbdebdb9fbf4d8e3ea157aca08ba16e16190a Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Sun, 16 Oct 2022 23:12:24 +0100
Subject: [PATCH] Add support for `FETCH` cluases in SQL `RETURN` statements

Closes #1323
---
 lib/src/sql/statements/output.rs | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/lib/src/sql/statements/output.rs b/lib/src/sql/statements/output.rs
index 64a8b9d5..eb8d9872 100644
--- a/lib/src/sql/statements/output.rs
+++ b/lib/src/sql/statements/output.rs
@@ -4,15 +4,19 @@ use crate::dbs::Transaction;
 use crate::err::Error;
 use crate::sql::comment::shouldbespace;
 use crate::sql::error::IResult;
+use crate::sql::fetch::{fetch, Fetchs};
 use crate::sql::value::{value, Value};
 use derive::Store;
 use nom::bytes::complete::tag_no_case;
+use nom::combinator::opt;
+use nom::sequence::preceded;
 use serde::{Deserialize, Serialize};
 use std::fmt;
 
 #[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize, Store)]
 pub struct OutputStatement {
 	pub what: Value,
+	pub fetch: Option<Fetchs>,
 }
 
 impl OutputStatement {
@@ -30,24 +34,38 @@ impl OutputStatement {
 		// Ensure futures are processed
 		let opt = &opt.futures(true);
 		// Process the output value
-		self.what.compute(ctx, opt, txn, doc).await
+		let mut val = self.what.compute(ctx, opt, txn, doc).await?;
+		// Fetch any
+		if let Some(fetchs) = &self.fetch {
+			for fetch in fetchs.iter() {
+				val.fetch(ctx, opt, txn, fetch).await?;
+			}
+		}
+		//
+		Ok(val)
 	}
 }
 
 impl fmt::Display for OutputStatement {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-		write!(f, "RETURN {}", self.what)
+		write!(f, "RETURN {}", self.what)?;
+		if let Some(ref v) = self.fetch {
+			write!(f, " {}", v)?
+		}
+		Ok(())
 	}
 }
 
 pub fn output(i: &str) -> IResult<&str, OutputStatement> {
 	let (i, _) = tag_no_case("RETURN")(i)?;
 	let (i, _) = shouldbespace(i)?;
-	let (i, v) = value(i)?;
+	let (i, what) = value(i)?;
+	let (i, fetch) = opt(preceded(shouldbespace, fetch))(i)?;
 	Ok((
 		i,
 		OutputStatement {
-			what: v,
+			what,
+			fetch,
 		},
 	))
 }