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