From a6ec9dbd02b663f9704d89521b3fcae25b4e419d Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@surrealdb.com>
Date: Thu, 24 Aug 2023 09:19:19 +0100
Subject: [PATCH] Ensure edge `in` and `out` fields are not overwritten (#2506)

---
 lib/src/doc/reset.rs | 9 ++++++++-
 lib/tests/relate.rs  | 6 ++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/lib/src/doc/reset.rs b/lib/src/doc/reset.rs
index 4c3b0367..050d56e2 100644
--- a/lib/src/doc/reset.rs
+++ b/lib/src/doc/reset.rs
@@ -2,6 +2,7 @@ use crate::ctx::Context;
 use crate::dbs::Options;
 use crate::dbs::Statement;
 use crate::dbs::Transaction;
+use crate::dbs::Workable;
 use crate::doc::Document;
 use crate::err::Error;
 use crate::sql::paths::EDGE;
@@ -21,7 +22,13 @@ impl<'a> Document<'a> {
 		let rid = self.id.as_ref().unwrap();
 		// Set default field values
 		self.current.doc.to_mut().def(rid);
-		// Ensure edge fields are reset
+		// This is a RELATE statement, so reset fields
+		if let Workable::Relate(l, r) = &self.extras {
+			self.current.doc.to_mut().put(&*EDGE, Value::Bool(true));
+			self.current.doc.to_mut().put(&*IN, l.clone().into());
+			self.current.doc.to_mut().put(&*OUT, r.clone().into());
+		}
+		// This is an UPDATE of a graph edge, so reset fields
 		if self.initial.doc.pick(&*EDGE).is_true() {
 			self.current.doc.to_mut().put(&*EDGE, Value::Bool(true));
 			self.current.doc.to_mut().put(&*IN, self.initial.doc.pick(&*IN));
diff --git a/lib/tests/relate.rs b/lib/tests/relate.rs
index 5a86a497..18562581 100644
--- a/lib/tests/relate.rs
+++ b/lib/tests/relate.rs
@@ -10,7 +10,7 @@ async fn relate_with_parameters() -> Result<(), Error> {
 	let sql = "
 		LET $tobie = person:tobie;
 		LET $jaime = person:jaime;
-		RELATE $tobie->knows->$jaime SET id = knows:test;
+		RELATE $tobie->knows->$jaime SET id = knows:test, brother = true;
 	";
 	let dbs = Datastore::new("memory").await?;
 	let ses = Session::owner().with_ns("test").with_db("test");
@@ -32,6 +32,7 @@ async fn relate_with_parameters() -> Result<(), Error> {
 				id: knows:test,
 				in: person:tobie,
 				out: person:jaime,
+				brother: true,
 			}
 		]",
 	);
@@ -45,7 +46,7 @@ async fn relate_and_overwrite() -> Result<(), Error> {
 	let sql = "
 		LET $tobie = person:tobie;
 		LET $jaime = person:jaime;
-		RELATE $tobie->knows->$jaime SET id = knows:test;
+		RELATE $tobie->knows->$jaime CONTENT { id: knows:test, brother: true };
 		UPDATE knows:test CONTENT { test: true };
 		SELECT * FROM knows:test;
 	";
@@ -69,6 +70,7 @@ async fn relate_and_overwrite() -> Result<(), Error> {
 				id: knows:test,
 				in: person:tobie,
 				out: person:jaime,
+				brother: true,
 			}
 		]",
 	);