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