entity_graph_mcp/
store.rs

1//! In-process entity store.
2//!
3//! Mirrors the shape of `@prometheus-ags/entity-graph-core`:
4//!   `entities[type][id] = json_value`
5//!
6//! Replace the `DashMap` internals with real DB calls for production use.
7
8use std::sync::Arc;
9
10use dashmap::DashMap;
11use serde_json::Value;
12use thiserror::Error;
13
14/// Errors returned by store operations.
15#[derive(Debug, Error)]
16#[allow(dead_code)]
17pub enum StoreError {
18    #[error("entity type '{0}' not found")]
19    TypeNotFound(String),
20
21    #[error("entity '{entity_type}/{id}' not found")]
22    EntityNotFound { entity_type: String, id: String },
23
24    #[error("invalid entity data: {0}")]
25    InvalidData(String),
26}
27
28/// A single entity record (JSON object keyed by string fields).
29pub type EntityData = serde_json::Map<String, Value>;
30
31/// Shareable, thread-safe entity store.
32///
33/// Structure: `types → (id → data)`.
34#[derive(Clone, Default)]
35pub struct EntityStore {
36    inner: Arc<DashMap<String, DashMap<String, EntityData>>>,
37}
38
39impl EntityStore {
40    /// Create an empty store.
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    // ── Read ─────────────────────────────────────────────────────────────────
46
47    /// Return all known entity type names.
48    pub fn list_types(&self) -> Vec<String> {
49        self.inner.iter().map(|e| e.key().clone()).collect()
50    }
51
52    /// Return every entity of the given type as `(id, data)` pairs.
53    pub fn list_entities(&self, entity_type: &str) -> Vec<(String, EntityData)> {
54        match self.inner.get(entity_type) {
55            None => vec![],
56            Some(bucket) => bucket
57                .iter()
58                .map(|e| (e.key().clone(), e.value().clone()))
59                .collect(),
60        }
61    }
62
63    /// Read a single entity.
64    pub fn get(&self, entity_type: &str, id: &str) -> Result<EntityData, StoreError> {
65        let bucket = self
66            .inner
67            .get(entity_type)
68            .ok_or_else(|| StoreError::TypeNotFound(entity_type.to_owned()))?;
69
70        bucket
71            .get(id)
72            .map(|e| e.value().clone())
73            .ok_or_else(|| StoreError::EntityNotFound {
74                entity_type: entity_type.to_owned(),
75                id: id.to_owned(),
76            })
77    }
78
79    // ── Write ─────────────────────────────────────────────────────────────────
80
81    /// Insert or merge-update an entity.
82    ///
83    /// If the entity already exists, the provided fields are merged (shallow
84    /// merge, matching the behaviour of `upsertEntity` in the TypeScript core).
85    pub fn upsert(&self, entity_type: &str, id: &str, data: EntityData) -> EntityData {
86        let bucket = self.inner.entry(entity_type.to_owned()).or_default();
87
88        let mut entry = bucket.entry(id.to_owned()).or_default();
89        // Shallow merge: incoming fields overwrite existing ones.
90        for (k, v) in data {
91            entry.insert(k, v);
92        }
93        entry.clone()
94    }
95
96    /// Delete an entity. Returns the removed data if present.
97    pub fn delete(&self, entity_type: &str, id: &str) -> Option<EntityData> {
98        let bucket = self.inner.get(entity_type)?;
99        bucket.remove(id).map(|(_, v)| v)
100    }
101
102    // ── Seeding ───────────────────────────────────────────────────────────────
103
104    /// Bulk-load entities (useful for testing or seeding demo data).
105    pub fn seed(&self, entity_type: &str, entities: Vec<(String, EntityData)>) {
106        let bucket = self.inner.entry(entity_type.to_owned()).or_default();
107        for (id, data) in entities {
108            bucket.insert(id, data);
109        }
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use serde_json::json;
117
118    fn map(v: Value) -> EntityData {
119        match v {
120            Value::Object(m) => m,
121            _ => panic!("expected object"),
122        }
123    }
124
125    #[test]
126    fn upsert_creates_new_entity() {
127        let store = EntityStore::new();
128        let data = map(json!({"name": "Acme", "status": "active"}));
129        store.upsert("Company", "co_1", data.clone());
130
131        let got = store.get("Company", "co_1").unwrap();
132        assert_eq!(got["name"], json!("Acme"));
133    }
134
135    #[test]
136    fn upsert_merges_existing_entity() {
137        let store = EntityStore::new();
138        store.upsert("Company", "co_1", map(json!({"name": "Old"})));
139        store.upsert(
140            "Company",
141            "co_1",
142            map(json!({"name": "New", "city": "NYC"})),
143        );
144
145        let got = store.get("Company", "co_1").unwrap();
146        assert_eq!(got["name"], json!("New"));
147        assert_eq!(got["city"], json!("NYC"));
148    }
149
150    #[test]
151    fn delete_removes_entity() {
152        let store = EntityStore::new();
153        store.upsert("Company", "co_1", map(json!({"name": "Acme"})));
154        let removed = store.delete("Company", "co_1");
155        assert!(removed.is_some());
156        assert!(store.get("Company", "co_1").is_err());
157    }
158
159    #[test]
160    fn list_types_returns_all_types() {
161        let store = EntityStore::new();
162        store.upsert("Company", "co_1", map(json!({"name": "A"})));
163        store.upsert("User", "u_1", map(json!({"email": "a@b.com"})));
164
165        let mut types = store.list_types();
166        types.sort();
167        assert_eq!(types, vec!["Company", "User"]);
168    }
169
170    #[test]
171    fn get_unknown_type_returns_error() {
172        let store = EntityStore::new();
173        let err = store.get("Ghost", "x").unwrap_err();
174        assert!(matches!(err, StoreError::TypeNotFound(_)));
175    }
176}
entity_graph_mcp/store.rs

entity_graph_mcp/
store.rs
