ensure and ensure_not

cozorocks/CMakeLists.txt

@@ -4,7 +4,7 @@ project(cozorocks)
 set(CMAKE_CXX_STANDARD 17)
 
 include_directories("bridge")
-include_directories("deps/rocksdb/include")
+include_directories("./rocksdb/include")
 include_directories("../target/cxxbridge")
 
 add_library(cozorocks "bridge/bridge.h" "bridge/common.h" "bridge/db.cpp" "bridge/db.h" "bridge/iter.h" "bridge/opts.h"

docs/source/stored.rst

@@ -34,6 +34,12 @@ To manipulate stored relations, use one of the following query options:
     Put data from the resulting relation into the named stored relation.
     If keys from the data exist beforehand, the rows are simply replaced with new ones.
 
+.. function:: :ensure <NAME> <SPEC>
+
+    Ensures that rows specified by the output relation and spec already exist in the database,
+    and that no other process has written to these rows at commit since the transaction starts.
+    Useful for ensuring read-write consistency.
+
 .. function:: :rm <NAME> <SPEC>
 
     Remove data from the resulting relation from the named stored relation.
@@ -41,6 +47,12 @@ To manipulate stored relations, use one of the following query options:
     If a row from the resulting relation does not match any keys, nothing happens for that row,
     and no error is raised.
 
+.. function:: :ensure_not <NAME> <SPEC>
+
+    Ensures that rows specified by the output relation and spec do not exist in the database,
+    and that no other process has written to these rows at commit since the transaction starts.
+    Useful for ensuring read-write consistency.
+
 You can rename and remove stored relations with the system ops ``::relation rename`` and ``::relation remove``,
 described in the system op chapter.
 
@@ -102,10 +114,10 @@ The expression is evaluated once for each row, so for example if you specified o
 you will get a different UUID for each row.
 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-Put and remove
+Put, remove, ensure and ensure-not
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-For ``:put`` and ``:remove``,
+For ``:put``, ``:remove``, ``:ensure`` and ``:ensure_not``,
 you do not need to specify all existing columns in the spec if the omitted columns have a default generator,
 in which case the generator will be used to generate a value,
 or the type of the column is nullable, in which case the value is ``null``.
@@ -113,8 +125,8 @@ Also, the order of the columns does not matter, and neither does whether a colum
 The spec specified when the relation was created will be consulted to know how to store data correctly.
 Specifying default values does not have any effect and will not replace existing ones.
 
-For ``:put``, the spec needs to contain enough bindings to generate all keys and values.
+For ``:put`` and ``:ensure``, the spec needs to contain enough bindings to generate all keys and values.
-For ``:rm``, it only needs to generate all keys.
+For ``:rm`` and ``:ensure_not``, it only needs to generate all keys.
 
 ------------------------------------------------------
 Chaining queries into a single transaction

src/query/stored.rs

@@ -1,7 +1,7 @@
 use std::collections::BTreeMap;
 
 use itertools::Itertools;
-use miette::{Diagnostic, Result};
+use miette::{bail, Diagnostic, Result};
 use smartstring::SmartString;
 use thiserror::Error;
 
@@ -145,13 +145,85 @@ impl SessionTx {
                         to_clear.extend(cleanups);
                     }
                 }
-            },
+            }
             RelationOp::Ensure => {
-                todo!()
+                let mut key_extractors = make_extractors(
-            },
+                    &relation_store.metadata.keys,
+                    &metadata.keys,
+                    key_bindings,
+                    headers,
+                )?;
+
+                let val_extractors = make_extractors(
+                    &relation_store.metadata.non_keys,
+                    &metadata.non_keys,
+                    dep_bindings,
+                    headers,
+                )?;
+                key_extractors.extend(val_extractors);
+
+                for tuple in res_iter {
+                    let tuple = tuple?;
+
+                    let extracted = Tuple(
+                        key_extractors
+                            .iter()
+                            .map(|ex| ex.extract_data(&tuple))
+                            .try_collect()?,
+                    );
+
+                    let key = relation_store.adhoc_encode_key(&extracted, *span)?;
+                    let val = relation_store.adhoc_encode_val(&extracted, *span)?;
+
+                    let existing = self.tx.get(&key, true)?;
+                    match existing {
+                        None => {
+                            bail!(TransactAssertionFailure {
+                                relation: relation_store.name.to_string(),
+                                key: extracted.0.clone(),
+                                notice: "key does not exist in database".to_string()
+                            })
+                        }
+                        Some(v) => {
+                            if &v as &[u8] != &val as &[u8] {
+                                bail!(TransactAssertionFailure {
+                                    relation: relation_store.name.to_string(),
+                                    key: extracted.0.clone(),
+                                    notice: "key exists in database, but value does not match"
+                                        .to_string()
+                                })
+                            }
+                        }
+                    }
+                }
+            }
             RelationOp::EnsureNot => {
-                todo!()
+                let key_extractors = make_extractors(
-            },
+                    &relation_store.metadata.keys,
+                    &metadata.keys,
+                    key_bindings,
+                    headers,
+                )?;
+
+                for tuple in res_iter {
+                    let tuple = tuple?;
+                    let extracted = Tuple(
+                        key_extractors
+                            .iter()
+                            .map(|ex| ex.extract_data(&tuple))
+                            .try_collect()?,
+                    );
+                    let key = relation_store.adhoc_encode_key(&extracted, *span)?;
+                    let existing = self.tx.get(&key, true)?;
+                    if existing.is_some() {
+                        bail!(TransactAssertionFailure {
+                            relation: relation_store.name.to_string(),
+                            key: extracted.0.clone(),
+                            notice: "key exists in database".to_string()
+                        })
+                    }
+                }
+            }
             RelationOp::Create | RelationOp::Replace | RelationOp::Put => {
                 let mut key_extractors = make_extractors(
                     &relation_store.metadata.keys,
@@ -241,6 +313,14 @@ impl SessionTx {
     }
 }
 
+#[derive(Debug, Error, Diagnostic)]
+#[error("Assertion failure for {key:?} of {relation}: {notice}")]
+struct TransactAssertionFailure {
+    relation: String,
+    key: Vec<DataValue>,
+    notice: String,
+}
+
 enum DataExtractor {
     DefaultExtractor(Expr, NullableColType),
     IndexExtractor(usize, NullableColType),