Skip to main content

teaql_runtime/
graph.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use teaql_core::{Record, Value};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum GraphOperation {
7    Upsert,
8    Reference,
9    Remove,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
13pub enum GraphMutationKind {
14    Create,
15    Update,
16    Delete,
17    Reference,
18}
19
20#[derive(Debug, Clone, PartialEq)]
21pub struct GraphMutationPlanItem {
22    pub entity: String,
23    pub kind: GraphMutationKind,
24    pub values: Record,
25    pub update_fields: Vec<String>,
26}
27
28#[derive(Debug, Clone, PartialEq)]
29pub struct GraphMutationBatch {
30    pub entity: String,
31    pub kind: GraphMutationKind,
32    pub update_fields: Vec<String>,
33    pub items: Vec<GraphMutationPlanItem>,
34}
35
36#[derive(Debug, Clone, PartialEq, Default)]
37pub struct GraphMutationPlan {
38    pub planned_root: Option<GraphNode>,
39    pub items: Vec<GraphMutationPlanItem>,
40    pub batches: Vec<GraphMutationBatch>,
41}
42
43impl GraphMutationPlan {
44    pub fn push(
45        &mut self,
46        entity: impl Into<String>,
47        kind: GraphMutationKind,
48        values: Record,
49        update_fields: Vec<String>,
50    ) {
51        self.items.push(GraphMutationPlanItem {
52            entity: entity.into(),
53            kind,
54            values,
55            update_fields,
56        });
57    }
58
59    pub fn rebuild_batches(&mut self) {
60        let mut grouped: BTreeMap<
61            (String, GraphMutationKind, Vec<String>),
62            Vec<GraphMutationPlanItem>,
63        > = BTreeMap::new();
64        for item in &self.items {
65            let update_fields = match item.kind {
66                GraphMutationKind::Update => item.update_fields.clone(),
67                _ => Vec::new(),
68            };
69            grouped
70                .entry((item.entity.clone(), item.kind, update_fields))
71                .or_default()
72                .push(item.clone());
73        }
74        self.batches = grouped
75            .into_iter()
76            .map(
77                |((entity, kind, update_fields), items)| GraphMutationBatch {
78                    entity,
79                    kind,
80                    update_fields,
81                    items,
82                },
83            )
84            .collect();
85    }
86
87    pub fn grouped_counts(&self) -> BTreeMap<(String, GraphMutationKind), usize> {
88        let mut counts = BTreeMap::new();
89        for batch in &self.batches {
90            *counts
91                .entry((batch.entity.clone(), batch.kind))
92                .or_insert(0) += batch.items.len();
93        }
94        counts
95    }
96
97    pub fn batch_count(&self) -> usize {
98        self.batches.len()
99    }
100
101    pub fn len(&self) -> usize {
102        self.items.len()
103    }
104
105    pub fn is_empty(&self) -> bool {
106        self.items.is_empty()
107    }
108}
109
110pub fn sorted_update_fields(
111    values: &Record,
112    excluded: impl IntoIterator<Item = String>,
113) -> Vec<String> {
114    let excluded = excluded.into_iter().collect::<BTreeSet<_>>();
115    values
116        .keys()
117        .filter(|field| !excluded.contains(*field))
118        .cloned()
119        .collect()
120}
121
122#[derive(Debug, Clone, PartialEq)]
123pub struct GraphNode {
124    pub entity: String,
125    pub values: Record,
126    pub relations: BTreeMap<String, Vec<GraphNode>>,
127    pub operation: GraphOperation,
128}
129
130impl GraphNode {
131    pub fn new(entity: impl Into<String>) -> Self {
132        Self {
133            entity: entity.into(),
134            values: Record::new(),
135            relations: BTreeMap::new(),
136            operation: GraphOperation::Upsert,
137        }
138    }
139
140    pub fn operation(mut self, operation: GraphOperation) -> Self {
141        self.operation = operation;
142        self
143    }
144
145    pub fn reference(mut self) -> Self {
146        self.operation = GraphOperation::Reference;
147        self
148    }
149
150    pub fn remove(mut self) -> Self {
151        self.operation = GraphOperation::Remove;
152        self
153    }
154
155    pub fn value(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
156        self.values.insert(field.into(), value.into());
157        self
158    }
159
160    pub fn relation(mut self, name: impl Into<String>, node: GraphNode) -> Self {
161        self.relations.entry(name.into()).or_default().push(node);
162        self
163    }
164
165    pub fn relations(
166        mut self,
167        name: impl Into<String>,
168        nodes: impl IntoIterator<Item = GraphNode>,
169    ) -> Self {
170        self.relations.entry(name.into()).or_default().extend(nodes);
171        self
172    }
173
174    pub fn id(&self) -> Option<&Value> {
175        self.values.get("id")
176    }
177}