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
691        // Drive the selection in a known order. Two traversals of the live
692        // store need not have the same order (or see the same compaction), so
693        // counting replacements in a separate traversal is not a stable bound.
694        let observe = |payload_bytes| {
695            let values: Vec<_> = (1..=64)
696                .rev()
697                .map(|id| (id, Value::Blob(vec![id as u8; payload_bytes])))
698                .collect();
699            let (latest, allocations) = measure_allocations(|| {
700                let mut latest = HashMap::new();
701                for (id, value) in &values {
702                    let mut shared = None;
703                    insert_latest_config_value(
704                        &mut latest,
705                        "fixture.copy".into(),
706                        *id,
707                        value,
708                        &mut shared,
709                    );
710                    insert_latest_config_value(
711                        &mut latest,
712                        "red.config/fixture.copy".into(),
713                        *id,
714                        value,
715                        &mut shared,
716                    );
717                }
718                latest
719            });
720            assert_eq!(latest["fixture.copy"].0, 64);
721            assert_eq!(
722                *latest["fixture.copy"].1,
723                Value::Blob(vec![64; payload_bytes])
724            );
725            assert!(Arc::ptr_eq(
726                &latest["fixture.copy"].1,
727                &latest["red.config/fixture.copy"].1
728            ));
729            allocations
730        };
731        let empty = observe(0);
732        let large = observe(64 * 1024);
733        assert!(
734            large.allocated_bytes <= empty.allocated_bytes + 2 * 64 * 1024,
735            "superseded config payload copies: {empty:?} -> {large:?}"
736        );
737    }
738
739    #[test]
740    fn config_snapshot_keeps_legacy_and_namespaced_aliases_after_collection_migration() {
741        use crate::storage::{EntityData, EntityId, RowData, UnifiedEntity};
742        let runtime = super::super::RedDBRuntime::in_memory().expect("runtime");
743        let store = runtime.db().store();
744        let manager = store.get_or_create_collection("red_config");
745        assert!(Arc::ptr_eq(
746            &manager,
747            &store.get_collection("red.config").unwrap()
748        ));
749        for (id, key, value, tombstone) in [
750            (1, "theme", "dark", false),
751            (2, "red.config.theme", "light", false),
752            (3, "other", "hidden", true),
753        ] {
754            let mut entity =
755                UnifiedEntity::table_row(EntityId::new(2_000_000 + id), "red_config", id, vec![]);
756            entity.data = EntityData::Row(RowData {
757                columns: vec![],
758                schema: None,
759                named: Some(HashMap::from([
760                    ("key".into(), Value::text(key)),
761                    ("value".into(), Value::text(value)),
762                    ("tombstone".into(), Value::Boolean(tombstone)),
763                ])),
764            });
765            manager.insert(entity).expect("config row");
766        }
767        let values = latest_config_snapshot(runtime.db().as_ref()).expect("snapshot");
768        assert_eq!(values.get("theme"), Some(&Value::text("dark")));
769        assert_eq!(values.get("red.config.theme"), Some(&Value::text("light")));
770        assert_eq!(values.get("red.config/theme"), Some(&Value::text("light")));
771        assert_eq!(values.get("other"), Some(&Value::text("hidden")));
772        assert!(!values.contains_key("red.config/other"));
773    }
774
775    fn with_kv_values<T>(values: HashMap<String, String>, f: impl FnOnce() -> T) -> T {
776        CURRENT_KV_RESOLVER.with(|cell| {
777            cell.replace(Some(KvResolver {
778                store: None,
779                values: Some(values),
780                identity: None,
781            }));
782        });
783        let result = f();
784        CURRENT_KV_RESOLVER.with(|cell| {
785            cell.replace(None);
786        });
787        result
788    }
789
790    #[test]
791    fn current_kv_value_resolves_namespaced_and_bare_keys() {
792        let values = HashMap::from([("acme.key".to_string(), "plain-value".to_string())]);
793
794        with_kv_values(values, || {
795            // `$kv.acme.key` desugars to `red.kv/acme.key`.
796            assert_eq!(
797                current_kv_value("red.kv/acme.key").as_deref(),
798                Some("plain-value")
799            );
800            // Bare-key lookup (post-`SET KV` snapshot update) also resolves.
801            assert_eq!(current_kv_value("acme.key").as_deref(), Some("plain-value"));
802            // Unknown keys resolve to None (→ SQL NULL), never an error.
803            assert_eq!(current_kv_value("red.kv/missing.key"), None);
804        });
805    }
806
807    #[test]
808    fn config_key_read_allowed_hard_blocks_reserved_namespace() {
809        // The reserved `red.config.*` system namespace is denied role- and
810        // store-independently, mirroring the `red.secret.*` block (#1743).
811        assert!(!config_key_read_allowed(
812            None,
813            None,
814            "red.config/red.config.aes_key"
815        ));
816        assert!(!config_key_read_allowed(None, None, "red.config.aes_key"));
817        // Ordinary config keys pass when no store is installed (non-IAM path).
818        assert!(config_key_read_allowed(
819            None,
820            None,
821            "red.config/ai.openrouter.default.key"
822        ));
823        assert!(config_key_read_allowed(None, None, "acme.flags.beta"));
824    }
825
826    #[test]
827    fn is_config_collection_matches_system_config_stores() {
828        assert!(is_config_collection("red_config"));
829        assert!(is_config_collection("RED_CONFIG"));
830        assert!(is_config_collection("red.config"));
831        assert!(!is_config_collection("users"));
832        assert!(!is_config_collection("red_kv"));
833    }
834
835    #[test]
836    fn update_current_kv_value_reflects_in_resolver() {
837        with_kv_values(HashMap::new(), || {
838            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
839            update_current_kv_value("feature.flag", Some("on".to_string()));
840            assert_eq!(
841                current_kv_value("red.kv/feature.flag").as_deref(),
842                Some("on")
843            );
844            update_current_kv_value("feature.flag", None);
845            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
846        });
847    }
848}
849
850fn latest_config_snapshot(db: &RedDB) -> RedDBResult<HashMap<String, Value>> {
851    let mut latest: HashMap<String, (u64, Arc<Value>)> = HashMap::new();
852
853    // The store canonicalizes the legacy `red_config` name to `red.config`.
854    // Scan it once, but preserve both legacy and namespaced resolver keys.
855    if let Some(manager) = db.store().get_collection("red.config") {
856        super::function_budget::try_scan(
857            |visit| manager.try_for_each_entity(usize::MAX, visit),
858            |entity| {
859                let Some(row) = entity.data.as_row() else {
860                    return true;
861                };
862                let Some(Value::Text(key)) = row.get_field("key") else {
863                    return true;
864                };
865                let value = row.get_field("value").unwrap_or(&Value::Null);
866                let id = entity.id.raw();
867                let key = key.to_ascii_lowercase();
868                // One value may be visible under several aliases. Clone the
869                // payload at most once for each advancing entity version.
870                let mut shared = None;
871                insert_latest_config_value(&mut latest, key.clone(), id, value, &mut shared);
872                if let Some(rest) = key.strip_prefix("red.config.") {
873                    insert_latest_config_value(
874                        &mut latest,
875                        format!("red.config/{rest}"),
876                        id,
877                        value,
878                        &mut shared,
879                    );
880                }
881                if !matches!(row.get_field("tombstone"), Some(Value::Boolean(true))) {
882                    insert_latest_config_value(
883                        &mut latest,
884                        format!("red.config/{key}"),
885                        id,
886                        value,
887                        &mut shared,
888                    );
889                }
890                true
891            },
892        )?;
893    }
894
895    Ok(latest
896        .into_iter()
897        .map(|(key, (_, value))| {
898            (
899                key,
900                Arc::try_unwrap(value).unwrap_or_else(|shared| (*shared).clone()),
901            )
902        })
903        .collect())
904}
905
906fn insert_latest_config_value(
907    latest: &mut HashMap<String, (u64, Arc<Value>)>,
908    key: String,
909    id: u64,
910    value: &Value,
911    shared: &mut Option<Arc<Value>>,
912) {
913    match latest.get(&key) {
914        Some((prev_id, _)) if *prev_id > id => {}
915        _ => {
916            let shared = shared.get_or_insert_with(|| Arc::new(value.clone()));
917            latest.insert(key, (id, Arc::clone(shared)));
918        }
919    }
920}
921
922struct ConfigResolver {
923    db: Arc<RedDB>,
924    store: Option<Arc<crate::auth::store::AuthStore>>,
925    identity: Option<(String, crate::auth::Role, Option<String>)>,
926    values: Option<HashMap<String, Value>>,
927}
928
929impl ConfigResolver {
930    fn can_read(&self, key: &str) -> bool {
931        config_key_read_allowed(self.store.as_ref(), self.identity.as_ref(), key)
932    }
933}
934
935/// The reserved `red.config.*` system-config namespace. Never expose it via
936/// the inline `$config`/`CONFIG()`/`KV()` resolver regardless of IAM role,
937/// mirroring the role-independent `red.secret.*` hard-block (#1743).
938fn is_reserved_config_key(key: &str) -> bool {
939    let key = key.strip_prefix("red.config/").unwrap_or(key);
940    key.starts_with("red.config.")
941        || key == "red.config"
942        || key == super::function_catalog::REGISTRY_KEY
943}
944
945/// The system config-store collections (`red_config` / `red.config`) reached
946/// by the inline resolver. A `KV(<collection>, key)` read of one of these is
947/// gated exactly like `$config`; any other collection is unaffected (#1743).
948fn is_config_collection(collection: &str) -> bool {
949    matches!(
950        collection.to_ascii_lowercase().as_str(),
951        "red_config" | "red.config"
952    )
953}
954
955fn config_resource_target(key: &str) -> String {
956    let bare = key.strip_prefix("red.config/").unwrap_or(key);
957    format!("red.config/{bare}")
958}
959
960/// Shared inline-config read gate: a role-independent hard-block on the
961/// reserved `red.config.*` namespace, then a `config:read` capability check
962/// mirroring `SHOW CONFIG` / the `$secret` sibling (#1743). Absent a store or
963/// with IAM authorization disabled, only the hard-block applies so existing
964/// non-IAM `$config` reads keep working.
965fn config_key_read_allowed(
966    store: Option<&Arc<crate::auth::store::AuthStore>>,
967    identity: Option<&(String, crate::auth::Role, Option<String>)>,
968    key: &str,
969) -> bool {
970    if is_reserved_config_key(key) {
971        return false;
972    }
973    let Some(store) = store else {
974        return true;
975    };
976    if !store.iam_authorization_enabled() {
977        return true;
978    }
979    let Some((username, role, tenant)) = identity else {
980        return true;
981    };
982    let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
983    let mut resource =
984        crate::auth::policies::ResourceRef::new("config".to_string(), config_resource_target(key));
985    if let Some(tenant) = tenant {
986        resource = resource.with_tenant(tenant.clone());
987    }
988    let ctx = crate::auth::policies::EvalContext {
989        principal_tenant: tenant.clone(),
990        current_tenant: tenant.clone(),
991        peer_ip: None,
992        mfa_present: false,
993        now_ms: crate::auth::now_ms(),
994        principal_is_admin_role: *role == crate::auth::Role::Admin,
995        principal_is_platform_scoped: tenant.is_none(),
996    };
997    store.check_policy_authz_with_role(&principal, "config:read", &resource, &ctx, *role)
998}
999
1000/// Whether the current principal may read config `key` through the inline
1001/// `$config`/`CONFIG()`/`KV('red_config', …)` resolver. Used by the db-fallback
1002/// lookup paths that bypass [`current_config_value`] (#1743).
1003pub(crate) fn config_read_permitted(key: &str) -> bool {
1004    let key = key.to_ascii_lowercase();
1005    CURRENT_CONFIG_RESOLVER.with(|cell| match cell.borrow().as_ref() {
1006        Some(resolver) => resolver.can_read(&key),
1007        None => !is_reserved_config_key(&key),
1008    })
1009}
1010
1011/// Gate a raw `KV(collection, key)` read: config system collections go through
1012/// the `$config` gate; every other collection is unaffected (#1743).
1013pub(crate) fn kv_read_permitted(collection: &str, key: &str) -> bool {
1014    if is_config_collection(collection) {
1015        config_read_permitted(key)
1016    } else {
1017        true
1018    }
1019}
1020
1021pub(crate) struct ConfigSnapshotGuard {
1022    previous: Option<ConfigResolver>,
1023}
1024
1025impl ConfigSnapshotGuard {
1026    pub(super) fn install(
1027        db: Arc<RedDB>,
1028        store: Option<Arc<crate::auth::store::AuthStore>>,
1029    ) -> Self {
1030        let identity = current_auth_identity().map(|(username, role)| {
1031            let tenant = current_tenant();
1032            (username, role, tenant)
1033        });
1034        let previous = CURRENT_CONFIG_RESOLVER.with(|cell| {
1035            cell.replace(Some(ConfigResolver {
1036                db,
1037                store,
1038                identity,
1039                values: None,
1040            }))
1041        });
1042        Self { previous }
1043    }
1044}
1045
1046impl Drop for ConfigSnapshotGuard {
1047    fn drop(&mut self) {
1048        let previous = self.previous.take();
1049        CURRENT_CONFIG_RESOLVER.with(|cell| {
1050            cell.replace(previous);
1051        });
1052    }
1053}
1054
1055/// Install the MVCC snapshot used by the current thread for the duration
1056/// of one statement. Paired with `clear_current_snapshot()` — callers
1057/// should prefer the `CurrentSnapshotGuard` RAII wrapper so early returns
1058/// still clean up.
1059pub fn set_current_snapshot(ctx: SnapshotContext) {
1060    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = Some(ctx));
1061    HAS_SNAPSHOT.with(|c| c.set(true));
1062}
1063
1064pub fn clear_current_snapshot() {
1065    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = None);
1066    HAS_SNAPSHOT.with(|c| c.set(false));
1067}
1068
1069/// Drop-guard that restores the previous snapshot on scope exit. Safe to
1070/// nest — each statement saves the caller's snapshot and puts it back
1071/// instead of blindly clearing, so a top-level `execute_query` called
1072/// from inside another statement dispatch (e.g. vector source subqueries)
1073/// doesn't strip visibility from the outer scan.
1074pub(crate) struct CurrentSnapshotGuard {
1075    previous: Option<SnapshotContext>,
1076}
1077
1078impl CurrentSnapshotGuard {
1079    pub(crate) fn install(ctx: SnapshotContext) -> Self {
1080        let previous = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1081        set_current_snapshot(ctx);
1082        Self { previous }
1083    }
1084}
1085
1086impl Drop for CurrentSnapshotGuard {
1087    fn drop(&mut self) {
1088        let prev = self.previous.take();
1089        let has = prev.is_some();
1090        CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = prev);
1091        HAS_SNAPSHOT.with(|c| c.set(has));
1092    }
1093}
1094
1095/// Is this entity visible under the current thread's MVCC snapshot?
1096///
1097/// Returns `true` (no filtering) when no snapshot is installed — that
1098/// path is used by embedded callers and by operations that intentionally
1099/// bypass MVCC (VACUUM, snapshot export, admin introspection).
1100///
1101/// When a snapshot is installed the result is
1102///   `snapshot.sees(xmin, xmax) && !mgr.is_aborted(xmin) && !xmax_half_abort`
1103/// where `xmax_half_abort` re-grants visibility for tuples whose
1104/// deleting transaction rolled back.
1105#[inline]
1106pub fn entity_visible_under_current_snapshot(
1107    entity: &crate::storage::unified::entity::UnifiedEntity,
1108) -> bool {
1109    // Moderation visibility gate (#1274, ADR 0057). A row carrying the
1110    // moderation status marker — quarantine-pending or rejected-tombstone
1111    // — is hidden from every normal read, on top of MVCC visibility. The
1112    // marker lives on the row itself, so the check is a single field probe
1113    // and rides the existing per-row visibility chokepoint rather than
1114    // adding a separate filter pass to each scan call-site.
1115    if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
1116        return false;
1117    }
1118    // Fast path — one `Cell<bool>` read, no RefCell borrow. Autocommit
1119    // reads (no active MVCC transaction) still hide superseded physical
1120    // versions while avoiding a full snapshot-context lookup.
1121    // This runs on every row of every scan; the slow path only fires
1122    // inside an explicit transaction.
1123    if !HAS_SNAPSHOT.with(|c| c.get()) {
1124        return entity.xmax == 0;
1125    }
1126    CURRENT_SNAPSHOT.with(|cell| {
1127        let guard = cell.borrow();
1128        let Some(ctx) = guard.as_ref() else {
1129            return true;
1130        };
1131        let visible = visibility_check(ctx, entity.xmin, entity.xmax);
1132        if visible {
1133            record_serializable_read(ctx, entity);
1134        }
1135        visible
1136    })
1137}
1138
1139/// Direct visibility check from raw `(xmin, xmax)` — bypasses the
1140/// entity borrow for callers that already decomposed the tuple (e.g.
1141/// pre-materialized scan caches). Same semantics as
1142/// `entity_visible_under_current_snapshot`.
1143#[inline]
1144pub(crate) fn xids_visible_under_current_snapshot(xmin: u64, xmax: u64) -> bool {
1145    if !HAS_SNAPSHOT.with(|c| c.get()) {
1146        return true;
1147    }
1148    CURRENT_SNAPSHOT.with(|cell| {
1149        let guard = cell.borrow();
1150        let Some(ctx) = guard.as_ref() else {
1151            return true;
1152        };
1153        visibility_check(ctx, xmin, xmax)
1154    })
1155}
1156
1157/// Clone the current thread's snapshot context. Parallel scan paths
1158/// (`query_all_zoned` with `std::thread::scope`) call this on the main
1159/// thread *before* spawning workers so the captured `SnapshotContext`
1160/// can be moved into every worker closure. Worker threads do not
1161/// inherit thread-locals, so calling `entity_visible_under_current_snapshot`
1162/// from inside a spawned closure would silently skip the filter.
1163///
1164/// This is the only sanctioned ambient read of the snapshot thread-local
1165/// (issue #2138). `CURRENT_SNAPSHOT` and `HAS_SNAPSHOT` are private to this
1166/// module and stay that way: they exist for statement-frame entry and for
1167/// detecting a frameless read, nothing else. Everything downstream must
1168/// take the resulting `Option<SnapshotContext>` as an explicit argument —
1169/// that is what makes `SegmentManager::scan*` thread-independent.
1170///
1171/// Statement-frame build also reads it, so a nested sub-execution inside
1172/// one statement (CTE prelude, materialized-view refresh, inline graph
1173/// TVF) inherits the enclosing snapshot instead of minting a fresh one
1174/// (#2183).
1175///
1176/// The `Option` is load-bearing: `None` means no statement frame is
1177/// installed — the low-level callers that intentionally bypass MVCC
1178/// (VACUUM, snapshot export, admin introspection) — and the scan API
1179/// keeps a documented fallback for that case rather than assuming a frame
1180/// always exists.
1181pub fn capture_current_snapshot() -> Option<SnapshotContext> {
1182    CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone())
1183}
1184
1185/// Whether the active read snapshot may need historical tuple versions
1186/// that the current secondary indexes cannot prove. Index paths can still
1187/// recheck visible candidates, but only a heap scan can discover versions
1188/// whose indexed value was changed or deleted after this snapshot.
1189pub(crate) fn current_snapshot_requires_index_fallback() -> bool {
1190    if !HAS_SNAPSHOT.with(|c| c.get()) {
1191        return false;
1192    }
1193    CURRENT_SNAPSHOT.with(|cell| {
1194        cell.borrow()
1195            .as_ref()
1196            .is_some_and(|ctx| ctx.requires_index_fallback)
1197    })
1198}
1199
1200/// Frozen MVCC + identity context for callers that need to reinstall
1201/// the same view across thread-local boundaries — long-lived cursors,
1202/// background batchers, anything that detaches from the dispatch path
1203/// and re-enters later.
1204///
1205/// The bundle bakes in the three thread-locals every read path
1206/// consults: `SnapshotContext` (MVCC visibility), the auth identity
1207/// (RLS policy gate), and the tenant id (RLS scope). A FETCH that
1208/// reinstalls the bundle sees exactly the same rows as the DECLARE
1209/// would have, regardless of writes that landed in between.
1210///
1211/// Cheap to clone — `SnapshotContext` is a clone of three
1212/// `Arc`-backed fields, identity is a `(String, Role)`, tenant is a
1213/// `String`. None of these contend with the read path.
1214#[derive(Clone, Default)]
1215pub struct SnapshotBundle {
1216    pub snapshot: Option<SnapshotContext>,
1217    pub auth: Option<(String, crate::auth::Role)>,
1218    pub auth_principal: Option<crate::auth::UserId>,
1219    pub tenant: Option<String>,
1220}
1221
1222/// Capture the three read-path thread-locals into a `SnapshotBundle`.
1223/// Pairs with `with_snapshot_bundle` for re-entry.
1224pub fn snapshot_bundle() -> SnapshotBundle {
1225    SnapshotBundle {
1226        snapshot: capture_current_snapshot(),
1227        auth: current_auth_identity(),
1228        auth_principal: current_auth_principal(),
1229        tenant: CURRENT_TENANT_ID.with(|cell| cell.borrow().clone()),
1230    }
1231}
1232
1233/// Reinstall a captured `SnapshotBundle` for the duration of `f`.
1234/// Restores the caller's previous thread-locals on exit (panic-safe via
1235/// the explicit guard struct so a panic in `f` cannot leak the
1236/// installed identity into the worker's next request).
1237pub fn with_snapshot_bundle<R>(bundle: &SnapshotBundle, f: impl FnOnce() -> R) -> R {
1238    struct Guard {
1239        prev_snapshot: Option<SnapshotContext>,
1240        prev_auth: Option<(String, crate::auth::Role)>,
1241        prev_auth_principal: Option<crate::auth::UserId>,
1242        prev_tenant: Option<String>,
1243    }
1244    impl Drop for Guard {
1245        fn drop(&mut self) {
1246            let snap = self.prev_snapshot.take();
1247            let has = snap.is_some();
1248            CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = snap);
1249            HAS_SNAPSHOT.with(|c| c.set(has));
1250            CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = self.prev_auth.take());
1251            CURRENT_AUTH_PRINCIPAL
1252                .with(|cell| *cell.borrow_mut() = self.prev_auth_principal.take());
1253            CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = self.prev_tenant.take());
1254        }
1255    }
1256
1257    let _guard = {
1258        let prev_snapshot = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1259        let prev_auth = CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone());
1260        let prev_auth_principal = CURRENT_AUTH_PRINCIPAL.with(|cell| cell.borrow().clone());
1261        let prev_tenant = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
1262
1263        match bundle.snapshot.clone() {
1264            Some(ctx) => set_current_snapshot(ctx),
1265            None => clear_current_snapshot(),
1266        }
1267        CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = bundle.auth.clone());
1268        CURRENT_AUTH_PRINCIPAL.with(|cell| *cell.borrow_mut() = bundle.auth_principal.clone());
1269        CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = bundle.tenant.clone());
1270
1271        Guard {
1272            prev_snapshot,
1273            prev_auth,
1274            prev_auth_principal,
1275            prev_tenant,
1276        }
1277    };
1278    f()
1279}
1280
1281/// Apply the same visibility rules used by the thread-local helpers
1282/// against a caller-provided context. Intended for parallel workers
1283/// that captured the snapshot with `capture_current_snapshot()`.
1284#[inline]
1285pub fn entity_visible_with_context(
1286    ctx: Option<&SnapshotContext>,
1287    entity: &crate::storage::unified::entity::UnifiedEntity,
1288) -> bool {
1289    reddb_engine::storage::visibility::entity_visible_with_context(ctx, entity)
1290}