Skip to main content

marsdb_graph/
store.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3
4use marsdb_storage::{ReadTransaction, ReadableMultimapTable, ReadableTable, StorageEngine, Txn, 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    /// Open a read transaction for a statement that never mutates
48    /// anything (`MATCH ... RETURN`) — a consistent point-in-time
49    /// snapshot that runs alongside any concurrent readers or a
50    /// concurrent writer without contending for redb's single-writer
51    /// lock. No commit/abort: a read transaction has nothing to roll
52    /// back, it just releases on drop.
53    pub fn begin_read(&self) -> Result<ReadTransaction, GraphError> {
54        Ok(self.storage.begin_read()?)
55    }
56
57    /// Commit a transaction obtained from [`begin_write`](Self::begin_write).
58    pub fn commit(write_txn: WriteTransaction) -> Result<(), GraphError> {
59        write_txn.commit()?;
60        Ok(())
61    }
62
63    /// Abort (roll back) a transaction obtained from
64    /// [`begin_write`](Self::begin_write), discarding any writes made
65    /// through it.
66    pub fn abort(write_txn: WriteTransaction) -> Result<(), GraphError> {
67        write_txn.abort()?;
68        Ok(())
69    }
70
71    pub fn create_node(
72        &self,
73        labels: &[&str],
74        props: BTreeMap<String, PropertyValue>,
75    ) -> Result<NodeId, GraphError> {
76        let write_txn = self.begin_write()?;
77        let id = Self::create_node_in_txn(&write_txn, labels, props)?;
78        write_txn.commit()?;
79        Ok(id)
80    }
81
82    pub fn create_node_in_txn(
83        write_txn: &WriteTransaction,
84        labels: &[&str],
85        props: BTreeMap<String, PropertyValue>,
86    ) -> Result<NodeId, GraphError> {
87        let label_ids = labels
88            .iter()
89            .map(|l| intern_label(write_txn, l))
90            .collect::<Result<Vec<_>, _>>()?;
91        let id = next_id(write_txn, "next_node_id")?;
92        let record = NodeRecord {
93            label_ids: label_ids.clone(),
94            props,
95        };
96        let bytes = encode(&record)?;
97        let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
98        nodes.insert(id, bytes.as_slice())?;
99        drop(nodes);
100        let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
101        for label_id in label_ids {
102            label_index.insert(label_id, id)?;
103        }
104        Ok(NodeId(id))
105    }
106
107    pub fn get_node(&self, id: NodeId) -> Result<Option<Node>, GraphError> {
108        let read_txn = self.begin_read()?;
109        Self::get_node_in_txn(Txn::Read(&read_txn), id)
110    }
111
112    pub fn get_node_in_txn(txn: Txn, id: NodeId) -> Result<Option<Node>, GraphError> {
113        let record: Option<NodeRecord> = {
114            let nodes = txn.open_table(marsdb_storage::tables::NODES)?;
115            let found = match nodes.get(id.0)? {
116                Some(guard) => Some(decode(guard.value())?),
117                None => None,
118            };
119            found
120        };
121        let Some(record) = record else { return Ok(None) };
122        let labels = record
123            .label_ids
124            .iter()
125            .map(|&lid| resolve_label(txn, lid))
126            .collect::<Result<Vec<_>, _>>()?;
127        Ok(Some(Node {
128            id,
129            labels,
130            props: record.props,
131        }))
132    }
133
134    pub fn create_edge(
135        &self,
136        label: &str,
137        src: NodeId,
138        dst: NodeId,
139        props: BTreeMap<String, PropertyValue>,
140    ) -> Result<EdgeId, GraphError> {
141        let write_txn = self.begin_write()?;
142        let id = Self::create_edge_in_txn(&write_txn, label, src, dst, props)?;
143        write_txn.commit()?;
144        Ok(id)
145    }
146
147    pub fn create_edge_in_txn(
148        write_txn: &WriteTransaction,
149        label: &str,
150        src: NodeId,
151        dst: NodeId,
152        props: BTreeMap<String, PropertyValue>,
153    ) -> Result<EdgeId, GraphError> {
154        let label_id = intern_label(write_txn, label)?;
155        let id = next_id(write_txn, "next_edge_id")?;
156        let record = EdgeRecord {
157            label_id,
158            src: src.0,
159            dst: dst.0,
160            props,
161        };
162        let bytes = encode(&record)?;
163        let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
164        edges.insert(id, bytes.as_slice())?;
165
166        let out_entry = AdjEntry {
167            edge_id: EdgeId(id),
168            other: dst,
169            label_id,
170        }
171        .encode();
172        let in_entry = AdjEntry {
173            edge_id: EdgeId(id),
174            other: src,
175            label_id,
176        }
177        .encode();
178        let mut adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
179        adj_out.insert(src.0, out_entry.as_slice())?;
180        let mut adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
181        adj_in.insert(dst.0, in_entry.as_slice())?;
182        Ok(EdgeId(id))
183    }
184
185    pub fn get_edge(&self, id: EdgeId) -> Result<Option<Edge>, GraphError> {
186        let read_txn = self.begin_read()?;
187        Self::get_edge_in_txn(Txn::Read(&read_txn), id)
188    }
189
190    pub fn get_edge_in_txn(txn: Txn, id: EdgeId) -> Result<Option<Edge>, GraphError> {
191        let record: Option<EdgeRecord> = {
192            let edges = txn.open_table(marsdb_storage::tables::EDGES)?;
193            let found = match edges.get(id.0)? {
194                Some(guard) => Some(decode(guard.value())?),
195                None => None,
196            };
197            found
198        };
199        let Some(record) = record else { return Ok(None) };
200        let label = resolve_label(txn, record.label_id)?;
201        Ok(Some(Edge {
202            id,
203            label,
204            src: NodeId(record.src),
205            dst: NodeId(record.dst),
206            props: record.props,
207        }))
208    }
209
210    /// Neighbors of `node` in `dir`, optionally filtered by edge label.
211    /// Reads directly from the adjacency multimap without touching `edges`.
212    pub fn neighbors(
213        &self,
214        node: NodeId,
215        dir: Direction,
216        label_filter: Option<&str>,
217    ) -> Result<Vec<AdjEntry>, GraphError> {
218        let read_txn = self.begin_read()?;
219        Self::neighbors_in_txn(Txn::Read(&read_txn), node, dir, label_filter)
220    }
221
222    pub fn neighbors_in_txn(
223        txn: Txn,
224        node: NodeId,
225        dir: Direction,
226        label_filter: Option<&str>,
227    ) -> Result<Vec<AdjEntry>, GraphError> {
228        let label_id_filter = match label_filter {
229            Some(l) => match lookup_label_id(txn, l)? {
230                Some(id) => Some(id),
231                None => return Ok(Vec::new()),
232            },
233            None => None,
234        };
235        let mut result = Vec::new();
236        let table_def = match dir {
237            Direction::Out => marsdb_storage::tables::ADJ_OUT,
238            Direction::In => marsdb_storage::tables::ADJ_IN,
239        };
240        let table = txn.open_multimap_table(table_def)?;
241        for item in table.get(node.0)? {
242            let entry = AdjEntry::decode(item?.value());
243            if label_id_filter.is_none_or(|lid| lid == entry.label_id) {
244                result.push(entry);
245            }
246        }
247        Ok(result)
248    }
249
250    pub fn delete_edge(&self, id: EdgeId) -> Result<bool, GraphError> {
251        let write_txn = self.begin_write()?;
252        let removed = Self::delete_edge_in_txn(&write_txn, id)?;
253        write_txn.commit()?;
254        Ok(removed)
255    }
256
257    pub fn delete_edge_in_txn(write_txn: &WriteTransaction, id: EdgeId) -> Result<bool, GraphError> {
258        let record_bytes: Option<Vec<u8>> = {
259            let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
260            let removed = edges.remove(id.0)?.map(|guard| guard.value().to_vec());
261            removed
262        };
263        let Some(record_bytes) = record_bytes else {
264            return Ok(false);
265        };
266        let record: EdgeRecord = decode(&record_bytes)?;
267        let out_entry = AdjEntry {
268            edge_id: id,
269            other: NodeId(record.dst),
270            label_id: record.label_id,
271        }
272        .encode();
273        let in_entry = AdjEntry {
274            edge_id: id,
275            other: NodeId(record.src),
276            label_id: record.label_id,
277        }
278        .encode();
279        {
280            let mut adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
281            adj_out.remove(record.src, out_entry.as_slice())?;
282        }
283        {
284            let mut adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
285            adj_in.remove(record.dst, in_entry.as_slice())?;
286        }
287        Ok(true)
288    }
289
290    /// Delete a node. If `detach` is false and the node has incident edges,
291    /// returns `GraphError::NodeHasEdges` instead of deleting anything.
292    pub fn delete_node(&self, id: NodeId, detach: bool) -> Result<bool, GraphError> {
293        let write_txn = self.begin_write()?;
294        let existed = Self::delete_node_in_txn(&write_txn, id, detach)?;
295        write_txn.commit()?;
296        Ok(existed)
297    }
298
299    pub fn delete_node_in_txn(write_txn: &WriteTransaction, id: NodeId, detach: bool) -> Result<bool, GraphError> {
300        let mut incident: Vec<EdgeId> = Vec::new();
301        {
302            let adj_out = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_OUT)?;
303            for item in adj_out.get(id.0)? {
304                incident.push(AdjEntry::decode(item?.value()).edge_id);
305            }
306            let adj_in = write_txn.open_multimap_table(marsdb_storage::tables::ADJ_IN)?;
307            for item in adj_in.get(id.0)? {
308                incident.push(AdjEntry::decode(item?.value()).edge_id);
309            }
310        }
311        if !incident.is_empty() && !detach {
312            return Err(GraphError::NodeHasEdges(id));
313        }
314        for edge_id in incident {
315            Self::delete_edge_in_txn(write_txn, edge_id)?;
316        }
317        let removed_bytes: Option<Vec<u8>> = {
318            let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
319            let removed = nodes.remove(id.0)?.map(|guard| guard.value().to_vec());
320            removed
321        };
322        let Some(removed_bytes) = removed_bytes else {
323            return Ok(false);
324        };
325        let record: NodeRecord = decode(&removed_bytes)?;
326        let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
327        for label_id in record.label_ids {
328            label_index.remove(label_id, id.0)?;
329        }
330        Ok(true)
331    }
332
333    pub fn set_node_prop(&self, id: NodeId, key: &str, value: PropertyValue) -> Result<bool, GraphError> {
334        let write_txn = self.begin_write()?;
335        let updated = Self::set_node_prop_in_txn(&write_txn, id, key, value)?;
336        if updated {
337            write_txn.commit()?;
338        } else {
339            write_txn.abort()?;
340        }
341        Ok(updated)
342    }
343
344    pub fn set_node_prop_in_txn(
345        write_txn: &WriteTransaction,
346        id: NodeId,
347        key: &str,
348        value: PropertyValue,
349    ) -> Result<bool, GraphError> {
350        let bytes_opt: Option<Vec<u8>> = {
351            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
352            let found = nodes.get(id.0)?.map(|g| g.value().to_vec());
353            found
354        };
355        let Some(bytes) = bytes_opt else { return Ok(false) };
356        let mut record: NodeRecord = decode(&bytes)?;
357        record.props.insert(key.to_string(), value);
358        let new_bytes = encode(&record)?;
359        let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
360        nodes.insert(id.0, new_bytes.as_slice())?;
361        Ok(true)
362    }
363
364    pub fn set_edge_prop(&self, id: EdgeId, key: &str, value: PropertyValue) -> Result<bool, GraphError> {
365        let write_txn = self.begin_write()?;
366        let updated = Self::set_edge_prop_in_txn(&write_txn, id, key, value)?;
367        if updated {
368            write_txn.commit()?;
369        } else {
370            write_txn.abort()?;
371        }
372        Ok(updated)
373    }
374
375    pub fn set_edge_prop_in_txn(
376        write_txn: &WriteTransaction,
377        id: EdgeId,
378        key: &str,
379        value: PropertyValue,
380    ) -> Result<bool, GraphError> {
381        let bytes_opt: Option<Vec<u8>> = {
382            let edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
383            let found = edges.get(id.0)?.map(|g| g.value().to_vec());
384            found
385        };
386        let Some(bytes) = bytes_opt else { return Ok(false) };
387        let mut record: EdgeRecord = decode(&bytes)?;
388        record.props.insert(key.to_string(), value);
389        let new_bytes = encode(&record)?;
390        let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
391        edges.insert(id.0, new_bytes.as_slice())?;
392        Ok(true)
393    }
394
395    pub fn remove_node_prop_in_txn(write_txn: &WriteTransaction, id: NodeId, key: &str) -> Result<bool, GraphError> {
396        let bytes_opt: Option<Vec<u8>> = {
397            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
398            let found = nodes.get(id.0)?.map(|g| g.value().to_vec());
399            found
400        };
401        let Some(bytes) = bytes_opt else { return Ok(false) };
402        let mut record: NodeRecord = decode(&bytes)?;
403        record.props.remove(key);
404        let new_bytes = encode(&record)?;
405        let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
406        nodes.insert(id.0, new_bytes.as_slice())?;
407        Ok(true)
408    }
409
410    pub fn remove_edge_prop_in_txn(write_txn: &WriteTransaction, id: EdgeId, key: &str) -> Result<bool, GraphError> {
411        let bytes_opt: Option<Vec<u8>> = {
412            let edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
413            let found = edges.get(id.0)?.map(|g| g.value().to_vec());
414            found
415        };
416        let Some(bytes) = bytes_opt else { return Ok(false) };
417        let mut record: EdgeRecord = decode(&bytes)?;
418        record.props.remove(key);
419        let new_bytes = encode(&record)?;
420        let mut edges = write_txn.open_table(marsdb_storage::tables::EDGES)?;
421        edges.insert(id.0, new_bytes.as_slice())?;
422        Ok(true)
423    }
424
425    /// Adds `label` to `id`'s label set -- a no-op (not an error) if it's
426    /// already there, same idempotent-add semantics real Cypher's `SET
427    /// n:Label` has.
428    pub fn add_node_label_in_txn(write_txn: &WriteTransaction, id: NodeId, label: &str) -> Result<bool, GraphError> {
429        let bytes_opt: Option<Vec<u8>> = {
430            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
431            let found = nodes.get(id.0)?.map(|g| g.value().to_vec());
432            found
433        };
434        let Some(bytes) = bytes_opt else { return Ok(false) };
435        let mut record: NodeRecord = decode(&bytes)?;
436        let label_id = intern_label(write_txn, label)?;
437        if !record.label_ids.contains(&label_id) {
438            record.label_ids.push(label_id);
439            let new_bytes = encode(&record)?;
440            let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
441            nodes.insert(id.0, new_bytes.as_slice())?;
442            let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
443            label_index.insert(label_id, id.0)?;
444        }
445        Ok(true)
446    }
447
448    /// Removes `label` from `id`'s label set -- a no-op (not an error) if
449    /// it's not there (label unknown entirely, or known but not on this
450    /// node), same as real Cypher's `REMOVE n:Label`.
451    pub fn remove_node_label_in_txn(write_txn: &WriteTransaction, id: NodeId, label: &str) -> Result<bool, GraphError> {
452        let bytes_opt: Option<Vec<u8>> = {
453            let nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
454            let found = nodes.get(id.0)?.map(|g| g.value().to_vec());
455            found
456        };
457        let Some(bytes) = bytes_opt else { return Ok(false) };
458        let Some(label_id) = lookup_label_id(Txn::Write(write_txn), label)? else { return Ok(true) };
459        let mut record: NodeRecord = decode(&bytes)?;
460        if let Some(pos) = record.label_ids.iter().position(|&l| l == label_id) {
461            record.label_ids.remove(pos);
462            let new_bytes = encode(&record)?;
463            let mut nodes = write_txn.open_table(marsdb_storage::tables::NODES)?;
464            nodes.insert(id.0, new_bytes.as_slice())?;
465            let mut label_index = write_txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
466            label_index.remove(label_id, id.0)?;
467        }
468        Ok(true)
469    }
470
471    /// Full scan of all nodes, optionally filtered by label. v1 has no
472    /// secondary index on label, so this is a linear scan of the table.
473    pub fn all_nodes(&self, label_filter: Option<&str>) -> Result<Vec<Node>, GraphError> {
474        let read_txn = self.begin_read()?;
475        Self::all_nodes_in_txn(Txn::Read(&read_txn), label_filter)
476    }
477
478    pub fn all_nodes_in_txn(txn: Txn, label_filter: Option<&str>) -> Result<Vec<Node>, GraphError> {
479        Self::all_nodes_limited_in_txn(txn, label_filter, usize::MAX)
480    }
481
482    /// Same as `all_nodes_in_txn`, but stops once `limit` nodes are found --
483    /// the storage-level end of `LIMIT` push-down (see the executor's
484    /// `scan()`/`eval_plan` docs for the query-level half): a query whose
485    /// entire plan is a bare scan feeding straight into a `LIMIT` doesn't
486    /// need to touch rows past the first `limit`, whether or not a label
487    /// filter narrows it first.
488    pub fn all_nodes_limited_in_txn(
489        txn: Txn,
490        label_filter: Option<&str>,
491        limit: usize,
492    ) -> Result<Vec<Node>, GraphError> {
493        // A label filter goes through NODE_LABEL_INDEX (label_id -> node_ids)
494        // plus a point lookup per match, instead of scanning every row in
495        // NODES — cost scales with the number of matching rows, not the
496        // table size. No filter means every row is wanted anyway, so a full
497        // scan is already optimal; the index wouldn't help.
498        let Some(label_filter) = label_filter else {
499            let mut result = Vec::new();
500            let nodes = txn.open_table(marsdb_storage::tables::NODES)?;
501            for item in nodes.iter()? {
502                if result.len() >= limit {
503                    break;
504                }
505                let (key, value) = item?;
506                let record: NodeRecord = decode(value.value())?;
507                let labels = record
508                    .label_ids
509                    .iter()
510                    .map(|&lid| resolve_label(txn, lid))
511                    .collect::<Result<Vec<_>, _>>()?;
512                result.push(Node {
513                    id: NodeId(key.value()),
514                    labels,
515                    props: record.props,
516                });
517            }
518            return Ok(result);
519        };
520        let Some(label_id) = lookup_label_id(txn, label_filter)? else {
521            return Ok(Vec::new());
522        };
523        let node_ids: Vec<u64> = {
524            let label_index = txn.open_multimap_table(marsdb_storage::tables::NODE_LABEL_INDEX)?;
525            // `.take(limit)` here, not a `.truncate()` after collecting --
526            // stops walking the multimap's own entries past `limit`, not
527            // just the (more expensive) per-id NODES point-reads below.
528            // Measured difference: without this, a labeled LIMIT query's
529            // cost still scaled with the *matching* row count, not `limit`
530            // (see BENCHMARKS.md's `execute_scan_limit_pushdown` numbers).
531            let ids: Vec<u64> = label_index
532                .get(label_id)?
533                .take(limit)
534                .map(|item| item.map(|g| g.value()))
535                .collect::<Result<_, _>>()?;
536            ids
537        };
538        let mut result = Vec::with_capacity(node_ids.len());
539        let nodes = txn.open_table(marsdb_storage::tables::NODES)?;
540        for id in node_ids {
541            let guard = nodes
542                .get(id)?
543                .expect("node_label_index entry must reference a live node");
544            let record: NodeRecord = decode(guard.value())?;
545            drop(guard);
546            let labels = record
547                .label_ids
548                .iter()
549                .map(|&lid| resolve_label(txn, lid))
550                .collect::<Result<Vec<_>, _>>()?;
551            result.push(Node {
552                id: NodeId(id),
553                labels,
554                props: record.props,
555            });
556        }
557        Ok(result)
558    }
559}