Merge pull request #3355 from albinsuresh/feat/3339/entity-store-query-api

feat: REST API for querying entity store

Author: albinsuresh

crates/core/tedge_agent/src/entity_manager/server.rs

@@ -21,6 +21,7 @@ pub enum EntityStoreRequest {
     Get(EntityTopicId),
     Create(EntityRegistrationMessage),
     Delete(EntityTopicId),
+    List(Option<EntityTopicId>),
     MqttMessage(MqttMessage),
 }
 
@@ -29,6 +30,7 @@ pub enum EntityStoreResponse {
     Get(Option<EntityMetadata>),
     Create(Result<Vec<RegisteredEntityData>, entity_store::Error>),
     Delete(Vec<EntityTopicId>),
+    List(Result<Vec<EntityMetadata>, entity_store::Error>),
     Ok,
 }
 
@@ -80,6 +82,12 @@ impl Server for EntityStoreServer {
                 let deleted_entities = self.deregister_entity(topic_id).await;
                 EntityStoreResponse::Delete(deleted_entities)
             }
+            EntityStoreRequest::List(topic_id) => {
+                let entities = self.entity_store.list_entity_tree(topic_id.as_ref());
+                EntityStoreResponse::List(
+                    entities.map(|entities| entities.into_iter().cloned().collect()),
+                )
+            }
             EntityStoreRequest::MqttMessage(mqtt_message) => {
                 self.process_mqtt_message(mqtt_message).await;
                 EntityStoreResponse::Ok

crates/core/tedge_agent/src/http_server/entity_store.rs

@@ -54,6 +54,7 @@ impl IntoResponse for Error {
             Error::InvalidEntityTopicId(_) => StatusCode::BAD_REQUEST,
             Error::EntityStoreError(err) => match err {
                 entity_store::Error::EntityAlreadyRegistered(_) => StatusCode::CONFLICT,
+                entity_store::Error::UnknownEntity(_) => StatusCode::NOT_FOUND,
                 _ => StatusCode::BAD_REQUEST,
             },
             Error::EntityNotFound(_) => StatusCode::NOT_FOUND,
@@ -68,7 +69,7 @@ impl IntoResponse for Error {
 
 pub(crate) fn entity_store_router(state: AgentState) -> Router {
     Router::new()
-        .route("/v1/entities", post(register_entity))
+        .route("/v1/entities", post(register_entity).get(list_entities))
         .route(
             "/v1/entities/*path",
             get(get_entity).delete(deregister_entity),
@@ -138,6 +139,22 @@ async fn deregister_entity(
     Ok(Json(deleted))
 }
 
+async fn list_entities(
+    State(state): State<AgentState>,
+) -> Result<Json<Vec<EntityMetadata>>, Error> {
+    let response = state
+        .entity_store_handle
+        .clone()
+        .await_response(EntityStoreRequest::List(None))
+        .await?;
+
+    let EntityStoreResponse::List(entities) = response else {
+        return Err(Error::InvalidEntityStoreResponse);
+    };
+
+    Ok(Json(entities?))
+}
+
 #[cfg(test)]
 mod tests {
     use super::AgentState;
@@ -153,6 +170,7 @@ mod tests {
     use serde_json::json;
     use serde_json::Map;
     use serde_json::Value;
+    use std::collections::HashSet;
     use tedge_actors::Builder;
     use tedge_actors::ClientMessageBox;
     use tedge_actors::MessageReceiver;
@@ -447,6 +465,81 @@ mod tests {
         assert_eq!(response.status(), StatusCode::OK);
     }
 
+    #[tokio::test]
+    async fn entity_list() {
+        let TestHandle {
+            mut app,
+            mut entity_store_box,
+        } = setup();
+
+        // Mock entity store actor response
+        tokio::spawn(async move {
+            if let Some(mut req) = entity_store_box.recv().await {
+                if let EntityStoreRequest::List(_) = req.request {
+                    req.reply_to
+                        .send(EntityStoreResponse::List(Ok(vec![
+                            EntityMetadata::main_device("main".to_string()),
+                            EntityMetadata::child_device("child0".to_string()).unwrap(),
+                            EntityMetadata::child_device("child1".to_string()).unwrap(),
+                        ])))
+                        .await
+                        .unwrap();
+                }
+            }
+        });
+
+        let req = Request::builder()
+            .method(Method::GET)
+            .uri("/v1/entities")
+            .body(Body::empty())
+            .expect("request builder");
+
+        let response = app.call(req).await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
+        let entities: Vec<EntityMetadata> = serde_json::from_slice(&body).unwrap();
+
+        let entity_set = entities
+            .iter()
+            .map(|e| e.topic_id.as_str())
+            .collect::<HashSet<_>>();
+        assert!(entity_set.contains("device/main//"));
+        assert!(entity_set.contains("device/child0//"));
+        assert!(entity_set.contains("device/child1//"));
+    }
+
+    #[tokio::test]
+    async fn entity_list_unknown_entity() {
+        let TestHandle {
+            mut app,
+            mut entity_store_box,
+        } = setup();
+
+        // Mock entity store actor response
+        tokio::spawn(async move {
+            if let Some(mut req) = entity_store_box.recv().await {
+                if let EntityStoreRequest::List(_) = req.request {
+                    req.reply_to
+                        .send(EntityStoreResponse::List(Err(
+                            entity_store::Error::UnknownEntity("unknown".to_string()),
+                        )))
+                        .await
+                        .unwrap();
+                }
+            }
+        });
+
+        let req = Request::builder()
+            .method(Method::GET)
+            .uri("/v1/entities")
+            .body(Body::empty())
+            .expect("request builder");
+
+        let response = app.call(req).await.unwrap();
+        assert_eq!(response.status(), StatusCode::NOT_FOUND);
+    }
+
     struct TestHandle {
         app: Router,
         entity_store_box: ServerMessageBox<EntityStoreRequest, EntityStoreResponse>,

crates/core/tedge_api/src/entity.rs

@@ -61,7 +61,7 @@ impl From<&EntityExternalId> for String {
 pub struct EntityMetadata {
     #[serde(rename = "@topic-id")]
     pub topic_id: EntityTopicId,
-    #[serde(rename = "@parent")]
+    #[serde(rename = "@parent", skip_serializing_if = "Option::is_none")]
     pub parent: Option<EntityTopicId>,
     #[serde(rename = "@type")]
     pub r#type: EntityType,

crates/core/tedge_api/src/entity_store.rs

@@ -33,6 +33,7 @@ use serde_json::Value as JsonValue;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::collections::VecDeque;
 use std::path::Path;
 
 // In the future, root will be read from config
@@ -514,6 +515,14 @@ impl EntityStore {
     pub fn cache_early_data_message(&mut self, message: MqttMessage) {
         self.pending_entity_store.cache_early_data_message(message)
     }
+
+    pub fn list_entity_tree<'a>(
+        &'a self,
+        root: Option<&'a EntityTopicId>,
+    ) -> Result<Vec<&'a EntityMetadata>, entity_store::Error> {
+        self.entities
+            .list_entity_tree(root.unwrap_or_else(|| self.main_device()))
+    }
 }
 
 /// In-memory representation of the entity tree
@@ -650,6 +659,26 @@ impl EntityTree {
             removed_entities.push(topic_id.clone())
         }
     }
+
+    pub fn list_entity_tree<'a>(
+        &'a self,
+        root: &'a EntityTopicId,
+    ) -> Result<Vec<&'a EntityMetadata>, Error> {
+        if let Some(metadata) = self.entities.get(root).map(|node| node.metadata()) {
+            let mut topic_ids = VecDeque::from(vec![root]);
+            let mut entities = vec![metadata];
+
+            while let Some(topic_id) = topic_ids.pop_front() {
+                let (child_topics, children): (Vec<_>, Vec<_>) =
+                    self.children(topic_id).into_iter().unzip();
+                topic_ids.extend(child_topics);
+                entities.extend(children);
+            }
+            Ok(entities)
+        } else {
+            Err(Error::UnknownEntity(root.to_string()))
+        }
+    }
 }
 
 /// Represents an error encountered while updating the store.
@@ -1020,6 +1049,124 @@ mod tests {
         assert!(children.iter().any(|&e| e == "device/child2//"));
     }
 
+    #[test]
+    fn list_entity_tree() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let mut store = new_entity_store(&temp_dir, true);
+
+        // Build the entity tree:
+        //
+        // main
+        // |-- child0
+        // |   |-- child00
+        // |   |   |-- child000
+        // |-- child1
+        // |-- child2
+        // |   |-- child20
+        // |   |-- child21
+        // |   |   |-- child210
+        // |   |   |   |-- child2100
+        // |   |   |-- child211
+        // |   |-- child22
+        build_entity_tree(
+            &mut store,
+            vec![
+                ("device/child0//", "child-device", None),
+                ("device/child1//", "child-device", None),
+                ("device/child2//", "child-device", None),
+                ("device/child00//", "child-device", Some("device/child0//")),
+                ("device/child20//", "child-device", Some("device/child2//")),
+                ("device/child21//", "child-device", Some("device/child2//")),
+                ("device/child22//", "child-device", Some("device/child2//")),
+                (
+                    "device/child000//",
+                    "child-device",
+                    Some("device/child00//"),
+                ),
+                (
+                    "device/child210//",
+                    "child-device",
+                    Some("device/child21//"),
+                ),
+                (
+                    "device/child211//",
+                    "child-device",
+                    Some("device/child21//"),
+                ),
+                (
+                    "device/child2100//",
+                    "child-device",
+                    Some("device/child210//"),
+                ),
+            ],
+        );
+
+        // List entity tree from root
+        let entities = store
+            .list_entity_tree(None)
+            .unwrap()
+            .iter()
+            .map(|e| e.topic_id.as_str())
+            .collect::<HashSet<_>>();
+        let expected_entities = [
+            "device/main//",
+            "device/child0//",
+            "device/child1//",
+            "device/child2//",
+            "device/child00//",
+            "device/child20//",
+            "device/child21//",
+            "device/child22//",
+            "device/child000//",
+            "device/child210//",
+            "device/child211//",
+            "device/child2100//",
+        ]
+        .into_iter()
+        .collect::<HashSet<_>>();
+        assert_eq!(entities, expected_entities);
+
+        // List entity tree from child2
+        let child_root = EntityTopicId::from_str("device/child2//").unwrap();
+        let entities = store
+            .list_entity_tree(Some(&child_root))
+            .unwrap()
+            .iter()
+            .map(|e| e.topic_id.as_str())
+            .collect::<HashSet<_>>();
+        let expected_entities = [
+            "device/child2//",
+            "device/child20//",
+            "device/child21//",
+            "device/child22//",
+            "device/child210//",
+            "device/child211//",
+            "device/child2100//",
+        ]
+        .into_iter()
+        .collect::<HashSet<_>>();
+        assert_eq!(entities, expected_entities);
+    }
+
+    // Each item in the vector represents (topic_id, type, parent)
+    fn build_entity_tree(store: &mut EntityStore, entity_tree: Vec<(&str, &str, Option<&str>)>) {
+        for entity in entity_tree {
+            let topic_id = EntityTopicId::from_str(entity.0).unwrap();
+            let r#type = EntityType::from_str(entity.1).unwrap();
+            let parent = entity.2.map(|p| EntityTopicId::from_str(p).unwrap());
+
+            store
+                .update(EntityRegistrationMessage {
+                    topic_id,
+                    r#type,
+                    external_id: None,
+                    parent,
+                    other: Map::new(),
+                })
+                .unwrap();
+        }
+    }
+
     #[test]
     fn lists_services() {
         let temp_dir = tempfile::tempdir().unwrap();