bee_graph/
lib.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2#[cfg(any(feature = "neo4j", feature = "nebulagraph", feature = "arangodb"))]
3mod http;
4
5#[cfg(feature = "arangodb")]
6pub mod arangodb;
7#[cfg(feature = "nebulagraph")]
8pub mod nebulagraph;
9#[cfg(feature = "neo4j")]
10pub mod neo4j;
11
12use async_trait::async_trait;
13use serde::{Deserialize, Serialize};
14use thiserror::Error;
15
16// ---------------------------------------------------------------------------
17// Public types
18// ---------------------------------------------------------------------------
19
20/// Unique identifier for a vertex (node).
21pub type VertexId = String;
22
23/// Unique identifier for an edge (relationship).
24pub type EdgeId = String;
25
26/// Key-value properties attached to vertices and edges.
27pub type Properties = serde_json::Map<String, serde_json::Value>;
28
29/// Named parameters for parameterised graph queries.
30pub type Params = serde_json::Map<String, serde_json::Value>;
31
32/// A vertex (node) in the graph.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct Vertex {
35    pub id: VertexId,
36    pub label: String,
37    pub properties: Properties,
38}
39
40/// An edge (relationship) between two vertices.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct Edge {
43    pub id: EdgeId,
44    pub label: String,
45    pub from: VertexId,
46    pub to: VertexId,
47    pub properties: Properties,
48}
49
50/// High-level traversal specification.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct Traversal {
53    /// Starting vertex ID.
54    pub start: VertexId,
55    /// Edge label to follow (`None` = any).
56    pub edge_label: Option<String>,
57    /// Maximum depth / hops.
58    pub max_depth: u32,
59    /// Direction in which to traverse edges.
60    pub direction: TraversalDirection,
61}
62
63/// Direction for graph traversals.
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub enum TraversalDirection {
66    Outgoing,
67    Incoming,
68    Both,
69}
70
71/// A single path result from a traversal.
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct PathResult {
74    /// Ordered vertices visited along the path.
75    pub vertices: Vec<Vertex>,
76    /// Ordered edges traversed.
77    pub edges: Vec<Edge>,
78}
79
80/// A generic result envelope for raw graph queries.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct QueryResult {
83    /// Column names returned by the query.
84    pub columns: Vec<String>,
85    /// Rows, each represented as a JSON value.
86    pub rows: Vec<serde_json::Value>,
87}
88
89// ---------------------------------------------------------------------------
90// Errors
91// ---------------------------------------------------------------------------
92
93/// Errors that can occur when interacting with a graph database.
94#[derive(Error, Debug)]
95pub enum GraphError {
96    /// The requested vertex does not exist.
97    #[error("vertex not found: {0}")]
98    VertexNotFound(VertexId),
99
100    /// The connection to the graph backend failed.
101    #[error("connection error: {0}")]
102    ConnectionError(String),
103
104    /// A query could not be executed.
105    #[error("query error: {0}")]
106    QueryError(String),
107}
108
109// ---------------------------------------------------------------------------
110// Trait
111// ---------------------------------------------------------------------------
112
113/// A generic async graph database trait.
114///
115/// Backends such as Neo4j, NebulaGraph, or ArangoDB can implement this to
116/// provide vertex/edge CRUD, traversals, and arbitrary query execution.
117#[async_trait]
118pub trait GraphDB: Send + Sync {
119    /// Insert a new vertex.  The `id` field may be assigned by the caller
120    /// or left empty for the backend to generate.
121    async fn add_vertex(&self, vertex: Vertex) -> Result<Vertex, GraphError>;
122
123    /// Return the vertex identified by `id`.
124    async fn get_vertex(&self, id: &VertexId) -> Result<Option<Vertex>, GraphError>;
125
126    /// Update properties of an existing vertex.  The provided `properties`
127    /// are merged into (not fully replaced by) the existing ones.
128    async fn update_vertex(
129        &self,
130        id: &VertexId,
131        properties: Properties,
132    ) -> Result<Vertex, GraphError>;
133
134    /// Delete the vertex identified by `id` and all incident edges.
135    async fn delete_vertex(&self, id: &VertexId) -> Result<(), GraphError>;
136
137    /// Insert a new edge.
138    async fn add_edge(&self, edge: Edge) -> Result<Edge, GraphError>;
139
140    /// Traverse the graph starting from `traversal.start`.
141    async fn traverse(&self, traversal: Traversal) -> Result<Vec<PathResult>, GraphError>;
142
143    /// Execute an arbitrary backend-specific query with optional parameters.
144    async fn query(&self, query: &str, params: Option<Params>) -> Result<QueryResult, GraphError>;
145}
146
147// ---------------------------------------------------------------------------
148// Tests
149// ---------------------------------------------------------------------------
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use std::collections::HashMap;
155    use std::sync::Mutex;
156
157    struct StubGraphDB {
158        vertices: Mutex<HashMap<VertexId, Vertex>>,
159        edges: Mutex<HashMap<EdgeId, Edge>>,
160    }
161
162    impl StubGraphDB {
163        fn new() -> Self {
164            Self { vertices: Mutex::new(HashMap::new()), edges: Mutex::new(HashMap::new()) }
165        }
166    }
167
168    #[allow(clippy::needless_lifetimes)]
169    #[async_trait]
170    impl GraphDB for StubGraphDB {
171        async fn add_vertex(&self, mut vertex: Vertex) -> Result<Vertex, GraphError> {
172            let mut map = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
173            if vertex.id.is_empty() {
174                vertex.id = format!("v{}", map.len());
175            }
176            map.insert(vertex.id.clone(), vertex.clone());
177            Ok(vertex)
178        }
179
180        async fn get_vertex(&self, id: &VertexId) -> Result<Option<Vertex>, GraphError> {
181            let map = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
182            Ok(map.get(id).cloned())
183        }
184
185        async fn update_vertex(
186            &self,
187            id: &VertexId,
188            properties: Properties,
189        ) -> Result<Vertex, GraphError> {
190            let mut map = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
191            let vertex = map.get_mut(id).ok_or_else(|| GraphError::VertexNotFound(id.clone()))?;
192            for (k, v) in properties {
193                vertex.properties.insert(k, v);
194            }
195            Ok(vertex.clone())
196        }
197
198        async fn delete_vertex(&self, id: &VertexId) -> Result<(), GraphError> {
199            let mut vmap = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
200            vmap.remove(id);
201            let mut emap = self.edges.lock().unwrap_or_else(|e| e.into_inner());
202            emap.retain(|_, e| e.from != *id && e.to != *id);
203            Ok(())
204        }
205
206        async fn add_edge(&self, mut edge: Edge) -> Result<Edge, GraphError> {
207            // Validate endpoints exist
208            let vmap = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
209            if !vmap.contains_key(&edge.from) {
210                return Err(GraphError::VertexNotFound(edge.from.clone()));
211            }
212            if !vmap.contains_key(&edge.to) {
213                return Err(GraphError::VertexNotFound(edge.to.clone()));
214            }
215            drop(vmap);
216
217            let mut map = self.edges.lock().unwrap_or_else(|e| e.into_inner());
218            if edge.id.is_empty() {
219                edge.id = format!("e{}", map.len());
220            }
221            map.insert(edge.id.clone(), edge.clone());
222            Ok(edge)
223        }
224
225        async fn traverse(&self, traversal: Traversal) -> Result<Vec<PathResult>, GraphError> {
226            // Acquire vmap before emap, matching delete_vertex, to avoid ABBA deadlock.
227            let vmap = self.vertices.lock().unwrap_or_else(|e| e.into_inner());
228            if !vmap.contains_key(&traversal.start) {
229                return Err(GraphError::VertexNotFound(traversal.start.clone()));
230            }
231            let emap = self.edges.lock().unwrap_or_else(|e| e.into_inner());
232
233            let mut results = Vec::new();
234            let mut frontier: Vec<(VertexId, Vec<Vertex>, Vec<Edge>)> = Vec::new();
235            if let Some(v) = vmap.get(&traversal.start).cloned() {
236                frontier.push((traversal.start.clone(), vec![v], Vec::new()));
237            }
238
239            // BFS, collecting every path of up to `max_depth` hops.
240            for _ in 0..traversal.max_depth {
241                let mut next = Vec::new();
242                for (cur, vs, es) in &frontier {
243                    for edge in emap.values() {
244                        let follows = match &traversal.direction {
245                            TraversalDirection::Outgoing => edge.from == *cur,
246                            TraversalDirection::Incoming => edge.to == *cur,
247                            TraversalDirection::Both => edge.from == *cur || edge.to == *cur,
248                        };
249                        if !follows {
250                            continue;
251                        }
252                        if let Some(label) = &traversal.edge_label
253                            && *label != edge.label
254                        {
255                            continue;
256                        }
257                        let next_id = match &traversal.direction {
258                            TraversalDirection::Outgoing => edge.to.clone(),
259                            TraversalDirection::Incoming => edge.from.clone(),
260                            TraversalDirection::Both => {
261                                if edge.from == *cur {
262                                    edge.to.clone()
263                                } else {
264                                    edge.from.clone()
265                                }
266                            }
267                        };
268                        if let Some(nv) = vmap.get(&next_id).cloned() {
269                            let mut nvs = vs.clone();
270                            let mut nes = es.clone();
271                            nvs.push(nv);
272                            nes.push(edge.clone());
273                            results.push(PathResult { vertices: nvs.clone(), edges: nes.clone() });
274                            next.push((next_id, nvs, nes));
275                        }
276                    }
277                    if results.len() >= 10 {
278                        break;
279                    }
280                }
281                if results.len() >= 10 {
282                    break;
283                }
284                frontier = next;
285                if frontier.is_empty() {
286                    break;
287                }
288            }
289            Ok(results)
290        }
291
292        async fn query(
293            &self,
294            _query: &str,
295            _params: Option<Params>,
296        ) -> Result<QueryResult, GraphError> {
297            Ok(QueryResult { columns: Vec::new(), rows: Vec::new() })
298        }
299    }
300
301    #[tokio::test]
302    async fn test_add_and_get_vertex() {
303        let db = StubGraphDB::new();
304        let v = Vertex { id: "v1".into(), label: "Person".into(), properties: Properties::new() };
305        let created = db.add_vertex(v).await.unwrap();
306        assert_eq!(created.id, "v1");
307        let found = db.get_vertex(&"v1".into()).await.unwrap();
308        assert!(found.is_some());
309    }
310
311    #[tokio::test]
312    async fn test_update_vertex() {
313        let db = StubGraphDB::new();
314        let v = Vertex {
315            id: "v1".into(),
316            label: "Person".into(),
317            properties: {
318                let mut m = Properties::new();
319                m.insert("name".into(), serde_json::json!("Alice"));
320                m
321            },
322        };
323        db.add_vertex(v).await.unwrap();
324
325        let mut props = Properties::new();
326        props.insert("age".into(), serde_json::json!(30));
327        let updated = db.update_vertex(&"v1".into(), props).await.unwrap();
328        assert_eq!(updated.properties.get("name").unwrap(), &serde_json::json!("Alice"));
329        assert_eq!(updated.properties.get("age").unwrap(), &serde_json::json!(30));
330    }
331
332    #[tokio::test]
333    async fn test_add_edge_and_traverse() {
334        let db = StubGraphDB::new();
335        let a = Vertex { id: "a".into(), label: "Person".into(), properties: Properties::new() };
336        let b = Vertex { id: "b".into(), label: "Person".into(), properties: Properties::new() };
337        db.add_vertex(a).await.unwrap();
338        db.add_vertex(b).await.unwrap();
339
340        let edge = Edge {
341            id: String::new(),
342            label: "KNOWS".into(),
343            from: "a".into(),
344            to: "b".into(),
345            properties: Properties::new(),
346        };
347        let created = db.add_edge(edge).await.unwrap();
348        assert!(!created.id.is_empty());
349
350        let paths = db
351            .traverse(Traversal {
352                start: "a".into(),
353                edge_label: Some("KNOWS".into()),
354                max_depth: 1,
355                direction: TraversalDirection::Outgoing,
356            })
357            .await
358            .unwrap();
359        assert_eq!(paths.len(), 1);
360        assert_eq!(paths[0].vertices.len(), 2);
361    }
362
363    #[tokio::test]
364    async fn test_delete_vertex_cascades_edges() {
365        let db = StubGraphDB::new();
366        let a = Vertex { id: "a".into(), label: "N".into(), properties: Properties::new() };
367        let b = Vertex { id: "b".into(), label: "N".into(), properties: Properties::new() };
368        db.add_vertex(a).await.unwrap();
369        db.add_vertex(b).await.unwrap();
370        db.add_edge(Edge {
371            id: "e1".into(),
372            label: "L".into(),
373            from: "a".into(),
374            to: "b".into(),
375            properties: Properties::new(),
376        })
377        .await
378        .unwrap();
379
380        db.delete_vertex(&"a".into()).await.unwrap();
381        assert!(db.get_vertex(&"a".into()).await.unwrap().is_none());
382        // traversal from a should now fail
383        let result = db
384            .traverse(Traversal {
385                start: "a".into(),
386                edge_label: None,
387                max_depth: 1,
388                direction: TraversalDirection::Both,
389            })
390            .await;
391        assert!(result.is_err());
392    }
393
394    #[tokio::test]
395    async fn test_query() {
396        let db = StubGraphDB::new();
397        let res = db.query("MATCH (n) RETURN n", None).await.unwrap();
398        assert!(res.columns.is_empty());
399        assert!(res.rows.is_empty());
400    }
401}
bee_graph/lib.rs

bee_graph/
lib.rs
