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}