Order query expression results (#32)

Expressions are scheduled based on dependencies and thusthe gathered
expression values may not be in the same order as in source.

Reorder to match the source.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2023-10-31 13:26:33 -07:00
committed by GitHub
parent 519cce5b33
commit 6228eaab4a
3 changed files with 66 additions and 35 deletions

View File

@@ -7,7 +7,7 @@ use clap::{Parser, Subcommand};
fn rego_eval(
files: &[String],
input: Option<String>,
query: Option<String>,
query: String,
enable_tracing: bool,
) -> Result<()> {
// User specified data.
@@ -83,16 +83,10 @@ fn rego_eval(
// Evaluate all the modules.
interpreter.eval(&Some(data), &input, false, Some(schedule))?;
// Fetch query string. If none specified, use "data".
let query = match &query {
Some(query) => query,
_ => "data",
};
// Parse the query.
let query_source = regorus::Source {
file: "<query.rego>",
contents: query,
contents: &query,
lines: query.split('\n').collect(),
};
let query_span = regorus::Span {
@@ -107,7 +101,7 @@ fn rego_eval(
let query_schedule = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?;
let results = interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?;
println!("eval results:\n{}", serde_json::to_string_pretty(&results)?);
println!("{}", serde_json::to_string_pretty(&results)?);
Ok(())
}
@@ -168,12 +162,7 @@ enum RegorusCommand {
/// Evaluate a Rego Query.
Eval {
/// Policy or data files. Rego, json or yaml.
#[arg(
required(true),
long,
short,
value_name = "policy.rego|data.json|data.yaml"
)]
#[arg(long, short, value_name = "policy.rego|data.json|data.yaml")]
data: Vec<String>,
/// Input file. json or yaml.
@@ -181,7 +170,7 @@ enum RegorusCommand {
input: Option<String>,
/// Query. Rego query block.
query: Option<String>,
query: String,
/// Enable tracing.
#[arg(long, short)]

View File

@@ -59,7 +59,7 @@ impl Default for QueryResult {
#[derive(Debug, Clone, Default, Serialize)]
pub struct QueryResults {
pub results: Vec<QueryResult>,
pub result: Vec<QueryResult>,
}
#[derive(Debug, Clone)]
@@ -782,6 +782,27 @@ impl<'source> Interpreter<'source> {
Ok(count > 0)
}
fn make_expression_result(span: &Span, v: &Value) -> Value {
let mut loc = BTreeMap::new();
loc.insert(
Value::String("row".to_string()),
Value::from_float(span.line as f64),
);
loc.insert(
Value::String("col".to_string()),
Value::from_float(span.col as f64),
);
let mut expr = BTreeMap::new();
expr.insert(Value::String("value".to_string()), v.clone());
expr.insert(Value::String("location".to_string()), Value::from_map(loc));
expr.insert(
Value::String("text".to_string()),
Value::String(span.text().to_string()),
);
Value::from_map(expr)
}
fn eval_stmt(
&mut self,
stmt: &'source LiteralStmt<'source>,
@@ -820,7 +841,7 @@ impl<'source> Interpreter<'source> {
}
let r = Ok(match &stmt.literal {
Literal::Expr { expr, .. } => {
Literal::Expr { span, expr, .. } => {
let value = match expr {
Expr::Call { span, fcn, params } => self.eval_call(
span,
@@ -834,7 +855,9 @@ impl<'source> Interpreter<'source> {
if let Some(ctx) = self.contexts.last_mut() {
if let Some(result) = &mut ctx.result {
result.expressions.push(value.clone());
result
.expressions
.push(Self::make_expression_result(span, &value))
}
}
@@ -847,7 +870,7 @@ impl<'source> Interpreter<'source> {
value != Value::Undefined
}
}
Literal::NotExpr { expr, .. } => {
Literal::NotExpr { span, expr, .. } => {
let value = match expr {
// Extra parameter is allowed; but a return argument is not allowed.
Expr::Call { span, fcn, params } => self.eval_call(
@@ -861,13 +884,15 @@ 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));
result
.expressions
.push(Self::make_expression_result(span, &Value::Bool(true)))
}
}
// https://github.com/open-policy-agent/opa/issues/1622#issuecomment-520547385
matches!(value, Value::Bool(false) | Value::Undefined)
}
Literal::SomeVars { vars, .. } => {
Literal::SomeVars { span, vars, .. } => {
for var in vars {
let name = var.text();
if let Ok(variable) = self.add_variable_or(name) {
@@ -881,7 +906,9 @@ 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));
result
.expressions
.push(Self::make_expression_result(span, &Value::Bool(true)))
}
}
true
@@ -894,7 +921,9 @@ 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));
result
.expressions
.push(Self::make_expression_result(span, &Value::Bool(true)))
}
}
self.eval_some_in(span, key, value, collection, stmts)?
@@ -908,7 +937,9 @@ 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));
result
.expressions
.push(Self::make_expression_result(span, &Value::Bool(true)))
}
}
self.eval_every(span, key, value, domain, query)?
@@ -1094,7 +1125,7 @@ impl<'source> Interpreter<'source> {
.insert(Value::String(name.to_string()), value.clone());
}
}
ctx.results.results.push(result);
ctx.results.result.push(result);
}
return Ok(true);
@@ -2300,22 +2331,33 @@ impl<'source> Interpreter<'source> {
// Eval the query.
let query_r = self.eval_query(query);
let mut results = match self.contexts.pop() {
Some(ctx) => ctx.results,
_ => bail!("internal error: no context"),
};
// Restore schedules.
if let Some(self_schedule) = &mut self.schedule {
for (k, _) in schedule.order.iter() {
for (k, ord) in schedule.order.iter() {
if k == &query {
for idx in 0..results.result.len() {
let mut ordered_expressions = vec![Value::Undefined; ord.len()];
for (expr_idx, value) in results.result[idx].expressions.iter().enumerate()
{
let orig_idx = ord[expr_idx] as usize;
ordered_expressions[orig_idx] = value.clone();
}
results.result[idx].expressions = ordered_expressions;
}
}
self_schedule.order.remove(k);
}
}
self.set_current_module(prev_module)?;
let r = match self.contexts.pop() {
Some(ctx) => Ok(ctx.results),
_ => bail!("internal error: no context"),
};
match query_r {
Ok(_) => r,
Ok(_) => Ok(results),
Err(e) => Err(e),
}
}

View File

@@ -182,12 +182,12 @@ pub fn check_output(computed_results: &[Value], expected_results: &[Value]) -> R
}
fn query_results_to_value(query_results: QueryResults) -> Result<Value> {
if let Some(query_result) = query_results.results.last() {
if let Some(query_result) = query_results.result.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()),
Some(v) => Ok(v["value"].clone()),
_ => bail!("no expressions in query results"),
};
}