From b2967a43dcd78915478f12209d16e2601d13679c Mon Sep 17 00:00:00 2001
From: Xianzhe Liang <liangxianzhe@gmail.com>
Date: Sun, 10 Mar 2024 20:41:12 -0700
Subject: [PATCH] NamedRow into Payload.

Once we have NamedRow (from previous query or construct manually),
convert it to Payload to make db changes.

This is similar to the insert/put/update/rm wrappers in pycozo.
---
 cozo-core/src/lib.rs           |  1 +
 cozo-core/src/runtime/db.rs    | 15 ++++++++++++++-
 cozo-core/src/runtime/tests.rs | 22 +++++++++++++++++++++-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/cozo-core/src/lib.rs b/cozo-core/src/lib.rs
index c0969a01..84a6297c 100644
--- a/cozo-core/src/lib.rs
+++ b/cozo-core/src/lib.rs
@@ -75,6 +75,7 @@ pub use crate::runtime::db::evaluate_expressions;
 pub use crate::runtime::db::get_variables;
 pub use crate::runtime::db::Poison;
 pub use crate::runtime::db::ScriptMutability;
+pub use crate::runtime::db::Payload;
 pub use crate::runtime::db::TransactionPayload;
 
 pub(crate) mod data;
diff --git a/cozo-core/src/runtime/db.rs b/cozo-core/src/runtime/db.rs
index 90996a57..a89248b6 100644
--- a/cozo-core/src/runtime/db.rs
+++ b/cozo-core/src/runtime/db.rs
@@ -231,11 +231,23 @@ impl NamedRows {
             next: None,
         })
     }
+
+    /// Create a query and parameters to apply an operation (insert, put, delete, rm) to a stored
+    /// relation with the named rows.
+    pub fn into_payload(self, relation: &str, op: &str) -> Payload {
+        let cols_str = self.headers.join(", ");
+        let query = format!("?[{cols_str}] <- $data :{op} {relation} {{ {cols_str} }}");
+        let data = DataValue::List(self.rows.into_iter().map(|r| DataValue::List(r)).collect());
+        (query, [("data".to_string(), data)].into())
+    }
 }
 
 const STATUS_STR: &str = "status";
 const OK_STR: &str = "OK";
 
+/// The query and parameters.
+pub type Payload = (String, BTreeMap<String, DataValue>);
+
 /// Commands to be sent to a multi-transaction
 #[derive(Eq, PartialEq, Debug)]
 pub enum TransactionPayload {
@@ -244,7 +256,7 @@ pub enum TransactionPayload {
     /// Abort the current transaction
     Abort,
     /// Run a query inside the transaction
-    Query((String, BTreeMap<String, DataValue>)),
+    Query(Payload),
 }
 
 impl<'s, S: Storage<'s>> Db<S> {
@@ -397,6 +409,7 @@ impl<'s, S: Storage<'s>> Db<S> {
             mutability == ScriptMutability::Immutable,
         )
     }
+
     /// Run the CozoScript passed in. The `params` argument is a map of parameters.
     pub fn run_script_read_only(
         &'s self,
diff --git a/cozo-core/src/runtime/tests.rs b/cozo-core/src/runtime/tests.rs
index 39e97d84..f2d3bb2f 100644
--- a/cozo-core/src/runtime/tests.rs
+++ b/cozo-core/src/runtime/tests.rs
@@ -1183,6 +1183,26 @@ fn deletion() {
     db.run_default(r"?[x] <- [[1]] :delete a {x}").unwrap();
 }
 
+#[test]
+fn into_payload() {
+    let db = DbInstance::new("mem", "", "").unwrap();
+    db.run_default(r":create a {x => y}").unwrap();
+    db.run_default(r"?[x, y] <- [[1, 2], [3, 4]] :insert a {x => y}",).unwrap();
+
+    let mut res = db.run_default(r"?[x, y] := *a[x, y]").unwrap();
+    assert_eq!(res.rows.len(), 2);
+
+    let delete = res.clone().into_payload("a", "rm");
+    db.run_script(delete.0.as_str(), delete.1, ScriptMutability::Mutable).unwrap();
+    assert_eq!(db.run_default(r"?[x, y] := *a[x, y]").unwrap().rows.len(), 0);
+
+    db.run_default(r":create b {m => n}").unwrap();
+    res.headers = vec!["m".into(), "n".into()];
+    let put = res.into_payload("b", "put");
+    db.run_script(put.0.as_str(), put.1, ScriptMutability::Mutable).unwrap();
+    assert_eq!(db.run_default(r"?[m, n] := *b[m, n]").unwrap().rows.len(), 2);
+}
+
 #[test]
 fn returning() {
     let db = DbInstance::new("mem", "", "").unwrap();
@@ -1553,4 +1573,4 @@ fn fts_drop() {
     db.run_default(r#"
         ::fts drop entity:fts_index
     "#).unwrap();
-}
\ No newline at end of file
+}