mirror of
https://github.com/microsoft/regorus.git
synced 2026-08-05 02:16:11 +00:00
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:
committed by
GitHub
parent
519cce5b33
commit
6228eaab4a
@@ -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)]
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"),
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user