From 78329abf977f28d7d5dc17b16fb17205e7731a89 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@surrealdb.com>
Date: Fri, 31 Mar 2023 19:46:13 +0100
Subject: [PATCH] Enable access to input data using `$input` variable

Closes #1543
---
 lib/src/doc/field.rs | 15 +++++++
 lib/tests/update.rs  | 97 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 lib/tests/update.rs

diff --git a/lib/src/doc/field.rs b/lib/src/doc/field.rs
index fdc5a771..783e8dcd 100644
--- a/lib/src/doc/field.rs
+++ b/lib/src/doc/field.rs
@@ -4,6 +4,7 @@ use crate::dbs::Statement;
 use crate::dbs::Transaction;
 use crate::doc::Document;
 use crate::err::Error;
+use crate::sql::data::Data;
 use crate::sql::permission::Permission;
 use crate::sql::value::Value;
 
@@ -27,10 +28,22 @@ impl<'a> Document<'a> {
 			for (k, mut val) in self.current.walk(&fd.name).into_iter() {
 				// Get the initial value
 				let old = self.initial.pick(&k);
+				// Get the input value
+				let inp = match stm.data() {
+					Some(Data::MergeExpression(v)) => v.pick(&k),
+					Some(Data::ContentExpression(v)) => v.pick(&k),
+					Some(Data::ReplaceExpression(v)) => v.pick(&k),
+					Some(Data::SetExpression(v)) => match v.iter().find(|f| f.0.eq(&k)) {
+						Some((_, _, v)) => v.to_owned(),
+						_ => Value::None,
+					},
+					_ => Value::None,
+				};
 				// Check for a VALUE clause
 				if let Some(expr) = &fd.value {
 					// Configure the context
 					let mut ctx = Context::new(ctx);
+					ctx.add_value("input".into(), &inp);
 					ctx.add_value("value".into(), &val);
 					ctx.add_value("after".into(), &val);
 					ctx.add_value("before".into(), &old);
@@ -49,6 +62,7 @@ impl<'a> Document<'a> {
 				if let Some(expr) = &fd.assert {
 					// Configure the context
 					let mut ctx = Context::new(ctx);
+					ctx.add_value("input".into(), &inp);
 					ctx.add_value("value".into(), &val);
 					ctx.add_value("after".into(), &val);
 					ctx.add_value("before".into(), &old);
@@ -81,6 +95,7 @@ impl<'a> Document<'a> {
 							let opt = &opt.perms(false);
 							// Configure the context
 							let mut ctx = Context::new(ctx);
+							ctx.add_value("input".into(), &inp);
 							ctx.add_value("value".into(), &val);
 							ctx.add_value("after".into(), &val);
 							ctx.add_value("before".into(), &old);
diff --git a/lib/tests/update.rs b/lib/tests/update.rs
new file mode 100644
index 00000000..d9305511
--- /dev/null
+++ b/lib/tests/update.rs
@@ -0,0 +1,97 @@
+mod parse;
+use parse::Parse;
+use surrealdb::dbs::Session;
+use surrealdb::err::Error;
+use surrealdb::kvs::Datastore;
+use surrealdb::sql::Value;
+
+#[tokio::test]
+async fn update_with_input() -> Result<(), Error> {
+	let sql = "
+		DEFINE FIELD name ON TABLE person
+			ASSERT
+				IF $input THEN
+					$input = /^[A-Z]{1}[a-z]+$/
+				ELSE
+					true
+				END
+			VALUE
+				IF $input THEN
+					'Name: ' + $input
+				ELSE
+					$value
+				END
+		;
+		UPDATE person:test CONTENT { name: 'Tobie' };
+		UPDATE person:test REPLACE { name: 'jaime' };
+		UPDATE person:test MERGE { name: 'Jaime' };
+		UPDATE person:test SET name = 'tobie';
+		UPDATE person:test SET name = 'Tobie';
+		SELECT * FROM person:test;
+	";
+	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(), 7);
+	//
+	let tmp = res.remove(0).result;
+	assert!(tmp.is_ok());
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:test,
+				name: 'Name: Tobie',
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result;
+	assert!(matches!(
+		tmp.err(),
+		Some(e) if e.to_string() == r#"Found 'Name: jaime' for field `name`, with record `person:test`, but field must conform to: IF $input THEN $input = /^[A-Z]{1}[a-z]+$/ ELSE true END"#
+	));
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:test,
+				name: 'Name: Jaime',
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result;
+	assert!(matches!(
+		tmp.err(),
+		Some(e) if e.to_string() == r#"Found 'Name: tobie' for field `name`, with record `person:test`, but field must conform to: IF $input THEN $input = /^[A-Z]{1}[a-z]+$/ ELSE true END"#
+	));
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:test,
+				name: 'Name: Tobie',
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	let tmp = res.remove(0).result?;
+	let val = Value::parse(
+		"[
+			{
+				id: person:test,
+				name: 'Name: Tobie',
+			}
+		]",
+	);
+	assert_eq!(tmp, val);
+	//
+	Ok(())
+}