fix: Merge data to init document (#293)

Init document is the aggregated data documen that the user has
specified using multiple `add_data` calls. Each query evaluation
starts of by initializing the current data to the init document.

Previously `add_data` was incorrectly added to the current document,
causing the added data to be lost if the addition happened after query
evaluation.

With this fix, scenarios where data addition may be interspersed with
query evaluation calls are supported.

Also provide a get_data method to obtain the (init) data document.

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

src/engine.rs

@@ -239,7 +239,7 @@ impl Engine {
     /// # }
     /// ```
     pub fn clear_data(&mut self) {
-        self.interpreter.set_data(Value::new_object());
+        self.interpreter.set_init_data(Value::new_object());
         self.prepared = false;
     }
 
@@ -276,7 +276,40 @@ impl Engine {
             bail!("data must be object");
         }
         self.prepared = false;
-        self.interpreter.get_data_mut().merge(data)
+        self.interpreter.get_init_data_mut().merge(data)
+    }
+
+    /// Get the data document.
+    ///
+    /// The returned value is the data document that has been constructed using
+    /// one or more calls to [`Engine::add_data`]. The values of policy rules are
+    /// not included in the returned document.
+    ///
+    ///
+    /// ```
+    /// # use regorus::*;
+    /// # fn main() -> anyhow::Result<()> {
+    /// let mut engine = Engine::new();
+    ///
+    /// // If not set, data document is empty.
+    /// assert_eq!(engine.get_data(), Value::new_object());
+    ///
+    /// // Merge { "x" : 1, "y" : {} }
+    /// assert!(engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?).is_ok());
+    ///
+    /// // Merge { "z" : 2 }
+    /// assert!(engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?).is_ok());
+    ///
+    /// let data = engine.get_data();
+    /// assert_eq!(data["x"], Value::from(1));
+    /// assert_eq!(data["y"], Value::new_object());
+    /// assert_eq!(data["z"], Value::from(2));
+    ///
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_data(&self) -> Value {
+        self.interpreter.get_init_data().clone()
     }
 
     pub fn add_data_json(&mut self, data_json: &str) -> Result<()> {
@@ -537,11 +570,7 @@ impl Engine {
             self.interpreter.set_modules(&self.modules);
 
             self.interpreter.clear_builtins_cache();
-            // when the interpreter is prepared the initial data is saved
-            // the data will be reset to init_data each time clean_internal_evaluation_state is called
-            let init_data = self.interpreter.get_data_mut().clone();
-            self.interpreter.set_init_data(init_data);
-
+            // clean_internal_evaluation_state will set data to an efficient clont of use supplied init_data
             // Initialize the with-document with initial data values.
             // with-modifiers will be applied to this document.
             self.interpreter.init_with_document()?;

src/interpreter.rs

@@ -216,18 +216,22 @@ impl Interpreter {
         self.modules = modules.to_vec();
     }
 
-    pub fn set_init_data(&mut self, init_data: Value) {
-        self.init_data = init_data;
-    }
-
-    pub fn set_data(&mut self, data: Value) {
-        self.data = data;
-    }
-
     pub fn get_data_mut(&mut self) -> &mut Value {
         &mut self.data
     }
 
+    pub fn set_init_data(&mut self, data: Value) {
+        self.init_data = data;
+    }
+
+    pub fn get_init_data(&self) -> &Value {
+        &self.init_data
+    }
+
+    pub fn get_init_data_mut(&mut self) -> &mut Value {
+        &mut self.init_data
+    }
+
     pub fn set_traces(&mut self, enable_tracing: bool) {
         self.traces = match enable_tracing {
             true => Some(vec![]),

src/tests/interpreter/mod.rs

@@ -422,3 +422,32 @@ fn one_yaml() -> Result<()> {
 fn run(path: &str) {
     yaml_test(path).unwrap()
 }
+
+#[test]
+fn test_get_data() -> Result<()> {
+    let mut engine = Engine::new();
+
+    // Merge { "x" : 1, "y" : {} }
+    engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?)?;
+
+    // Merge { "z" : 2 }
+    engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?)?;
+
+    // Add a policy
+    engine.add_policy("policy.rego".to_string(), "package a".to_string())?;
+
+    // Evaluate virtual data document. The virtual document includes all rules as well.
+    let v_data = engine.eval_query("data".to_string(), false)?.result[0].expressions[0]
+        .value
+        .clone();
+    // There must be an empty package.
+    assert_eq!(v_data["a"], Value::new_object());
+
+    // Get the data document.
+    let data = engine.get_data();
+
+    // There must NOT be any value of `a`.
+    assert_eq!(data["a"], Value::Undefined);
+
+    Ok(())
+}

tests/kata/main.rs

@@ -172,6 +172,64 @@ struct Cli {
     generate: bool,
 }
 
+fn stateful_policy_test() -> Result<()> {
+    // Create an engine for evaluating Rego policies.
+    let mut engine = regorus::Engine::new();
+
+    let policy = String::from(
+        r#"
+       package example
+       import rego.v1
+
+       default allow := false
+
+       allow if {
+         print("data.allowed_actions = ", data.allowed_actions)
+         input.action in data.allowed_actions["user1"]
+         print("This rule should be allowed")
+       }
+	"#,
+    );
+
+    // Add policy to the engine.
+    engine.add_policy(String::from("policy.rego"), policy)?;
+
+    // Evaluate first input. Expect to evaluate to false, since state is not set
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_bool_query(String::from("data.example.allow"), false)?;
+    println!("Received result: {:?}", r);
+    assert_eq!(r, false);
+
+    // Add data to engine. Set state
+    engine.add_data(regorus::Value::from_json_str(
+        r#"{
+     "allowed_actions": {
+        "user1" : ["read", "write"]
+     }}"#,
+    )?)?;
+
+    // Evaluate second input. Expect to evaluate to true, since state has been set now
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_bool_query(String::from("data.example.allow"), false)?;
+    println!("Received result: {:?}", r);
+    assert_eq!(
+        r, true,
+        "expect result to be true since rule evaluates to true after state has been updated, per rego logs"
+    );
+
+    Ok(())
+}
+
 fn main() -> Result<()> {
     let cli = Cli::parse();
     run_kata_tests(
@@ -179,5 +237,6 @@ fn main() -> Result<()> {
         &cli.name,
         cli.coverage,
         cli.generate,
-    )
+    )?;
+    stateful_policy_test()
 }