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