eval_user_query for OPA style results (#29)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

Cargo.toml

@@ -16,6 +16,7 @@ lazy_static = "1.4.0"
 rand = "0.8.5"
 
 [dev-dependencies]
+clap = { version = "4.4.7", features = ["derive"] }
 serde_yaml = "0.9.16"
 test-generator = "0.3.1"
 walkdir = "2.3.2"

examples/regorus.rs

@@ -0,0 +1,111 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use anyhow::{bail, Context, Result};
+use clap::Parser;
+
+#[derive(clap::Parser)]
+#[command(author, version, about, long_about = None)]
+struct Cli {
+    /// Policy or data files. Rego, json or yaml.
+    #[arg(required(true), long, short, value_name = "policy.rego")]
+    data: Vec<String>,
+
+    /// Input file. json or yaml.
+    #[arg(long, short, value_name = "input.rego")]
+    input: Option<String>,
+
+    // Query. Rego expression.
+    #[arg(long, short)]
+    query: Option<String>,
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+    let enable_tracing = false;
+
+    // User specified data.
+    let mut data = regorus::Value::new_object();
+
+    // Read all policy files.
+    let mut policies = vec![];
+    for file in cli.data.iter() {
+        let contents =
+            std::fs::read_to_string(file).with_context(|| format!("Failed to read {file}"))?;
+
+        if file.ends_with(".rego") {
+            policies.push(contents);
+        } else {
+            let value: regorus::Value = if file.ends_with(".json") {
+                serde_json::from_str(&contents)?
+            } else if file.ends_with(".yaml") {
+                serde_yaml::from_str(&contents)?
+            } else {
+                bail!("Unsupported data file `{file}`. Must be rego, json or yaml.")
+            };
+
+            if let Err(err) = data.merge(value) {
+                bail!("Error processing {file}. {err}");
+            }
+        }
+    }
+
+    // Create source objects.
+    let mut sources = vec![];
+    for (idx, rego) in policies.iter().enumerate() {
+        sources.push(regorus::Source {
+            file: &cli.data[idx],
+            contents: rego.as_str(),
+            lines: rego.split('\n').collect(),
+        });
+    }
+
+    // Parse the policy files.
+    let mut modules = vec![];
+    for source in &sources {
+        let mut parser = regorus::Parser::new(source)?;
+        modules.push(parser.parse()?);
+    }
+
+    // Analyze the modules and determine how statements must be schedules.
+    let analyzer = regorus::Analyzer::new();
+    let schedule = analyzer.analyze(&modules)?;
+
+    // Create interpreter object.
+    let modules_ref: Vec<&regorus::Module> = modules.iter().collect();
+    let mut interpreter = regorus::Interpreter::new(modules_ref)?;
+
+    // Prepare for evalution.
+    interpreter.prepare_for_eval(Some(&schedule), &Some(data))?;
+
+    // Evaluate all the modules.
+    interpreter.eval(&None, &None, false, Some(&schedule))?;
+
+    // Fetch query string. If none specified, use "data".
+    let query = match &cli.query {
+        Some(query) => query,
+        _ => "data",
+    };
+
+    // Parse the query.
+    let query_source = regorus::Source {
+        file: "<query.rego>",
+        contents: query,
+        lines: query.split('\n').collect(),
+    };
+    let query_span = regorus::Span {
+        source: &query_source,
+        line: 1,
+        col: 1,
+        start: 0,
+        end: query.len() as u16,
+    };
+    let mut parser = regorus::Parser::new(&query_source)?;
+    let query_node = parser.parse_query(query_span, "")?;
+    let stmt_order = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?;
+
+    let results = interpreter.eval_user_query(&query_node, &stmt_order, enable_tracing)?;
+    println!("eval results:\n{}", serde_json::to_string_pretty(&results)?);
+
+    Ok(())
+}

src/builtins/aggregates.rs

@@ -27,7 +27,7 @@ fn count(span: &Span, params: &[Expr], args: &[Value]) -> Result<Value> {
         Value::Array(a) => a.len() as Float,
         Value::Set(a) => a.len() as Float,
         Value::Object(a) => a.len() as Float,
-        Value::String(a) => a.len() as Float,
+        Value::String(a) => a.encode_utf16().count() as Float,
         a => {
             let span = params[0].span();
             bail!(span.error(

src/interpreter.rs

@@ -11,6 +11,7 @@ use crate::value::*;
 
 use anyhow::{anyhow, bail, Result};
 use log::info;
+use serde::Serialize;
 use std::collections::{hash_map::Entry, BTreeMap, BTreeSet, HashMap};
 use std::rc::Rc;
 
@@ -39,11 +40,35 @@ pub struct Interpreter<'source> {
     traces: Option<Vec<String>>,
 }
 
+#[derive(Debug, Clone, Serialize)]
+pub struct QueryResult {
+    // Expressions is shown first to match OPA.
+    pub expressions: Vec<Value>,
+    #[serde(skip_serializing_if = "Value::is_empty_object")]
+    pub bindings: Value,
+}
+
+impl Default for QueryResult {
+    fn default() -> Self {
+        Self {
+            bindings: Value::new_object(),
+            expressions: vec![],
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct QueryResults {
+    pub results: Vec<QueryResult>,
+}
+
 #[derive(Debug, Clone)]
 struct Context<'source> {
     key_expr: Option<&'source Expr<'source>>,
     output_expr: Option<&'source Expr<'source>>,
     value: Value,
+    result: Option<QueryResult>,
+    results: QueryResults,
 }
 
 #[derive(Debug)]
@@ -123,9 +148,6 @@ impl<'source> Interpreter<'source> {
     }
 
     fn current_scope(&mut self) -> Result<&Scope> {
-        if self.scopes.is_empty() {
-            println!("here");
-        }
         self.scopes
             .last()
             .ok_or_else(|| anyhow!("internal error: no active scope"))
@@ -484,6 +506,8 @@ impl<'source> Interpreter<'source> {
             key_expr: None,
             output_expr: None,
             value: Value::new_set(),
+            result: None,
+            results: QueryResults::default(),
         });
         let mut r = true;
         match domain {
@@ -585,8 +609,9 @@ impl<'source> Interpreter<'source> {
                 }
                 type_match.insert(expr);
 
+                let mut r = false;
                 for (idx, item) in items.iter().enumerate() {
-                    self.make_bindings(is_last, type_match, cache, item, &a[idx])?;
+                    r = self.make_bindings(is_last, type_match, cache, item, &a[idx])? || r;
                 }
 
                 Ok(true)
@@ -594,6 +619,7 @@ impl<'source> Interpreter<'source> {
 
             // Destructure objects
             (Expr::Object { fields, .. }, Value::Object(_)) => {
+                let mut r = true;
                 for (_, key_expr, value_expr) in fields.iter() {
                     // Rego does not support bindings in keys.
                     // Therefore, just eval key_expr.
@@ -608,9 +634,18 @@ impl<'source> Interpreter<'source> {
                     }
 
                     // Match patterns in value_expr
-                    self.make_bindings(is_last, type_match, cache, value_expr, field_value)?;
+                    r = r
+                        && self.make_bindings(
+                            is_last,
+                            type_match,
+                            cache,
+                            value_expr,
+                            field_value,
+                        )?;
                 }
-                Ok(true)
+                type_match.insert(expr);
+
+                Ok(r)
             }
             _ => {
                 let expr_value = self.lookup_or_eval_expr(cache, expr)?;
@@ -797,6 +832,12 @@ impl<'source> Interpreter<'source> {
                     _ => self.eval_expr(expr)?,
                 };
 
+                if let Some(ctx) = self.contexts.last_mut() {
+                    if let Some(result) = &mut ctx.result {
+                        result.expressions.push(value.clone());
+                    }
+                }
+
                 if let Value::Bool(bool) = value {
                     bool
                 } else {
@@ -818,6 +859,11 @@ impl<'source> Interpreter<'source> {
                     )?,
                     _ => self.eval_expr(expr)?,
                 };
+                if let Some(ctx) = self.contexts.last_mut() {
+                    if let Some(result) = &mut ctx.result {
+                        result.expressions.push(Value::Bool(true));
+                    }
+                }
                 // https://github.com/open-policy-agent/opa/issues/1622#issuecomment-520547385
                 matches!(value, Value::Bool(false) | Value::Undefined)
             }
@@ -833,6 +879,11 @@ impl<'source> Interpreter<'source> {
                         }
                     }
                 }
+                if let Some(ctx) = self.contexts.last_mut() {
+                    if let Some(result) = &mut ctx.result {
+                        result.expressions.push(Value::Bool(true));
+                    }
+                }
                 true
             }
             Literal::SomeIn {
@@ -840,14 +891,28 @@ impl<'source> Interpreter<'source> {
                 key,
                 value,
                 collection,
-            } => self.eval_some_in(span, key, value, collection, stmts)?,
+            } => {
+                if let Some(ctx) = self.contexts.last_mut() {
+                    if let Some(result) = &mut ctx.result {
+                        result.expressions.push(Value::Bool(true));
+                    }
+                }
+                self.eval_some_in(span, key, value, collection, stmts)?
+            }
             Literal::Every {
                 span,
                 key,
                 value,
                 domain,
                 query,
-            } => self.eval_every(span, key, value, domain, query)?,
+            } => {
+                if let Some(ctx) = self.contexts.last_mut() {
+                    if let Some(result) = &mut ctx.result {
+                        result.expressions.push(Value::Bool(true));
+                    }
+                }
+                self.eval_every(span, key, value, domain, query)?
+            }
         });
 
         for (path, value) in to_restore.into_iter().rev() {
@@ -901,6 +966,7 @@ impl<'source> Interpreter<'source> {
             // that the effects of the current loop iteration are cleared.
             let scope_saved = self.current_scope()?.clone();
 
+            let query_result = self.get_current_context()?.result.clone();
             match loop_expr_value {
                 Value::Array(items) => {
                     for (idx, v) in items.iter().enumerate() {
@@ -910,6 +976,9 @@ impl<'source> Interpreter<'source> {
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
                         self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
+                        if let Some(ctx) = self.contexts.last_mut() {
+                            ctx.result = query_result.clone();
+                        }
                     }
                 }
                 Value::Set(items) => {
@@ -920,6 +989,9 @@ impl<'source> Interpreter<'source> {
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
                         self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
+                        if let Some(ctx) = self.contexts.last_mut() {
+                            ctx.result = query_result.clone();
+                        }
                     }
                 }
                 Value::Object(obj) => {
@@ -930,6 +1002,9 @@ impl<'source> Interpreter<'source> {
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
                         self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
+                        if let Some(ctx) = self.contexts.last_mut() {
+                            ctx.result = query_result.clone();
+                        }
                     }
                 }
                 _ => {
@@ -1007,8 +1082,21 @@ impl<'source> Interpreter<'source> {
                 _ => (),
             }
 
-            // Push the context back so that it is available to the caller.
-            //            self.contexts.push(ctx);
+            // If a query snippet is being run, gather results.
+            let ctx = self.contexts.last_mut().expect("no current context");
+            if let Some(result) = &ctx.result {
+                let mut result = result.clone();
+                if let Some(scope) = self.scopes.last() {
+                    for (name, value) in scope.iter() {
+                        result
+                            .bindings
+                            .as_object_mut()?
+                            .insert(Value::String(name.to_string()), value.clone());
+                    }
+                }
+                ctx.results.results.push(result);
+            }
+
             return Ok(true);
         }
 
@@ -1214,7 +1302,8 @@ impl<'source> Interpreter<'source> {
                 }
             }
             _ => {
-                return Err(anyhow!("\"{}\" must be array, object, or set", collection));
+                false
+                //bail!(collection_expr.span().error("collection must be array, object or set"));
             }
         };
 
@@ -1231,6 +1320,8 @@ impl<'source> Interpreter<'source> {
             key_expr: None,
             output_expr: Some(term),
             value: Value::new_array(),
+            result: None,
+            results: QueryResults::default(),
         });
 
         // Evaluate body first.
@@ -1252,6 +1343,8 @@ impl<'source> Interpreter<'source> {
             key_expr: None,
             output_expr: Some(term),
             value: Value::new_set(),
+            result: None,
+            results: QueryResults::default(),
         });
 
         self.eval_query(query)?;
@@ -1273,6 +1366,8 @@ impl<'source> Interpreter<'source> {
             key_expr: Some(key),
             output_expr: Some(value),
             value: Value::new_object(),
+            result: None,
+            results: QueryResults::default(),
         });
 
         self.eval_query(query)?;
@@ -1400,6 +1495,8 @@ impl<'source> Interpreter<'source> {
                 key_expr: None,
                 output_expr,
                 value: Value::new_set(),
+                result: None,
+                results: QueryResults::default(),
             };
 
             // Back up local variables of current function and empty
@@ -1646,6 +1743,8 @@ impl<'source> Interpreter<'source> {
                         key_expr,
                         output_expr,
                         value,
+                        result: None,
+                        results: QueryResults::default(),
                     },
                     path,
                 ))
@@ -1657,6 +1756,8 @@ impl<'source> Interpreter<'source> {
                         key_expr: None,
                         output_expr: key.as_ref(),
                         value: Value::new_set(),
+                        result: None,
+                        results: QueryResults::default(),
                     },
                     path,
                 ))
@@ -1771,35 +1872,11 @@ impl<'source> Interpreter<'source> {
         }
     }
 
-    pub fn merge_value(span: &Span<'source>, value: &mut Value, mut new: Value) -> Result<()> {
-        match (value, &mut new) {
-            (v @ Value::Undefined, _) => *v = new,
-            (Value::Set(ref mut set), Value::Set(new)) => {
-                Rc::make_mut(set).append(Rc::make_mut(new))
-            }
-            (Value::Object(map), Value::Object(new)) => {
-                for (k, v) in new.iter() {
-                    match map.get(k) {
-                        Some(pv) if *pv != *v => {
-                            return Err(span.source.error(
-                                span.line,
-                                span.col,
-                                format!(
-                                    "value for key `{}` generated multiple times: `{}` and `{}`",
-                                    serde_json::to_string_pretty(&k)?,
-                                    serde_json::to_string_pretty(&pv)?,
-                                    serde_json::to_string_pretty(&v)?,
-                                )
-                                .as_str(),
-                            ));
-                        }
-                        _ => Rc::make_mut(map).insert(k.clone(), v.clone()),
-                    };
-                }
-            }
-            _ => bail!("internal error: could not merge value"),
-        };
-        Ok(())
+    pub fn merge_value(span: &Span<'source>, value: &mut Value, new: Value) -> Result<()> {
+        match value.merge(new) {
+            Ok(()) => Ok(()),
+            Err(err) => return Err(span.error(format!("{err}").as_str())),
+        }
     }
 
     pub fn get_path_string(refr: &Expr, document: Option<&str>) -> Result<String> {
@@ -2187,11 +2264,12 @@ impl<'source> Interpreter<'source> {
         self.eval_modules(input, enable_tracing)
     }
 
-    pub fn eval_query_snippet(
+    pub fn eval_user_query(
         &mut self,
-        snippet: &'source Expr<'source>,
+        query: &'source Query<'source>,
+        order: &[u16],
         enable_tracing: bool,
-    ) -> Result<Value> {
+    ) -> Result<QueryResults> {
         self.traces = match enable_tracing {
             true => Some(vec![]),
             false => None,
@@ -2200,40 +2278,28 @@ impl<'source> Interpreter<'source> {
         // Create a new scope for evaluating the expression.
         self.scopes.push(Scope::new());
         let prev_module = self.set_current_module(self.modules.last().copied())?;
-        let value = self.eval_expr(snippet)?;
-        // Pop the scope.
 
-        let scope = self.scopes.pop();
-        let r = match scope {
-            Some(scope) if !scope.is_empty() => {
-                let mut r = Value::new_object();
-                let map = r.as_object_mut()?;
-                // Capture each binding.
-                for (name, v) in scope {
-                    map.insert(Value::String(name), v);
-                }
-                Ok(r)
-            }
-            _ => Ok(value),
-        };
-        /*        let r = match snippet {
-            Expr::AssignExpr { .. } => {
-                if let Some(scope) = scope {
-                    let mut r = Value::new_object();
-                    let map = r.as_object_mut()?;
-                    // Capture each binding.
-                    for (name, v) in scope {
-                        map.insert(Value::String(name), v);
-                    }
-                    Ok(r)
-                } else {
-                    bail!("internal error: expression scope not found");
-                }
-            }
-            _ => Ok(value),
-        };*/
+        // Push new context.
+        self.contexts.push(Context {
+            key_expr: None,
+            output_expr: None,
+            value: Value::new_set(),
+            // Request that results be gathered.
+            result: Some(QueryResult::default()),
+            results: QueryResults::default(),
+        });
+
+        let ordered_stmts: Vec<&'source LiteralStmt<'source>> =
+            order.iter().map(|i| &query.stmts[*i as usize]).collect();
+        let _value = self.eval_stmts(&ordered_stmts);
+
+        // Pop the scope.
+        let _scope = self.scopes.pop();
         self.set_current_module(prev_module)?;
-        r
+        match self.contexts.pop() {
+            Some(ctx) => Ok(ctx.results),
+            _ => bail!("internal error: no context"),
+        }
     }
 
     fn gather_rules(&mut self) -> Result<()> {

src/parser.rs

@@ -972,7 +972,9 @@ impl<'source> Parser<'source> {
             literals.push(stmt);
         }
 
-        self.expect(end_delim, "while parsing query")?;
+        if !end_delim.is_empty() {
+            self.expect(end_delim, "while parsing query")?;
+        }
         span.end = self.end;
         Ok(Query {
             span,

src/scheduler.rs

@@ -404,6 +404,33 @@ impl<'a> Analyzer<'a> {
     }
 
     pub fn analyze(mut self, modules: &'a [Module<'a>]) -> Result<Schedule> {
+        self.add_rules(modules)?;
+
+        for m in modules {
+            self.analyze_module(m)?;
+        }
+
+        Ok(Schedule {
+            scopes: self.locals,
+            order: self.order,
+        })
+    }
+
+    pub fn analyze_query_snippet(
+        mut self,
+        modules: &'a [Module<'a>],
+        query: &'a Query<'a>,
+    ) -> Result<Vec<u16>> {
+        self.add_rules(modules)?;
+        self.analyze_query(None, None, query, Scope::default())?;
+        Ok(self
+            .order
+            .get(query)
+            .expect("could not schedule user query")
+            .clone())
+    }
+
+    fn add_rules(&mut self, modules: &'a [Module<'a>]) -> Result<()> {
         for m in modules {
             let path = utils::get_path_string(&m.package.refr, Some("data"))?;
             let scope: &mut Scope = self.packages.entry(path).or_default();
@@ -422,14 +449,7 @@ impl<'a> Analyzer<'a> {
             }
         }
 
-        for m in modules {
-            self.analyze_module(m)?;
-        }
-
-        Ok(Schedule {
-            scopes: self.locals,
-            order: self.order,
-        })
+        Ok(())
     }
 
     fn analyze_module(&mut self, m: &'a Module<'a>) -> Result<()> {

src/value.rs

@@ -192,6 +192,10 @@ impl Value {
         matches!(self, Value::Null)
     }
 
+    pub fn is_empty_object(&self) -> bool {
+        self == &Value::new_object()
+    }
+
     pub fn as_bool(&self) -> Result<&bool> {
         match self {
             Value::Bool(b) => Ok(b),
@@ -307,6 +311,32 @@ impl Value {
             _ => bail!("internal error: make: not an selfect {self:?}"),
         }
     }
+
+    pub fn merge(&mut self, mut new: Value) -> Result<()> {
+        match (self, &mut new) {
+            (v @ Value::Undefined, _) => *v = new,
+            (Value::Set(ref mut set), Value::Set(new)) => {
+                Rc::make_mut(set).append(Rc::make_mut(new))
+            }
+            (Value::Object(map), Value::Object(new)) => {
+                for (k, v) in new.iter() {
+                    match map.get(k) {
+                        Some(pv) if *pv != *v => {
+                            bail!(
+                                "value for key `{}` generated multiple times: `{}` and `{}`",
+                                serde_json::to_string_pretty(&k)?,
+                                serde_json::to_string_pretty(&pv)?,
+                                serde_json::to_string_pretty(&v)?,
+                            )
+                        }
+                        _ => Rc::make_mut(map).insert(k.clone(), v.clone()),
+                    };
+                }
+            }
+            _ => bail!("internal error: could not merge value"),
+        };
+        Ok(())
+    }
 }
 impl ops::Index<usize> for Value {
     type Output = Value;

tests/interpreter/mod.rs

@@ -181,6 +181,20 @@ pub fn check_output(computed_results: &[Value], expected_results: &[Value]) -> R
     Ok(())
 }
 
+fn query_results_to_value(query_results: QueryResults) -> Result<Value> {
+    if let Some(query_result) = query_results.results.last() {
+        if !query_result.bindings.is_empty_object() {
+            return Ok(query_result.bindings.clone());
+        } else {
+            return match query_result.expressions.last() {
+                Some(v) => Ok(v.clone()),
+                _ => bail!("no expressions in query results"),
+            };
+        }
+    }
+    bail!("query result incomplete")
+}
+
 pub fn eval_file_first_rule(
     regos: &[String],
     data_opt: Option<Value>,
@@ -194,15 +208,21 @@ pub fn eval_file_first_rule(
     let mut modules = vec![];
     let mut modules_ref = vec![];
 
-    // the query is parsed for later
-    let source = Source {
+    let query_source = regorus::Source {
         file: "<query.rego>",
         contents: query,
         lines: query.split('\n').collect(),
     };
-    let mut parser = Parser::new(&source)?;
-    let expr = parser.parse_assign_expr()?;
-
+    let query_span = regorus::Span {
+        source: &query_source,
+        line: 1,
+        col: 1,
+        start: 0,
+        end: query.len() as u16,
+    };
+    let mut parser = regorus::Parser::new(&query_source)?;
+    let query_node = parser.parse_query(query_span, "")?;
+    let query_stmt_order = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?;
     for (idx, _) in regos.iter().enumerate() {
         files.push(format!("rego_{idx}"));
     }
@@ -248,14 +268,22 @@ pub fn eval_file_first_rule(
             }
 
             // Now eval the query.
-            results.push(interpreter.eval_query_snippet(&expr, enable_tracing)?);
+            results.push(query_results_to_value(interpreter.eval_user_query(
+                &query_node,
+                &query_stmt_order,
+                enable_tracing,
+            )?)?);
         }
     } else {
         // it no input is defined then one evaluation of all modules is performed
         interpreter.eval(&data_opt, &None, enable_tracing, Some(&schedule))?;
 
         // Now eval the query.
-        results.push(interpreter.eval_query_snippet(&expr, enable_tracing)?);
+        results.push(query_results_to_value(interpreter.eval_user_query(
+            &query_node,
+            &query_stmt_order,
+            enable_tracing,
+        )?)?);
     }
 
     Ok(results)
@@ -274,14 +302,21 @@ pub fn eval_file(
     let mut modules = vec![];
     let mut modules_ref = vec![];
 
-    // the query is parsed for later
-    let source = Source {
+    let query_source = regorus::Source {
         file: "<query.rego>",
         contents: query,
         lines: query.split('\n').collect(),
     };
-    let mut parser = Parser::new(&source)?;
-    let expr = parser.parse_assign_expr()?;
+    let query_span = regorus::Span {
+        source: &query_source,
+        line: 1,
+        col: 1,
+        start: 0,
+        end: query.len() as u16,
+    };
+    let mut parser = regorus::Parser::new(&query_source)?;
+    let query_node = parser.parse_query(query_span, "")?;
+    let query_stmt_order = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?;
 
     for (idx, _) in regos.iter().enumerate() {
         files.push(format!("rego_{idx}"));
@@ -324,14 +359,22 @@ pub fn eval_file(
             interpreter.eval_modules(&Some(input), enable_tracing)?;
 
             // Now eval the query.
-            results.push(interpreter.eval_query_snippet(&expr, enable_tracing)?);
+            results.push(query_results_to_value(interpreter.eval_user_query(
+                &query_node,
+                &query_stmt_order,
+                enable_tracing,
+            )?)?);
         }
     } else {
         // it no input is defined then one evaluation of all modules is performed
         interpreter.eval(&data_opt, &None, enable_tracing, Some(&schedule))?;
 
         // Now eval the query.
-        results.push(interpreter.eval_query_snippet(&expr, enable_tracing)?);
+        results.push(query_results_to_value(interpreter.eval_user_query(
+            &query_node,
+            &query_stmt_order,
+            enable_tracing,
+        )?)?);
     }
 
     Ok(results)
@@ -380,6 +423,7 @@ fn one_file() -> Result<()> {
     interpreter.prepare_for_eval(Some(&schedule), &None)?;
     let results = interpreter.eval_modules(&input, true)?;
     println!("eval results:\n{}", serde_json::to_string_pretty(&results)?);
+
     Ok(())
 }
 
@@ -421,7 +465,7 @@ impl<'de> Deserialize<'de> for ValueOrVec {
 
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
 struct TestCase {
-    data: Value,
+    data: Option<Value>,
     input: Option<ValueOrVec>,
     modules: Vec<String>,
     note: String,
@@ -431,6 +475,8 @@ struct TestCase {
     skip: Option<bool>,
     error: Option<String>,
     traces: Option<bool>,
+    want_error: Option<String>,
+    want_error_code: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
@@ -453,6 +499,7 @@ fn yaml_test_impl(file: &str, is_opa_test: bool) -> Result<()> {
 
         match (&case.want_result, &case.error) {
             (Some(_), None) | (None, Some(_)) => (),
+            _ if is_opa_test => (),
             _ => panic!("either want_result or error must be specified in test case."),
         }
 
@@ -460,7 +507,7 @@ fn yaml_test_impl(file: &str, is_opa_test: bool) -> Result<()> {
 
         match eval_file(
             &case.modules,
-            Some(case.data),
+            case.data,
             case.input,
             case.query.as_str(),
             enable_tracing,
@@ -479,13 +526,20 @@ fn yaml_test_impl(file: &str, is_opa_test: bool) -> Result<()> {
                         // Convert value to json compatible representation.
                         let results =
                             Value::from_json_str(serde_json::to_string(&results)?.as_str())?;
+                        dbg!((&results, &expected_results[0]));
                         match_values(&results, &expected_results[0])?;
                     } else {
                         check_output(&results, &expected_results)?;
                     }
                 }
-                _ => panic!("eval succeeded and did not produce any errors"),
+                _ => bail!("eval succeeded and did not produce any errors"),
             },
+            Err(actual) if is_opa_test => {
+                if case.want_error.is_none() && case.want_error_code.is_none() {
+                    return Err(actual);
+                }
+                // opa test expects execution to fail and it did.
+            }
             Err(actual) => match &case.error {
                 Some(expected) => {
                     let actual = actual.to_string();
@@ -564,7 +618,7 @@ fn run_opa_tests() -> Result<()> {
             .filter_map(|e| e.ok())
         {
             let path = entry.path().to_string_lossy().to_string();
-            if Path::new(&path).is_dir() {
+            if !Path::new(&path).is_file() || !path.ends_with(".yaml") {
                 continue;
             }
             let yaml = path;