Skip to main content

grust_core/
lib.rs

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;