Skip to main content

marsdb_graph/
store.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3
4use marsdb_storage::{ReadableMultimapTable, ReadableTable, StorageEngine, WriteTransaction};
5
6use crate::encode::{decode, encode, EdgeRecord, NodeRecord};
7use crate::error::GraphError;
8use crate::id::next_id;
9use crate::labels::{intern_label, lookup_label_id, resolve_label};
10use crate::model::{AdjEntry, Direction, Edge, EdgeId, Node, NodeId, PropertyValue};
11
12pub struct GraphStore {
13    storage: StorageEngine,
14}
15
16impl GraphStore {
17    pub fn open_file(path: impl AsRef<Path>) -> Result<Self, GraphError> {
18        Ok(Self {
19            storage: StorageEngine::open_file(path)?,
20        })
21    }
22
23    pub fn open_memory() -> Result<Self, GraphError> {
24        Ok(Self {
25            storage: StorageEngine::open_memory()?,
26        })
27    }
28
29    /// Open a write transaction spanning multiple graph operations. Callers
30    /// (e.g. the query executor) drive an entire Cypher statement through
31    /// the `*_in_txn` methods below using this one transaction, then call
32    /// `write_txn.commit()` themselves — this is the crash-safety boundary
33    /// from the plan: one statement = one transaction, not one transaction
34    /// per individual node/edge write.
35    ///
36    /// v1 uses a write transaction even for pure-read statements (rather
37    /// than a separate read-only path) to keep one code path and guarantee
38    /// every statement — reads included — sees one consistent snapshot.
39    /// Trade-off: this serializes concurrent readers behind redb's
40    /// single-writer lock instead of allowing true concurrent reads; a
41    /// read-only transaction path is the natural follow-up if read
42    /// concurrency becomes a bottleneck.
43    pub fn begin_write(&self) -> Result<WriteTransaction, GraphError> {
44        Ok(self.storage.begin_write()?)
45    }
46
47    /// Commit a transaction obtained from [`begin_write`](Self::begin_write).
48    pub fn commit(write_txn: WriteTransaction) -> Result<(), GraphError> {
49        write_txn.commit()?;
50        Ok(())
51    }
52
53    /// Abort (roll back) a transaction obtained from
54    /// [`begin_write`](Self::begin_write), discarding any writes made
55    /// through it.
56    pub fn abort(write_txn: WriteTransaction) -> Result<(), GraphError> {
57        write_txn.abort()?;
58        Ok(())
59    }
60
61    pub fn create_node(
62        &self,
63        labels: &[&str],
64        props: BTreeMap<String, PropertyValue>,
65    ) -> Result<NodeId, GraphError> {
66        let write_txn = self.begin_write()?;
67        let id = Self::create_node_in_txn(&write_txn, labels, props)?;
68        write_txn.commit()?;
69        Ok(id)
70    }
71
72    pub fn create_node_in_txn(
73        write_txn: &WriteTransaction,
74        labels: &[&str],
75        props: BTreeMap<String, PropertyValue>,
76    ) -> Result<NodeId, GraphError> {
77        let label_ids = labels
78            .iter()
79            .map(|l| intern_label(write_txn, l))
80            .collect::<Result<Vec<_>, _>>()?;
81        let id = next_id(write_txn, "next_node_id")?;
82        let record = NodeRecord {
83            label_ids: label_ids.clone(),
84            props,
85        };
86        let bytes = encode(&record)?;
87        let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
88        nodes.insert(id, bytes.as_slice())?;
89        drop(nodes);
90        let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
91        for label_id in label_ids {
92            label_index.insert(label_id, id)?;
93        }
94        Ok(NodeId(id))
95    }
96
97    pub fn get_node(&self, id: NodeId) -> Result<Option<Node>, GraphError> {
98        let write_txn = self.begin_write()?;
99        let node = Self::get_node_in_txn(&write_txn, id)?;
100        write_txn.abort()?;
101        Ok(node)
102    }
103
104    pub fn get_node_in_txn(write_txn: &WriteTransaction, id: NodeId) -> Result<Option<Node>, GraphError> {
105        let record: Option<NodeRecord> = {
106            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
107            let found = match nodes.get(id.0)? {
108                Some(guard) => Some(decode(guard.value())?),
109                None => None,
110            };
111            found
112        };
113        let Some(record) = record else { return Ok(None) };
114        let labels = record
115            .label_ids
116            .iter()
117            .map(|&lid| resolve_label(write_txn, lid))
118            .collect::<Result<Vec<_>, _>>()?;
119        Ok(Some(Node {
120            id,
121            labels,
122            props: record.props,
123        }))
124    }
125
126    pub fn create_edge(
127        &self,
128        label: &str,
129        src: NodeId,
130        dst: NodeId,
131        props: BTreeMap<String, PropertyValue>,
132    ) -> Result<EdgeId, GraphError> {
133        let write_txn = self.begin_write()?;
134        let id = Self::create_edge_in_txn(&write_txn, label, src, dst, props)?;
135        write_txn.commit()?;
136        Ok(id)
137    }
138
139    pub fn create_edge_in_txn(
140        write_txn: &WriteTransaction,
141        label: &str,
142        src: NodeId,
143        dst: NodeId,
144        props: BTreeMap<String, PropertyValue>,
145    ) -> Result<EdgeId, GraphError> {
146        let label_id = intern_label(write_txn, label)?;
147        let id = next_id(write_txn, "next_edge_id")?;
148        let record = EdgeRecord {
149            label_id,
150            src: src.0,
151            dst: dst.0,
152            props,
153        };
154        let bytes = encode(&record)?;
155        let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
156        edges.insert(id, bytes.as_slice())?;
157
158        let out_entry = AdjEntry {
159            edge_id: EdgeId(id),
160            other: dst,
161            label_id,
162        }
163        .encode();
164        let in_entry = AdjEntry {
165            edge_id: EdgeId(id),
166            other: src,
167            label_id,
168        }
169        .encode();
170        let mut adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
171        adj_out.insert(src.0, out_entry.as_slice())?;
172        let mut adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
173        adj_in.insert(dst.0, in_entry.as_slice())?;
174        Ok(EdgeId(id))
175    }
176
177    pub fn get_edge(&self, id: EdgeId) -> Result<Option<Edge>, GraphError> {
178        let write_txn = self.begin_write()?;
179        let edge = Self::get_edge_in_txn(&write_txn, id)?;
180        write_txn.abort()?;
181        Ok(edge)
182    }
183
184    pub fn get_edge_in_txn(write_txn: &WriteTransaction, id: EdgeId) -> Result<Option<Edge>, GraphError> {
185        let record: Option<EdgeRecord> = {
186            let edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
187            let found = match edges.get(id.0)? {
188                Some(guard) => Some(decode(guard.value())?),
189                None => None,
190            };
191            found
192        };
193        let Some(record) = record else { return Ok(None) };
194        let label = resolve_label(write_txn, record.label_id)?;
195        Ok(Some(Edge {
196            id,
197            label,
198            src: NodeId(record.src),
199            dst: NodeId(record.dst),
200            props: record.props,
201        }))
202    }
203
204    /// Neighbors of `node` in `dir`, optionally filtered by edge label.
205    /// Reads directly from the adjacency multimap without touching `edges`.
206    pub fn neighbors(
207        &self,
208        node: NodeId,
209        dir: Direction,
210        label_filter: Option<&str>,
211    ) -> Result<Vec<AdjEntry>, GraphError> {
212        let write_txn = self.begin_write()?;
213        let result = Self::neighbors_in_txn(&write_txn, node, dir, label_filter)?;
214        write_txn.abort()?;
215        Ok(result)
216    }
217
218    pub fn neighbors_in_txn(
219        write_txn: &WriteTransaction,
220        node: NodeId,
221        dir: Direction,
222        label_filter: Option<&str>,
223    ) -> Result<Vec<AdjEntry>, GraphError> {
224        let label_id_filter = match label_filter {
225            Some(l) => match lookup_label_id(write_txn, l)? {
226                Some(id) => Some(id),
227                None => return Ok(Vec::new()),
228            },
229            None => None,
230        };
231        let mut result = Vec::new();
232        let table_def = match dir {
233            Direction::Out => marsdb_storage::tables::ADJ_OUT,
234            Direction::In => marsdb_storage::tables::ADJ_IN,
235        };
236        let table = write_txn.open_multimap_table(table_def)?;
237        for item in table.get(node.0)? {
238            let entry = AdjEntry::decode(item?.value());
239            if label_id_filter.is_none_or(|lid| lid == entry.label_id) {
240                result.push(entry);
241            }
242        }
243        Ok(result)
244    }
245
246    pub fn delete_edge(&self, id: EdgeId) -> Result<bool, GraphError> {
247        let write_txn = self.begin_write()?;
248        let removed = Self::delete_edge_in_txn(&write_txn, id)?;
249        write_txn.commit()?;
250        Ok(removed)
251    }
252
253    pub fn delete_edge_in_txn(write_txn: &WriteTransaction, id: EdgeId) -> Result<bool, GraphError> {
254        let record_bytes: Option<Vec<u8>> = {
255            let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
256            let removed = edges.remove(id.0)?.map(|guard| guard.value().to_vec());
257            removed
258        };
259        let Some(record_bytes) = record_bytes else {
260            return Ok(false);
261        };
262        let record: EdgeRecord = decode(&record_bytes)?;
263        let out_entry = AdjEntry {
264            edge_id: id,
265            other: NodeId(record.dst),
266            label_id: record.label_id,
267        }
268        .encode();
269        let in_entry = AdjEntry {
270            edge_id: id,
271            other: NodeId(record.src),
272            label_id: record.label_id,
273        }
274        .encode();
275        {
276            let mut adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
277            adj_out.remove(record.src, out_entry.as_slice())?;
278        }
279        {
280            let mut adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
281            adj_in.remove(record.dst, in_entry.as_slice())?;
282        }
283        Ok(true)
284    }
285
286    /// Delete a node. If `detach` is false and the node has incident edges,
287    /// returns `GraphError::NodeHasEdges` instead of deleting anything.
288    pub fn delete_node(&self, id: NodeId, detach: bool) -> Result<bool, GraphError> {
289        let write_txn = self.begin_write()?;
290        let existed = Self::delete_node_in_txn(&write_txn, id, detach)?;
291        write_txn.commit()?;
292        Ok(existed)
293    }
294
295    pub fn delete_node_in_txn(write_txn: &WriteTransaction, id: NodeId, detach: bool) -> Result<bool, GraphError> {
296        let mut incident: Vec<EdgeId> = Vec::new();
297        {
298            let adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
299            for item in adj_out.get(id.0)? {
300                incident.push(AdjEntry::decode(item?.value()).edge_id);
301            }
302            let adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
303            for item in adj_in.get(id.0)? {
304                incident.push(AdjEntry::decode(item?.value()).edge_id);
305            }
306        }
307        if !incident.is_empty() && !detach {
308            return Err(GraphError::NodeHasEdges(id));
309        }
310        for edge_id in incident {
311            Self::delete_edge_in_txn(write_txn, edge_id)?;
312        }
313        let removed_bytes: Option<Vec<u8>> = {
314            let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
315            let removed = nodes.remove(id.0)?.map(|guard| guard.value().to_vec());
316            removed
317        };
318        let Some(removed_bytes) = removed_bytes else {
319            return Ok(false);
320        };
321        let record: NodeRecord = decode(&removed_bytes)?;
322        let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
323        for label_id in record.label_ids {
324            label_index.remove(label_id, id.0)?;
325        }
326        Ok(true)
327    }
328
329    pub fn set_node_prop(&self, id: NodeId, key: &str, value: PropertyValue) -> Result<bool, GraphError> {
330        let write_txn = self.begin_write()?;
331        let updated = Self::set_node_prop_in_txn(&write_txn, id, key, value)?;
332        if updated {
333            write_txn.commit()?;
334        } else {
335            write_txn.abort()?;
336        }
337        Ok(updated)
338    }
339
340    pub fn set_node_prop_in_txn(
341        write_txn: &WriteTransaction,
342        id: NodeId,
343        key: &str,
344        value: PropertyValue,
345    ) -> Result<bool, GraphError> {
346        let bytes_opt: Option<Vec<u8>> = {
347            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
348            let found = nodes.get(id.0)?.map(|g| g.value().to_vec());
349            found
350        };
351        let Some(bytes) = bytes_opt else { return Ok(false) };
352        let mut record: NodeRecord = decode(&bytes)?;
353        record.props.insert(key.to_string(), value);
354        let new_bytes = encode(&record)?;
355        let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
356        nodes.insert(id.0, new_bytes.as_slice())?;
357        Ok(true)
358    }
359
360    pub fn set_edge_prop(&self, id: EdgeId, key: &str, value: PropertyValue) -> Result<bool, GraphError> {
361        let write_txn = self.begin_write()?;
362        let updated = Self::set_edge_prop_in_txn(&write_txn, id, key, value)?;
363        if updated {
364            write_txn.commit()?;
365        } else {
366            write_txn.abort()?;
367        }
368        Ok(updated)
369    }
370
371    pub fn set_edge_prop_in_txn(
372        write_txn: &WriteTransaction,
373        id: EdgeId,
374        key: &str,
375        value: PropertyValue,
376    ) -> Result<bool, GraphError> {
377        let bytes_opt: Option<Vec<u8>> = {
378            let edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
379            let found = edges.get(id.0)?.map(|g| g.value().to_vec());
380            found
381        };
382        let Some(bytes) = bytes_opt else { return Ok(false) };
383        let mut record: EdgeRecord = decode(&bytes)?;
384        record.props.insert(key.to_string(), value);
385        let new_bytes = encode(&record)?;
386        let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
387        edges.insert(id.0, new_bytes.as_slice())?;
388        Ok(true)
389    }
390
391    /// Full scan of all nodes, optionally filtered by label. v1 has no
392    /// secondary index on label, so this is a linear scan of the table.
393    pub fn all_nodes(&self, label_filter: Option<&str>) -> Result<Vec<Node>, GraphError> {
394        let write_txn = self.begin_write()?;
395        let result = Self::all_nodes_in_txn(&write_txn, label_filter)?;
396        write_txn.abort()?;
397        Ok(result)
398    }
399
400    pub fn all_nodes_in_txn(write_txn: &WriteTransaction, label_filter: Option<&str>) -> Result<Vec<Node>, GraphError> {
401        // A label filter goes through NODE_LABEL_INDEX (label_id -> node_ids)
402        // plus a point lookup per match, instead of scanning every row in
403        // NODES — cost scales with the number of matching rows, not the
404        // table size. No filter means every row is wanted anyway, so a full
405        // scan is already optimal; the index wouldn't help.
406        let Some(label_filter) = label_filter else {
407            let mut result = Vec::new();
408            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
409            for item in nodes.iter()? {
410                let (key, value) = item?;
411                let record: NodeRecord = decode(value.value())?;
412                let labels = record
413                    .label_ids
414                    .iter()
415                    .map(|&lid| resolve_label(write_txn, lid))
416                    .collect::<Result<Vec<_>, _>>()?;
417                result.push(Node {
418                    id: NodeId(key.value()),
419                    labels,
420                    props: record.props,
421                });
422            }
423            return Ok(result);
424        };
425        let Some(label_id) = lookup_label_id(write_txn, label_filter)? else {
426            return Ok(Vec::new());
427        };
428        let node_ids: Vec<u64> = {
429            let label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
430            let ids: Vec<u64> = label_index
431                .get(label_id)?
432                .map(|item| item.map(|g| g.value()))
433                .collect::<Result<_, _>>()?;
434            ids
435        };
436        let mut result = Vec::with_capacity(node_ids.len());
437        let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
438        for id in node_ids {
439            let guard = nodes
440                .get(id)?
441                .expect("node_label_index entry must reference a live node");
442            let record: NodeRecord = decode(guard.value())?;
443            drop(guard);
444            let labels = record
445                .label_ids
446                .iter()
447                .map(|&lid| resolve_label(write_txn, lid))
448                .collect::<Result<Vec<_>, _>>()?;
449            result.push(Node {
450                id: NodeId(id),
451                labels,
452                props: record.props,
453            });
454        }
455        Ok(result)
456    }
457}