Skip to main content

nova_boot_graphdb/
memory.rs

1use crate::{error::GraphDbError, traits::GraphStore, types::*};
2use async_trait::async_trait;
3use serde_json::Value as JsonValue;
4use std::collections::{HashMap, HashSet, VecDeque};
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8#[derive(Default)]
9pub struct InMemoryGraphStore {
10    nodes: Arc<RwLock<HashMap<String, GraphNode>>>,
11    edges: Arc<RwLock<HashMap<String, GraphEdge>>>,
12}
13
14#[async_trait]
15impl GraphStore for InMemoryGraphStore {
16    async fn execute(&self, _query: GraphQuery) -> Result<JsonValue, GraphDbError> {
17        Err(GraphDbError::NotImplemented(
18            "in-memory store does not parse free-form query text",
19        ))
20    }
21
22    async fn upsert_node(&self, node: GraphNode) -> Result<(), GraphDbError> {
23        if node.id.is_empty() {
24            return Err(GraphDbError::InvalidInput(
25                "node id cannot be empty".to_string(),
26            ));
27        }
28        self.nodes.write().await.insert(node.id.clone(), node);
29        Ok(())
30    }
31
32    async fn upsert_edge(&self, edge: GraphEdge) -> Result<(), GraphDbError> {
33        if edge.id.is_empty() {
34            return Err(GraphDbError::InvalidInput(
35                "edge id cannot be empty".to_string(),
36            ));
37        }
38
39        let nodes = self.nodes.read().await;
40        if !nodes.contains_key(&edge.from) || !nodes.contains_key(&edge.to) {
41            return Err(GraphDbError::InvalidInput(
42                "edge endpoints must exist before edge upsert".to_string(),
43            ));
44        }
45        drop(nodes);
46
47        self.edges.write().await.insert(edge.id.clone(), edge);
48        Ok(())
49    }
50
51    async fn get_node(&self, node_id: &str) -> Result<Option<GraphNode>, GraphDbError> {
52        Ok(self.nodes.read().await.get(node_id).cloned())
53    }
54
55    async fn neighbors(&self, node_id: &str) -> Result<Vec<GraphNode>, GraphDbError> {
56        let edges = self.edges.read().await;
57        let nodes = self.nodes.read().await;
58        let mut out = Vec::new();
59
60        for edge in edges.values() {
61            if edge.from == node_id
62                && let Some(node) = nodes.get(&edge.to)
63            {
64                out.push(node.clone());
65            }
66        }
67
68        Ok(out)
69    }
70
71    async fn traverse(&self, start: &str, max_depth: usize) -> Result<GraphSubgraph, GraphDbError> {
72        let nodes_map = self.nodes.read().await;
73        if !nodes_map.contains_key(start) {
74            return Ok(GraphSubgraph::default());
75        }
76        let edges_map = self.edges.read().await;
77
78        let mut visited: HashSet<String> = HashSet::new();
79        let mut q: VecDeque<(String, usize)> = VecDeque::new();
80        let mut nodes = Vec::new();
81        let mut edges = Vec::new();
82
83        visited.insert(start.to_string());
84        q.push_back((start.to_string(), 0));
85
86        while let Some((current, depth)) = q.pop_front() {
87            if let Some(node) = nodes_map.get(&current) {
88                nodes.push(node.clone());
89            }
90
91            if depth >= max_depth {
92                continue;
93            }
94
95            for edge in edges_map.values() {
96                if edge.from == current {
97                    edges.push(edge.clone());
98
99                    if !visited.contains(&edge.to) {
100                        visited.insert(edge.to.clone());
101                        q.push_back((edge.to.clone(), depth + 1));
102                    }
103                }
104            }
105        }
106
107        Ok(GraphSubgraph { nodes, edges })
108    }
109}