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