1use std::collections::{HashMap, HashSet};
7use std::sync::Arc;
8
9use crate::api::{RedDBError, RedDBResult};
10use crate::storage::RedDB;
11use reddb_types::Value;
12
13thread_local! {
14 static CURRENT_CONN_ID: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
18
19 static CURRENT_AUTH_IDENTITY: std::cell::RefCell<Option<(String, crate::auth::Role)>> =
27 const { std::cell::RefCell::new(None) };
28
29 static CURRENT_SNAPSHOT: std::cell::RefCell<Option<SnapshotContext>> =
39 const { std::cell::RefCell::new(None) };
40
41 static HAS_SNAPSHOT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
47
48 static CURRENT_TENANT_ID: std::cell::RefCell<Option<String>> =
58 const { std::cell::RefCell::new(None) };
59
60 static CURRENT_CONFIG_RESOLVER: std::cell::RefCell<Option<ConfigResolver>> =
64 const { std::cell::RefCell::new(None) };
65
66 static CURRENT_SECRET_RESOLVER: std::cell::RefCell<Option<SecretResolver>> =
70 const { std::cell::RefCell::new(None) };
71
72 static CURRENT_KV_RESOLVER: std::cell::RefCell<Option<KvResolver>> =
76 const { std::cell::RefCell::new(None) };
77}
78
79pub use reddb_engine::transaction::read::SnapshotContext;
80use reddb_engine::transaction::read::{record_serializable_read, visibility_check};
81
82pub fn set_current_connection_id(id: u64) {
91 CURRENT_CONN_ID.with(|c| c.set(id));
92}
93
94pub fn clear_current_connection_id() {
96 CURRENT_CONN_ID.with(|c| c.set(0));
97}
98
99pub fn current_connection_id() -> u64 {
102 CURRENT_CONN_ID.with(|c| c.get())
103}
104
105pub fn set_current_auth_identity(username: String, role: crate::auth::Role) {
109 CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = Some((username, role)));
110}
111
112pub fn clear_current_auth_identity() {
116 CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = None);
117}
118
119pub(crate) fn current_auth_identity() -> Option<(String, crate::auth::Role)> {
122 CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone())
123}
124
125pub fn current_auth_identity_for_audit() -> Option<(String, crate::auth::Role)> {
129 current_auth_identity()
130}
131
132pub fn set_current_tenant(tenant_id: String) {
137 CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = Some(tenant_id));
138}
139
140pub fn clear_current_tenant() {
143 CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = None);
144}
145
146pub fn current_tenant() -> Option<String> {
157 let inherited = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
158 if let Some(over) = current_scope_override() {
159 if over.tenant.is_active() {
160 return over.tenant.resolve(inherited);
161 }
162 }
163 if let Some(tx_local) = current_tx_local_tenant() {
164 return tx_local;
165 }
166 inherited
167}
168
169thread_local! {
170 static TX_LOCAL_TENANT: std::cell::RefCell<Option<Option<String>>> =
179 const { std::cell::RefCell::new(None) };
180}
181
182fn current_tx_local_tenant() -> Option<Option<String>> {
183 TX_LOCAL_TENANT.with(|cell| cell.borrow().clone())
184}
185
186pub(crate) fn parse_set_local_tenant(query: &str) -> RedDBResult<Option<Option<String>>> {
192 let mut tokens = query.split_ascii_whitespace();
193 let Some(w1) = tokens.next() else {
194 return Ok(None);
195 };
196 if !w1.eq_ignore_ascii_case("SET") {
197 return Ok(None);
198 }
199 let Some(w2) = tokens.next() else {
200 return Ok(None);
201 };
202 if !w2.eq_ignore_ascii_case("LOCAL") {
203 return Ok(None);
204 }
205 let Some(w3) = tokens.next() else {
206 return Ok(None);
207 };
208 if !w3.eq_ignore_ascii_case("TENANT") {
209 return Ok(None);
210 }
211 let rest: String = tokens.collect::<Vec<_>>().join(" ");
212 let rest = rest.trim().trim_end_matches(';').trim();
213 let value_str = rest.strip_prefix('=').map(|s| s.trim()).unwrap_or(rest);
214 if value_str.is_empty() {
215 return Err(RedDBError::Query(
216 "SET LOCAL TENANT expects a string literal or NULL".to_string(),
217 ));
218 }
219 if value_str.eq_ignore_ascii_case("NULL") {
220 return Ok(Some(None));
221 }
222 if value_str.starts_with('\'') && value_str.ends_with('\'') && value_str.len() >= 2 {
223 let inner = &value_str[1..value_str.len() - 1];
224 return Ok(Some(Some(inner.to_string())));
225 }
226 Err(RedDBError::Query(format!(
227 "SET LOCAL TENANT expects a string literal or NULL, got `{value_str}`"
228 )))
229}
230
231pub(crate) struct TxLocalTenantGuard;
232
233impl TxLocalTenantGuard {
234 pub fn install(value: Option<Option<String>>) -> Self {
235 TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = value);
236 Self
237 }
238}
239
240impl Drop for TxLocalTenantGuard {
241 fn drop(&mut self) {
242 TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = None);
243 }
244}
245
246thread_local! {
247 static SCOPE_OVERRIDES: std::cell::RefCell<Vec<crate::runtime::within_clause::ScopeOverride>> =
254 const { std::cell::RefCell::new(Vec::new()) };
255}
256
257pub(crate) fn push_scope_override(over: crate::runtime::within_clause::ScopeOverride) {
258 SCOPE_OVERRIDES.with(|cell| cell.borrow_mut().push(over));
259}
260
261pub(crate) fn pop_scope_override() {
262 SCOPE_OVERRIDES.with(|cell| {
263 cell.borrow_mut().pop();
264 });
265}
266
267pub(crate) fn current_scope_override() -> Option<crate::runtime::within_clause::ScopeOverride> {
268 SCOPE_OVERRIDES.with(|cell| cell.borrow().last().cloned())
269}
270
271pub(crate) fn has_scope_override_active() -> bool {
275 SCOPE_OVERRIDES.with(|cell| !cell.borrow().is_empty())
276}
277
278pub(crate) struct ScopeOverrideGuard;
282
283impl ScopeOverrideGuard {
284 pub fn install(over: crate::runtime::within_clause::ScopeOverride) -> Self {
285 push_scope_override(over);
286 Self
287 }
288}
289
290impl Drop for ScopeOverrideGuard {
291 fn drop(&mut self) {
292 pop_scope_override();
293 }
294}
295
296pub(crate) fn current_user_projected() -> Option<String> {
302 let inherited = current_auth_identity().map(|(u, _)| u);
303 if let Some(over) = current_scope_override() {
304 if over.user.is_active() {
305 return over.user.resolve(inherited);
306 }
307 }
308 inherited
309}
310
311pub(crate) fn current_role_projected() -> Option<String> {
312 let inherited = current_auth_identity().map(|(_, r)| format!("{r:?}").to_lowercase());
313 if let Some(over) = current_scope_override() {
314 if over.role.is_active() {
315 return over.role.resolve(inherited);
316 }
317 }
318 inherited
319}
320
321pub(crate) fn current_secret_value(path: &str) -> Option<Value> {
322 let key = path.to_ascii_lowercase();
323 CURRENT_SECRET_RESOLVER.with(|cell| {
324 let mut resolver = cell.borrow_mut();
325 let resolver = resolver.as_mut()?;
326 if let Some(value) = resolver.values.get(&key) {
327 return value.clone();
328 }
329 let value = resolver.runtime.query_secret_value(&key).ok().flatten();
330 resolver.values.insert(key, value.clone());
331 value
332 })
333}
334
335pub(crate) fn secret_query_input(value: Value) -> Option<Value> {
338 if !matches!(value, Value::Secret(_)) {
339 return Some(value);
340 }
341 CURRENT_SECRET_RESOLVER.with(|cell| {
342 let resolver = cell.borrow();
343 resolver.as_ref()?.runtime.secret_query_input(value)
344 })
345}
346
347pub(crate) fn secret_query_output(value: Value) -> crate::RedDBResult<Value> {
348 if value.is_null() || matches!(value, Value::Secret(_)) {
349 return Ok(value);
350 }
351 CURRENT_SECRET_RESOLVER.with(|cell| {
352 let resolver = cell.borrow();
353 let resolver = resolver.as_ref().ok_or_else(|| {
354 RedDBError::Query("secret expression requires an active statement".into())
355 })?;
356 resolver.runtime.secret_query_output(value)
357 })
358}
359
360#[cfg(test)]
361fn secret_value_from_snapshot(values: &HashMap<String, String>, key: &str) -> Option<String> {
362 if key.starts_with("red.secret.") {
363 return None;
364 }
365 values.get(key).cloned()
366}
367
368struct SecretResolver {
369 runtime: super::RedDBRuntime,
370 values: HashMap<String, Option<Value>>,
371}
372
373pub(crate) struct SecretStoreGuard {
374 previous: Option<SecretResolver>,
375}
376
377impl SecretStoreGuard {
378 pub(super) fn install(runtime: &super::RedDBRuntime) -> Self {
379 let previous = CURRENT_SECRET_RESOLVER.with(|cell| {
380 cell.replace(Some(SecretResolver {
381 runtime: runtime.clone(),
382 values: HashMap::new(),
383 }))
384 });
385 Self { previous }
386 }
387}
388
389impl Drop for SecretStoreGuard {
390 fn drop(&mut self) {
391 let previous = self.previous.take();
392 CURRENT_SECRET_RESOLVER.with(|cell| {
393 cell.replace(previous);
394 });
395 }
396}
397
398pub(crate) fn current_kv_value(path: &str) -> Option<String> {
406 let key = path.to_ascii_lowercase();
407 CURRENT_KV_RESOLVER.with(|cell| {
408 let mut resolver = cell.borrow_mut();
409 let resolver = resolver.as_mut()?;
410 if resolver.values.is_none() {
411 resolver.values = resolver
412 .store
413 .as_ref()
414 .map(|store| store.plain_kv_snapshot());
415 }
416 let values = resolver.values.as_ref()?;
417 let found = values
421 .get(&key)
422 .map(|value| (key.as_str(), value))
423 .or_else(|| {
424 key.strip_prefix("red.kv/")
425 .and_then(|rest| values.get(rest).map(|value| (rest, value)))
426 })?;
427 if !resolver.can_read(found.0) {
428 return None;
429 }
430 Some(found.1.clone())
431 })
432}
433
434struct KvResolver {
435 store: Option<Arc<crate::auth::store::AuthStore>>,
436 values: Option<HashMap<String, String>>,
437 identity: Option<(String, crate::auth::Role, Option<String>)>,
438}
439
440impl KvResolver {
441 fn can_read(&self, key: &str) -> bool {
442 let Some(store) = &self.store else {
443 return true;
444 };
445 let Some((username, role, tenant)) = &self.identity else {
446 return true;
447 };
448 let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
449 let mut resource =
450 crate::auth::policies::ResourceRef::new("kv".to_string(), key.to_string());
451 if let Some(tenant) = tenant {
452 resource = resource.with_tenant(tenant.clone());
453 }
454 let ctx = crate::auth::policies::EvalContext {
455 principal_tenant: tenant.clone(),
456 current_tenant: tenant.clone(),
457 peer_ip: None,
458 mfa_present: false,
459 now_ms: crate::auth::now_ms(),
460 principal_is_admin_role: *role == crate::auth::Role::Admin,
461 principal_is_platform_scoped: tenant.is_none(),
462 };
463 store.check_policy_authz_with_role(&principal, "kv:read", &resource, &ctx, *role)
464 }
465}
466
467pub(crate) struct KvStoreGuard {
468 previous: Option<Box<KvResolver>>,
473}
474
475impl KvStoreGuard {
476 pub(super) fn install(store: Option<Arc<crate::auth::store::AuthStore>>) -> Self {
477 let previous = CURRENT_KV_RESOLVER.with(|cell| {
478 cell.replace(Some(KvResolver {
479 store,
480 values: None,
481 identity: current_auth_identity().map(|(username, role)| {
482 let tenant = current_tenant();
483 (username, role, tenant)
484 }),
485 }))
486 });
487 Self {
488 previous: previous.map(Box::new),
489 }
490 }
491}
492
493impl Drop for KvStoreGuard {
494 fn drop(&mut self) {
495 let previous = self.previous.take().map(|b| *b);
496 CURRENT_KV_RESOLVER.with(|cell| {
497 cell.replace(previous);
498 });
499 }
500}
501
502pub(crate) fn update_current_kv_value(path: &str, value: Option<String>) {
505 let key = path.to_ascii_lowercase();
506 CURRENT_KV_RESOLVER.with(|cell| {
507 if let Some(resolver) = cell.borrow_mut().as_mut() {
508 let Some(values) = resolver.values.as_mut() else {
509 return;
510 };
511 match value {
512 Some(value) => {
513 values.insert(key, value);
514 }
515 None => {
516 values.remove(&key);
517 }
518 }
519 }
520 });
521}
522
523pub(crate) fn current_config_value(path: &str) -> Option<Value> {
524 let key = path.to_ascii_lowercase();
525 CURRENT_CONFIG_RESOLVER.with(|cell| {
526 let mut resolver = cell.borrow_mut();
527 let resolver = resolver.as_mut()?;
528 if resolver.values.is_none() {
529 resolver.values = Some(latest_config_snapshot(&resolver.db));
530 }
531 let found = {
537 let values = resolver.values.as_ref()?;
538 values
539 .get(&key)
540 .map(|value| (key.clone(), value.clone()))
541 .or_else(|| {
542 key.strip_prefix("red.config/").and_then(|rest| {
543 values
544 .get(rest)
545 .map(|value| (rest.to_string(), value.clone()))
546 .or_else(|| {
547 let legacy = format!("red.config.{rest}");
548 values.get(&legacy).map(|value| (legacy, value.clone()))
549 })
550 })
551 })
552 }?;
553 if !resolver.can_read(&found.0) {
557 return None;
558 }
559 Some(found.1)
560 })
561}
562
563pub(crate) fn update_current_config_value(path: &str, value: Value) {
564 let key = path.to_ascii_lowercase();
565 CURRENT_CONFIG_RESOLVER.with(|cell| {
566 if let Some(resolver) = cell.borrow_mut().as_mut() {
567 if let Some(values) = resolver.values.as_mut() {
568 values.insert(key, value);
569 }
570 }
571 });
572}
573
574pub(crate) fn update_current_secret_value(_path: &str, _value: Option<String>) {
575 CURRENT_SECRET_RESOLVER.with(|cell| {
576 if let Some(resolver) = cell.borrow_mut().as_mut() {
577 resolver.values.clear();
578 }
579 });
580}
581
582#[cfg(test)]
583mod tests {
584 use super::*;
585
586 #[test]
593 fn snapshot_thread_local_stays_private_to_the_statement_frame() {
594 let src = include_str!("execution_context.rs");
595
596 for name in ["CURRENT_SNAPSHOT", "HAS_SNAPSHOT"] {
597 assert!(
598 !src.contains(&format!("pub static {name}"))
599 && !src.contains(&format!("pub(crate) static {name}")),
600 "{name} must stay private to execution_context; expose an explicit \
601 Option<SnapshotContext> via capture_current_snapshot instead"
602 );
603 }
604
605 assert!(
606 src.contains("the only sanctioned ambient read"),
607 "capture_current_snapshot must keep documenting that it is the single \
608 sanctioned ambient read of the snapshot thread-local"
609 );
610 }
611
612 #[test]
613 fn secret_snapshot_hard_blocks_system_keys() {
614 let values = HashMap::from([
615 ("red.secret.aes_key".into(), "internal".into()),
616 ("acme.key".into(), "user-value".into()),
617 ]);
618 assert_eq!(
619 secret_value_from_snapshot(&values, "red.secret.aes_key"),
620 None
621 );
622 assert_eq!(
623 secret_value_from_snapshot(&values, "acme.key").as_deref(),
624 Some("user-value")
625 );
626 }
627
628 fn with_kv_values<T>(values: HashMap<String, String>, f: impl FnOnce() -> T) -> T {
629 CURRENT_KV_RESOLVER.with(|cell| {
630 cell.replace(Some(KvResolver {
631 store: None,
632 values: Some(values),
633 identity: None,
634 }));
635 });
636 let result = f();
637 CURRENT_KV_RESOLVER.with(|cell| {
638 cell.replace(None);
639 });
640 result
641 }
642
643 #[test]
644 fn current_kv_value_resolves_namespaced_and_bare_keys() {
645 let values = HashMap::from([("acme.key".to_string(), "plain-value".to_string())]);
646
647 with_kv_values(values, || {
648 assert_eq!(
650 current_kv_value("red.kv/acme.key").as_deref(),
651 Some("plain-value")
652 );
653 assert_eq!(current_kv_value("acme.key").as_deref(), Some("plain-value"));
655 assert_eq!(current_kv_value("red.kv/missing.key"), None);
657 });
658 }
659
660 #[test]
661 fn config_key_read_allowed_hard_blocks_reserved_namespace() {
662 assert!(!config_key_read_allowed(
665 None,
666 None,
667 "red.config/red.config.aes_key"
668 ));
669 assert!(!config_key_read_allowed(None, None, "red.config.aes_key"));
670 assert!(config_key_read_allowed(
672 None,
673 None,
674 "red.config/ai.openrouter.default.key"
675 ));
676 assert!(config_key_read_allowed(None, None, "acme.flags.beta"));
677 }
678
679 #[test]
680 fn is_config_collection_matches_system_config_stores() {
681 assert!(is_config_collection("red_config"));
682 assert!(is_config_collection("RED_CONFIG"));
683 assert!(is_config_collection("red.config"));
684 assert!(!is_config_collection("users"));
685 assert!(!is_config_collection("red_kv"));
686 }
687
688 #[test]
689 fn update_current_kv_value_reflects_in_resolver() {
690 with_kv_values(HashMap::new(), || {
691 assert_eq!(current_kv_value("red.kv/feature.flag"), None);
692 update_current_kv_value("feature.flag", Some("on".to_string()));
693 assert_eq!(
694 current_kv_value("red.kv/feature.flag").as_deref(),
695 Some("on")
696 );
697 update_current_kv_value("feature.flag", None);
698 assert_eq!(current_kv_value("red.kv/feature.flag"), None);
699 });
700 }
701}
702
703fn latest_config_snapshot(db: &RedDB) -> HashMap<String, Value> {
704 let mut latest: HashMap<String, (u64, Value)> = HashMap::new();
705
706 if let Some(manager) = db.store().get_collection("red_config") {
707 manager.for_each_entity(|entity| {
708 let Some(row) = entity.data.as_row() else {
709 return true;
710 };
711 let Some(Value::Text(key)) = row.get_field("key") else {
712 return true;
713 };
714 let value = row.get_field("value").cloned().unwrap_or(Value::Null);
715 let id = entity.id.raw();
716 let key = key.to_ascii_lowercase();
717 insert_latest_config_value(&mut latest, key.clone(), id, value.clone());
718 if let Some(rest) = key.strip_prefix("red.config.") {
719 insert_latest_config_value(&mut latest, format!("red.config/{rest}"), id, value);
720 }
721 true
722 });
723 }
724
725 if let Some(manager) = db.store().get_collection("red.config") {
726 manager.for_each_entity(|entity| {
727 let Some(row) = entity.data.as_row() else {
728 return true;
729 };
730 if matches!(row.get_field("tombstone"), Some(Value::Boolean(true))) {
731 return true;
732 }
733 let Some(Value::Text(key)) = row.get_field("key") else {
734 return true;
735 };
736 let value = row.get_field("value").cloned().unwrap_or(Value::Null);
737 insert_latest_config_value(
738 &mut latest,
739 format!("red.config/{}", key.to_ascii_lowercase()),
740 entity.id.raw(),
741 value,
742 );
743 true
744 });
745 }
746
747 latest
748 .into_iter()
749 .map(|(key, (_, value))| (key, value))
750 .collect()
751}
752
753fn insert_latest_config_value(
754 latest: &mut HashMap<String, (u64, Value)>,
755 key: String,
756 id: u64,
757 value: Value,
758) {
759 match latest.get(&key) {
760 Some((prev_id, _)) if *prev_id > id => {}
761 _ => {
762 latest.insert(key, (id, value));
763 }
764 }
765}
766
767struct ConfigResolver {
768 db: Arc<RedDB>,
769 store: Option<Arc<crate::auth::store::AuthStore>>,
770 identity: Option<(String, crate::auth::Role, Option<String>)>,
771 values: Option<HashMap<String, Value>>,
772}
773
774impl ConfigResolver {
775 fn can_read(&self, key: &str) -> bool {
776 config_key_read_allowed(self.store.as_ref(), self.identity.as_ref(), key)
777 }
778}
779
780fn is_reserved_config_key(key: &str) -> bool {
784 let key = key.strip_prefix("red.config/").unwrap_or(key);
785 key.starts_with("red.config.")
786 || key == "red.config"
787 || key == super::function_catalog::REGISTRY_KEY
788}
789
790fn is_config_collection(collection: &str) -> bool {
794 matches!(
795 collection.to_ascii_lowercase().as_str(),
796 "red_config" | "red.config"
797 )
798}
799
800fn config_resource_target(key: &str) -> String {
801 let bare = key.strip_prefix("red.config/").unwrap_or(key);
802 format!("red.config/{bare}")
803}
804
805fn config_key_read_allowed(
811 store: Option<&Arc<crate::auth::store::AuthStore>>,
812 identity: Option<&(String, crate::auth::Role, Option<String>)>,
813 key: &str,
814) -> bool {
815 if is_reserved_config_key(key) {
816 return false;
817 }
818 let Some(store) = store else {
819 return true;
820 };
821 if !store.iam_authorization_enabled() {
822 return true;
823 }
824 let Some((username, role, tenant)) = identity else {
825 return true;
826 };
827 let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
828 let mut resource =
829 crate::auth::policies::ResourceRef::new("config".to_string(), config_resource_target(key));
830 if let Some(tenant) = tenant {
831 resource = resource.with_tenant(tenant.clone());
832 }
833 let ctx = crate::auth::policies::EvalContext {
834 principal_tenant: tenant.clone(),
835 current_tenant: tenant.clone(),
836 peer_ip: None,
837 mfa_present: false,
838 now_ms: crate::auth::now_ms(),
839 principal_is_admin_role: *role == crate::auth::Role::Admin,
840 principal_is_platform_scoped: tenant.is_none(),
841 };
842 store.check_policy_authz_with_role(&principal, "config:read", &resource, &ctx, *role)
843}
844
845pub(crate) fn config_read_permitted(key: &str) -> bool {
849 let key = key.to_ascii_lowercase();
850 CURRENT_CONFIG_RESOLVER.with(|cell| match cell.borrow().as_ref() {
851 Some(resolver) => resolver.can_read(&key),
852 None => !is_reserved_config_key(&key),
853 })
854}
855
856pub(crate) fn kv_read_permitted(collection: &str, key: &str) -> bool {
859 if is_config_collection(collection) {
860 config_read_permitted(key)
861 } else {
862 true
863 }
864}
865
866pub(crate) struct ConfigSnapshotGuard {
867 previous: Option<ConfigResolver>,
868}
869
870impl ConfigSnapshotGuard {
871 pub(super) fn install(
872 db: Arc<RedDB>,
873 store: Option<Arc<crate::auth::store::AuthStore>>,
874 ) -> Self {
875 let identity = current_auth_identity().map(|(username, role)| {
876 let tenant = current_tenant();
877 (username, role, tenant)
878 });
879 let previous = CURRENT_CONFIG_RESOLVER.with(|cell| {
880 cell.replace(Some(ConfigResolver {
881 db,
882 store,
883 identity,
884 values: None,
885 }))
886 });
887 Self { previous }
888 }
889}
890
891impl Drop for ConfigSnapshotGuard {
892 fn drop(&mut self) {
893 let previous = self.previous.take();
894 CURRENT_CONFIG_RESOLVER.with(|cell| {
895 cell.replace(previous);
896 });
897 }
898}
899
900pub fn set_current_snapshot(ctx: SnapshotContext) {
905 CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = Some(ctx));
906 HAS_SNAPSHOT.with(|c| c.set(true));
907}
908
909pub fn clear_current_snapshot() {
910 CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = None);
911 HAS_SNAPSHOT.with(|c| c.set(false));
912}
913
914pub(crate) struct CurrentSnapshotGuard {
920 previous: Option<SnapshotContext>,
921}
922
923impl CurrentSnapshotGuard {
924 pub(crate) fn install(ctx: SnapshotContext) -> Self {
925 let previous = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
926 set_current_snapshot(ctx);
927 Self { previous }
928 }
929}
930
931impl Drop for CurrentSnapshotGuard {
932 fn drop(&mut self) {
933 let prev = self.previous.take();
934 let has = prev.is_some();
935 CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = prev);
936 HAS_SNAPSHOT.with(|c| c.set(has));
937 }
938}
939
940#[inline]
951pub fn entity_visible_under_current_snapshot(
952 entity: &crate::storage::unified::entity::UnifiedEntity,
953) -> bool {
954 if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
961 return false;
962 }
963 if !HAS_SNAPSHOT.with(|c| c.get()) {
969 return entity.xmax == 0;
970 }
971 CURRENT_SNAPSHOT.with(|cell| {
972 let guard = cell.borrow();
973 let Some(ctx) = guard.as_ref() else {
974 return true;
975 };
976 let visible = visibility_check(ctx, entity.xmin, entity.xmax);
977 if visible {
978 record_serializable_read(ctx, entity);
979 }
980 visible
981 })
982}
983
984#[inline]
989pub(crate) fn xids_visible_under_current_snapshot(xmin: u64, xmax: u64) -> bool {
990 if !HAS_SNAPSHOT.with(|c| c.get()) {
991 return true;
992 }
993 CURRENT_SNAPSHOT.with(|cell| {
994 let guard = cell.borrow();
995 let Some(ctx) = guard.as_ref() else {
996 return true;
997 };
998 visibility_check(ctx, xmin, xmax)
999 })
1000}
1001
1002pub fn capture_current_snapshot() -> Option<SnapshotContext> {
1027 CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone())
1028}
1029
1030pub(crate) fn current_snapshot_requires_index_fallback() -> bool {
1035 if !HAS_SNAPSHOT.with(|c| c.get()) {
1036 return false;
1037 }
1038 CURRENT_SNAPSHOT.with(|cell| {
1039 cell.borrow()
1040 .as_ref()
1041 .is_some_and(|ctx| ctx.requires_index_fallback)
1042 })
1043}
1044
1045#[derive(Clone, Default)]
1060pub struct SnapshotBundle {
1061 pub snapshot: Option<SnapshotContext>,
1062 pub auth: Option<(String, crate::auth::Role)>,
1063 pub tenant: Option<String>,
1064}
1065
1066pub fn snapshot_bundle() -> SnapshotBundle {
1069 SnapshotBundle {
1070 snapshot: capture_current_snapshot(),
1071 auth: current_auth_identity(),
1072 tenant: CURRENT_TENANT_ID.with(|cell| cell.borrow().clone()),
1073 }
1074}
1075
1076pub fn with_snapshot_bundle<R>(bundle: &SnapshotBundle, f: impl FnOnce() -> R) -> R {
1081 struct Guard {
1082 prev_snapshot: Option<SnapshotContext>,
1083 prev_auth: Option<(String, crate::auth::Role)>,
1084 prev_tenant: Option<String>,
1085 }
1086 impl Drop for Guard {
1087 fn drop(&mut self) {
1088 let snap = self.prev_snapshot.take();
1089 let has = snap.is_some();
1090 CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = snap);
1091 HAS_SNAPSHOT.with(|c| c.set(has));
1092 CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = self.prev_auth.take());
1093 CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = self.prev_tenant.take());
1094 }
1095 }
1096
1097 let _guard = {
1098 let prev_snapshot = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1099 let prev_auth = CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone());
1100 let prev_tenant = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
1101
1102 match bundle.snapshot.clone() {
1103 Some(ctx) => set_current_snapshot(ctx),
1104 None => clear_current_snapshot(),
1105 }
1106 CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = bundle.auth.clone());
1107 CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = bundle.tenant.clone());
1108
1109 Guard {
1110 prev_snapshot,
1111 prev_auth,
1112 prev_tenant,
1113 }
1114 };
1115 f()
1116}
1117
1118#[inline]
1122pub fn entity_visible_with_context(
1123 ctx: Option<&SnapshotContext>,
1124 entity: &crate::storage::unified::entity::UnifiedEntity,
1125) -> bool {
1126 if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
1130 return false;
1131 }
1132 match ctx {
1133 Some(ctx) => {
1134 let visible = visibility_check(ctx, entity.xmin, entity.xmax);
1135 if visible {
1136 record_serializable_read(ctx, entity);
1137 }
1138 visible
1139 }
1140 None => true,
1141 }
1142}