Skip to main content

reddb_server/runtime/
execution_context.rs

1//! Per-thread runtime execution context.
2//!
3//! This module owns connection identity, auth/tenant scope, statement-local
4//! config/secret resolvers, and MVCC snapshot visibility helpers.
5
6use 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    /// Current connection id for the executing statement. Set by the
15    /// per-connection wrapper (stdio/gRPC handlers) before dispatching
16    /// into `execute_query`; falls back to `0` for embedded callers.
17    static CURRENT_CONN_ID: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
18
19    /// Authenticated user + role for the executing statement (Phase 2.5.2
20    /// RLS enforcement). Set by the transport middleware after validating
21    /// credentials (password / cert / oauth); unset means "anonymous" /
22    /// "embedded" — RLS policies degrade to the role-agnostic subset.
23    ///
24    /// `None` skips RLS injection entirely; `Some((username, role))`
25    /// passes `role` to `matching_rls_policies(table, Some(role), action)`.
26    static CURRENT_AUTH_IDENTITY: std::cell::RefCell<Option<(String, crate::auth::Role)>> =
27        const { std::cell::RefCell::new(None) };
28
29    /// Immutable principal produced by credential validation. Unlike
30    /// `CURRENT_TENANT_ID`, this value is never changed by routing/scope SQL.
31    static CURRENT_AUTH_PRINCIPAL: std::cell::RefCell<Option<crate::auth::UserId>> =
32        const { std::cell::RefCell::new(None) };
33
34    /// MVCC snapshot scoped to the currently-executing statement (Phase
35    /// 2.3.2d PG parity). `execute_query` captures it on entry and drops
36    /// it on exit; every scan consults it via
37    /// `entity_visible_under_current_snapshot` to hide tuples whose xmin
38    /// hasn't committed or whose xmax already has.
39    ///
40    /// `None` means "pre-MVCC semantics" — the read path returns every
41    /// tuple regardless of xmin/xmax. All embedded callers that bypass
42    /// `execute_query` see this default.
43    static CURRENT_SNAPSHOT: std::cell::RefCell<Option<SnapshotContext>> =
44        const { std::cell::RefCell::new(None) };
45
46    /// Cheap presence flag for `CURRENT_SNAPSHOT`. Scan hot paths
47    /// poll this instead of `borrow()`-ing the RefCell on every
48    /// row — the common case (autocommit / no MVCC session) reads
49    /// one atomic `Cell<bool>` and short-circuits, saving ~10ns × N
50    /// rows on aggregate_group / select_range scans.
51    static HAS_SNAPSHOT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
52
53    /// Session-scoped tenant id for the current connection (Phase 2.5.3
54    /// multi-tenancy). Populated by `SET TENANT 'id'` or by transport
55    /// middleware after resolving tenant from auth claims. Read by the
56    /// `CURRENT_TENANT()` scalar function — RLS policies typically
57    /// combine it as `USING (tenant_id = CURRENT_TENANT())` to scope
58    /// every query to one tenant.
59    ///
60    /// `None` means "no tenant bound" — `CURRENT_TENANT()` returns
61    /// NULL, and RLS policies that gate on it hide every row.
62    static CURRENT_TENANT_ID: std::cell::RefCell<Option<String>> =
63        const { std::cell::RefCell::new(None) };
64
65    /// Statement-local config resolver. SQL expressions materialize the
66    /// `red_config` snapshot lazily on the first `$config.*`/`CONFIG()`
67    /// access, keeping ordinary statements on the zero-scan path.
68    static CURRENT_CONFIG_RESOLVER: std::cell::RefCell<Option<ConfigResolver>> =
69        const { std::cell::RefCell::new(None) };
70
71    /// Statement-local secret resolver. SQL expressions materialize the
72    /// vault KV snapshot lazily on first `$secret.*` access, then use
73    /// lock-free map reads for the rest of the statement.
74    static CURRENT_SECRET_RESOLVER: std::cell::RefCell<Option<SecretResolver>> =
75        const { std::cell::RefCell::new(None) };
76
77    /// Statement-local plain-KV resolver (#1602). Mirrors the secret
78    /// resolver but reads the non-encrypted `plain_kv` store and gates on
79    /// `kv:read`. Materialized lazily on the first `$kv.*` access.
80    static CURRENT_KV_RESOLVER: std::cell::RefCell<Option<KvResolver>> =
81        const { std::cell::RefCell::new(None) };
82}
83
84pub use reddb_engine::transaction::read::SnapshotContext;
85use reddb_engine::transaction::read::{record_serializable_read, visibility_check};
86
87/// Install a connection id on the current thread for the duration of a
88/// statement. The Transaction State module is keyed
89/// by this id so different connections can hold independent BEGINs.
90///
91/// Pub so transports (PG wire, gRPC, HTTP per-request spawners) and
92/// tests can emulate per-connection isolation. Call it once when
93/// binding the connection's worker thread; pair with
94/// `clear_current_connection_id` on teardown.
95pub fn set_current_connection_id(id: u64) {
96    CURRENT_CONN_ID.with(|c| c.set(id));
97}
98
99/// Reset the thread's connection id back to `0` (autocommit).
100pub fn clear_current_connection_id() {
101    CURRENT_CONN_ID.with(|c| c.set(0));
102}
103
104/// Read the connection id set by `set_current_connection_id`. Returns
105/// `0` when no wrapper installed one — auto-commit path.
106pub fn current_connection_id() -> u64 {
107    CURRENT_CONN_ID.with(|c| c.get())
108}
109
110/// Install the authenticated identity for the current thread (Phase 2.5.2
111/// RLS enforcement). Transport layers call this right after resolving
112/// auth so the query dispatch can fold RLS policies into the filter.
113pub fn set_current_auth_identity(username: String, role: crate::auth::Role) {
114    set_current_auth_principal(crate::auth::UserId::platform(username), role);
115}
116
117/// Install the structured principal returned by credential validation.
118pub fn set_current_auth_principal(principal: crate::auth::UserId, role: crate::auth::Role) {
119    CURRENT_AUTH_IDENTITY.with(|cell| {
120        *cell.borrow_mut() = Some((principal.username.clone(), role));
121    });
122    CURRENT_AUTH_PRINCIPAL.with(|cell| *cell.borrow_mut() = Some(principal));
123}
124
125/// Replace only the SQL/RLS projection while preserving the immutable
126/// credential principal used by authorization.
127pub(crate) fn set_current_auth_projection(username: String, role: crate::auth::Role) {
128    CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = Some((username, role)));
129}
130
131/// Clear the thread-local auth identity. Transports call this after the
132/// statement completes so pooled threads don't leak identities across
133/// requests.
134pub fn clear_current_auth_identity() {
135    CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = None);
136    CURRENT_AUTH_PRINCIPAL.with(|cell| *cell.borrow_mut() = None);
137}
138
139/// Read the current-thread auth identity. `None` when no transport
140/// installed one (embedded mode / anonymous access).
141pub(crate) fn current_auth_identity() -> Option<(String, crate::auth::Role)> {
142    CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone())
143}
144
145/// Read the immutable credential principal, independently of mutable scope.
146pub(crate) fn current_auth_principal() -> Option<crate::auth::UserId> {
147    CURRENT_AUTH_PRINCIPAL.with(|cell| cell.borrow().clone())
148}
149
150/// Public probe of the thread-local auth identity for callers outside
151/// the `runtime` module (e.g. the AI credential resolver, which audits
152/// who triggered a secret read on behalf of a query).
153pub fn current_auth_identity_for_audit() -> Option<(String, crate::auth::Role)> {
154    current_auth_identity()
155}
156
157/// Install the session tenant id for the current thread (Phase 2.5.3
158/// multi-tenancy). Called by `SET TENANT 'id'` dispatch and by
159/// transport middleware that resolves tenant from auth claims (e.g.
160/// JWT `tenant` claim, HTTP header, subdomain).
161pub fn set_current_tenant(tenant_id: String) {
162    CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = Some(tenant_id));
163}
164
165/// Clear the current-thread tenant — `CURRENT_TENANT()` will then
166/// return NULL and any RLS policy gated on it will hide every row.
167pub fn clear_current_tenant() {
168    CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = None);
169}
170
171/// Read the current-thread tenant id, applying overrides in priority order:
172///   1. `WITHIN TENANT '<id>' …` per-statement override (highest)
173///   2. `SET LOCAL TENANT '<id>'` transaction-local override (consulted
174///      only when the current connection has an open transaction)
175///   3. `SET TENANT '<id>'` session-level thread-local
176///   4. `None` (deny-default for RLS).
177///
178/// The transaction-local layer is read through the runtime; an embedded
179/// helper crate that has no `RedDBRuntime` access still gets correct
180/// behaviour for layers 1, 3, and 4.
181pub fn current_tenant() -> Option<String> {
182    let inherited = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
183    if let Some(over) = current_scope_override() {
184        if over.tenant.is_active() {
185            return over.tenant.resolve(inherited);
186        }
187    }
188    if let Some(tx_local) = current_tx_local_tenant() {
189        return tx_local;
190    }
191    inherited
192}
193
194thread_local! {
195    /// Snapshot of the active connection's transaction-local tenant for
196    /// the current `execute_query` call. Outer `Some(_)` means "a
197    /// transaction-local tenant override is active for this call";
198    /// inner is the override's value (`Some(s)` overrides to `s`,
199    /// `None` overrides to NULL/cleared). Refreshed at the top of every
200    /// `execute_query` invocation and cleared by the RAII guard on
201    /// return so pooled connections cannot leak the override past the
202    /// statement that owns it.
203    static TX_LOCAL_TENANT: std::cell::RefCell<Option<Option<String>>> =
204        const { std::cell::RefCell::new(None) };
205}
206
207fn current_tx_local_tenant() -> Option<Option<String>> {
208    TX_LOCAL_TENANT.with(|cell| cell.borrow().clone())
209}
210
211/// Recognise `SET LOCAL TENANT '<id>'` / `SET LOCAL TENANT NULL` —
212/// returns `Ok(Some(Some(id)))` for an explicit value, `Ok(Some(None))`
213/// for an explicit NULL clear, `Ok(None)` when the input is not a
214/// `SET LOCAL TENANT` statement at all, and `Err` when the prefix
215/// matches but the value is malformed.
216pub(crate) fn parse_set_local_tenant(query: &str) -> RedDBResult<Option<Option<String>>> {
217    let mut tokens = query.split_ascii_whitespace();
218    let Some(w1) = tokens.next() else {
219        return Ok(None);
220    };
221    if !w1.eq_ignore_ascii_case("SET") {
222        return Ok(None);
223    }
224    let Some(w2) = tokens.next() else {
225        return Ok(None);
226    };
227    if !w2.eq_ignore_ascii_case("LOCAL") {
228        return Ok(None);
229    }
230    let Some(w3) = tokens.next() else {
231        return Ok(None);
232    };
233    if !w3.eq_ignore_ascii_case("TENANT") {
234        return Ok(None);
235    }
236    let rest: String = tokens.collect::<Vec<_>>().join(" ");
237    let rest = rest.trim().trim_end_matches(';').trim();
238    let value_str = rest.strip_prefix('=').map(|s| s.trim()).unwrap_or(rest);
239    if value_str.is_empty() {
240        return Err(RedDBError::Query(
241            "SET LOCAL TENANT expects a string literal or NULL".to_string(),
242        ));
243    }
244    if value_str.eq_ignore_ascii_case("NULL") {
245        return Ok(Some(None));
246    }
247    if value_str.starts_with('\'') && value_str.ends_with('\'') && value_str.len() >= 2 {
248        let inner = &value_str[1..value_str.len() - 1];
249        return Ok(Some(Some(inner.to_string())));
250    }
251    Err(RedDBError::Query(format!(
252        "SET LOCAL TENANT expects a string literal or NULL, got `{value_str}`"
253    )))
254}
255
256pub(crate) struct TxLocalTenantGuard;
257
258impl TxLocalTenantGuard {
259    pub fn install(value: Option<Option<String>>) -> Self {
260        TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = value);
261        Self
262    }
263}
264
265impl Drop for TxLocalTenantGuard {
266    fn drop(&mut self) {
267        TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = None);
268    }
269}
270
271thread_local! {
272    /// Stack of `WITHIN ... <stmt>` overrides active on the current
273    /// thread. Every entry corresponds to one in-flight `execute_query`
274    /// call that started with a `WITHIN` prefix; the entry is pushed
275    /// before dispatch and popped before the call returns. The stack
276    /// shape supports nested invocations (e.g. a view body that itself
277    /// re-enters execute_query).
278    static SCOPE_OVERRIDES: std::cell::RefCell<Vec<crate::runtime::within_clause::ScopeOverride>> =
279        const { std::cell::RefCell::new(Vec::new()) };
280}
281
282pub(crate) fn push_scope_override(over: crate::runtime::within_clause::ScopeOverride) {
283    SCOPE_OVERRIDES.with(|cell| cell.borrow_mut().push(over));
284}
285
286pub(crate) fn pop_scope_override() {
287    SCOPE_OVERRIDES.with(|cell| {
288        cell.borrow_mut().pop();
289    });
290}
291
292pub(crate) fn current_scope_override() -> Option<crate::runtime::within_clause::ScopeOverride> {
293    SCOPE_OVERRIDES.with(|cell| cell.borrow().last().cloned())
294}
295
296/// Cheap probe: is any `WITHIN …` scope override active on this
297/// thread? The fast-path needs to know without paying for the full
298/// `.last().cloned()` allocation — just peek at stack length.
299pub(crate) fn has_scope_override_active() -> bool {
300    SCOPE_OVERRIDES.with(|cell| !cell.borrow().is_empty())
301}
302
303/// RAII guard pairing `push_scope_override` with the matching pop, so
304/// the stack stays balanced even when the inner `execute_query` returns
305/// early via `?`.
306pub(crate) struct ScopeOverrideGuard;
307
308impl ScopeOverrideGuard {
309    pub fn install(over: crate::runtime::within_clause::ScopeOverride) -> Self {
310        push_scope_override(over);
311        Self
312    }
313}
314
315impl Drop for ScopeOverrideGuard {
316    fn drop(&mut self) {
317        pop_scope_override();
318    }
319}
320
321/// Read the current-thread auth identity, honouring per-statement
322/// `WITHIN ... USER '<u>' AS ROLE '<r>'` overrides. The override only
323/// supplies projected strings — it never grants additional privilege —
324/// so callers that need to make authorisation decisions must read from
325/// the underlying `current_auth_identity()` directly.
326pub(crate) fn current_user_projected() -> Option<String> {
327    let inherited = current_auth_identity().map(|(u, _)| u);
328    if let Some(over) = current_scope_override() {
329        if over.user.is_active() {
330            return over.user.resolve(inherited);
331        }
332    }
333    inherited
334}
335
336pub(crate) fn current_role_projected() -> Option<String> {
337    let inherited = current_auth_identity().map(|(_, r)| format!("{r:?}").to_lowercase());
338    if let Some(over) = current_scope_override() {
339        if over.role.is_active() {
340            return over.role.resolve(inherited);
341        }
342    }
343    inherited
344}
345
346pub(crate) fn current_secret_value(path: &str) -> Option<Value> {
347    try_current_secret_value(path).ok().flatten()
348}
349
350pub(crate) fn try_current_secret_value(path: &str) -> RedDBResult<Option<Value>> {
351    let key = path.to_ascii_lowercase();
352    CURRENT_SECRET_RESOLVER.with(|cell| {
353        let mut resolver = cell.borrow_mut();
354        let Some(resolver) = resolver.as_mut() else {
355            return Ok(None);
356        };
357        if let Some(value) = resolver.values.get(&key) {
358            return Ok(value.clone());
359        }
360        let value = resolver.runtime.query_secret_value(&key)?;
361        resolver.values.insert(key, value.clone());
362        Ok(value)
363    })
364}
365
366/// Decrypt only while evaluating a query. Stored and projected values retain
367/// their Secret representation, so formatters cannot expose plaintext.
368pub(crate) fn secret_query_input(value: Value) -> Option<Value> {
369    if !matches!(value, Value::Secret(_)) {
370        return Some(value);
371    }
372    CURRENT_SECRET_RESOLVER.with(|cell| {
373        let resolver = cell.borrow();
374        resolver.as_ref()?.runtime.secret_query_input(value)
375    })
376}
377
378pub(crate) fn try_secret_query_input(
379    value: Value,
380) -> Result<Option<Value>, super::impl_dml_crypto::SecretInputReadLimit> {
381    if !matches!(value, Value::Secret(_)) {
382        return Ok(Some(value));
383    }
384    CURRENT_SECRET_RESOLVER.with(|cell| {
385        let resolver = cell.borrow();
386        match resolver.as_ref() {
387            Some(resolver) => resolver.runtime.try_secret_query_input(value),
388            None => Ok(None),
389        }
390    })
391}
392
393pub(crate) fn secret_query_output(value: Value) -> crate::RedDBResult<Value> {
394    if value.is_null() || matches!(value, Value::Secret(_)) {
395        return Ok(value);
396    }
397    CURRENT_SECRET_RESOLVER.with(|cell| {
398        let resolver = cell.borrow();
399        let resolver = resolver.as_ref().ok_or_else(|| {
400            RedDBError::Query("secret expression requires an active statement".into())
401        })?;
402        resolver.runtime.secret_query_output(value)
403    })
404}
405
406#[cfg(test)]
407fn secret_value_from_snapshot(values: &HashMap<String, String>, key: &str) -> Option<String> {
408    if key.starts_with("red.secret.") {
409        return None;
410    }
411    values.get(key).cloned()
412}
413
414struct SecretResolver {
415    runtime: super::RedDBRuntime,
416    values: HashMap<String, Option<Value>>,
417}
418
419pub(crate) struct SecretStoreGuard {
420    previous: Option<SecretResolver>,
421}
422
423impl SecretStoreGuard {
424    pub(super) fn install(runtime: &super::RedDBRuntime) -> Self {
425        let previous = CURRENT_SECRET_RESOLVER.with(|cell| {
426            cell.replace(Some(SecretResolver {
427                runtime: runtime.clone(),
428                values: HashMap::new(),
429            }))
430        });
431        Self { previous }
432    }
433}
434
435impl Drop for SecretStoreGuard {
436    fn drop(&mut self) {
437        let previous = self.previous.take();
438        CURRENT_SECRET_RESOLVER.with(|cell| {
439            cell.replace(previous);
440        });
441    }
442}
443
444/// Resolve `$kv.<path>` from the plain (non-encrypted) KV store (#1602).
445///
446/// The parser desugars `$kv.<path>` to `__KV_REF("red.kv/<path>")`; this
447/// strips the `red.kv/` prefix and looks up the bare key in the
448/// statement-local snapshot, gating each hit on `kv:read`. Returns `None`
449/// (→ SQL NULL) when the key is absent or the principal is denied — never
450/// an error, mirroring the secret resolver's deny-to-NULL behaviour.
451pub(crate) fn current_kv_value(path: &str) -> Option<String> {
452    let key = path.to_ascii_lowercase();
453    CURRENT_KV_RESOLVER.with(|cell| {
454        let mut resolver = cell.borrow_mut();
455        let resolver = resolver.as_mut()?;
456        if resolver.values.is_none() {
457            resolver.values = resolver
458                .store
459                .as_ref()
460                .map(|store| store.plain_kv_snapshot());
461        }
462        let values = resolver.values.as_ref()?;
463        // Accept both the namespaced `red.kv/<key>` form (from the
464        // `$kv.*` desugar) and the bare key (from `update_current_kv_value`
465        // after a `SET KV`).
466        let found = values
467            .get(&key)
468            .map(|value| (key.as_str(), value))
469            .or_else(|| {
470                key.strip_prefix("red.kv/")
471                    .and_then(|rest| values.get(rest).map(|value| (rest, value)))
472            })?;
473        if !resolver.can_read(found.0) {
474            return None;
475        }
476        Some(found.1.clone())
477    })
478}
479
480struct KvResolver {
481    store: Option<Arc<crate::auth::store::AuthStore>>,
482    values: Option<HashMap<String, String>>,
483    identity: Option<(String, crate::auth::Role, Option<String>)>,
484}
485
486impl KvResolver {
487    fn can_read(&self, key: &str) -> bool {
488        let Some(store) = &self.store else {
489            return true;
490        };
491        let Some((username, role, tenant)) = &self.identity else {
492            return true;
493        };
494        let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
495        let mut resource =
496            crate::auth::policies::ResourceRef::new("kv".to_string(), key.to_string());
497        if let Some(tenant) = tenant {
498            resource = resource.with_tenant(tenant.clone());
499        }
500        let ctx = crate::auth::policies::EvalContext {
501            principal_tenant: tenant.clone(),
502            current_tenant: tenant.clone(),
503            peer_ip: None,
504            mfa_present: false,
505            now_ms: crate::auth::now_ms(),
506            principal_is_admin_role: *role == crate::auth::Role::Admin,
507            principal_is_platform_scoped: tenant.is_none(),
508        };
509        store.check_policy_authz_with_role(&principal, "kv:read", &resource, &ctx, *role)
510    }
511}
512
513pub(crate) struct KvStoreGuard {
514    // Boxed to keep the stack footprint of nested guards small. Each
515    // StatementFrameGuards is live for the full recursive call depth of a
516    // materialized-view refresh chain; unboxed KvResolver (≈120 B) would
517    // push those deep paths past the 8 MB stack limit.
518    previous: Option<Box<KvResolver>>,
519}
520
521impl KvStoreGuard {
522    pub(super) fn install(store: Option<Arc<crate::auth::store::AuthStore>>) -> Self {
523        let previous = CURRENT_KV_RESOLVER.with(|cell| {
524            cell.replace(Some(KvResolver {
525                store,
526                values: None,
527                identity: current_auth_identity().map(|(username, role)| {
528                    let tenant = current_tenant();
529                    (username, role, tenant)
530                }),
531            }))
532        });
533        Self {
534            previous: previous.map(Box::new),
535        }
536    }
537}
538
539impl Drop for KvStoreGuard {
540    fn drop(&mut self) {
541        let previous = self.previous.take().map(|b| *b);
542        CURRENT_KV_RESOLVER.with(|cell| {
543            cell.replace(previous);
544        });
545    }
546}
547
548/// Update the statement-local KV snapshot after a `SET KV` / `DELETE KV`
549/// so a subsequent `$kv.*` read in the same statement observes the write.
550pub(crate) fn update_current_kv_value(path: &str, value: Option<String>) {
551    let key = path.to_ascii_lowercase();
552    CURRENT_KV_RESOLVER.with(|cell| {
553        if let Some(resolver) = cell.borrow_mut().as_mut() {
554            let Some(values) = resolver.values.as_mut() else {
555                return;
556            };
557            match value {
558                Some(value) => {
559                    values.insert(key, value);
560                }
561                None => {
562                    values.remove(&key);
563                }
564            }
565        }
566    });
567}
568
569pub(crate) fn current_config_value(path: &str) -> Option<Value> {
570    try_current_config_value(path).ok().flatten()
571}
572
573pub(crate) fn try_current_config_value(path: &str) -> RedDBResult<Option<Value>> {
574    let key = path.to_ascii_lowercase();
575    CURRENT_CONFIG_RESOLVER.with(|cell| {
576        let mut resolver = cell.borrow_mut();
577        let Some(resolver) = resolver.as_mut() else {
578            return Ok(None);
579        };
580        if !resolver.can_read(&key) {
581            return Ok(None);
582        }
583        if resolver.values.is_none() {
584            resolver.values = Some(latest_config_snapshot(&resolver.db)?);
585        }
586        // #1370 — `$config.<path>` desugars to `CONFIG("red.config/<path>")`,
587        // but `SET CONFIG <path>` stores under the bare key. Mirror the secret
588        // resolver: after the namespaced key, fall back to the stripped bare
589        // key, then the dotted `red.config.<rest>` legacy form. Track the key
590        // that matched so the `config:read` gate authorizes the real target.
591        let found = {
592            let values = resolver
593                .values
594                .as_ref()
595                .expect("invariant: successful snapshot is installed");
596            values
597                .get(&key)
598                .map(|value| (key.clone(), value.clone()))
599                .or_else(|| {
600                    key.strip_prefix("red.config/").and_then(|rest| {
601                        values
602                            .get(rest)
603                            .map(|value| (rest.to_string(), value.clone()))
604                            .or_else(|| {
605                                let legacy = format!("red.config.{rest}");
606                                values.get(&legacy).map(|value| (legacy, value.clone()))
607                            })
608                    })
609                })
610        };
611        let Some(found) = found else { return Ok(None) };
612        // #1743 — hard-block the reserved `red.config.*` namespace and gate the
613        // read on `config:read`, mirroring the `$secret` sibling. Denied reads
614        // resolve to `None` (→ SQL NULL), never an error.
615        if !resolver.can_read(&found.0) {
616            return Ok(None);
617        }
618        Ok(Some(found.1))
619    })
620}
621
622pub(crate) fn update_current_config_value(path: &str, value: Value) {
623    let key = path.to_ascii_lowercase();
624    CURRENT_CONFIG_RESOLVER.with(|cell| {
625        if let Some(resolver) = cell.borrow_mut().as_mut() {
626            if let Some(values) = resolver.values.as_mut() {
627                values.insert(key, value);
628            }
629        }
630    });
631}
632
633pub(crate) fn update_current_secret_value(_path: &str, _value: Option<String>) {
634    CURRENT_SECRET_RESOLVER.with(|cell| {
635        if let Some(resolver) = cell.borrow_mut().as_mut() {
636            resolver.values.clear();
637        }
638    });
639}
640
641#[cfg(test)]
642mod tests {
643    use super::*;
644
645    /// Companion to `raw_iteration_confined_to_maintenance_allowlist`:
646    /// the snapshot thread-local is an implementation detail of the
647    /// statement frame (issue #2138). It must stay private to this module,
648    /// so the only way any other module can observe ambient snapshot state
649    /// is `capture_current_snapshot`, which hands back an owned `Option`
650    /// the caller then passes explicitly into the scan API.
651    #[test]
652    fn snapshot_thread_local_stays_private_to_the_statement_frame() {
653        let src = include_str!("execution_context.rs");
654
655        for name in ["CURRENT_SNAPSHOT", "HAS_SNAPSHOT"] {
656            assert!(
657                !src.contains(&format!("pub static {name}"))
658                    && !src.contains(&format!("pub(crate) static {name}")),
659                "{name} must stay private to execution_context; expose an explicit \
660                 Option<SnapshotContext> via capture_current_snapshot instead"
661            );
662        }
663
664        assert!(
665            src.contains("the only sanctioned ambient read"),
666            "capture_current_snapshot must keep documenting that it is the single \
667             sanctioned ambient read of the snapshot thread-local"
668        );
669    }
670
671    #[test]
672    fn secret_snapshot_hard_blocks_system_keys() {
673        let values = HashMap::from([
674            ("red.secret.aes_key".into(), "internal".into()),
675            ("acme.key".into(), "user-value".into()),
676        ]);
677        assert_eq!(
678            secret_value_from_snapshot(&values, "red.secret.aes_key"),
679            None
680        );
681        assert_eq!(
682            secret_value_from_snapshot(&values, "acme.key").as_deref(),
683            Some("user-value")
684        );
685    }
686
687    #[test]
688    fn config_snapshot_does_not_copy_superseded_values() {
689        use crate::storage::no_malloc_hot_paths::measure_allocations;
690        use crate::storage::{EntityData, EntityId, RowData, UnifiedEntity};
691        let observe = |payload_bytes| {
692            let runtime = super::super::RedDBRuntime::in_memory().expect("runtime");
693            let manager = runtime.db().store().get_or_create_collection("red_config");
694            // Insertion and segment traversal can disagree. Only candidates
695            // that advance the selected version may copy the value.
696            for id in (1..=64).rev() {
697                let mut entity = UnifiedEntity::table_row(
698                    EntityId::new(1_000_000 + id),
699                    "red_config",
700                    id,
701                    vec![],
702                );
703                entity.data = EntityData::Row(RowData {
704                    columns: vec![],
705                    schema: None,
706                    named: Some(HashMap::from([
707                        ("key".into(), Value::text("fixture.copy")),
708                        ("value".into(), Value::Blob(vec![id as u8; payload_bytes])),
709                    ])),
710                });
711                manager.insert(entity).expect("config history");
712            }
713            let mut newest = 0;
714            let mut replacements = 0;
715            manager.try_for_each_entity(usize::MAX, |entity| {
716                if matches!(entity.data.as_row().and_then(|row| row.get_field("key")), Some(Value::Text(key)) if key.as_ref() == "fixture.copy")
717                    && entity.id.raw() >= newest {
718                    newest = entity.id.raw();
719                    replacements += 1;
720                }
721                true
722            }).expect("scan order");
723            let (snapshot, allocations) =
724                measure_allocations(|| latest_config_snapshot(runtime.db().as_ref()));
725            assert_eq!(
726                snapshot.expect("snapshot").get("fixture.copy"),
727                Some(&Value::Blob(vec![64; payload_bytes]))
728            );
729            (allocations, replacements)
730        };
731        let (empty, _) = observe(0);
732        let (large, replacements) = observe(64 * 1024);
733        assert!(
734            large.allocated_bytes <= empty.allocated_bytes + replacements * 64 * 1024 + 64 * 1024,
735            "superseded config payload copies: {empty:?} -> {large:?}"
736        );
737        eprintln!(
738            "config snapshot allocated bytes, empty/64 KiB values: {} / {}",
739            empty.allocated_bytes, large.allocated_bytes
740        );
741    }
742
743    fn with_kv_values<T>(values: HashMap<String, String>, f: impl FnOnce() -> T) -> T {
744        CURRENT_KV_RESOLVER.with(|cell| {
745            cell.replace(Some(KvResolver {
746                store: None,
747                values: Some(values),
748                identity: None,
749            }));
750        });
751        let result = f();
752        CURRENT_KV_RESOLVER.with(|cell| {
753            cell.replace(None);
754        });
755        result
756    }
757
758    #[test]
759    fn current_kv_value_resolves_namespaced_and_bare_keys() {
760        let values = HashMap::from([("acme.key".to_string(), "plain-value".to_string())]);
761
762        with_kv_values(values, || {
763            // `$kv.acme.key` desugars to `red.kv/acme.key`.
764            assert_eq!(
765                current_kv_value("red.kv/acme.key").as_deref(),
766                Some("plain-value")
767            );
768            // Bare-key lookup (post-`SET KV` snapshot update) also resolves.
769            assert_eq!(current_kv_value("acme.key").as_deref(), Some("plain-value"));
770            // Unknown keys resolve to None (→ SQL NULL), never an error.
771            assert_eq!(current_kv_value("red.kv/missing.key"), None);
772        });
773    }
774
775    #[test]
776    fn config_key_read_allowed_hard_blocks_reserved_namespace() {
777        // The reserved `red.config.*` system namespace is denied role- and
778        // store-independently, mirroring the `red.secret.*` block (#1743).
779        assert!(!config_key_read_allowed(
780            None,
781            None,
782            "red.config/red.config.aes_key"
783        ));
784        assert!(!config_key_read_allowed(None, None, "red.config.aes_key"));
785        // Ordinary config keys pass when no store is installed (non-IAM path).
786        assert!(config_key_read_allowed(
787            None,
788            None,
789            "red.config/ai.openrouter.default.key"
790        ));
791        assert!(config_key_read_allowed(None, None, "acme.flags.beta"));
792    }
793
794    #[test]
795    fn is_config_collection_matches_system_config_stores() {
796        assert!(is_config_collection("red_config"));
797        assert!(is_config_collection("RED_CONFIG"));
798        assert!(is_config_collection("red.config"));
799        assert!(!is_config_collection("users"));
800        assert!(!is_config_collection("red_kv"));
801    }
802
803    #[test]
804    fn update_current_kv_value_reflects_in_resolver() {
805        with_kv_values(HashMap::new(), || {
806            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
807            update_current_kv_value("feature.flag", Some("on".to_string()));
808            assert_eq!(
809                current_kv_value("red.kv/feature.flag").as_deref(),
810                Some("on")
811            );
812            update_current_kv_value("feature.flag", None);
813            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
814        });
815    }
816}
817
818fn latest_config_snapshot(db: &RedDB) -> RedDBResult<HashMap<String, Value>> {
819    let mut latest: HashMap<String, (u64, Value)> = HashMap::new();
820
821    if let Some(manager) = db.store().get_collection("red_config") {
822        super::function_budget::try_scan(
823            |visit| manager.try_for_each_entity(usize::MAX, visit),
824            |entity| {
825                let Some(row) = entity.data.as_row() else {
826                    return true;
827                };
828                let Some(Value::Text(key)) = row.get_field("key") else {
829                    return true;
830                };
831                let value = row.get_field("value").unwrap_or(&Value::Null);
832                let id = entity.id.raw();
833                let key = key.to_ascii_lowercase();
834                insert_latest_config_value(&mut latest, key.clone(), id, value);
835                if let Some(rest) = key.strip_prefix("red.config.") {
836                    insert_latest_config_value(
837                        &mut latest,
838                        format!("red.config/{rest}"),
839                        id,
840                        value,
841                    );
842                }
843                true
844            },
845        )?;
846    }
847
848    if let Some(manager) = db.store().get_collection("red.config") {
849        super::function_budget::try_scan(
850            |visit| manager.try_for_each_entity(usize::MAX, visit),
851            |entity| {
852                let Some(row) = entity.data.as_row() else {
853                    return true;
854                };
855                if matches!(row.get_field("tombstone"), Some(Value::Boolean(true))) {
856                    return true;
857                }
858                let Some(Value::Text(key)) = row.get_field("key") else {
859                    return true;
860                };
861                let value = row.get_field("value").unwrap_or(&Value::Null);
862                insert_latest_config_value(
863                    &mut latest,
864                    format!("red.config/{}", key.to_ascii_lowercase()),
865                    entity.id.raw(),
866                    value,
867                );
868                true
869            },
870        )?;
871    }
872
873    Ok(latest
874        .into_iter()
875        .map(|(key, (_, value))| (key, value))
876        .collect())
877}
878
879fn insert_latest_config_value(
880    latest: &mut HashMap<String, (u64, Value)>,
881    key: String,
882    id: u64,
883    value: &Value,
884) {
885    match latest.get(&key) {
886        Some((prev_id, _)) if *prev_id > id => {}
887        _ => {
888            latest.insert(key, (id, value.clone()));
889        }
890    }
891}
892
893struct ConfigResolver {
894    db: Arc<RedDB>,
895    store: Option<Arc<crate::auth::store::AuthStore>>,
896    identity: Option<(String, crate::auth::Role, Option<String>)>,
897    values: Option<HashMap<String, Value>>,
898}
899
900impl ConfigResolver {
901    fn can_read(&self, key: &str) -> bool {
902        config_key_read_allowed(self.store.as_ref(), self.identity.as_ref(), key)
903    }
904}
905
906/// The reserved `red.config.*` system-config namespace. Never expose it via
907/// the inline `$config`/`CONFIG()`/`KV()` resolver regardless of IAM role,
908/// mirroring the role-independent `red.secret.*` hard-block (#1743).
909fn is_reserved_config_key(key: &str) -> bool {
910    let key = key.strip_prefix("red.config/").unwrap_or(key);
911    key.starts_with("red.config.")
912        || key == "red.config"
913        || key == super::function_catalog::REGISTRY_KEY
914}
915
916/// The system config-store collections (`red_config` / `red.config`) reached
917/// by the inline resolver. A `KV(<collection>, key)` read of one of these is
918/// gated exactly like `$config`; any other collection is unaffected (#1743).
919fn is_config_collection(collection: &str) -> bool {
920    matches!(
921        collection.to_ascii_lowercase().as_str(),
922        "red_config" | "red.config"
923    )
924}
925
926fn config_resource_target(key: &str) -> String {
927    let bare = key.strip_prefix("red.config/").unwrap_or(key);
928    format!("red.config/{bare}")
929}
930
931/// Shared inline-config read gate: a role-independent hard-block on the
932/// reserved `red.config.*` namespace, then a `config:read` capability check
933/// mirroring `SHOW CONFIG` / the `$secret` sibling (#1743). Absent a store or
934/// with IAM authorization disabled, only the hard-block applies so existing
935/// non-IAM `$config` reads keep working.
936fn config_key_read_allowed(
937    store: Option<&Arc<crate::auth::store::AuthStore>>,
938    identity: Option<&(String, crate::auth::Role, Option<String>)>,
939    key: &str,
940) -> bool {
941    if is_reserved_config_key(key) {
942        return false;
943    }
944    let Some(store) = store else {
945        return true;
946    };
947    if !store.iam_authorization_enabled() {
948        return true;
949    }
950    let Some((username, role, tenant)) = identity else {
951        return true;
952    };
953    let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
954    let mut resource =
955        crate::auth::policies::ResourceRef::new("config".to_string(), config_resource_target(key));
956    if let Some(tenant) = tenant {
957        resource = resource.with_tenant(tenant.clone());
958    }
959    let ctx = crate::auth::policies::EvalContext {
960        principal_tenant: tenant.clone(),
961        current_tenant: tenant.clone(),
962        peer_ip: None,
963        mfa_present: false,
964        now_ms: crate::auth::now_ms(),
965        principal_is_admin_role: *role == crate::auth::Role::Admin,
966        principal_is_platform_scoped: tenant.is_none(),
967    };
968    store.check_policy_authz_with_role(&principal, "config:read", &resource, &ctx, *role)
969}
970
971/// Whether the current principal may read config `key` through the inline
972/// `$config`/`CONFIG()`/`KV('red_config', …)` resolver. Used by the db-fallback
973/// lookup paths that bypass [`current_config_value`] (#1743).
974pub(crate) fn config_read_permitted(key: &str) -> bool {
975    let key = key.to_ascii_lowercase();
976    CURRENT_CONFIG_RESOLVER.with(|cell| match cell.borrow().as_ref() {
977        Some(resolver) => resolver.can_read(&key),
978        None => !is_reserved_config_key(&key),
979    })
980}
981
982/// Gate a raw `KV(collection, key)` read: config system collections go through
983/// the `$config` gate; every other collection is unaffected (#1743).
984pub(crate) fn kv_read_permitted(collection: &str, key: &str) -> bool {
985    if is_config_collection(collection) {
986        config_read_permitted(key)
987    } else {
988        true
989    }
990}
991
992pub(crate) struct ConfigSnapshotGuard {
993    previous: Option<ConfigResolver>,
994}
995
996impl ConfigSnapshotGuard {
997    pub(super) fn install(
998        db: Arc<RedDB>,
999        store: Option<Arc<crate::auth::store::AuthStore>>,
1000    ) -> Self {
1001        let identity = current_auth_identity().map(|(username, role)| {
1002            let tenant = current_tenant();
1003            (username, role, tenant)
1004        });
1005        let previous = CURRENT_CONFIG_RESOLVER.with(|cell| {
1006            cell.replace(Some(ConfigResolver {
1007                db,
1008                store,
1009                identity,
1010                values: None,
1011            }))
1012        });
1013        Self { previous }
1014    }
1015}
1016
1017impl Drop for ConfigSnapshotGuard {
1018    fn drop(&mut self) {
1019        let previous = self.previous.take();
1020        CURRENT_CONFIG_RESOLVER.with(|cell| {
1021            cell.replace(previous);
1022        });
1023    }
1024}
1025
1026/// Install the MVCC snapshot used by the current thread for the duration
1027/// of one statement. Paired with `clear_current_snapshot()` — callers
1028/// should prefer the `CurrentSnapshotGuard` RAII wrapper so early returns
1029/// still clean up.
1030pub fn set_current_snapshot(ctx: SnapshotContext) {
1031    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = Some(ctx));
1032    HAS_SNAPSHOT.with(|c| c.set(true));
1033}
1034
1035pub fn clear_current_snapshot() {
1036    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = None);
1037    HAS_SNAPSHOT.with(|c| c.set(false));
1038}
1039
1040/// Drop-guard that restores the previous snapshot on scope exit. Safe to
1041/// nest — each statement saves the caller's snapshot and puts it back
1042/// instead of blindly clearing, so a top-level `execute_query` called
1043/// from inside another statement dispatch (e.g. vector source subqueries)
1044/// doesn't strip visibility from the outer scan.
1045pub(crate) struct CurrentSnapshotGuard {
1046    previous: Option<SnapshotContext>,
1047}
1048
1049impl CurrentSnapshotGuard {
1050    pub(crate) fn install(ctx: SnapshotContext) -> Self {
1051        let previous = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1052        set_current_snapshot(ctx);
1053        Self { previous }
1054    }
1055}
1056
1057impl Drop for CurrentSnapshotGuard {
1058    fn drop(&mut self) {
1059        let prev = self.previous.take();
1060        let has = prev.is_some();
1061        CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = prev);
1062        HAS_SNAPSHOT.with(|c| c.set(has));
1063    }
1064}
1065
1066/// Is this entity visible under the current thread's MVCC snapshot?
1067///
1068/// Returns `true` (no filtering) when no snapshot is installed — that
1069/// path is used by embedded callers and by operations that intentionally
1070/// bypass MVCC (VACUUM, snapshot export, admin introspection).
1071///
1072/// When a snapshot is installed the result is
1073///   `snapshot.sees(xmin, xmax) && !mgr.is_aborted(xmin) && !xmax_half_abort`
1074/// where `xmax_half_abort` re-grants visibility for tuples whose
1075/// deleting transaction rolled back.
1076#[inline]
1077pub fn entity_visible_under_current_snapshot(
1078    entity: &crate::storage::unified::entity::UnifiedEntity,
1079) -> bool {
1080    // Moderation visibility gate (#1274, ADR 0057). A row carrying the
1081    // moderation status marker — quarantine-pending or rejected-tombstone
1082    // — is hidden from every normal read, on top of MVCC visibility. The
1083    // marker lives on the row itself, so the check is a single field probe
1084    // and rides the existing per-row visibility chokepoint rather than
1085    // adding a separate filter pass to each scan call-site.
1086    if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
1087        return false;
1088    }
1089    // Fast path — one `Cell<bool>` read, no RefCell borrow. Autocommit
1090    // reads (no active MVCC transaction) still hide superseded physical
1091    // versions while avoiding a full snapshot-context lookup.
1092    // This runs on every row of every scan; the slow path only fires
1093    // inside an explicit transaction.
1094    if !HAS_SNAPSHOT.with(|c| c.get()) {
1095        return entity.xmax == 0;
1096    }
1097    CURRENT_SNAPSHOT.with(|cell| {
1098        let guard = cell.borrow();
1099        let Some(ctx) = guard.as_ref() else {
1100            return true;
1101        };
1102        let visible = visibility_check(ctx, entity.xmin, entity.xmax);
1103        if visible {
1104            record_serializable_read(ctx, entity);
1105        }
1106        visible
1107    })
1108}
1109
1110/// Direct visibility check from raw `(xmin, xmax)` — bypasses the
1111/// entity borrow for callers that already decomposed the tuple (e.g.
1112/// pre-materialized scan caches). Same semantics as
1113/// `entity_visible_under_current_snapshot`.
1114#[inline]
1115pub(crate) fn xids_visible_under_current_snapshot(xmin: u64, xmax: u64) -> bool {
1116    if !HAS_SNAPSHOT.with(|c| c.get()) {
1117        return true;
1118    }
1119    CURRENT_SNAPSHOT.with(|cell| {
1120        let guard = cell.borrow();
1121        let Some(ctx) = guard.as_ref() else {
1122            return true;
1123        };
1124        visibility_check(ctx, xmin, xmax)
1125    })
1126}
1127
1128/// Clone the current thread's snapshot context. Parallel scan paths
1129/// (`query_all_zoned` with `std::thread::scope`) call this on the main
1130/// thread *before* spawning workers so the captured `SnapshotContext`
1131/// can be moved into every worker closure. Worker threads do not
1132/// inherit thread-locals, so calling `entity_visible_under_current_snapshot`
1133/// from inside a spawned closure would silently skip the filter.
1134///
1135/// This is the only sanctioned ambient read of the snapshot thread-local
1136/// (issue #2138). `CURRENT_SNAPSHOT` and `HAS_SNAPSHOT` are private to this
1137/// module and stay that way: they exist for statement-frame entry and for
1138/// detecting a frameless read, nothing else. Everything downstream must
1139/// take the resulting `Option<SnapshotContext>` as an explicit argument —
1140/// that is what makes `SegmentManager::scan*` thread-independent.
1141///
1142/// Statement-frame build also reads it, so a nested sub-execution inside
1143/// one statement (CTE prelude, materialized-view refresh, inline graph
1144/// TVF) inherits the enclosing snapshot instead of minting a fresh one
1145/// (#2183).
1146///
1147/// The `Option` is load-bearing: `None` means no statement frame is
1148/// installed — the low-level callers that intentionally bypass MVCC
1149/// (VACUUM, snapshot export, admin introspection) — and the scan API
1150/// keeps a documented fallback for that case rather than assuming a frame
1151/// always exists.
1152pub fn capture_current_snapshot() -> Option<SnapshotContext> {
1153    CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone())
1154}
1155
1156/// Whether the active read snapshot may need historical tuple versions
1157/// that the current secondary indexes cannot prove. Index paths can still
1158/// recheck visible candidates, but only a heap scan can discover versions
1159/// whose indexed value was changed or deleted after this snapshot.
1160pub(crate) fn current_snapshot_requires_index_fallback() -> bool {
1161    if !HAS_SNAPSHOT.with(|c| c.get()) {
1162        return false;
1163    }
1164    CURRENT_SNAPSHOT.with(|cell| {
1165        cell.borrow()
1166            .as_ref()
1167            .is_some_and(|ctx| ctx.requires_index_fallback)
1168    })
1169}
1170
1171/// Frozen MVCC + identity context for callers that need to reinstall
1172/// the same view across thread-local boundaries — long-lived cursors,
1173/// background batchers, anything that detaches from the dispatch path
1174/// and re-enters later.
1175///
1176/// The bundle bakes in the three thread-locals every read path
1177/// consults: `SnapshotContext` (MVCC visibility), the auth identity
1178/// (RLS policy gate), and the tenant id (RLS scope). A FETCH that
1179/// reinstalls the bundle sees exactly the same rows as the DECLARE
1180/// would have, regardless of writes that landed in between.
1181///
1182/// Cheap to clone — `SnapshotContext` is a clone of three
1183/// `Arc`-backed fields, identity is a `(String, Role)`, tenant is a
1184/// `String`. None of these contend with the read path.
1185#[derive(Clone, Default)]
1186pub struct SnapshotBundle {
1187    pub snapshot: Option<SnapshotContext>,
1188    pub auth: Option<(String, crate::auth::Role)>,
1189    pub auth_principal: Option<crate::auth::UserId>,
1190    pub tenant: Option<String>,
1191}
1192
1193/// Capture the three read-path thread-locals into a `SnapshotBundle`.
1194/// Pairs with `with_snapshot_bundle` for re-entry.
1195pub fn snapshot_bundle() -> SnapshotBundle {
1196    SnapshotBundle {
1197        snapshot: capture_current_snapshot(),
1198        auth: current_auth_identity(),
1199        auth_principal: current_auth_principal(),
1200        tenant: CURRENT_TENANT_ID.with(|cell| cell.borrow().clone()),
1201    }
1202}
1203
1204/// Reinstall a captured `SnapshotBundle` for the duration of `f`.
1205/// Restores the caller's previous thread-locals on exit (panic-safe via
1206/// the explicit guard struct so a panic in `f` cannot leak the
1207/// installed identity into the worker's next request).
1208pub fn with_snapshot_bundle<R>(bundle: &SnapshotBundle, f: impl FnOnce() -> R) -> R {
1209    struct Guard {
1210        prev_snapshot: Option<SnapshotContext>,
1211        prev_auth: Option<(String, crate::auth::Role)>,
1212        prev_auth_principal: Option<crate::auth::UserId>,
1213        prev_tenant: Option<String>,
1214    }
1215    impl Drop for Guard {
1216        fn drop(&mut self) {
1217            let snap = self.prev_snapshot.take();
1218            let has = snap.is_some();
1219            CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = snap);
1220            HAS_SNAPSHOT.with(|c| c.set(has));
1221            CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = self.prev_auth.take());
1222            CURRENT_AUTH_PRINCIPAL
1223                .with(|cell| *cell.borrow_mut() = self.prev_auth_principal.take());
1224            CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = self.prev_tenant.take());
1225        }
1226    }
1227
1228    let _guard = {
1229        let prev_snapshot = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1230        let prev_auth = CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone());
1231        let prev_auth_principal = CURRENT_AUTH_PRINCIPAL.with(|cell| cell.borrow().clone());
1232        let prev_tenant = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
1233
1234        match bundle.snapshot.clone() {
1235            Some(ctx) => set_current_snapshot(ctx),
1236            None => clear_current_snapshot(),
1237        }
1238        CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = bundle.auth.clone());
1239        CURRENT_AUTH_PRINCIPAL.with(|cell| *cell.borrow_mut() = bundle.auth_principal.clone());
1240        CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = bundle.tenant.clone());
1241
1242        Guard {
1243            prev_snapshot,
1244            prev_auth,
1245            prev_auth_principal,
1246            prev_tenant,
1247        }
1248    };
1249    f()
1250}
1251
1252/// Apply the same visibility rules used by the thread-local helpers
1253/// against a caller-provided context. Intended for parallel workers
1254/// that captured the snapshot with `capture_current_snapshot()`.
1255#[inline]
1256pub fn entity_visible_with_context(
1257    ctx: Option<&SnapshotContext>,
1258    entity: &crate::storage::unified::entity::UnifiedEntity,
1259) -> bool {
1260    reddb_engine::storage::visibility::entity_visible_with_context(ctx, entity)
1261}