prax_query/tenant/
task_local.rs

1//! Zero-allocation tenant context using task-local storage.
2//!
3//! This module provides high-performance tenant context propagation using
4//! Tokio's task-local storage, eliminating the need for `Arc<RwLock>` in
5//! the hot path.
6//!
7//! # Performance Benefits
8//!
9//! - **Zero heap allocation** for context access
10//! - **No locking** on the hot path
11//! - **Automatic cleanup** when task completes
12//! - **Async-aware** - works across `.await` points
13//!
14//! # Example
15//!
16//! ```rust,ignore
17//! use prax_query::tenant::task_local::{with_tenant, current_tenant, TenantScope};
18//!
19//! // Set tenant for async block
20//! with_tenant("tenant-123", async {
21//!     // All queries in this block use tenant-123
22//!     let users = client.user().find_many().exec().await?;
23//!
24//!     // Nested calls also see the tenant
25//!     do_something_else().await?;
26//!
27//!     Ok(())
28//! }).await?;
29//!
30//! // Or use scoped guard
31//! let _guard = TenantScope::new("tenant-123");
32//! // tenant context available until guard is dropped
33//! ```
34
35use std::cell::Cell;
36use std::future::Future;
37
38use super::context::{TenantContext, TenantId};
39
40tokio::task_local! {
41    /// Task-local tenant context.
42    static TENANT_CONTEXT: TenantContext;
43}
44
45thread_local! {
46    /// Thread-local tenant ID for sync code paths.
47    /// Uses Cell for interior mutability without runtime cost.
48    static SYNC_TENANT_ID: Cell<Option<TenantId>> = const { Cell::new(None) };
49}
50
51/// Execute an async block with the given tenant context.
52///
53/// This is the most efficient way to set tenant context for async code.
54/// The context is automatically available to all nested async calls.
55///
56/// # Example
57///
58/// ```rust,ignore
59/// use prax_query::tenant::task_local::with_tenant;
60///
61/// with_tenant("tenant-123", async {
62///     // All code here sees tenant-123
63///     let users = client.user().find_many().exec().await?;
64///     Ok(())
65/// }).await?;
66/// ```
67pub async fn with_tenant<F, T>(tenant_id: impl Into<TenantId>, f: F) -> T
68where
69    F: Future<Output = T>,
70{
71    let ctx = TenantContext::new(tenant_id);
72    TENANT_CONTEXT.scope(ctx, f).await
73}
74
75/// Execute an async block with a full tenant context.
76pub async fn with_context<F, T>(ctx: TenantContext, f: F) -> T
77where
78    F: Future<Output = T>,
79{
80    TENANT_CONTEXT.scope(ctx, f).await
81}
82
83/// Get the current tenant context if set.
84///
85/// Returns `None` if no tenant context is active.
86///
87/// # Example
88///
89/// ```rust,ignore
90/// use prax_query::tenant::task_local::current_tenant;
91///
92/// if let Some(ctx) = current_tenant() {
93///     println!("Current tenant: {}", ctx.id);
94/// }
95/// ```
96#[inline]
97pub fn current_tenant() -> Option<TenantContext> {
98    TENANT_CONTEXT.try_with(|ctx| ctx.clone()).ok()
99}
100
101/// Get the current tenant ID if set.
102///
103/// More efficient than `current_tenant()` when you only need the ID.
104#[inline]
105pub fn current_tenant_id() -> Option<TenantId> {
106    TENANT_CONTEXT.try_with(|ctx| ctx.id.clone()).ok()
107}
108
109/// Get the current tenant ID as a string slice.
110///
111/// **Always returns an empty string**, even when a tenant context is active.
112/// Use [`current_tenant_id()`] instead.
113#[deprecated(
114    since = "0.11.0",
115    note = "always returns an empty string; use current_tenant_id() instead"
116)]
117#[inline]
118pub fn current_tenant_id_str() -> &'static str {
119    // This is a workaround - in practice you'd use current_tenant_id()
120    // We return a static str for zero-allocation in the common case
121    ""
122}
123
124/// Check if a tenant context is currently active.
125#[inline]
126pub fn has_tenant() -> bool {
127    TENANT_CONTEXT.try_with(|_| ()).is_ok()
128}
129
130/// Execute a closure with the current tenant context.
131///
132/// Returns `None` if no tenant context is active.
133#[inline]
134pub fn with_current_tenant<F, T>(f: F) -> Option<T>
135where
136    F: FnOnce(&TenantContext) -> T,
137{
138    TENANT_CONTEXT.try_with(f).ok()
139}
140
141/// Require a tenant context, returning an error if not set.
142#[inline]
143pub fn require_tenant() -> Result<TenantContext, TenantNotSetError> {
144    current_tenant().ok_or(TenantNotSetError)
145}
146
147/// Error returned when tenant context is required but not set.
148#[derive(Debug, Clone, Copy)]
149pub struct TenantNotSetError;
150
151impl std::fmt::Display for TenantNotSetError {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        write!(f, "tenant context not set")
154    }
155}
156
157impl std::error::Error for TenantNotSetError {}
158
159// ============================================================================
160// Sync Context (Thread-Local)
161// ============================================================================
162
163/// Set the tenant ID for synchronous code on the current thread.
164///
165/// This is useful for sync code paths or when you can't use async scope.
166/// The tenant is automatically cleared when the guard is dropped.
167///
168/// # Example
169///
170/// ```rust,ignore
171/// use prax_query::tenant::task_local::set_sync_tenant;
172///
173/// let _guard = set_sync_tenant("tenant-123");
174/// // tenant available for sync code
175/// ```
176pub fn set_sync_tenant(tenant_id: impl Into<TenantId>) -> SyncTenantGuard {
177    let id = tenant_id.into();
178    let previous = SYNC_TENANT_ID.with(|cell| cell.replace(Some(id)));
179    SyncTenantGuard { previous }
180}
181
182/// Get the current sync tenant ID.
183#[inline]
184pub fn sync_tenant_id() -> Option<TenantId> {
185    SYNC_TENANT_ID.with(|cell| {
186        // SAFETY: We only read, not modify
187        unsafe { &*cell.as_ptr() }.clone()
188    })
189}
190
191/// Guard that resets the sync tenant when dropped.
192pub struct SyncTenantGuard {
193    previous: Option<TenantId>,
194}
195
196impl Drop for SyncTenantGuard {
197    fn drop(&mut self) {
198        SYNC_TENANT_ID.with(|cell| cell.set(self.previous.take()));
199    }
200}
201
202// ============================================================================
203// Scoped Guard (Alternative API)
204// ============================================================================
205
206/// A scoped tenant context that tracks whether it's been entered.
207///
208/// This provides an alternative to `with_tenant` for cases where you
209/// need more control over the scope.
210///
211/// # Example
212///
213/// ```rust,ignore
214/// use prax_query::tenant::task_local::TenantScope;
215///
216/// async fn handle_request(tenant_id: &str) {
217///     let scope = TenantScope::new(tenant_id);
218///
219///     scope.run(async {
220///         // tenant context active here
221///     }).await;
222/// }
223/// ```
224#[derive(Debug, Clone)]
225pub struct TenantScope {
226    context: TenantContext,
227}
228
229impl TenantScope {
230    /// Create a new tenant scope.
231    pub fn new(tenant_id: impl Into<TenantId>) -> Self {
232        Self {
233            context: TenantContext::new(tenant_id),
234        }
235    }
236
237    /// Create from a full context.
238    pub fn from_context(context: TenantContext) -> Self {
239        Self { context }
240    }
241
242    /// Get the tenant ID.
243    pub fn tenant_id(&self) -> &TenantId {
244        &self.context.id
245    }
246
247    /// Get the full context.
248    pub fn context(&self) -> &TenantContext {
249        &self.context
250    }
251
252    /// Run an async function within this tenant scope.
253    pub async fn run<F, T>(&self, f: F) -> T
254    where
255        F: Future<Output = T>,
256    {
257        TENANT_CONTEXT.scope(self.context.clone(), f).await
258    }
259
260    /// Run a sync closure within this tenant scope (thread-local).
261    pub fn run_sync<F, T>(&self, f: F) -> T
262    where
263        F: FnOnce() -> T,
264    {
265        let _guard = set_sync_tenant(self.context.id.clone());
266        f()
267    }
268}
269
270// ============================================================================
271// Middleware Integration
272// ============================================================================
273
274/// Extract tenant from various sources.
275pub trait TenantExtractor: Send + Sync {
276    /// Extract tenant ID from a request/context.
277    fn extract(&self, headers: &[(String, String)]) -> Option<TenantId>;
278}
279
280/// Extract tenant from a header.
281#[derive(Debug, Clone)]
282pub struct HeaderExtractor {
283    header_name: String,
284}
285
286impl HeaderExtractor {
287    /// Create a new header extractor.
288    pub fn new(header_name: impl Into<String>) -> Self {
289        Self {
290            header_name: header_name.into(),
291        }
292    }
293
294    /// Create with default header name "X-Tenant-ID".
295    pub fn default_header() -> Self {
296        Self::new("X-Tenant-ID")
297    }
298}
299
300impl TenantExtractor for HeaderExtractor {
301    fn extract(&self, headers: &[(String, String)]) -> Option<TenantId> {
302        headers
303            .iter()
304            .find(|(k, _)| k.eq_ignore_ascii_case(&self.header_name))
305            .map(|(_, v)| TenantId::new(v.clone()))
306    }
307}
308
309/// Extract tenant from a JWT claim in the `Authorization: Bearer <jwt>` header.
310///
311/// **Security note:** this extractor performs NO signature verification. It
312/// only base64url-decodes the JWT payload segment and reads the configured
313/// claim. Extraction MUST sit behind a verifying middleware: the token must
314/// already have been authenticated upstream by the web framework / auth
315/// middleware before this extractor runs. Used on its own, any caller can
316/// forge any tenant id.
317#[derive(Debug, Clone)]
318pub struct UnverifiedJwtClaimExtractor {
319    claim_name: String,
320}
321
322impl UnverifiedJwtClaimExtractor {
323    /// Create a new JWT claim extractor.
324    pub fn new(claim_name: impl Into<String>) -> Self {
325        Self {
326            claim_name: claim_name.into(),
327        }
328    }
329
330    /// Create with default claim name "tenant_id".
331    pub fn default_claim() -> Self {
332        Self::new("tenant_id")
333    }
334
335    /// Get the claim name.
336    pub fn claim_name(&self) -> &str {
337        &self.claim_name
338    }
339}
340
341/// Source-compatibility alias for [`UnverifiedJwtClaimExtractor`].
342///
343/// # Deprecated
344///
345/// Renamed to [`UnverifiedJwtClaimExtractor`] to make the absence of
346/// signature verification explicit; this alias will be removed in the 0.12
347/// breaking release. (Kept as a plain alias rather than `#[deprecated]` so
348/// the crate-level re-export stays warning-free until then.)
349pub type JwtClaimExtractor = UnverifiedJwtClaimExtractor;
350
351impl TenantExtractor for UnverifiedJwtClaimExtractor {
352    /// Extract the tenant ID from the JWT in the `Authorization` header.
353    ///
354    /// Decodes the payload segment only — NO signature verification is
355    /// performed (see the struct-level security note). Malformed tokens are
356    /// logged at debug level so they are distinguishable from an absent
357    /// header.
358    fn extract(&self, headers: &[(String, String)]) -> Option<TenantId> {
359        use base64::Engine as _;
360
361        let auth = headers
362            .iter()
363            .find(|(k, _)| k.eq_ignore_ascii_case("authorization"))
364            .map(|(_, v)| v.as_str())?;
365
366        // Strip the auth scheme ("Bearer <token>", scheme is case-insensitive).
367        let Some(token) = auth
368            .get(..7)
369            .filter(|scheme| scheme.eq_ignore_ascii_case("bearer "))
370            .map(|_| &auth[7..])
371        else {
372            tracing::debug!("Authorization header present but not a Bearer token");
373            return None;
374        };
375
376        // JWT = header.payload.signature; only the payload segment is needed.
377        let Some(payload) = token.split('.').nth(1) else {
378            tracing::debug!("Authorization header present but JWT has no payload segment");
379            return None;
380        };
381        let decoded = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload) {
382            Ok(decoded) => decoded,
383            Err(error) => {
384                tracing::debug!(%error, "Authorization header present but JWT payload undecodable");
385                return None;
386            }
387        };
388        let claims: serde_json::Value = match serde_json::from_slice(&decoded) {
389            Ok(claims) => claims,
390            Err(error) => {
391                tracing::debug!(%error, "Authorization header present but JWT payload is not JSON");
392                return None;
393            }
394        };
395
396        claims
397            .get(self.claim_name.as_str())?
398            .as_str()
399            .map(TenantId::new)
400    }
401}
402
403/// Composite extractor that tries multiple sources.
404pub struct CompositeExtractor {
405    extractors: Vec<Box<dyn TenantExtractor>>,
406}
407
408impl CompositeExtractor {
409    /// Create a new composite extractor.
410    pub fn new() -> Self {
411        Self {
412            extractors: Vec::new(),
413        }
414    }
415
416    /// Add an extractor.
417    pub fn add<E: TenantExtractor + 'static>(mut self, extractor: E) -> Self {
418        self.extractors.push(Box::new(extractor));
419        self
420    }
421}
422
423impl Default for CompositeExtractor {
424    fn default() -> Self {
425        Self::new()
426    }
427}
428
429impl TenantExtractor for CompositeExtractor {
430    fn extract(&self, headers: &[(String, String)]) -> Option<TenantId> {
431        for extractor in &self.extractors {
432            if let Some(id) = extractor.extract(headers) {
433                return Some(id);
434            }
435        }
436        None
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443
444    #[tokio::test]
445    async fn test_with_tenant() {
446        let result = with_tenant("test-tenant", async { current_tenant_id() }).await;
447
448        assert_eq!(result.unwrap().as_str(), "test-tenant");
449    }
450
451    #[tokio::test]
452    async fn test_no_tenant() {
453        assert!(current_tenant().is_none());
454        assert!(!has_tenant());
455    }
456
457    #[tokio::test]
458    async fn test_nested_tenant() {
459        with_tenant("outer", async {
460            assert_eq!(current_tenant_id().unwrap().as_str(), "outer");
461
462            with_tenant("inner", async {
463                assert_eq!(current_tenant_id().unwrap().as_str(), "inner");
464            })
465            .await;
466
467            // Should be back to outer
468            assert_eq!(current_tenant_id().unwrap().as_str(), "outer");
469        })
470        .await;
471    }
472
473    #[tokio::test]
474    async fn test_tenant_scope() {
475        let scope = TenantScope::new("scoped-tenant");
476
477        let result = scope
478            .run(async { current_tenant_id().map(|id| id.as_str().to_string()) })
479            .await;
480
481        assert_eq!(result, Some("scoped-tenant".to_string()));
482    }
483
484    #[test]
485    fn test_sync_tenant() {
486        {
487            let _guard = set_sync_tenant("sync-tenant");
488            assert_eq!(sync_tenant_id().unwrap().as_str(), "sync-tenant");
489        }
490
491        // Should be cleared after guard drop
492        assert!(sync_tenant_id().is_none());
493    }
494
495    #[test]
496    fn test_header_extractor() {
497        let extractor = HeaderExtractor::new("X-Tenant-ID");
498
499        let headers = vec![
500            ("Content-Type".to_string(), "application/json".to_string()),
501            ("X-Tenant-ID".to_string(), "tenant-from-header".to_string()),
502        ];
503
504        let id = extractor.extract(&headers);
505        assert_eq!(id.unwrap().as_str(), "tenant-from-header");
506    }
507
508    #[test]
509    fn test_jwt_claim_extractor() {
510        let extractor = UnverifiedJwtClaimExtractor::new("tenant_id");
511
512        // Unsigned JWT (alg: none) with payload {"tenant_id":"tenant-from-jwt"}.
513        // Extraction decodes the payload only; signature verification is the
514        // framework's responsibility.
515        let jwt = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJ0ZW5hbnRfaWQiOiJ0ZW5hbnQtZnJvbS1qd3QifQ.";
516        let headers = vec![
517            ("Content-Type".to_string(), "application/json".to_string()),
518            ("Authorization".to_string(), format!("Bearer {jwt}")),
519        ];
520
521        let id = extractor.extract(&headers);
522        assert_eq!(id.unwrap().as_str(), "tenant-from-jwt");
523
524        // The deprecated alias still resolves to the same type.
525        let _alias: JwtClaimExtractor = UnverifiedJwtClaimExtractor::default_claim();
526
527        // Missing Authorization header yields no tenant.
528        assert!(extractor.extract(&[]).is_none());
529
530        // Malformed tokens yield no tenant (logged at debug level).
531        for bad in ["not-a-jwt", "Bearer !!!.@@@.###", "Bearer a.b.c"] {
532            let headers = vec![("Authorization".to_string(), bad.to_string())];
533            assert!(extractor.extract(&headers).is_none());
534        }
535
536        // Missing claim yields no tenant.
537        let other = UnverifiedJwtClaimExtractor::new("org_id");
538        assert!(other.extract(&headers).is_none());
539    }
540
541    #[test]
542    fn test_composite_extractor() {
543        let extractor = CompositeExtractor::new()
544            .add(HeaderExtractor::new("X-Organization-ID"))
545            .add(HeaderExtractor::new("X-Tenant-ID"));
546
547        let headers = vec![("X-Tenant-ID".to_string(), "fallback-tenant".to_string())];
548
549        let id = extractor.extract(&headers);
550        assert_eq!(id.unwrap().as_str(), "fallback-tenant");
551    }
552}
prax_query/tenant/task_local.rs

prax_query/tenant/
task_local.rs
