Add APIs to get schema annotations (#1389)

Signed-off-by: Shaobo He <shaobohe@amazon.com>
Co-authored-by: John Kastner <130772734+john-h-kastner-aws@users.noreply.github.com>
Co-authored-by: Craig Disselkoen <cdiss@amazon.com>

cedar-policy/src/api.rs

@@ -30,6 +30,7 @@ use cedar_policy_validator::entity_manifest;
 pub use cedar_policy_validator::entity_manifest::{
     AccessTrie, EntityManifest, EntityRoot, Fields, RootAccessTrie,
 };
+use cedar_policy_validator::json_schema;
 use cedar_policy_validator::typecheck::{PolicyCheck, Typechecker};
 pub use id::*;
 
@@ -1397,7 +1398,172 @@ pub struct SchemaFragment {
     lossless: cedar_policy_validator::json_schema::Fragment<cedar_policy_validator::RawName>,
 }
 
+fn get_annotation_by_key(
+    annotations: &est::Annotations,
+    annotation_key: impl AsRef<str>,
+) -> Option<&str> {
+    annotations
+        .0
+        .get(&annotation_key.as_ref().parse().ok()?)
+        .map(|value| annotation_value_to_str_ref(value.as_ref()))
+}
+
+fn annotation_value_to_str_ref(value: Option<&ast::Annotation>) -> &str {
+    value.map_or("", |a| a.as_ref())
+}
+
+fn annotations_to_pairs(annotations: &est::Annotations) -> impl Iterator<Item = (&str, &str)> {
+    annotations
+        .0
+        .iter()
+        .map(|(key, value)| (key.as_ref(), annotation_value_to_str_ref(value.as_ref())))
+}
+
 impl SchemaFragment {
+    /// Get annotations of a non-empty namespace.
+    ///
+    /// We do not allow namespace-level annotations on the empty namespace.
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`]
+    pub fn namespace_annotations(
+        &self,
+        namespace: EntityNamespace,
+    ) -> Option<impl Iterator<Item = (&str, &str)>> {
+        self.lossless
+            .0
+            .get(&Some(namespace.0))
+            .map(|ns_def| annotations_to_pairs(&ns_def.annotations))
+    }
+
+    /// Get annotation value of a non-empty namespace by annotation key
+    /// `annotation_key`
+    ///
+    /// We do not allow namespace-level annotations on the empty namespace.
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`]
+    /// or `annotation_key` is not a valid annotation key
+    /// or it does not exist
+    pub fn namespace_annotation(
+        &self,
+        namespace: EntityNamespace,
+        annotation_key: impl AsRef<str>,
+    ) -> Option<&str> {
+        let ns = self.lossless.0.get(&Some(namespace.0))?;
+        get_annotation_by_key(&ns.annotations, annotation_key)
+    }
+
+    /// Get annotations of a common type declaration
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`] or
+    /// `ty` is not a valid common type ID or `ty` is not found in the
+    /// corresponding namespace definition
+    pub fn common_type_annotations(
+        &self,
+        namespace: Option<EntityNamespace>,
+        ty: &str,
+    ) -> Option<impl Iterator<Item = (&str, &str)>> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        let ty = json_schema::CommonTypeId::new(ast::UnreservedId::from_normalized_str(ty).ok()?)
+            .ok()?;
+        ns_def
+            .common_types
+            .get(&ty)
+            .map(|ty| annotations_to_pairs(&ty.annotations))
+    }
+
+    /// Get annotation value of a common type declaration by annotation key
+    /// `annotation_key`
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`]
+    /// or `ty` is not a valid common type ID
+    /// or `ty` is not found in the corresponding namespace definition
+    /// or `annotation_key` is not a valid annotation key
+    /// or it does not exist
+    pub fn common_type_annotation(
+        &self,
+        namespace: Option<EntityNamespace>,
+        ty: &str,
+        annotation_key: impl AsRef<str>,
+    ) -> Option<&str> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        let ty = json_schema::CommonTypeId::new(ast::UnreservedId::from_normalized_str(ty).ok()?)
+            .ok()?;
+        get_annotation_by_key(&ns_def.common_types.get(&ty)?.annotations, annotation_key)
+    }
+
+    /// Get annotations of an entity type declaration
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`] or
+    /// `ty` is not a valid entity type name or `ty` is not found in the
+    /// corresponding namespace definition
+    pub fn entity_type_annotations(
+        &self,
+        namespace: Option<EntityNamespace>,
+        ty: &str,
+    ) -> Option<impl Iterator<Item = (&str, &str)>> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        let ty = ast::UnreservedId::from_normalized_str(ty).ok()?;
+        ns_def
+            .entity_types
+            .get(&ty)
+            .map(|ty| annotations_to_pairs(&ty.annotations))
+    }
+
+    /// Get annotation value of an entity type declaration by annotation key
+    /// `annotation_key`
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`]
+    /// or `ty` is not a valid entity type name
+    /// or `ty` is not found in the corresponding namespace definition
+    /// or `annotation_key` is not a valid annotation key
+    /// or it does not exist
+    pub fn entity_type_annotation(
+        &self,
+        namespace: Option<EntityNamespace>,
+        ty: &str,
+        annotation_key: impl AsRef<str>,
+    ) -> Option<&str> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        let ty = ast::UnreservedId::from_normalized_str(ty).ok()?;
+        get_annotation_by_key(&ns_def.entity_types.get(&ty)?.annotations, annotation_key)
+    }
+
+    /// Get annotations of an action declaration
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`] or
+    /// `id` is not found in the corresponding namespace definition
+    pub fn action_annotations(
+        &self,
+        namespace: Option<EntityNamespace>,
+        id: EntityId,
+    ) -> Option<impl Iterator<Item = (&str, &str)>> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        ns_def
+            .actions
+            .get(id.as_ref())
+            .map(|a| annotations_to_pairs(&a.annotations))
+    }
+
+    /// Get annotation value of an action declaration by annotation key
+    /// `annotation_key`
+    ///
+    /// Returns `None` if `namespace` is not found in the [`SchemaFragment`]
+    /// or `id` is not found in the corresponding namespace definition
+    /// or `annotation_key` is not a valid annotation key
+    /// or it does not exist
+    pub fn action_annotation(
+        &self,
+        namespace: Option<EntityNamespace>,
+        id: EntityId,
+        annotation_key: impl AsRef<str>,
+    ) -> Option<&str> {
+        let ns_def = self.lossless.0.get(&namespace.map(|n| n.0))?;
+        get_annotation_by_key(
+            &ns_def.actions.get(id.as_ref())?.annotations,
+            annotation_key,
+        )
+    }
+
     /// Extract namespaces defined in this [`SchemaFragment`].
     ///
     /// `None` indicates the empty namespace.