Skip to main content

teaql_runtime/
context.rs

1use std::any::{Any, TypeId};
2use std::collections::{BTreeMap, HashMap};
3use std::future::Future;
4
5use std::pin::Pin;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Condvar, Mutex, OnceLock};
8use std::time::{Duration, Instant, SystemTime};
9
10use teaql_core::{EntityDescriptor, Value};
11use teaql_sql::{CompiledQuery, DatabaseKind};
12
13use crate::EntityRuntimeState;
14use crate::{
15    CheckObjectStatus, CheckResult, CheckResults, CheckerRegistry, ContextError,
16    EntityDataServiceBehavior, EntityDataServiceBehaviorRegistry, EntityGraphBuilder,
17    EntityRegistry, GraphNode, InMemoryEntityGraphDecoderRegistry, InternalIdGenerator, Language,
18    MetadataStore, ObjectLocation, RawAuditEvent, RawAuditEventSink, RequestPolicy, RuntimeError,
19    local_id_generator,
20};
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct ContextEntityRef {
24    pub entity_type: String,
25    pub id: u64,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct ContextRootError {
30    pub expected_entity_type: String,
31    pub actual_root: Option<ContextEntityRef>,
32}
33
34impl std::fmt::Display for ContextRootError {
35    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        match &self.actual_root {
37            None => write!(
38                formatter,
39                "active root {} is missing from UserContext",
40                self.expected_entity_type
41            ),
42            Some(actual) => write!(
43                formatter,
44                "active root type is {}, expected {}",
45                actual.entity_type, self.expected_entity_type
46            ),
47        }
48    }
49}
50
51impl std::error::Error for ContextRootError {}
52
53#[cfg(test)]
54mod active_root_tests {
55    use super::UserContext;
56
57    #[test]
58    fn active_root_is_typed_and_fails_closed() {
59        let context = UserContext::new().with_active_root("Tenant", 42);
60        assert_eq!(context.require_active_root("Tenant").unwrap().id, 42);
61        assert!(context.require_active_root("Organization").is_err());
62        assert!(UserContext::new().require_active_root("Tenant").is_err());
63    }
64}
65
66#[derive(Debug, Clone, PartialEq)]
67pub struct ContinuousPageCursor {
68    pub cursor_id: String,
69    pub query_key: String,
70    pub entity: String,
71    pub direction: teaql_core::SortDirection,
72    pub boundary: Value,
73    pub page_size: u64,
74    pub next_offset: u64,
75    pub expires_at: SystemTime,
76}
77
78#[async_trait::async_trait]
79pub trait ContinuousPageCursorStore: Send + Sync + 'static {
80    async fn get(
81        &self,
82        query_key: &str,
83        target_offset: u64,
84    ) -> Result<Option<ContinuousPageCursor>, String>;
85    async fn put(&self, cursor: ContinuousPageCursor) -> Result<(), String>;
86    async fn invalidate(&self, query_key: &str) -> Result<(), String>;
87}
88
89pub struct InMemoryContinuousPageCursorStore {
90    cursors: Mutex<HashMap<String, ContinuousPageCursor>>,
91    max_entries: usize,
92}
93
94#[derive(Debug, Clone)]
95pub struct RetainedIdSet {
96    pub query_key: String,
97    pub ids: Arc<Vec<u64>>,
98    pub expires_at: SystemTime,
99}
100
101#[async_trait::async_trait]
102pub trait IdSetStore: Send + Sync + 'static {
103    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String>;
104    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String>;
105    async fn invalidate(&self, query_key: &str) -> Result<(), String>;
106}
107
108pub struct InMemoryIdSetStore {
109    sets: Mutex<HashMap<String, RetainedIdSet>>,
110    max_entries: usize,
111    max_bytes: usize,
112}
113
114impl Default for InMemoryIdSetStore {
115    fn default() -> Self {
116        Self {
117            sets: Mutex::new(HashMap::new()),
118            max_entries: 64,
119            max_bytes: 256 * 1024 * 1024,
120        }
121    }
122}
123
124impl InMemoryIdSetStore {
125    fn retained_bytes(sets: &HashMap<String, RetainedIdSet>) -> usize {
126        sets.values()
127            .map(|value| value.ids.len().saturating_mul(std::mem::size_of::<u64>()))
128            .sum()
129    }
130}
131
132#[async_trait::async_trait]
133impl IdSetStore for InMemoryIdSetStore {
134    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String> {
135        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
136        if sets
137            .get(query_key)
138            .is_some_and(|value| value.expires_at <= SystemTime::now())
139        {
140            sets.remove(query_key);
141        }
142        Ok(sets.get(query_key).cloned())
143    }
144
145    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String> {
146        let incoming_bytes = id_set.ids.len().saturating_mul(std::mem::size_of::<u64>());
147        if incoming_bytes > self.max_bytes {
148            return Err("ID set exceeds the process-local store memory ceiling".to_owned());
149        }
150        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
151        sets.retain(|_, value| value.expires_at > SystemTime::now());
152        while sets.len() >= self.max_entries
153            || Self::retained_bytes(&sets).saturating_add(incoming_bytes) > self.max_bytes
154        {
155            let Some(oldest) = sets
156                .iter()
157                .min_by_key(|(_, value)| value.expires_at)
158                .map(|(key, _)| key.clone())
159            else {
160                break;
161            };
162            sets.remove(&oldest);
163        }
164        sets.insert(id_set.query_key.clone(), id_set);
165        Ok(())
166    }
167
168    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
169        self.sets
170            .lock()
171            .map_err(|error| error.to_string())?
172            .remove(query_key);
173        Ok(())
174    }
175}
176
177fn id_set_build_lock(query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
178    static LOCKS: OnceLock<Mutex<HashMap<String, std::sync::Weak<futures_util::lock::Mutex<()>>>>> =
179        OnceLock::new();
180    let mut locks = LOCKS
181        .get_or_init(|| Mutex::new(HashMap::new()))
182        .lock()
183        .expect("ID set build lock registry poisoned");
184    locks.retain(|_, lock| lock.strong_count() > 0);
185    if let Some(lock) = locks.get(query_key).and_then(std::sync::Weak::upgrade) {
186        return lock;
187    }
188    let lock = Arc::new(futures_util::lock::Mutex::new(()));
189    locks.insert(query_key.to_owned(), Arc::downgrade(&lock));
190    lock
191}
192
193impl Default for InMemoryContinuousPageCursorStore {
194    fn default() -> Self {
195        Self {
196            cursors: Mutex::new(HashMap::new()),
197            max_entries: 4096,
198        }
199    }
200}
201
202#[async_trait::async_trait]
203impl ContinuousPageCursorStore for InMemoryContinuousPageCursorStore {
204    async fn get(
205        &self,
206        query_key: &str,
207        target_offset: u64,
208    ) -> Result<Option<ContinuousPageCursor>, String> {
209        let key = format!("{query_key}:{target_offset}");
210        let mut cursors = self.cursors.lock().map_err(|e| e.to_string())?;
211        if cursors
212            .get(&key)
213            .is_some_and(|cursor| cursor.expires_at <= SystemTime::now())
214        {
215            cursors.remove(&key);
216        }
217        Ok(cursors.get(&key).cloned())
218    }
219
220    async fn put(&self, cursor: ContinuousPageCursor) -> Result<(), String> {
221        let key = format!("{}:{}", cursor.query_key, cursor.next_offset);
222        let mut cursors = self.cursors.lock().map_err(|e| e.to_string())?;
223        if cursors.len() >= self.max_entries {
224            if let Some(expired_or_oldest) = cursors
225                .iter()
226                .min_by_key(|(_, value)| value.expires_at)
227                .map(|(key, _)| key.clone())
228            {
229                cursors.remove(&expired_or_oldest);
230            }
231        }
232        cursors.insert(key, cursor);
233        Ok(())
234    }
235
236    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
237        let prefix = format!("{query_key}:");
238        self.cursors
239            .lock()
240            .map_err(|e| e.to_string())?
241            .retain(|key, _| !key.starts_with(&prefix));
242        Ok(())
243    }
244}
245
246#[derive(Debug, Clone, Copy, PartialEq, Eq)]
247pub enum SqlLogOperation {
248    Select,
249    Insert,
250    Update,
251    Delete,
252    Recover,
253}
254
255impl SqlLogOperation {
256    pub fn is_select(self) -> bool {
257        matches!(self, Self::Select)
258    }
259
260    pub fn is_mutation(self) -> bool {
261        !self.is_select()
262    }
263}
264
265#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
266pub struct SqlLogOptions {
267    pub select: bool,
268    pub mutation: bool,
269}
270
271impl SqlLogOptions {
272    pub fn disabled() -> Self {
273        Self {
274            select: false,
275            mutation: false,
276        }
277    }
278
279    pub fn select_only() -> Self {
280        Self {
281            select: true,
282            mutation: false,
283        }
284    }
285
286    pub fn mutation_only() -> Self {
287        Self {
288            select: false,
289            mutation: true,
290        }
291    }
292
293    pub fn all() -> Self {
294        Self {
295            select: true,
296            mutation: true,
297        }
298    }
299
300    pub fn enabled_for(self, operation: SqlLogOperation) -> bool {
301        match operation.is_select() {
302            true => self.select,
303            false => self.mutation,
304        }
305    }
306}
307
308#[derive(Debug, Clone, PartialEq)]
309pub struct SqlLogEntry {
310    pub operation: SqlLogOperation,
311    pub sql: String,
312    pub params: Vec<Value>,
313    pub debug_sql: String,
314    pub pretty_sql: String,
315    pub started_at: SystemTime,
316    pub ended_at: SystemTime,
317    pub elapsed: Duration,
318    pub result_count: Option<usize>,
319    pub result_type: Option<String>,
320    pub affected_rows: Option<u64>,
321    pub result_summary: String,
322}
323
324#[derive(Debug, Clone, PartialEq)]
325pub struct UnifiedLogEntry {
326    pub timestamp: SystemTime,
327    pub user_identifier: Option<String>,
328    pub trace_chain: Vec<teaql_core::TraceNode>,
329    pub payload: LogPayload,
330}
331
332#[derive(Debug, Clone, PartialEq)]
333pub enum LogPayload {
334    Sql(SqlLogEntry),
335    Info(InfoLogEntry),
336}
337
338#[derive(Debug, Clone, PartialEq)]
339pub struct InfoLogEntry {
340    pub message: String,
341}
342
343#[derive(Clone, Default)]
344pub struct UnifiedLogBuffer {
345    pub entries: std::sync::Arc<Mutex<Vec<UnifiedLogEntry>>>,
346}
347
348pub trait SchemaProvider: Send + Sync {
349    fn ensure_schema<'a>(
350        &'a self,
351        context: &'a UserContext,
352    ) -> Pin<Box<dyn Future<Output = Result<(), RuntimeError>> + Send + 'a>>;
353}
354
355pub struct UserContext {
356    active_root: Option<ContextEntityRef>,
357    pub(crate) metadata: Option<Box<dyn MetadataStore>>,
358    pub(crate) entity_registry: Option<Box<dyn EntityRegistry>>,
359    pub(crate) entity_graph_decoders: InMemoryEntityGraphDecoderRegistry,
360    pub(crate) entity_data_service_behavior_registry:
361        Option<Box<dyn EntityDataServiceBehaviorRegistry>>,
362    pub(crate) request_policy: Option<Box<dyn RequestPolicy>>,
363    pub(crate) checker_registry: Option<Box<dyn CheckerRegistry>>,
364    pub(crate) event_sink: Option<Box<dyn RawAuditEventSink>>,
365    pub(crate) custom_event_sink: Option<Box<dyn crate::SafeAuditEventSink>>,
366    pub(crate) internal_id_generator: Option<Box<dyn InternalIdGenerator>>,
367    schema_provider: Option<Box<dyn SchemaProvider>>,
368    language: Language,
369    i18n_catalog: Arc<crate::I18nCatalog>,
370    typed_resources: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
371    named_resources: BTreeMap<String, Box<dyn Any + Send + Sync>>,
372    locals: BTreeMap<String, Value>,
373    pub(crate) initial_graphs: Vec<GraphNode>,
374    pub(crate) root_graphs: Vec<GraphNode>,
375    entity_runtime_state: EntityRuntimeState,
376    sql_log_options: SqlLogOptions,
377    sql_log_entries: Mutex<Vec<SqlLogEntry>>,
378    user_identifier: Option<String>,
379    timezone: Option<String>,
380    trace_id: String,
381    continuous_page_cursor_store: std::sync::Arc<dyn ContinuousPageCursorStore>,
382    continuous_page_observation: Mutex<(String, Option<String>)>,
383    id_set_store: Arc<dyn IdSetStore>,
384    id_set_observation: Mutex<(String, Option<u64>)>,
385    local_lock_owner: u64,
386    remote_lock_owner: String,
387    runtime_telemetry: Arc<dyn crate::RuntimeTelemetry>,
388}
389
390#[derive(Clone, Copy)]
391struct LocalLockEntry {
392    owner: u64,
393    expires_at: Option<Instant>,
394}
395
396#[derive(Default)]
397struct ProcessLocalLocks {
398    entries: Mutex<HashMap<String, LocalLockEntry>>,
399    changed: Condvar,
400}
401
402static PROCESS_LOCAL_LOCKS: OnceLock<ProcessLocalLocks> = OnceLock::new();
403static NEXT_LOCAL_LOCK_OWNER: AtomicU64 = AtomicU64::new(1);
404
405impl Default for UserContext {
406    fn default() -> Self {
407        let pid = std::process::id();
408        let thread_id_str = format!("{:?}", std::thread::current().id());
409        let numeric_thread_id = thread_id_str
410            .strip_prefix("ThreadId(")
411            .and_then(|s| s.strip_suffix(")"))
412            .unwrap_or(&thread_id_str);
413        let os_user = std::env::var("USER")
414            .or_else(|_| std::env::var("USERNAME"))
415            .unwrap_or_else(|_| "main".to_owned());
416        let user_id = format!("{os_user}@pid-{pid}.tid-{numeric_thread_id}");
417        let owner_sequence = NEXT_LOCAL_LOCK_OWNER.fetch_add(1, Ordering::Relaxed);
418        Self {
419            active_root: None,
420            metadata: None,
421            entity_registry: None,
422            entity_graph_decoders: InMemoryEntityGraphDecoderRegistry::default(),
423            entity_data_service_behavior_registry: None,
424            request_policy: None,
425            checker_registry: None,
426            event_sink: None,
427            custom_event_sink: None,
428            internal_id_generator: None,
429            schema_provider: None,
430            language: Language::default(),
431            i18n_catalog: crate::I18nCatalog::builtin().clone(),
432            typed_resources: HashMap::new(),
433            named_resources: BTreeMap::new(),
434            locals: BTreeMap::new(),
435            initial_graphs: Vec::new(),
436            root_graphs: Vec::new(),
437            entity_runtime_state: EntityRuntimeState::default(),
438            sql_log_options: SqlLogOptions::all(),
439            sql_log_entries: Mutex::new(Vec::new()),
440            user_identifier: Some(user_id),
441            timezone: Some("UTC".to_owned()),
442            trace_id: format!(
443                "req-{pid}-{numeric_thread_id}-{:x}",
444                std::time::SystemTime::now()
445                    .duration_since(std::time::UNIX_EPOCH)
446                    .unwrap_or_default()
447                    .as_micros()
448            ),
449            continuous_page_cursor_store: std::sync::Arc::new(
450                InMemoryContinuousPageCursorStore::default(),
451            ),
452            continuous_page_observation: Mutex::new(("DISABLED".to_owned(), None)),
453            id_set_store: Arc::new(InMemoryIdSetStore::default()),
454            id_set_observation: Mutex::new(("ID_SET_DISABLED".to_owned(), None)),
455            local_lock_owner: owner_sequence,
456            remote_lock_owner: format!(
457                "teaql:{pid}:{owner_sequence}:{}",
458                SystemTime::now()
459                    .duration_since(SystemTime::UNIX_EPOCH)
460                    .unwrap_or_default()
461                    .as_nanos()
462            ),
463            runtime_telemetry: Arc::new(crate::NoopRuntimeTelemetry),
464        }
465    }
466}
467
468#[async_trait::async_trait]
469pub trait DataStore: Send + Sync + 'static {
470    async fn get(&self, key: &str) -> Option<Value>;
471    async fn put(&self, key: &str, value: Value, timeout_seconds: Option<u64>);
472    async fn remove(&self, key: &str);
473}
474
475/// Provider-neutral distributed lock boundary.
476///
477/// Implementations must associate an acquired lock with `owner_token` and
478/// release it only while that token still owns the key. A zero timeout is one
479/// non-blocking attempt; a zero expiry means no automatic lease expiry.
480#[async_trait::async_trait]
481pub trait RemoteLockProvider: Send + Sync + 'static {
482    async fn try_remote_lock(
483        &self,
484        key: &str,
485        owner_token: &str,
486        timeout_millis: u64,
487        expire_millis: u64,
488    ) -> bool;
489
490    async fn unlock_remote(&self, key: &str, owner_token: &str) -> bool;
491}
492
493#[derive(Default)]
494pub struct InMemoryDataStore {
495    cache: std::sync::RwLock<HashMap<String, (Value, Option<std::time::Instant>)>>,
496}
497
498#[async_trait::async_trait]
499impl DataStore for InMemoryDataStore {
500    async fn get(&self, key: &str) -> Option<Value> {
501        let lock = self.cache.read().unwrap();
502        if let Some((val, expires_at)) = lock.get(key) {
503            if let Some(exp) = expires_at {
504                if std::time::Instant::now() > *exp {
505                    return None;
506                }
507            }
508            return Some(val.clone());
509        }
510        None
511    }
512
513    async fn put(&self, key: &str, value: Value, timeout_seconds: Option<u64>) {
514        let mut lock = self.cache.write().unwrap();
515        let expires_at = timeout_seconds
516            .map(|secs| std::time::Instant::now() + std::time::Duration::from_secs(secs));
517        lock.insert(key.to_string(), (value, expires_at));
518    }
519
520    async fn remove(&self, key: &str) {
521        let mut lock = self.cache.write().unwrap();
522        lock.remove(key);
523    }
524}
525
526impl UserContext {
527    pub fn new() -> Self {
528        Self::default()
529    }
530
531    pub fn with_active_root(mut self, entity_type: impl Into<String>, id: u64) -> Self {
532        let entity_type = entity_type.into();
533        assert!(
534            !entity_type.trim().is_empty(),
535            "active root entity type is required"
536        );
537        assert!(id > 0, "active root id must be positive");
538        self.active_root = Some(ContextEntityRef { entity_type, id });
539        self
540    }
541
542    pub fn require_active_root(
543        &self,
544        expected_entity_type: &str,
545    ) -> Result<&ContextEntityRef, ContextRootError> {
546        match &self.active_root {
547            Some(root) if root.entity_type == expected_entity_type => Ok(root),
548            actual_root => Err(ContextRootError {
549                expected_entity_type: expected_entity_type.to_owned(),
550                actual_root: actual_root.clone(),
551            }),
552        }
553    }
554
555    pub(crate) fn active_root_ref(&self) -> Option<&ContextEntityRef> {
556        self.active_root.as_ref()
557    }
558
559    pub fn with_runtime_telemetry(mut self, telemetry: Arc<dyn crate::RuntimeTelemetry>) -> Self {
560        self.runtime_telemetry = telemetry;
561        self
562    }
563
564    pub fn set_runtime_telemetry(&mut self, telemetry: Arc<dyn crate::RuntimeTelemetry>) {
565        self.runtime_telemetry = telemetry;
566    }
567
568    pub fn runtime_telemetry(&self) -> &Arc<dyn crate::RuntimeTelemetry> {
569        &self.runtime_telemetry
570    }
571
572    pub(crate) fn runtime_telemetry_is_noop(&self) -> bool {
573        self.runtime_telemetry.is_noop()
574    }
575
576    pub fn start_runtime_operation(
577        &self,
578        operation: crate::RuntimeOperation,
579    ) -> crate::FailOpenRuntimeTelemetryScope {
580        crate::start_runtime_operation(&self.runtime_telemetry, operation)
581    }
582
583    pub fn try_local_lock(&self, key: &str, timeout_millis: u64, expire_millis: u64) -> bool {
584        let locks = PROCESS_LOCAL_LOCKS.get_or_init(ProcessLocalLocks::default);
585        let deadline = Instant::now() + Duration::from_millis(timeout_millis);
586        let mut entries = locks.entries.lock().expect("local lock state poisoned");
587        loop {
588            let now = Instant::now();
589            match entries.get(key).copied() {
590                None => {
591                    entries.insert(
592                        key.to_owned(),
593                        LocalLockEntry {
594                            owner: self.local_lock_owner,
595                            expires_at: (expire_millis > 0)
596                                .then(|| now + Duration::from_millis(expire_millis)),
597                        },
598                    );
599                    return true;
600                }
601                Some(current)
602                    if current.owner == self.local_lock_owner
603                        || current.expires_at.is_some_and(|expiry| now >= expiry) =>
604                {
605                    entries.insert(
606                        key.to_owned(),
607                        LocalLockEntry {
608                            owner: self.local_lock_owner,
609                            expires_at: (expire_millis > 0)
610                                .then(|| now + Duration::from_millis(expire_millis)),
611                        },
612                    );
613                    return true;
614                }
615                Some(current) => {
616                    if timeout_millis == 0 || now >= deadline {
617                        return false;
618                    }
619                    let wake_after = current
620                        .expires_at
621                        .map(|expiry| expiry.saturating_duration_since(now))
622                        .unwrap_or_else(|| deadline.saturating_duration_since(now))
623                        .min(deadline.saturating_duration_since(now));
624                    let waited = locks
625                        .changed
626                        .wait_timeout(entries, wake_after)
627                        .expect("local lock state poisoned");
628                    entries = waited.0;
629                }
630            }
631        }
632    }
633
634    pub fn unlock_local(&self, key: &str) {
635        let locks = PROCESS_LOCAL_LOCKS.get_or_init(ProcessLocalLocks::default);
636        let mut entries = locks.entries.lock().expect("local lock state poisoned");
637        if entries
638            .get(key)
639            .is_some_and(|entry| entry.owner == self.local_lock_owner)
640        {
641            entries.remove(key);
642            locks.changed.notify_all();
643        }
644    }
645
646    /// Attempts to acquire a provider-backed distributed lock.
647    ///
648    /// A missing provider remains a no-op success, matching the optional
649    /// Remote Lock boundary in the other TeaQL runtimes. Install an
650    /// `Arc<dyn RemoteLockProvider>` resource to enable distributed exclusion.
651    pub async fn try_remote_lock(
652        &self,
653        key: &str,
654        timeout_millis: u64,
655        expire_millis: u64,
656    ) -> bool {
657        match self.get_resource::<Arc<dyn RemoteLockProvider>>() {
658            Some(provider) => {
659                provider
660                    .try_remote_lock(key, &self.remote_lock_owner, timeout_millis, expire_millis)
661                    .await
662            }
663            None => true,
664        }
665    }
666
667    /// Releases a distributed lock only when this context still owns it.
668    pub async fn unlock_remote(&self, key: &str) -> bool {
669        match self.get_resource::<Arc<dyn RemoteLockProvider>>() {
670            Some(provider) => provider.unlock_remote(key, &self.remote_lock_owner).await,
671            None => true,
672        }
673    }
674
675    pub fn user_identifier(&self) -> Option<&str> {
676        self.user_identifier.as_deref()
677    }
678
679    pub fn set_user_identifier(&mut self, user_identifier: impl Into<String>) {
680        self.user_identifier = Some(user_identifier.into());
681    }
682
683    pub fn set_continuous_page_cursor_store(
684        &mut self,
685        store: std::sync::Arc<dyn ContinuousPageCursorStore>,
686    ) {
687        self.continuous_page_cursor_store = store;
688    }
689
690    pub fn continuous_page_plan(&self) -> Option<String> {
691        self.continuous_page_observation
692            .lock()
693            .ok()
694            .map(|value| value.0.clone())
695    }
696
697    pub fn continuous_page_cursor_id(&self) -> Option<String> {
698        self.continuous_page_observation
699            .lock()
700            .ok()
701            .and_then(|value| value.1.clone())
702    }
703
704    pub(crate) fn observe_continuous_page(
705        &self,
706        plan: impl Into<String>,
707        cursor_id: Option<String>,
708    ) {
709        if let Ok(mut observation) = self.continuous_page_observation.lock() {
710            *observation = (plan.into(), cursor_id);
711        }
712    }
713
714    pub(crate) fn continuous_page_cursor_store(&self) -> &dyn ContinuousPageCursorStore {
715        self.continuous_page_cursor_store.as_ref()
716    }
717
718    pub fn set_id_set_store(&mut self, store: Arc<dyn IdSetStore>) {
719        self.id_set_store = store;
720    }
721
722    pub fn id_set_plan(&self) -> Option<String> {
723        self.id_set_observation
724            .lock()
725            .ok()
726            .map(|observation| observation.0.clone())
727    }
728
729    pub fn id_set_count(&self) -> Option<u64> {
730        self.id_set_observation
731            .lock()
732            .ok()
733            .and_then(|observation| observation.1)
734    }
735
736    pub(crate) fn observe_id_set(&self, plan: impl Into<String>, count: Option<u64>) {
737        if let Ok(mut observation) = self.id_set_observation.lock() {
738            *observation = (plan.into(), count);
739        }
740    }
741
742    pub(crate) fn id_set_store(&self) -> &dyn IdSetStore {
743        self.id_set_store.as_ref()
744    }
745
746    pub(crate) fn id_set_build_lock(&self, query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
747        id_set_build_lock(query_key)
748    }
749
750    pub fn with_user_identifier(mut self, user_identifier: impl Into<String>) -> Self {
751        self.user_identifier = Some(user_identifier.into());
752        self
753    }
754
755    pub fn set_user_identifier_option(&mut self, user_identifier: Option<String>) {
756        self.user_identifier = user_identifier;
757    }
758
759    pub fn with_user_identifier_option(mut self, user_identifier: Option<String>) -> Self {
760        self.user_identifier = user_identifier;
761        self
762    }
763
764    pub fn timezone(&self) -> Option<&str> {
765        self.timezone.as_deref()
766    }
767
768    pub fn set_timezone(&mut self, timezone: impl Into<String>) {
769        self.timezone = Some(timezone.into());
770    }
771
772    pub fn with_timezone(mut self, timezone: impl Into<String>) -> Self {
773        self.timezone = Some(timezone.into());
774        self
775    }
776
777    pub fn trace_id(&self) -> &str {
778        &self.trace_id
779    }
780
781    pub fn set_trace_id(&mut self, trace_id: impl Into<String>) {
782        self.trace_id = trace_id.into();
783    }
784
785    pub fn with_trace_id(mut self, trace_id: impl Into<String>) -> Self {
786        self.trace_id = trace_id.into();
787        self
788    }
789
790    pub fn with_module(mut self, module: crate::RuntimeModule) -> Self {
791        module.apply_to(&mut self);
792        self
793    }
794
795    pub fn entity_runtime_state(&self) -> EntityRuntimeState {
796        // UserContext owns only the immutable identity-graph anchor. Every query/new-entity
797        // operation receives fresh mutation state, even when the same context is reused.
798        EntityRuntimeState::fresh_with_shared_graph(&self.entity_runtime_state)
799    }
800
801    pub fn initial_graphs(&self) -> &[GraphNode] {
802        &self.initial_graphs
803    }
804
805    pub fn set_initial_graphs(&mut self, graphs: Vec<GraphNode>) {
806        self.initial_graphs = graphs;
807    }
808
809    pub fn root_graphs(&self) -> &[GraphNode] {
810        &self.root_graphs
811    }
812
813    pub fn set_root_graphs(&mut self, graphs: Vec<GraphNode>) {
814        self.root_graphs = graphs;
815    }
816
817    pub fn with_metadata(mut self, metadata: impl MetadataStore + 'static) -> Self {
818        self.metadata = Some(Box::new(metadata));
819        self
820    }
821
822    pub fn set_metadata(&mut self, metadata: impl MetadataStore + 'static) {
823        self.metadata = Some(Box::new(metadata));
824    }
825
826    pub fn with_entity_registry(mut self, registry: impl EntityRegistry + 'static) -> Self {
827        self.entity_registry = Some(Box::new(registry));
828        self
829    }
830
831    pub fn set_entity_registry(&mut self, registry: impl EntityRegistry + 'static) {
832        self.entity_registry = Some(Box::new(registry));
833    }
834
835    pub fn set_entity_graph_decoder_registry(
836        &mut self,
837        registry: InMemoryEntityGraphDecoderRegistry,
838    ) {
839        self.entity_graph_decoders = registry;
840    }
841
842    pub(crate) fn has_entity_graph_decoder(&self, entity: &str) -> bool {
843        self.entity_graph_decoders.contains(entity)
844    }
845
846    pub(crate) fn decode_compact_entity_into_graph(
847        &self,
848        entity: &str,
849        row: teaql_core::CompactRow,
850        root: &EntityRuntimeState,
851        graph: &mut EntityGraphBuilder,
852    ) -> Result<(), teaql_core::EntityError> {
853        self.entity_graph_decoders
854            .decode_compact(entity, row, root, graph)
855    }
856
857    pub(crate) fn decode_compact_entity_list_into_graph(
858        &self,
859        entity: &str,
860        rows: Vec<teaql_core::CompactRow>,
861        root: &EntityRuntimeState,
862        graph: &mut EntityGraphBuilder,
863        owner_entity: &str,
864        owner_id: u64,
865        relation: &str,
866    ) -> Result<(), teaql_core::EntityError> {
867        self.entity_graph_decoders.decode_compact_list(
868            entity,
869            rows,
870            root,
871            graph,
872            owner_entity,
873            owner_id,
874            relation,
875        )
876    }
877
878    pub(crate) fn decode_compact_entity_batch_into_graph(
879        &self,
880        entity: &str,
881        rows: Vec<teaql_core::CompactRow>,
882        root: &EntityRuntimeState,
883        graph: &mut EntityGraphBuilder,
884    ) -> Result<(), teaql_core::EntityError> {
885        self.entity_graph_decoders
886            .decode_compact_batch(entity, rows, root, graph)
887    }
888
889    pub(crate) fn decode_compact_entity_option_into_graph(
890        &self,
891        entity: &str,
892        rows: Vec<teaql_core::CompactRow>,
893        root: &EntityRuntimeState,
894        graph: &mut EntityGraphBuilder,
895        owner_entity: &str,
896        owner_id: u64,
897        relation: &str,
898    ) -> Result<(), teaql_core::EntityError> {
899        self.entity_graph_decoders.decode_compact_option(
900            entity,
901            rows,
902            root,
903            graph,
904            owner_entity,
905            owner_id,
906            relation,
907        )
908    }
909
910    pub fn with_entity_data_service_behavior_registry(
911        mut self,
912        registry: impl EntityDataServiceBehaviorRegistry + 'static,
913    ) -> Self {
914        self.entity_data_service_behavior_registry = Some(Box::new(registry));
915        self
916    }
917
918    pub fn set_entity_data_service_behavior_registry(
919        &mut self,
920        registry: impl EntityDataServiceBehaviorRegistry + 'static,
921    ) {
922        self.entity_data_service_behavior_registry = Some(Box::new(registry));
923    }
924
925    pub fn with_request_policy(mut self, policy: impl RequestPolicy + 'static) -> Self {
926        self.request_policy = Some(Box::new(policy));
927        self
928    }
929
930    pub fn set_request_policy(&mut self, policy: impl RequestPolicy + 'static) {
931        self.request_policy = Some(Box::new(policy));
932    }
933
934    pub fn clear_request_policy(&mut self) {
935        self.request_policy = None;
936    }
937
938    pub fn with_checker_registry(mut self, registry: impl CheckerRegistry + 'static) -> Self {
939        self.checker_registry = Some(Box::new(registry));
940        self
941    }
942
943    pub fn set_checker_registry(&mut self, registry: impl CheckerRegistry + 'static) {
944        self.checker_registry = Some(Box::new(registry));
945    }
946
947    pub(crate) fn with_event_sink(mut self, sink: impl RawAuditEventSink + 'static) -> Self {
948        self.event_sink = Some(Box::new(sink));
949        self
950    }
951
952    pub(crate) fn set_event_sink(&mut self, sink: impl RawAuditEventSink + 'static) {
953        self.event_sink = Some(Box::new(sink));
954    }
955
956    pub fn with_custom_event_sink(
957        mut self,
958        sink: impl crate::SafeAuditEventSink + 'static,
959    ) -> Self {
960        self.custom_event_sink = Some(Box::new(sink));
961        self
962    }
963
964    pub fn set_custom_event_sink(&mut self, sink: impl crate::SafeAuditEventSink + 'static) {
965        self.custom_event_sink = Some(Box::new(sink));
966    }
967
968    pub fn with_internal_id_generator(
969        mut self,
970        generator: impl InternalIdGenerator + 'static,
971    ) -> Self {
972        self.internal_id_generator = Some(Box::new(generator));
973        self
974    }
975
976    pub fn set_internal_id_generator(&mut self, generator: impl InternalIdGenerator + 'static) {
977        self.internal_id_generator = Some(Box::new(generator));
978    }
979
980    pub fn with_schema_provider(mut self, provider: impl SchemaProvider + 'static) -> Self {
981        self.schema_provider = Some(Box::new(provider));
982        self
983    }
984
985    pub fn set_schema_provider(&mut self, provider: impl SchemaProvider + 'static) {
986        self.schema_provider = Some(Box::new(provider));
987    }
988
989    pub async fn ensure_schema(&self) -> Result<(), RuntimeError> {
990        let provider = self
991            .schema_provider
992            .as_ref()
993            .ok_or_else(|| RuntimeError::Schema("missing schema provider".to_owned()))?;
994        provider.ensure_schema(self).await
995    }
996
997    pub fn with_language(mut self, language: Language) -> Self {
998        self.language = language;
999        self
1000    }
1001
1002    pub fn set_language(&mut self, language: Language) {
1003        self.language = language;
1004    }
1005
1006    pub fn with_i18n_catalog(mut self, catalog: Arc<crate::I18nCatalog>) -> Self {
1007        self.i18n_catalog = catalog;
1008        self
1009    }
1010
1011    pub fn set_i18n_catalog(&mut self, catalog: Arc<crate::I18nCatalog>) {
1012        self.i18n_catalog = catalog;
1013    }
1014
1015    pub fn with_sql_log_options(mut self, options: SqlLogOptions) -> Self {
1016        self.sql_log_options = options;
1017        self
1018    }
1019
1020    pub fn set_sql_log_options(&mut self, options: SqlLogOptions) {
1021        self.sql_log_options = options;
1022    }
1023
1024    pub fn enable_select_sql_log(&mut self) {
1025        self.sql_log_options.select = true;
1026    }
1027
1028    pub fn enable_mutation_sql_log(&mut self) {
1029        self.sql_log_options.mutation = true;
1030    }
1031
1032    pub fn enable_all_sql_log(&mut self) {
1033        self.sql_log_options = SqlLogOptions::all();
1034    }
1035
1036    pub fn disable_sql_log(&mut self) {
1037        self.sql_log_options = SqlLogOptions::disabled();
1038        self.clear_sql_logs();
1039    }
1040
1041    pub fn sql_log_options(&self) -> SqlLogOptions {
1042        self.sql_log_options
1043    }
1044
1045    pub fn sql_logs(&self) -> Vec<SqlLogEntry> {
1046        self.sql_log_entries
1047            .lock()
1048            .map(|entries| entries.clone())
1049            .unwrap_or_default()
1050    }
1051
1052    pub fn clear_sql_logs(&self) {
1053        if let Ok(mut entries) = self.sql_log_entries.lock() {
1054            entries.clear();
1055        }
1056    }
1057
1058    pub(crate) fn record_sql_log(
1059        &self,
1060        operation: SqlLogOperation,
1061        query: &CompiledQuery,
1062        database_kind: DatabaseKind,
1063        started_at: SystemTime,
1064        ended_at: SystemTime,
1065        elapsed: Duration,
1066        result_count: Option<usize>,
1067        result_type: Option<String>,
1068        affected_rows: Option<u64>,
1069        trace_chain: Vec<teaql_core::TraceNode>,
1070    ) {
1071        if !self.sql_log_options.enabled_for(operation) {
1072            return;
1073        }
1074        let debug_sql = query.debug_sql(database_kind);
1075        let result_summary = sql_result_summary(
1076            operation,
1077            result_count,
1078            result_type.as_deref(),
1079            affected_rows,
1080            &debug_sql,
1081        );
1082
1083        let sql_log_entry = SqlLogEntry {
1084            operation,
1085            sql: query.sql.clone(),
1086            params: query.params.clone(),
1087            pretty_sql: pretty_sql(&debug_sql),
1088            debug_sql: debug_sql.clone(),
1089            started_at,
1090            ended_at,
1091            elapsed,
1092            result_summary: result_summary.clone(),
1093            result_count,
1094            result_type,
1095            affected_rows,
1096        };
1097
1098        if let Ok(mut entries) = self.sql_log_entries.lock() {
1099            // Keep sql_log_entries backwards-compatible for now if needed,
1100            // wait, we modified SqlLogEntry. We can just push it directly since we removed comment.
1101            // Wait, we need to push a cloned SqlLogEntry since it doesn't have comment.
1102            entries.push(sql_log_entry.clone());
1103        }
1104
1105        if let Some(buf) = self.get_resource::<UnifiedLogBuffer>() {
1106            if let Ok(mut entries) = buf.entries.lock() {
1107                entries.push(UnifiedLogEntry {
1108                    timestamp: started_at,
1109                    user_identifier: self.user_identifier.clone(),
1110                    trace_chain: trace_chain.clone(),
1111                    payload: LogPayload::Sql(sql_log_entry.clone()),
1112                });
1113            }
1114        }
1115
1116        crate::log_formatter::LogManager::write_sql_log(&trace_chain, &sql_log_entry);
1117    }
1118
1119    pub(crate) fn record_metadata_log(&self, metadata: &teaql_data_service::ExecutionMetadata) {
1120        let operation = match metadata.operation {
1121            teaql_data_service::DataServiceOperation::Query => SqlLogOperation::Select,
1122            teaql_data_service::DataServiceOperation::Insert => SqlLogOperation::Insert,
1123            teaql_data_service::DataServiceOperation::Update => SqlLogOperation::Update,
1124            teaql_data_service::DataServiceOperation::Delete => SqlLogOperation::Delete,
1125            teaql_data_service::DataServiceOperation::Recover => SqlLogOperation::Update,
1126            teaql_data_service::DataServiceOperation::Batch => SqlLogOperation::Update,
1127            teaql_data_service::DataServiceOperation::Schema => SqlLogOperation::Update,
1128        };
1129        if !self.sql_log_options.enabled_for(operation) {
1130            return;
1131        }
1132        if let Some(debug_sql) = &metadata.debug_query {
1133            let sql_log_entry = SqlLogEntry {
1134                operation,
1135                sql: metadata.parameterized_query.clone().unwrap_or_default(),
1136                params: metadata.params.clone(),
1137                pretty_sql: pretty_sql(debug_sql),
1138                debug_sql: debug_sql.clone(),
1139                started_at: metadata.started_at,
1140                ended_at: metadata.ended_at,
1141                elapsed: metadata
1142                    .ended_at
1143                    .duration_since(metadata.started_at)
1144                    .unwrap_or_default(),
1145                result_count: metadata.result_count,
1146                result_type: None, // Not directly available
1147                affected_rows: metadata.affected_rows,
1148                result_summary: String::new(), // We can synthesize this if needed, or leave it empty/basic
1149            };
1150
1151            // synthesize a summary for the log
1152            let mut summary = String::new();
1153            if let Some(c) = metadata.result_count {
1154                summary = format!("{} rows returned", c);
1155            } else if let Some(a) = metadata.affected_rows {
1156                summary = format!("{} rows affected", a);
1157            }
1158
1159            let mut final_entry = sql_log_entry;
1160            final_entry.result_summary = summary;
1161
1162            if let Ok(mut entries) = self.sql_log_entries.lock() {
1163                entries.push(final_entry.clone());
1164            }
1165
1166            if let Some(buf) = self.get_resource::<UnifiedLogBuffer>() {
1167                if let Ok(mut entries) = buf.entries.lock() {
1168                    entries.push(UnifiedLogEntry {
1169                        timestamp: metadata.started_at,
1170                        user_identifier: self.user_identifier.clone(),
1171                        trace_chain: metadata.trace_chain.clone(),
1172                        payload: LogPayload::Sql(final_entry.clone()),
1173                    });
1174                }
1175            }
1176
1177            crate::log_formatter::LogManager::write_sql_log(&metadata.trace_chain, &final_entry);
1178        }
1179    }
1180
1181    pub fn language(&self) -> Language {
1182        self.language
1183    }
1184
1185    pub fn set_language_code(&mut self, code: &str) -> Result<(), RuntimeError> {
1186        let Some(language) = Language::from_code(code) else {
1187            return Err(RuntimeError::UnsupportedLocale(code.to_owned()));
1188        };
1189        self.language = language;
1190        Ok(())
1191    }
1192
1193    pub fn set_locale_code(&mut self, code: &str) -> Result<(), RuntimeError> {
1194        self.set_language_code(code)
1195    }
1196
1197    pub fn generate_id(&self, entity: &str) -> Result<Option<u64>, RuntimeError> {
1198        self.internal_id_generator
1199            .as_ref()
1200            .map(|generator| generator.generate_id(entity))
1201            .transpose()
1202    }
1203
1204    pub fn next_id(&self, entity: &str) -> Result<u64, RuntimeError> {
1205        match self.generate_id(entity)? {
1206            Some(id) => Ok(id),
1207            None => local_id_generator().generate_id(entity),
1208        }
1209    }
1210
1211    pub fn entity(&self, name: &str) -> Option<&EntityDescriptor> {
1212        self.metadata
1213            .as_ref()
1214            .and_then(|metadata| metadata.entity(name))
1215    }
1216
1217    pub fn all_entities(&self) -> Vec<&EntityDescriptor> {
1218        self.metadata
1219            .as_ref()
1220            .map(|metadata| metadata.all_entities())
1221            .unwrap_or_default()
1222    }
1223
1224    pub fn require_entity(&self, name: &str) -> Result<&EntityDescriptor, RuntimeError> {
1225        self.entity(name)
1226            .ok_or_else(|| RuntimeError::MissingEntity(name.to_owned()))
1227    }
1228
1229    pub fn insert_resource<T>(&mut self, resource: T)
1230    where
1231        T: Send + Sync + 'static,
1232    {
1233        self.typed_resources
1234            .insert(TypeId::of::<T>(), Box::new(resource));
1235    }
1236
1237    pub fn get_resource<T>(&self) -> Option<&T>
1238    where
1239        T: Send + Sync + 'static,
1240    {
1241        self.typed_resources
1242            .get(&TypeId::of::<T>())
1243            .and_then(|value| value.downcast_ref::<T>())
1244    }
1245
1246    pub fn require_resource<T>(&self) -> Result<&T, ContextError>
1247    where
1248        T: Send + Sync + 'static,
1249    {
1250        self.get_resource::<T>()
1251            .ok_or(ContextError::MissingTypedResource(
1252                std::any::type_name::<T>(),
1253            ))
1254    }
1255
1256    pub fn insert_named_resource<T>(&mut self, name: impl Into<String>, resource: T)
1257    where
1258        T: Send + Sync + 'static,
1259    {
1260        self.named_resources.insert(name.into(), Box::new(resource));
1261    }
1262
1263    pub fn get_named_resource<T>(&self, name: &str) -> Option<&T>
1264    where
1265        T: Send + Sync + 'static,
1266    {
1267        self.named_resources
1268            .get(name)
1269            .and_then(|value| value.downcast_ref::<T>())
1270    }
1271
1272    pub fn require_named_resource<T>(&self, name: &str) -> Result<&T, ContextError>
1273    where
1274        T: Send + Sync + 'static,
1275    {
1276        self.get_named_resource::<T>(name)
1277            .ok_or_else(|| ContextError::MissingResource(name.to_owned()))
1278    }
1279
1280    pub fn put_local(&mut self, key: impl Into<String>, value: impl Into<Value>) {
1281        self.locals.insert(key.into(), value.into());
1282    }
1283
1284    pub fn local(&self, key: &str) -> Option<&Value> {
1285        self.locals.get(key)
1286    }
1287
1288    pub fn remove_local(&mut self, key: &str) -> Option<Value> {
1289        self.locals.remove(key)
1290    }
1291
1292    pub fn has_entity_data_service(&self, entity: &str) -> bool {
1293        let in_registry = self
1294            .entity_registry
1295            .as_ref()
1296            .map(|registry| registry.contains(entity))
1297            .unwrap_or(false);
1298        in_registry || self.entity(entity).is_some()
1299    }
1300
1301    pub fn entity_data_service_behavior(
1302        &self,
1303        entity: &str,
1304    ) -> Option<std::sync::Arc<dyn EntityDataServiceBehavior>> {
1305        self.entity_data_service_behavior_registry
1306            .as_ref()
1307            .and_then(|registry| registry.behavior(entity))
1308    }
1309
1310    pub fn has_checker(&self, entity: &str) -> bool {
1311        self.checker_registry
1312            .as_ref()
1313            .and_then(|registry| registry.checker(entity))
1314            .is_some()
1315    }
1316
1317    pub fn check_and_fix_values(
1318        &self,
1319        entity: &str,
1320        values: &mut crate::EntityValues,
1321    ) -> Result<(), RuntimeError> {
1322        self.check_and_fix_values_at(entity, values, &ObjectLocation::root())
1323    }
1324
1325    pub fn check_and_fix_values_at(
1326        &self,
1327        entity: &str,
1328        values: &mut crate::EntityValues,
1329        location: &ObjectLocation,
1330    ) -> Result<(), RuntimeError> {
1331        let status = CheckObjectStatus::from_values(values);
1332        let checker = self
1333            .checker_registry
1334            .as_ref()
1335            .and_then(|registry| registry.checker(entity));
1336        let mut results = CheckResults::new();
1337        if let Some(checker) = checker {
1338            checker.check_and_fix(self, values, location, &mut results);
1339        }
1340
1341        // Keep runtime validation aligned with the schema generated from the
1342        // same metadata. Custom checkers get the first chance to supply or fix
1343        // a value; afterwards every NOT NULL property must be present on a
1344        // create, and an update must not explicitly clear one.
1345        if let Some(descriptor) = self
1346            .metadata
1347            .as_ref()
1348            .and_then(|metadata| metadata.entity(entity))
1349        {
1350            for property in descriptor
1351                .properties
1352                .iter()
1353                // The optimistic-lock version is runtime-managed. Insert
1354                // preparation assigns its initial value after check/fix, so a
1355                // create caller must never be required to provide it.
1356                .filter(|property| !property.nullable && !property.is_version)
1357            {
1358                let missing = !values.contains_key(&property.name);
1359                let null = matches!(values.get(&property.name), Some(Value::Null));
1360                let property_location = location.clone().member(&property.name);
1361                let already_reported = results.iter().any(|result| {
1362                    result.rule == crate::CheckRule::Required
1363                        && result.location == property_location
1364                });
1365                if ((status.is_create() && missing) || null) && !already_reported {
1366                    results.push(CheckResult::required(property_location));
1367                }
1368            }
1369        }
1370        if results.is_empty() {
1371            return Ok(());
1372        }
1373        self.translate_check_results(&mut results);
1374        Err(RuntimeError::Check(results))
1375    }
1376
1377    pub fn translate_check_results(&self, results: &mut CheckResults) {
1378        for result in results {
1379            if result.message.is_none() {
1380                result.message = Some(
1381                    self.i18n_catalog
1382                        .translate_check_result(self.language, result),
1383                );
1384            }
1385        }
1386    }
1387
1388    pub fn send_event(&self, event: RawAuditEvent) -> Result<(), RuntimeError> {
1389        let scope = self.start_runtime_operation(
1390            crate::RuntimeOperation::new("audit", format!("{}.event", event.entity))
1391                .attribute("teaql.entity.type", event.entity.clone()),
1392        );
1393        let result = self.send_event_inner(event);
1394        match &result {
1395            Ok(()) => scope.success(std::collections::BTreeMap::new()),
1396            Err(_) => scope.failure("audit_error"),
1397        }
1398        result
1399    }
1400
1401    fn send_event_inner(&self, event: RawAuditEvent) -> Result<(), RuntimeError> {
1402        if let Some(sink) = self.event_sink.as_ref() {
1403            sink.on_event(self, &event)?;
1404        }
1405        if let Some(sink) = self.custom_event_sink.as_ref() {
1406            let (mask_fields, max_len) = self
1407                .metadata
1408                .as_ref()
1409                .and_then(|metadata| metadata.entity(&event.entity))
1410                .map(|desc| (desc.audit_mask_fields.clone(), desc.audit_value_max_len))
1411                .unwrap_or_else(|| (vec![], None));
1412
1413            let safe_event = event.build_safe_event(&mask_fields, max_len);
1414            sink.on_safe_event(self, &safe_event)?;
1415        }
1416
1417        crate::log_formatter::LogManager::write_audit_log(&event);
1418
1419        Ok(())
1420    }
1421
1422    pub async fn get_in_store(&self, key: &str) -> Option<Value> {
1423        let store = self.get_resource::<Box<dyn DataStore>>()?;
1424        store.get(key).await
1425    }
1426
1427    pub async fn put_in_store(
1428        &self,
1429        key: &str,
1430        value: impl Into<Value>,
1431        timeout_seconds: Option<u64>,
1432    ) {
1433        if let Some(store) = self.get_resource::<Box<dyn DataStore>>() {
1434            store.put(key, value.into(), timeout_seconds).await;
1435        }
1436    }
1437
1438    pub async fn clear_in_store(&self, key: &str) {
1439        if let Some(store) = self.get_resource::<Box<dyn DataStore>>() {
1440            store.remove(key).await;
1441        }
1442    }
1443}
1444
1445fn extract_id_from_sql(sql: &str) -> Option<String> {
1446    let sql_lower = sql.to_lowercase();
1447    let where_idx = sql_lower.find("where")?;
1448    let where_clause = &sql_lower[where_idx + 5..];
1449
1450    let bytes = where_clause.as_bytes();
1451    let mut i = 0;
1452    while i < bytes.len() {
1453        if i + 1 < bytes.len() && &bytes[i..i + 2] == b"id" {
1454            // Check boundary before
1455            let prev_ok = i == 0 || {
1456                let prev_char = bytes[i - 1] as char;
1457                !prev_char.is_ascii_alphanumeric() && prev_char != '_' && prev_char != '.'
1458            };
1459            // Check boundary after
1460            let next_ok = i + 2 == bytes.len() || {
1461                let next_char = bytes[i + 2] as char;
1462                !next_char.is_ascii_alphanumeric() && next_char != '_'
1463            };
1464
1465            if prev_ok && next_ok {
1466                // Found the standalone "id" word!
1467                // Now look for "=" after it
1468                let mut j = i + 2;
1469                while j < bytes.len() && (bytes[j] as char).is_whitespace() {
1470                    j += 1;
1471                }
1472                if j < bytes.len() && bytes[j] == b'=' {
1473                    j += 1;
1474                    while j < bytes.len() && (bytes[j] as char).is_whitespace() {
1475                        j += 1;
1476                    }
1477                    // Now extract the value
1478                    let mut val_str = String::new();
1479                    if j < bytes.len() && bytes[j] == b'\'' {
1480                        j += 1; // consume single quote
1481                        while j < bytes.len() && bytes[j] != b'\'' {
1482                            val_str.push(bytes[j] as char);
1483                            j += 1;
1484                        }
1485                        return Some(val_str);
1486                    }
1487                    // No else needed — falls through to unquoted parsing
1488                    while j < bytes.len() {
1489                        let c = bytes[j] as char;
1490                        if !c.is_ascii_alphanumeric() && c != '_' && c != '-' {
1491                            break;
1492                        }
1493                        val_str.push(c);
1494                        j += 1;
1495                    }
1496                    if !val_str.is_empty() {
1497                        return Some(val_str);
1498                    }
1499                }
1500            }
1501        }
1502        i += 1;
1503    }
1504    None
1505}
1506
1507fn sql_result_summary(
1508    operation: SqlLogOperation,
1509    result_count: Option<usize>,
1510    result_type: Option<&str>,
1511    affected_rows: Option<u64>,
1512    debug_sql: &str,
1513) -> String {
1514    match operation {
1515        SqlLogOperation::Select => {
1516            let count = result_count.unwrap_or(0);
1517            match count {
1518                0 => "MISS".to_owned(),
1519                1 => match result_type {
1520                    Some(result_type) => extract_id_from_sql(debug_sql)
1521                        .map(|id| format!("{result_type}({id})"))
1522                        .unwrap_or_else(|| result_type.to_owned()),
1523                    None => "row".to_owned(),
1524                },
1525                _ => match result_type {
1526                    Some(result_type) => format!("{count}*{result_type}"),
1527                    None => format!("{count}*rows"),
1528                },
1529            }
1530        }
1531        _ => {
1532            let affected = affected_rows.unwrap_or(0);
1533            format!("{affected} UPDATED")
1534        }
1535    }
1536}
1537
1538fn pretty_sql(sql: &str) -> String {
1539    let mut pretty = sql.to_owned();
1540    for keyword in [
1541        " FROM ",
1542        " WHERE ",
1543        " GROUP BY ",
1544        " HAVING ",
1545        " ORDER BY ",
1546        " LIMIT ",
1547        " OFFSET ",
1548        " RETURNING ",
1549    ] {
1550        pretty = pretty.replace(keyword, &format!("\n{}", keyword.trim_start()));
1551    }
1552    pretty.replace(" AND ", "\n  AND ")
1553}
1554
1555#[cfg(test)]
1556mod sql_log_option_tests {
1557    use super::*;
1558
1559    #[test]
1560    fn disabled_sql_log_rejects_executor_metadata_before_recording() {
1561        let mut context = UserContext::default();
1562        context.disable_sql_log();
1563        let now = SystemTime::now();
1564        context.record_metadata_log(&teaql_data_service::ExecutionMetadata {
1565            backend: "sql".to_owned(),
1566            operation: teaql_data_service::DataServiceOperation::Query,
1567            started_at: now,
1568            ended_at: now,
1569            affected_rows: None,
1570            result_count: Some(1),
1571            trace_chain: Vec::new(),
1572            comment: Some("disabled log test".to_owned()),
1573            backend_request_id: None,
1574            parameterized_query: Some("SELECT id FROM sample WHERE id = $1".to_owned()),
1575            params: vec![Value::I64(1)],
1576            debug_query: Some("SELECT id FROM sample WHERE id = 1".to_owned()),
1577        });
1578        assert!(context.sql_logs().is_empty());
1579    }
1580}
1581
1582#[cfg(test)]
1583mod entity_runtime_state_tests {
1584    use super::*;
1585    use crate::EntityKey;
1586
1587    #[test]
1588    fn reused_user_context_returns_independent_mutation_ledgers() {
1589        let context = UserContext::default();
1590        let first = context.entity_runtime_state();
1591        let key = EntityKey::new("School", 1_u64);
1592        first.set(key.clone(), "name", "First");
1593
1594        let second = context.entity_runtime_state();
1595
1596        assert_eq!(first.changed_field_names(&key).len(), 1);
1597        assert!(second.changed_field_names(&key).is_empty());
1598    }
1599}