1use std::{
2 collections::{BTreeMap, BTreeSet},
3 fmt,
4};
5
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8
9pub type Result<T> = std::result::Result<T, GrustError>;
10pub type Props = BTreeMap<String, Value>;
11
12#[derive(Debug, thiserror::Error)]
13pub enum GrustError {
14 #[error("backend error: {0}")]
15 Backend(String),
16 #[error("schema error: {0}")]
17 Schema(String),
18 #[error("unsupported graph feature: {0}")]
19 Unsupported(String),
20 #[error("serialization error: {0}")]
21 Serialization(String),
22}
23
24#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
25pub struct NodeId(String);
26
27impl NodeId {
28 pub fn new(value: impl Into<String>) -> Self {
29 Self(value.into())
30 }
31
32 pub fn as_str(&self) -> &str {
33 &self.0
34 }
35
36 pub fn into_string(self) -> String {
37 self.0
38 }
39}
40
41impl fmt::Display for NodeId {
42 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43 f.write_str(&self.0)
44 }
45}
46
47impl From<String> for NodeId {
48 fn from(value: String) -> Self {
49 Self::new(value)
50 }
51}
52
53impl From<&str> for NodeId {
54 fn from(value: &str) -> Self {
55 Self::new(value)
56 }
57}
58
59impl From<&String> for NodeId {
60 fn from(value: &String) -> Self {
61 Self::new(value.clone())
62 }
63}
64
65impl From<&NodeId> for NodeId {
66 fn from(value: &NodeId) -> Self {
67 value.clone()
68 }
69}
70
71#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
72pub struct EdgeId(String);
73
74impl EdgeId {
75 pub fn new(value: impl Into<String>) -> Self {
76 Self(value.into())
77 }
78
79 pub fn as_str(&self) -> &str {
80 &self.0
81 }
82
83 pub fn into_string(self) -> String {
84 self.0
85 }
86}
87
88impl fmt::Display for EdgeId {
89 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90 f.write_str(&self.0)
91 }
92}
93
94impl From<String> for EdgeId {
95 fn from(value: String) -> Self {
96 Self::new(value)
97 }
98}
99
100impl From<&str> for EdgeId {
101 fn from(value: &str) -> Self {
102 Self::new(value)
103 }
104}
105
106impl From<&String> for EdgeId {
107 fn from(value: &String) -> Self {
108 Self::new(value.clone())
109 }
110}
111
112impl From<&EdgeId> for EdgeId {
113 fn from(value: &EdgeId) -> Self {
114 value.clone()
115 }
116}
117
118#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
119pub struct Label(String);
120
121impl Label {
122 pub fn new(value: impl Into<String>) -> Self {
123 Self(value.into())
124 }
125
126 pub fn as_str(&self) -> &str {
127 &self.0
128 }
129
130 pub fn into_string(self) -> String {
131 self.0
132 }
133}
134
135impl fmt::Display for Label {
136 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137 f.write_str(&self.0)
138 }
139}
140
141impl From<String> for Label {
142 fn from(value: String) -> Self {
143 Self::new(value)
144 }
145}
146
147impl From<&str> for Label {
148 fn from(value: &str) -> Self {
149 Self::new(value)
150 }
151}
152
153impl From<&String> for Label {
154 fn from(value: &String) -> Self {
155 Self::new(value.clone())
156 }
157}
158
159impl From<&Label> for Label {
160 fn from(value: &Label) -> Self {
161 value.clone()
162 }
163}
164
165#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
166#[serde(tag = "type", content = "value", rename_all = "snake_case")]
167pub enum Value {
168 Null,
169 Bool(bool),
170 Int(i64),
171 Float(f64),
172 String(String),
173 StringArray(Vec<String>),
174 Json(serde_json::Value),
175}
176
177impl Value {
178 pub fn as_str(&self) -> Option<&str> {
179 match self {
180 Self::String(value) => Some(value),
181 _ => None,
182 }
183 }
184
185 pub fn as_string_array(&self) -> Option<&[String]> {
186 match self {
187 Self::StringArray(values) => Some(values),
188 _ => None,
189 }
190 }
191}
192
193impl From<String> for Value {
194 fn from(value: String) -> Self {
195 Self::String(value)
196 }
197}
198
199impl From<&str> for Value {
200 fn from(value: &str) -> Self {
201 Self::String(value.to_string())
202 }
203}
204
205impl From<&String> for Value {
206 fn from(value: &String) -> Self {
207 Self::String(value.clone())
208 }
209}
210
211impl From<Vec<String>> for Value {
212 fn from(value: Vec<String>) -> Self {
213 Self::StringArray(value)
214 }
215}
216
217impl From<bool> for Value {
218 fn from(value: bool) -> Self {
219 Self::Bool(value)
220 }
221}
222
223impl From<i64> for Value {
224 fn from(value: i64) -> Self {
225 Self::Int(value)
226 }
227}
228
229impl From<i32> for Value {
230 fn from(value: i32) -> Self {
231 Self::Int(i64::from(value))
232 }
233}
234
235impl From<usize> for Value {
236 fn from(value: usize) -> Self {
237 Self::Int(value as i64)
238 }
239}
240
241impl From<f64> for Value {
242 fn from(value: f64) -> Self {
243 Self::Float(value)
244 }
245}
246
247impl From<serde_json::Value> for Value {
248 fn from(value: serde_json::Value) -> Self {
249 match value {
250 serde_json::Value::Null => Self::Null,
251 serde_json::Value::Bool(value) => Self::Bool(value),
252 serde_json::Value::Number(value) => {
253 if let Some(value) = value.as_i64() {
254 Self::Int(value)
255 } else if let Some(value) = value.as_f64() {
256 Self::Float(value)
257 } else {
258 Self::Json(serde_json::Value::Number(value))
259 }
260 }
261 serde_json::Value::String(value) => Self::String(value),
262 serde_json::Value::Array(values) => {
263 let strings = values
264 .iter()
265 .filter_map(|value| value.as_str().map(ToString::to_string))
266 .collect::<Vec<_>>();
267 if strings.len() == values.len() {
268 Self::StringArray(strings)
269 } else {
270 Self::Json(serde_json::Value::Array(values))
271 }
272 }
273 serde_json::Value::Object(value) => Self::Json(serde_json::Value::Object(value)),
274 }
275 }
276}
277
278#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
279pub struct Node {
280 pub id: NodeId,
281 pub label: Label,
282 pub props: Props,
283}
284
285impl Node {
286 pub fn new(label: impl Into<Label>, id: impl Into<NodeId>, props: impl Into<Props>) -> Self {
287 let id = id.into();
288 let mut props = props.into();
289 props
290 .entry("id".to_string())
291 .or_insert_with(|| Value::from(id.as_str()));
292 Self {
293 id,
294 label: label.into(),
295 props,
296 }
297 }
298}
299
300#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
301pub struct Edge {
302 pub id: Option<EdgeId>,
303 pub from: NodeId,
304 pub to: NodeId,
305 pub label: Label,
306 pub props: Props,
307}
308
309impl Edge {
310 pub fn new(
311 label: impl Into<Label>,
312 from: impl Into<NodeId>,
313 to: impl Into<NodeId>,
314 props: impl Into<Props>,
315 ) -> Self {
316 Self {
317 id: None,
318 from: from.into(),
319 to: to.into(),
320 label: label.into(),
321 props: props.into(),
322 }
323 }
324
325 pub fn with_id(mut self, id: impl Into<EdgeId>) -> Self {
326 self.id = Some(id.into());
327 self
328 }
329}
330
331#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
332pub struct Graph {
333 pub nodes: Vec<Node>,
334 pub edges: Vec<Edge>,
335}
336
337impl Graph {
338 pub fn new(nodes: Vec<Node>, edges: Vec<Edge>) -> Self {
339 Self { nodes, edges }
340 }
341
342 pub fn builder() -> GraphBuilder {
343 GraphBuilder::new()
344 }
345}
346
347#[derive(Clone, Debug, Eq, PartialEq)]
348pub enum EdgePolicy {
349 AllowDuplicates,
350 DedupeByFromLabelTo,
351}
352
353impl Default for EdgePolicy {
354 fn default() -> Self {
355 Self::DedupeByFromLabelTo
356 }
357}
358
359#[derive(Clone, Debug, Default)]
360pub struct GraphBuilder {
361 nodes: BTreeMap<NodeId, Node>,
362 edges: Vec<Edge>,
363 edge_keys: BTreeSet<(NodeId, Label, NodeId)>,
364 edge_policy: EdgePolicy,
365}
366
367impl GraphBuilder {
368 pub fn new() -> Self {
369 Self::default()
370 }
371
372 pub fn edge_policy(mut self, edge_policy: EdgePolicy) -> Self {
373 self.edge_policy = edge_policy;
374 self
375 }
376
377 pub fn node<'a>(
378 &'a mut self,
379 label: impl Into<Label>,
380 id: impl Into<NodeId>,
381 ) -> NodeBuilder<'a> {
382 NodeBuilder {
383 builder: self,
384 label: label.into(),
385 id: id.into(),
386 props: Props::new(),
387 }
388 }
389
390 pub fn edge<'a>(
391 &'a mut self,
392 label: impl Into<Label>,
393 from: impl Into<NodeId>,
394 to: impl Into<NodeId>,
395 ) -> EdgeBuilder<'a> {
396 EdgeBuilder {
397 builder: self,
398 id: None,
399 label: label.into(),
400 from: from.into(),
401 to: to.into(),
402 props: Props::new(),
403 }
404 }
405
406 pub fn add_node(&mut self, node: Node) -> NodeId {
407 let id = node.id.clone();
408 self.nodes
409 .entry(id.clone())
410 .and_modify(|existing| {
411 if existing.label == node.label {
412 existing.props.extend(node.props.clone());
413 }
414 })
415 .or_insert(node);
416 id
417 }
418
419 pub fn add_edge(&mut self, edge: Edge) -> Option<EdgeId> {
420 let id = edge.id.clone();
421 match self.edge_policy {
422 EdgePolicy::AllowDuplicates => self.edges.push(edge),
423 EdgePolicy::DedupeByFromLabelTo => {
424 let key = (edge.from.clone(), edge.label.clone(), edge.to.clone());
425 if self.edge_keys.insert(key) {
426 self.edges.push(edge);
427 }
428 }
429 }
430 id
431 }
432
433 pub fn build(self) -> Graph {
434 Graph {
435 nodes: self.nodes.into_values().collect(),
436 edges: self.edges,
437 }
438 }
439}
440
441pub struct NodeBuilder<'a> {
442 builder: &'a mut GraphBuilder,
443 label: Label,
444 id: NodeId,
445 props: Props,
446}
447
448impl<'a> NodeBuilder<'a> {
449 pub fn prop(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
450 self.props.insert(key.into(), value.into());
451 self
452 }
453
454 pub fn props(mut self, props: Props) -> Self {
455 self.props.extend(props);
456 self
457 }
458
459 pub fn finish(self) -> NodeId {
460 let node = Node::new(self.label, self.id, self.props);
461 self.builder.add_node(node)
462 }
463}
464
465pub struct EdgeBuilder<'a> {
466 builder: &'a mut GraphBuilder,
467 id: Option<EdgeId>,
468 label: Label,
469 from: NodeId,
470 to: NodeId,
471 props: Props,
472}
473
474impl<'a> EdgeBuilder<'a> {
475 pub fn id(mut self, id: impl Into<EdgeId>) -> Self {
476 self.id = Some(id.into());
477 self
478 }
479
480 pub fn prop(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
481 self.props.insert(key.into(), value.into());
482 self
483 }
484
485 pub fn props(mut self, props: Props) -> Self {
486 self.props.extend(props);
487 self
488 }
489
490 pub fn finish(self) -> Option<EdgeId> {
491 let mut edge = Edge::new(self.label, self.from, self.to, self.props);
492 edge.id = self.id;
493 self.builder.add_edge(edge)
494 }
495}
496
497#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
498pub struct GraphSchema {
499 pub nodes: Vec<NodeType>,
500 pub edges: Vec<EdgeType>,
501}
502
503#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
504pub struct NodeType {
505 pub label: Label,
506 pub fields: Vec<Field>,
507}
508
509#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
510pub struct EdgeType {
511 pub label: Label,
512 pub from: Vec<Label>,
513 pub to: Vec<Label>,
514 pub fields: Vec<Field>,
515 pub directed: bool,
516 pub uniqueness: EdgeUniqueness,
517}
518
519#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
520pub struct Field {
521 pub name: String,
522 pub ty: FieldType,
523 pub required: bool,
524}
525
526#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
527pub enum FieldType {
528 String,
529 Int,
530 Float,
531 Bool,
532 DateTime,
533 StringArray,
534 Json,
535}
536
537#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
538pub enum EdgeUniqueness {
539 None,
540 FromTo,
541 FromLabelTo,
542}
543
544#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
545pub struct Traversal {
546 pub start: Start,
547 pub steps: Vec<Step>,
548 pub limit: Option<u32>,
549}
550
551impl Traversal {
552 pub fn from_node(id: impl Into<NodeId>) -> Self {
553 Self {
554 start: Start::Node(id.into()),
555 steps: Vec::new(),
556 limit: None,
557 }
558 }
559
560 pub fn out(mut self, edge: impl Into<Label>) -> Self {
561 self.steps.push(Step {
562 direction: Direction::Out,
563 edge: Some(edge.into()),
564 node: None,
565 });
566 self
567 }
568
569 pub fn in_(mut self, edge: impl Into<Label>) -> Self {
570 self.steps.push(Step {
571 direction: Direction::In,
572 edge: Some(edge.into()),
573 node: None,
574 });
575 self
576 }
577
578 pub fn both(mut self, edge: impl Into<Label>) -> Self {
579 self.steps.push(Step {
580 direction: Direction::Both,
581 edge: Some(edge.into()),
582 node: None,
583 });
584 self
585 }
586
587 pub fn to(mut self, node: impl Into<Label>) -> Self {
588 if let Some(step) = self.steps.last_mut() {
589 step.node = Some(node.into());
590 }
591 self
592 }
593
594 pub fn limit(mut self, limit: u32) -> Self {
595 self.limit = Some(limit);
596 self
597 }
598}
599
600#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
601pub enum Start {
602 Node(NodeId),
603 NodesByLabel(Label),
604 NodesByProperty {
605 label: Label,
606 key: String,
607 value: Value,
608 },
609}
610
611#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
612pub struct Step {
613 pub direction: Direction,
614 pub edge: Option<Label>,
615 pub node: Option<Label>,
616}
617
618#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
619pub enum Direction {
620 Out,
621 In,
622 Both,
623}
624
625#[derive(Clone, Debug, Default, PartialEq)]
626pub struct EdgeQuery {
627 pub from: Option<NodeId>,
628 pub to: Option<NodeId>,
629 pub label: Option<Label>,
630}
631
632#[derive(Clone, Debug, Default, Eq, PartialEq)]
633pub struct LoadReport {
634 pub nodes: usize,
635 pub edges: usize,
636}
637
638#[async_trait]
639pub trait GraphStore: Send + Sync {
640 async fn apply_schema(&self, _schema: &GraphSchema) -> Result<()> {
641 Ok(())
642 }
643
644 async fn put_node(&self, node: &Node) -> Result<NodeId>;
645 async fn put_edge(&self, edge: &Edge) -> Result<Option<EdgeId>>;
646
647 async fn put_graph(&self, graph: &Graph) -> Result<LoadReport> {
648 let mut report = LoadReport::default();
649 for node in &graph.nodes {
650 self.put_node(node).await?;
651 report.nodes += 1;
652 }
653 for edge in &graph.edges {
654 self.put_edge(edge).await?;
655 report.edges += 1;
656 }
657 Ok(report)
658 }
659
660 async fn get_node(&self, id: &NodeId) -> Result<Option<Node>>;
661 async fn get_edges(&self, query: EdgeQuery) -> Result<Vec<Edge>>;
662 async fn traverse(&self, traversal: Traversal) -> Result<Vec<Node>>;
663}
664
665#[async_trait]
666pub trait GraphAdminStore: GraphStore {
667 async fn bootstrap(&self) -> Result<()> {
668 Ok(())
669 }
670
671 async fn clear(&self) -> Result<()>;
672}
673
674pub mod prelude {
675 pub use crate::{
676 Direction, Edge, EdgeId, EdgePolicy, EdgeQuery, EdgeType, Field, FieldType, Graph,
677 GraphAdminStore, GraphBuilder, GraphSchema, GraphStore, GrustError, Label, LoadReport,
678 Node, NodeId, NodeType, Props, Result, Start, Step, Traversal, Value,
679 };
680}
681
682#[cfg(test)]
683mod tests;