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 pub fn begin_write(&self) -> Result<WriteTransaction, GraphError> {
44 Ok(self.storage.begin_write()?)
45 }
46
47 pub fn commit(write_txn: WriteTransaction) -> Result<(), GraphError> {
49 write_txn.commit()?;
50 Ok(())
51 }
52
53 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 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 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 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 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}