From 6017d424eb9fd4683a8b6487f87c471cd9fbfc0b Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Wed, 23 Mar 2022 14:01:28 +0000
Subject: [PATCH] Implement FETCH clauses

---
 lib/src/dbs/iterator.rs   | 75 +++++++++++++++++++++++++++++++++++----
 lib/src/dbs/statement.rs  |  8 +++++
 lib/src/sql/value/get.rs  |  1 +
 lib/src/sql/value/pick.rs |  1 +
 4 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/lib/src/dbs/iterator.rs b/lib/src/dbs/iterator.rs
index 7750c054..ee82c51e 100644
--- a/lib/src/dbs/iterator.rs
+++ b/lib/src/dbs/iterator.rs
@@ -8,6 +8,7 @@ use crate::dbs::Transaction;
 use crate::doc::Document;
 use crate::err::Error;
 use crate::sql::id::Id;
+use crate::sql::part::Part;
 use crate::sql::statements::create::CreateStatement;
 use crate::sql::statements::delete::DeleteStatement;
 use crate::sql::statements::insert::InsertStatement;
@@ -137,13 +138,11 @@ impl Iterator {
 		// Process any ORDER clause
 		self.output_order(&ctx, opt, txn).await?;
 		// Process any START clause
-		if let Some(v) = self.stm.start() {
-			self.results = mem::take(&mut self.results).into_iter().skip(v.0).collect();
-		}
+		self.output_start(&ctx, opt, txn).await?;
 		// Process any LIMIT clause
-		if let Some(v) = self.stm.limit() {
-			self.results = mem::take(&mut self.results).into_iter().take(v.0).collect();
-		}
+		self.output_limit(&ctx, opt, txn).await?;
+		// Process any FETCH clause
+		self.output_fetch(&ctx, opt, txn).await?;
 		// Output the results
 		Ok(mem::take(&mut self.results).into())
 	}
@@ -230,6 +229,70 @@ impl Iterator {
 		Ok(())
 	}
 
+	#[inline]
+	async fn output_start(
+		&mut self,
+		_ctx: &Runtime,
+		_opt: &Options,
+		_txn: &Transaction,
+	) -> Result<(), Error> {
+		if let Some(v) = self.stm.start() {
+			self.results = mem::take(&mut self.results).into_iter().skip(v.0).collect();
+		}
+		Ok(())
+	}
+
+	#[inline]
+	async fn output_limit(
+		&mut self,
+		_ctx: &Runtime,
+		_opt: &Options,
+		_txn: &Transaction,
+	) -> Result<(), Error> {
+		if let Some(v) = self.stm.limit() {
+			self.results = mem::take(&mut self.results).into_iter().take(v.0).collect();
+		}
+		Ok(())
+	}
+
+	#[inline]
+	async fn output_fetch(
+		&mut self,
+		ctx: &Runtime,
+		opt: &Options,
+		txn: &Transaction,
+	) -> Result<(), Error> {
+		if let Some(fetchs) = self.stm.fetch() {
+			for fetch in &fetchs.0 {
+				// Loop over each value
+				for obj in &mut self.results {
+					// Get the value at the path
+					let val = obj.get(ctx, opt, txn, &fetch.fetch).await?;
+					// Set the value at the path
+					match val {
+						Value::Array(v) => {
+							// Fetch all remote records
+							let val = Value::Array(v).get(ctx, opt, txn, &[Part::All]).await?;
+							// Set the value at the path
+							obj.set(ctx, opt, txn, &fetch.fetch, val).await?;
+						}
+						Value::Thing(v) => {
+							// Fetch all remote records
+							let val = Value::Thing(v).get(ctx, opt, txn, &[Part::All]).await?;
+							// Set the value at the path
+							obj.set(ctx, opt, txn, &fetch.fetch, val).await?;
+						}
+						_ => {
+							// Set the value at the path
+							obj.set(ctx, opt, txn, &fetch.fetch, val).await?;
+						}
+					}
+				}
+			}
+		}
+		Ok(())
+	}
+
 	#[cfg(not(feature = "parallel"))]
 	async fn iterate(
 		&mut self,
diff --git a/lib/src/dbs/statement.rs b/lib/src/dbs/statement.rs
index 94569c7e..d5d587bd 100644
--- a/lib/src/dbs/statement.rs
+++ b/lib/src/dbs/statement.rs
@@ -1,4 +1,5 @@
 use crate::sql::cond::Cond;
+use crate::sql::fetch::Fetchs;
 use crate::sql::group::Groups;
 use crate::sql::limit::Limit;
 use crate::sql::order::Orders;
@@ -112,6 +113,13 @@ impl Statement {
 			_ => None,
 		}
 	}
+	// Returns any FETCH clause if specified
+	pub fn fetch(self: &Statement) -> Option<&Fetchs> {
+		match self {
+			Statement::Select(v) => v.fetch.as_ref(),
+			_ => None,
+		}
+	}
 	// Returns any START clause if specified
 	pub fn start(self: &Statement) -> Option<&Start> {
 		match self {
diff --git a/lib/src/sql/value/get.rs b/lib/src/sql/value/get.rs
index f6b6c3b3..251ae9ad 100644
--- a/lib/src/sql/value/get.rs
+++ b/lib/src/sql/value/get.rs
@@ -30,6 +30,7 @@ impl Value {
 						Some(v) => v.get(ctx, opt, txn, path.next()).await,
 						None => Ok(Value::None),
 					},
+					Part::All => Ok(self.clone()),
 					_ => Ok(Value::None),
 				},
 				// Current path part is an array
diff --git a/lib/src/sql/value/pick.rs b/lib/src/sql/value/pick.rs
index c14f1872..30c88c21 100644
--- a/lib/src/sql/value/pick.rs
+++ b/lib/src/sql/value/pick.rs
@@ -13,6 +13,7 @@ impl Value {
 						Some(v) => v.pick(path.next()),
 						None => Value::None,
 					},
+					Part::All => self.clone(),
 					_ => Value::None,
 				},
 				// Current path part is an array