query output assertions

src/cozoscript.pest

@@ -91,7 +91,8 @@ term = _{ literal | param | grouping | apply | var | list }
 list = { "[" ~ (expr ~ ",")* ~ expr? ~ "]" }
 grouping = { "(" ~ expr ~ ")" }
 
-option = _{(limit_option|offset_option|out_option|sort_option|relation_option|timeout_option) ~ ";"?}
+option = _{(limit_option|offset_option|out_option|sort_option|relation_option|timeout_option|
+            assert_none_option|assert_some_option) ~ ";"?}
 limit_option = {":limit"  ~ expr}
 offset_option = {":offset" ~ expr}
 out_option = {":pull" ~ var ~ ("@" ~ expr)? ~ out_spec}
@@ -107,6 +108,8 @@ sort_arg = { sort_dir? ~ var }
 sort_dir = _{ sort_asc | sort_desc }
 sort_asc = {"+"}
 sort_desc = {"-"}
+assert_none_option = {":assert" ~ "none"}
+assert_some_option = {":assert" ~ "some"}
 
 out_spec = {"{" ~ (pull_field ~ ",")* ~ pull_field? ~ "}"}
 pull_field = {rev_pull_marker? ~ compound_ident ~ (":" ~ out_spec)?}

src/data/expr.rs

@@ -15,7 +15,7 @@ use crate::data::tuple::Tuple;
 use crate::data::value::{DataValue, LARGEST_UTF_CHAR};
 use crate::parse::SourceSpan;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) enum Expr {
     Binding {
         var: Symbol,

src/data/program.rs

@@ -31,6 +31,12 @@ pub(crate) struct ConstRule {
     pub(crate) span: SourceSpan,
 }
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) enum QueryAssertion {
+    AssertNone(SourceSpan),
+    AssertSome(SourceSpan),
+}
+
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub(crate) struct QueryOutOptions {
     pub(crate) out_spec: BTreeMap<Symbol, (Vec<OutPullSpec>, Option<Validity>)>,
@@ -40,6 +46,7 @@ pub(crate) struct QueryOutOptions {
     pub(crate) timeout: Option<u64>,
     pub(crate) sorters: Vec<(Symbol, SortDir)>,
     pub(crate) store_relation: Option<(RelationMetadata, RelationOp)>,
+    pub(crate) assertion: Option<QueryAssertion>,
 }
 
 impl Default for QueryOutOptions {
@@ -52,6 +59,7 @@ impl Default for QueryOutOptions {
             timeout: None,
             sorters: vec![],
             store_relation: None,
+            assertion: None,
         }
     }
 }
@@ -502,7 +510,7 @@ impl InputProgram {
         if let Some(ConstRule { data, .. }) = self.const_rules.get(&MagicSymbol::Muggle {
             inner: Symbol::new(PROG_ENTRY, SourceSpan(0, 0)),
         }) {
-            return Ok(data.get(0).map(|row| row.0.len()).unwrap_or(0))
+            return Ok(data.get(0).map(|row| row.0.len()).unwrap_or(0));
         }
 
         Err(NoEntryError.into())

src/parse/query.rs

@@ -17,7 +17,7 @@ use crate::data::id::Validity;
 use crate::data::program::{
     AlgoApply, AlgoRuleArg, ConstRule, ConstRules, InputAtom, InputAttrTripleAtom, InputProgram,
     InputRelationApplyAtom, InputRule, InputRuleApplyAtom, InputRulesOrAlgo, InputTerm,
-    MagicSymbol, QueryOutOptions, RelationOp, SortDir, TripleDir, Unification,
+    MagicSymbol, QueryAssertion, QueryOutOptions, RelationOp, SortDir, TripleDir, Unification,
 };
 use crate::data::symb::Symbol;
 use crate::data::tuple::Tuple;
@@ -45,6 +45,11 @@ struct OptionNotPosIntError(&'static str, #[label] SourceSpan);
 #[derive(Debug)]
 struct MultipleRuleDefinitionError(String, Vec<SourceSpan>);
 
+#[derive(Debug, Error, Diagnostic)]
+#[error("Multiple query output assertions defined")]
+#[diagnostic(code(parser::multiple_out_assert))]
+struct DuplicateQueryAssertion(#[label] SourceSpan);
+
 impl Error for MultipleRuleDefinitionError {}
 
 impl Display for MultipleRuleDefinitionError {
@@ -380,6 +385,20 @@ pub(crate) fn parse_query(
                 };
                 out_opts.store_relation = Some((meta, op));
             }
+            Rule::assert_none_option => {
+                ensure!(
+                    out_opts.assertion.is_none(),
+                    DuplicateQueryAssertion(pair.extract_span())
+                );
+                out_opts.assertion = Some(QueryAssertion::AssertNone(pair.extract_span()))
+            }
+            Rule::assert_some_option => {
+                ensure!(
+                    out_opts.assertion.is_none(),
+                    DuplicateQueryAssertion(pair.extract_span())
+                );
+                out_opts.assertion = Some(QueryAssertion::AssertSome(pair.extract_span()))
+            }
             Rule::EOI => break,
             r => unreachable!("{:?}", r),
         }

src/runtime/db.rs

@@ -24,14 +24,14 @@ use crate::data::encode::{
 };
 use crate::data::id::{EntityId, TxId, Validity};
 use crate::data::json::JsonValue;
-use crate::data::program::{InputProgram, RelationOp};
+use crate::data::program::{InputProgram, QueryAssertion, RelationOp};
 use crate::data::symb::Symbol;
 use crate::data::tuple::{rusty_scratch_cmp, EncodedTuple, Tuple, SCRATCH_DB_KEY_PREFIX_LEN};
 use crate::data::value::{DataValue, LARGEST_UTF_CHAR};
 use crate::parse::schema::AttrTxItem;
 use crate::parse::sys::{CompactTarget, SysOp};
 use crate::parse::tx::TripleTx;
-use crate::parse::{parse_script, CozoScript};
+use crate::parse::{parse_script, CozoScript, SourceSpan};
 use crate::runtime::relation::{RelationId, RelationMetadata};
 use crate::runtime::transact::SessionTx;
 use crate::transact::meta::AttrNotFoundError;
@@ -425,6 +425,34 @@ impl Db {
             },
             poison,
         )?;
+        if let Some(assertion) = &input_program.out_opts.assertion {
+            match assertion {
+                QueryAssertion::AssertNone(span) => {
+                    if let Some(tuple) = result.scan_all().next() {
+                        let tuple = tuple?;
+
+                        #[derive(Debug, Error, Diagnostic)]
+                        #[error(
+                            "The query is asserted to return no result, but a tuple {0:?} is found"
+                        )]
+                        #[diagnostic(code(eval::assert_none_failure))]
+                        struct AssertNoneFailure(Tuple, #[label] SourceSpan);
+                        bail!(AssertNoneFailure(tuple, *span))
+                    }
+                }
+                QueryAssertion::AssertSome(span) => {
+                    if let Some(tuple) = result.scan_all().next() {
+                        let _ = tuple?;
+                    } else {
+                        #[derive(Debug, Error, Diagnostic)]
+                        #[error("The query is asserted to return some results, but returned none")]
+                        #[diagnostic(code(eval::assert_some_failure))]
+                        struct AssertSomeFailure(#[label] SourceSpan);
+                        bail!(AssertSomeFailure(*span))
+                    }
+                }
+            }
+        }
         let json_headers = match input_program.get_entry_head() {
             Err(_) => JsonValue::Null,
             Ok(headers) => headers.iter().map(|v| json!(v.name)).collect(),