Skip to main content

teaql_runtime/
entity_runtime.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::{Arc, Mutex};
3
4use teaql_core::{Record, Value};
5
6#[derive(Debug, Clone)]
7pub struct EntityKey {
8    pub entity: String,
9    pub id: Value,
10    id_key: String,
11}
12
13impl EntityKey {
14    pub fn new(entity: impl Into<String>, id: impl Into<Value>) -> Self {
15        let id = id.into();
16        Self {
17            entity: entity.into(),
18            id_key: value_key(&id),
19            id,
20        }
21    }
22}
23
24impl PartialEq for EntityKey {
25    fn eq(&self, other: &Self) -> bool {
26        self.entity == other.entity && self.id_key == other.id_key
27    }
28}
29
30impl Eq for EntityKey {}
31
32impl PartialOrd for EntityKey {
33    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
34        Some(self.cmp(other))
35    }
36}
37
38impl Ord for EntityKey {
39    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
40        self.entity
41            .cmp(&other.entity)
42            .then_with(|| self.id_key.cmp(&other.id_key))
43    }
44}
45
46fn value_key(value: &Value) -> String {
47    match value {
48        Value::Null => "null".to_owned(),
49        Value::Bool(value) => format!("bool:{value}"),
50        Value::I64(value) => format!("i64:{value}"),
51        Value::U64(value) => format!("u64:{value}"),
52        Value::F64(value) => format!("f64:{value}"),
53        Value::Decimal(value) => format!("decimal:{value}"),
54        Value::Text(value) => format!("text:{value}"),
55        Value::Json(value) => format!("json:{value}"),
56        Value::Date(value) => format!("date:{value}"),
57        Value::Timestamp(value) => format!("timestamp:{}", value.to_rfc3339()),
58        Value::Object(_) => "object".to_owned(),
59        Value::List(_) => "list".to_owned(),
60    }
61}
62
63#[derive(Debug, Clone, Default, PartialEq)]
64pub struct EntityChangeSet {
65    changes: BTreeMap<EntityKey, Record>,
66}
67
68impl EntityChangeSet {
69    pub fn is_empty(&self) -> bool {
70        self.changes.is_empty()
71    }
72
73    pub fn set(&mut self, key: EntityKey, field: impl Into<String>, value: Value) {
74        self.changes
75            .entry(key)
76            .or_default()
77            .insert(field.into(), value);
78    }
79
80    pub fn get(&self, key: &EntityKey, field: &str) -> Option<&Value> {
81        self.changes.get(key).and_then(|changes| changes.get(field))
82    }
83
84    pub fn changes(&self) -> &BTreeMap<EntityKey, Record> {
85        &self.changes
86    }
87
88    /// Remove all pending changes for a specific entity key.
89    pub fn clear_entity(&mut self, key: &EntityKey) {
90        self.changes.remove(key);
91    }
92
93    /// Get the set of field names that have been modified for a given entity key.
94    pub fn field_names(&self, key: &EntityKey) -> BTreeSet<String> {
95        self.changes
96            .get(key)
97            .map(|record| record.keys().cloned().collect())
98            .unwrap_or_default()
99    }
100}
101
102#[derive(Debug, Clone, Default, PartialEq)]
103pub struct ChangeSetStack {
104    stack: Vec<EntityChangeSet>,
105}
106
107impl ChangeSetStack {
108    pub fn current_mut(&mut self) -> &mut EntityChangeSet {
109        if self.stack.is_empty() {
110            self.stack.push(EntityChangeSet::default());
111        }
112        self.stack.last_mut().expect("change set stack has current")
113    }
114
115    pub fn current(&self) -> Option<&EntityChangeSet> {
116        self.stack.last()
117    }
118
119    pub fn push(&mut self) {
120        self.stack.push(EntityChangeSet::default());
121    }
122
123    pub fn pop(&mut self) -> Option<EntityChangeSet> {
124        self.stack.pop()
125    }
126
127    pub fn get(&self, key: &EntityKey, field: &str) -> Option<Value> {
128        self.stack
129            .iter()
130            .rev()
131            .find_map(|change_set| change_set.get(key, field).cloned())
132    }
133
134    pub fn set(&mut self, key: EntityKey, field: impl Into<String>, value: Value) {
135        self.current_mut().set(key, field, value);
136    }
137
138    pub fn clear_current(&mut self) {
139        if let Some(current) = self.stack.last_mut() {
140            *current = EntityChangeSet::default();
141        }
142    }
143
144    /// Remove all pending changes for a specific entity key across all stack levels.
145    pub fn clear_entity(&mut self, key: &EntityKey) {
146        for change_set in &mut self.stack {
147            change_set.clear_entity(key);
148        }
149    }
150
151    /// Get the union of all changed field names for a given entity key across all stack levels.
152    /// This is the Rust equivalent of Java's `entity.getUpdatedProperties()`.
153    pub fn changed_field_names(&self, key: &EntityKey) -> BTreeSet<String> {
154        let mut fields = BTreeSet::new();
155        for change_set in &self.stack {
156            fields.extend(change_set.field_names(key));
157        }
158        fields
159    }
160}
161
162#[derive(Debug, Default)]
163pub struct RootContext {
164    change_sets: ChangeSetStack,
165    /// Annotation comment for observability during graph save.
166    comment: Option<String>,
167    /// Entity keys that have been marked for deletion.
168    /// When the entity is saved, the graph save pipeline will treat these as Remove operations.
169    deleted_keys: std::collections::BTreeSet<EntityKey>,
170    /// Entity keys that have been marked as newly inserted.
171    new_keys: std::collections::BTreeSet<EntityKey>,
172    /// The original loaded snapshot record, used to avoid redundant fetching during save.
173    original_record: Option<Record>,
174    /// Trace chains associated with each entity key.
175    trace_chains: std::collections::BTreeMap<EntityKey, Vec<teaql_core::TraceNode>>,
176    /// Original versions of entities to perform optimistic concurrency control.
177    original_versions: std::collections::BTreeMap<EntityKey, i64>,
178    /// Indicates if this entity root is entirely new.
179    is_new: bool,
180}
181
182#[derive(Debug, Clone, Default)]
183pub struct EntityRoot {
184    inner: Arc<Mutex<RootContext>>,
185}
186
187impl PartialEq for EntityRoot {
188    fn eq(&self, other: &Self) -> bool {
189        Arc::ptr_eq(&self.inner, &other.inner)
190    }
191}
192
193impl EntityRoot {
194    pub fn push_change_set(&self) {
195        self.inner
196            .lock()
197            .unwrap_or_else(|e| e.into_inner())
198            .change_sets
199            .push();
200    }
201
202    pub fn pop_change_set(&self) -> Option<EntityChangeSet> {
203        self.inner
204            .lock()
205            .unwrap_or_else(|e| e.into_inner())
206            .change_sets
207            .pop()
208    }
209
210    pub fn clear_current_change_set(&self) {
211        self.inner
212            .lock()
213            .unwrap_or_else(|e| e.into_inner())
214            .change_sets
215            .clear_current();
216    }
217
218    pub fn set(&self, key: EntityKey, field: impl Into<String>, value: impl Into<Value>) {
219        self.inner
220            .lock()
221            .unwrap_or_else(|e| e.into_inner())
222            .change_sets
223            .set(key, field, value.into());
224    }
225
226    pub fn get(&self, key: &EntityKey, field: &str) -> Option<Value> {
227        self.inner
228            .lock()
229            .unwrap_or_else(|e| e.into_inner())
230            .change_sets
231            .get(key, field)
232    }
233
234    pub fn current_change_set(&self) -> EntityChangeSet {
235        self.inner
236            .lock()
237            .unwrap_or_else(|e| e.into_inner())
238            .change_sets
239            .current()
240            .cloned()
241            .unwrap_or_default()
242    }
243
244    /// Set an annotation comment on this entity root.
245    /// The comment propagates through the graph save process for observability.
246    pub fn set_comment(&self, comment: impl Into<String>) {
247        self.inner.lock().unwrap_or_else(|e| e.into_inner()).comment = Some(comment.into());
248    }
249
250    /// Get the annotation comment, if any.
251    pub fn get_comment(&self) -> Option<String> {
252        self.inner
253            .lock()
254            .unwrap_or_else(|e| e.into_inner())
255            .comment
256            .clone()
257    }
258
259    /// Mark this entity root as a newly created entity in memory.
260    pub fn mark_as_new(&self, key: EntityKey) {
261        self.inner
262            .lock()
263            .unwrap_or_else(|e| e.into_inner())
264            .new_keys
265            .insert(key);
266    }
267
268    /// Check if this entity root is marked as newly created.
269    pub fn is_new(&self, key: &EntityKey) -> bool {
270        self.inner
271            .lock()
272            .unwrap_or_else(|e| e.into_inner())
273            .new_keys
274            .contains(key)
275    }
276
277    /// Store the original record when loaded from DB.
278    pub fn set_original_record(&self, record: Record) {
279        self.inner
280            .lock()
281            .unwrap_or_else(|e| e.into_inner())
282            .original_record = Some(record);
283    }
284
285    /// Retrieve the original record.
286    pub fn original_record(&self) -> Option<Record> {
287        self.inner
288            .lock()
289            .unwrap_or_else(|e| e.into_inner())
290            .original_record
291            .clone()
292    }
293
294    /// Mark an entity as deleted. The next `save()` call will treat this entity
295    /// as a Remove operation in the graph save pipeline.
296    /// Any pending field changes for this entity are cleared — they are irrelevant
297    /// when the entity is being deleted.
298    pub fn mark_as_delete(&self, key: EntityKey) {
299        let mut ctx = self.inner.lock().unwrap_or_else(|e| e.into_inner());
300        ctx.change_sets.clear_entity(&key);
301        ctx.deleted_keys.insert(key);
302    }
303
304    /// Check whether an entity has been marked for deletion.
305    pub fn is_marked_as_delete(&self, key: &EntityKey) -> bool {
306        self.inner
307            .lock()
308            .unwrap_or_else(|e| e.into_inner())
309            .deleted_keys
310            .contains(key)
311    }
312
313    /// Get the set of field names that have been modified for the given entity key.
314    /// This is the Rust equivalent of Java's `entity.getUpdatedProperties()`.
315    pub fn changed_field_names(&self, key: &EntityKey) -> BTreeSet<String> {
316        self.inner
317            .lock()
318            .unwrap_or_else(|e| e.into_inner())
319            .change_sets
320            .changed_field_names(key)
321    }
322    pub fn deleted_keys(&self) -> std::collections::BTreeSet<EntityKey> {
323        self.inner
324            .lock()
325            .unwrap_or_else(|e| e.into_inner())
326            .deleted_keys
327            .clone()
328    }
329
330    pub fn new_keys(&self) -> std::collections::BTreeSet<EntityKey> {
331        self.inner
332            .lock()
333            .unwrap_or_else(|e| e.into_inner())
334            .new_keys
335            .clone()
336    }
337
338    pub fn get_original_version(&self, key: &EntityKey) -> Option<i64> {
339        self.inner
340            .lock()
341            .unwrap_or_else(|e| e.into_inner())
342            .original_versions
343            .get(key)
344            .cloned()
345    }
346
347    pub fn get_trace_chain(&self, key: &EntityKey) -> Vec<teaql_core::TraceNode> {
348        self.inner
349            .lock()
350            .unwrap_or_else(|e| e.into_inner())
351            .trace_chains
352            .get(key)
353            .cloned()
354            .unwrap_or_default()
355    }
356
357    pub fn set_original_version(&self, key: EntityKey, version: i64) {
358        self.inner
359            .lock()
360            .unwrap_or_else(|e| e.into_inner())
361            .original_versions
362            .insert(key, version);
363    }
364}
365
366pub trait LedgerEntity: teaql_core::Entity {
367    fn entity_root(&self) -> Option<EntityRoot>;
368}