uni_plugin/host/
principal.rs

1//! Authenticated-principal task-local plumbing.
2//!
3//! Hosts (Session / Transaction execute boundaries) install the current
4//! [`Principal`] into a Tokio task-local at the top of each query so deeper
5//! procedure / UDF invocation sites can read it without threading the
6//! principal through every internal API.
7//!
8//! This module also exposes [`maybe_scope_with_principal`] — the convenience
9//! wrapper that collapses the `match principal { Some(p) => scoped..., None
10//! => fut }.await` pattern previously duplicated across
11//! `uni::api::{session,transaction}` and `uni-query::df_udfs`.
12//!
13//! # Stability
14//!
15//! Moved here in Phase 5 of the §1.2 consolidation pass. `uni-query`
16//! re-exports the items below for backwards compatibility.
17
18// Rust guideline compliant
19
20use std::future::Future;
21use std::sync::Arc;
22
23use crate::traits::connector::Principal;
24
25tokio::task_local! {
26    /// Tokio task-local carrying the **authenticated principal** for
27    /// the in-flight query.
28    ///
29    /// Set by the host-crate execute boundaries (`Session::query`,
30    /// `Transaction::query`, `Transaction::execute`) so procedure
31    /// invocation sites can populate `ProcedureContext::with_principal`
32    /// without threading the principal through every internal API.
33    /// Read at `uni-query`'s procedure-call paths immediately before
34    /// calling `plugin.invoke(ctx, ...)`.
35    ///
36    /// Propagates across `.await` points within the same task tree;
37    /// does NOT propagate across `tokio::spawn`. The synthetic
38    /// procedure body bridge (`block_in_place` + `Handle::block_on`)
39    /// stays on the same task so the principal remains visible to
40    /// any nested `execute_inner_query` calls.
41    pub static CURRENT_PRINCIPAL: Arc<Principal>;
42}
43
44/// Run `fut` inside a scope where [`current_principal`] resolves to
45/// `principal`.
46///
47/// Use this at every host-crate boundary where a `Session` or
48/// `Transaction` dispatches into the executor.
49///
50/// # Examples
51///
52/// ```no_run
53/// use std::sync::Arc;
54/// use uni_plugin::host::principal::{scoped_with_principal, current_principal};
55/// use uni_plugin::traits::connector::Principal;
56///
57/// # async fn demo(principal: Arc<Principal>) {
58/// scoped_with_principal(principal.clone(), async {
59///     assert!(current_principal().is_some());
60/// })
61/// .await;
62/// # }
63/// ```
64pub fn scoped_with_principal<F: Future>(
65    principal: Arc<Principal>,
66    fut: F,
67) -> tokio::task::futures::TaskLocalFuture<Arc<Principal>, F> {
68    CURRENT_PRINCIPAL.scope(principal, fut)
69}
70
71/// Borrow the principal active for the current execute scope, if any.
72///
73/// Returns `None` outside a [`scoped_with_principal`] scope (e.g.,
74/// low-level unit tests that bypass `Session`).
75#[must_use]
76pub fn current_principal() -> Option<Arc<Principal>> {
77    CURRENT_PRINCIPAL.try_with(|p| p.clone()).ok()
78}
79
80/// Run `fut` either inside a principal task-local scope or unwrapped,
81/// depending on whether a principal was supplied.
82///
83/// Replaces the duplicated
84/// `match principal { Some(p) => scoped_with_principal(p, fut).await, None => fut.await }`
85/// pattern across `uni::api::{session,transaction}` and
86/// `uni-query::df_udfs`.
87///
88/// # Examples
89///
90/// ```no_run
91/// use std::sync::Arc;
92/// use uni_plugin::host::principal::maybe_scope_with_principal;
93/// use uni_plugin::traits::connector::Principal;
94///
95/// # async fn demo(principal: Option<Arc<Principal>>) -> u32 {
96/// maybe_scope_with_principal(principal, async { 42 }).await
97/// # }
98/// ```
99pub async fn maybe_scope_with_principal<F>(principal: Option<Arc<Principal>>, fut: F) -> F::Output
100where
101    F: Future,
102{
103    match principal {
104        Some(p) => scoped_with_principal(p, fut).await,
105        None => fut.await,
106    }
107}
uni_plugin/host/principal.rs

uni_plugin/host/
principal.rs
