adds upsert_entities

cedar-policy/CHANGELOG.md

@@ -22,6 +22,7 @@ Cedar Language Version: TBD
 - Added `Entities::remove_entities()` to remove `Entity`s from an `Entities` struct (resolving #701)
 - Added `PolicySet::merge_policyset()` to merge a `PolicySet` into another `PolicySet` struct (resolving #610)
 - Implemented [RFC 53 (enumerated entity types)](https://github.com/cedar-policy/rfcs/blob/main/text/0053-enum-entities.md)  (#1377)
+- Added `Entities::upsert_entities()` to add or update `Entity`s in an `Entities` struct (resolving #1479)
 
 ### Fixed
 

cedar-policy/src/api.rs

@@ -465,6 +465,43 @@ impl Entities {
         )?))
     }
 
+    /// Updates or adds all of the [`Entity`]s in the collection to this [`Entities`]
+    /// structure, re-computing the transitive closure.
+    ///
+    /// If a `schema` is provided, this method will ensure that the added
+    /// entities fully conform to the schema -- for instance, it will error if
+    /// attributes have the wrong types (e.g., string instead of integer), or if
+    /// required attributes are missing or superfluous attributes are provided.
+    /// (This method will not add action entities from the `schema`.)
+    ///
+    /// Re-computing the transitive closure can be expensive, so it is advised
+    /// to not call this method in a loop.
+    /// ## Errors
+    /// - [`EntitiesError::InvalidEntity`] if `schema` is not none and any entities do not conform
+    ///   to the schema
+    pub fn upsert_entities(
+        self,
+        entities: impl IntoIterator<Item = Entity>,
+        schema: Option<&Schema>,
+    ) -> Result<Self, EntitiesError> {
+        let entities: Vec<_> = entities.into_iter().collect();
+
+        let removed = self.0.remove_entities(
+            entities.iter().map(|e| e.uid().0),
+            cedar_policy_core::entities::TCComputation::AssumeAlreadyComputed
+        )?;
+
+        Ok(Self(
+            removed.add_entities(
+                entities.into_iter().map(|e| Arc::new(e.0)), // Second use
+                schema.map(|s| cedar_policy_validator::CoreSchema::new(&s.0))
+                    .as_ref(),
+                cedar_policy_core::entities::TCComputation::ComputeNow,
+                Extensions::all_available(),
+            )?
+        ))
+    }
+
     /// Parse an entities JSON file (in [&str] form) and add them into this
     /// [`Entities`] structure, re-computing the transitive closure
     ///