fix: Merge data to init document

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,16 +216,20 @@ impl Interpreter {
         self.modules = modules.to_vec();
     }
 
-    pub fn set_init_data(&mut self, init_data: Value) {
-        self.init_data = init_data;
+    pub fn get_data_mut(&mut self) -> &mut Value {
+        &mut self.data
     }
 
-    pub fn set_data(&mut self, data: Value) {
-        self.data = data;
+    pub fn set_init_data(&mut self, data: Value) {
+        self.init_data = data;
     }
 
-    pub fn get_data_mut(&mut self) -> &mut Value {
-        &mut self.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) {

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,12 +172,71 @@ 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(
         &Path::new(&cli.test_dir),
         &cli.name,
         cli.coverage,
         cli.generate,
-    )
+    )?;
+    stateful_policy_test()
 }