prax_query/tenant/
context.rs

1//! Tenant context for tracking the current tenant.
2
3use std::any::Any;
4use std::collections::HashMap;
5use std::fmt;
6use std::sync::Arc;
7
8/// A unique identifier for a tenant.
9#[derive(Debug, Clone, PartialEq, Eq, Hash)]
10pub struct TenantId(String);
11
12impl TenantId {
13    /// Create a new tenant ID.
14    pub fn new(id: impl Into<String>) -> Self {
15        Self(id.into())
16    }
17
18    /// Get the tenant ID as a string slice.
19    pub fn as_str(&self) -> &str {
20        &self.0
21    }
22
23    /// Convert to the inner string.
24    pub fn into_inner(self) -> String {
25        self.0
26    }
27}
28
29impl fmt::Display for TenantId {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        write!(f, "{}", self.0)
32    }
33}
34
35impl From<&str> for TenantId {
36    fn from(s: &str) -> Self {
37        Self::new(s)
38    }
39}
40
41impl From<String> for TenantId {
42    fn from(s: String) -> Self {
43        Self::new(s)
44    }
45}
46
47impl From<uuid::Uuid> for TenantId {
48    fn from(u: uuid::Uuid) -> Self {
49        Self::new(u.to_string())
50    }
51}
52
53impl From<i64> for TenantId {
54    fn from(i: i64) -> Self {
55        Self::new(i.to_string())
56    }
57}
58
59impl From<i32> for TenantId {
60    fn from(i: i32) -> Self {
61        Self::new(i.to_string())
62    }
63}
64
65/// Additional information about a tenant.
66#[derive(Debug, Clone, Default)]
67pub struct TenantInfo {
68    /// Display name for the tenant.
69    pub name: Option<String>,
70    /// Schema name (for schema-based isolation).
71    pub schema: Option<String>,
72    /// Database name (for database-based isolation).
73    pub database: Option<String>,
74    /// Whether this tenant has superuser privileges (bypasses filters).
75    pub is_superuser: bool,
76    /// Whether this is the system/default tenant.
77    pub is_system: bool,
78    /// Custom metadata.
79    metadata: HashMap<String, Arc<dyn Any + Send + Sync>>,
80}
81
82impl TenantInfo {
83    /// Create a new tenant info.
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// Set the tenant name.
89    pub fn with_name(mut self, name: impl Into<String>) -> Self {
90        self.name = Some(name.into());
91        self
92    }
93
94    /// Set the schema name.
95    pub fn with_schema(mut self, schema: impl Into<String>) -> Self {
96        self.schema = Some(schema.into());
97        self
98    }
99
100    /// Set the database name.
101    pub fn with_database(mut self, database: impl Into<String>) -> Self {
102        self.database = Some(database.into());
103        self
104    }
105
106    /// Mark as superuser.
107    pub fn as_superuser(mut self) -> Self {
108        self.is_superuser = true;
109        self
110    }
111
112    /// Mark as system tenant.
113    pub fn as_system(mut self) -> Self {
114        self.is_system = true;
115        self
116    }
117
118    /// Add custom metadata.
119    pub fn with_metadata<T: Any + Send + Sync>(mut self, key: impl Into<String>, value: T) -> Self {
120        self.metadata.insert(key.into(), Arc::new(value));
121        self
122    }
123
124    /// Get custom metadata.
125    pub fn get_metadata<T: Any + Send + Sync>(&self, key: &str) -> Option<&T> {
126        self.metadata.get(key).and_then(|v| v.downcast_ref())
127    }
128}
129
130/// Context for the current tenant.
131///
132/// This context is passed through the query pipeline to ensure all operations
133/// are scoped to the correct tenant.
134#[derive(Debug, Clone)]
135pub struct TenantContext {
136    /// The tenant identifier.
137    pub id: TenantId,
138    /// Additional tenant information.
139    pub info: TenantInfo,
140}
141
142impl TenantContext {
143    /// Create a new tenant context with just an ID.
144    pub fn new(id: impl Into<TenantId>) -> Self {
145        Self {
146            id: id.into(),
147            info: TenantInfo::default(),
148        }
149    }
150
151    /// Create a tenant context with additional info.
152    pub fn with_info(id: impl Into<TenantId>, info: TenantInfo) -> Self {
153        Self {
154            id: id.into(),
155            info,
156        }
157    }
158
159    /// Create a system/superuser context that bypasses tenant filters.
160    pub fn system() -> Self {
161        Self {
162            id: TenantId::new("__system__"),
163            info: TenantInfo::new().as_system().as_superuser(),
164        }
165    }
166
167    /// Check if this context should bypass tenant filters.
168    pub fn should_bypass(&self) -> bool {
169        self.info.is_superuser || self.info.is_system
170    }
171
172    /// Get the schema for this tenant (schema-based isolation).
173    pub fn schema(&self) -> Option<&str> {
174        self.info.schema.as_deref()
175    }
176
177    /// Get the database for this tenant (database-based isolation).
178    pub fn database(&self) -> Option<&str> {
179        self.info.database.as_deref()
180    }
181}
182
183/// Thread-local tenant context storage.
184#[cfg(feature = "thread-local-tenant")]
185mod thread_local {
186    use super::TenantContext;
187    use std::cell::RefCell;
188
189    thread_local! {
190        static CURRENT_TENANT: RefCell<Option<TenantContext>> = const { RefCell::new(None) };
191    }
192
193    /// Set the current tenant context for this thread.
194    pub fn set_current_tenant(ctx: TenantContext) {
195        CURRENT_TENANT.with(|t| {
196            *t.borrow_mut() = Some(ctx);
197        });
198    }
199
200    /// Get the current tenant context for this thread.
201    pub fn get_current_tenant() -> Option<TenantContext> {
202        CURRENT_TENANT.with(|t| t.borrow().clone())
203    }
204
205    /// Clear the current tenant context.
206    pub fn clear_current_tenant() {
207        CURRENT_TENANT.with(|t| {
208            *t.borrow_mut() = None;
209        });
210    }
211
212    /// Execute a closure with a specific tenant context.
213    pub fn with_tenant<F, R>(ctx: TenantContext, f: F) -> R
214    where
215        F: FnOnce() -> R,
216    {
217        let previous = get_current_tenant();
218        set_current_tenant(ctx);
219        let result = f();
220        if let Some(prev) = previous {
221            set_current_tenant(prev);
222        } else {
223            clear_current_tenant();
224        }
225        result
226    }
227}
228
229#[cfg(feature = "thread-local-tenant")]
230#[allow(unused_imports)]
231pub use thread_local::{clear_current_tenant, get_current_tenant, set_current_tenant, with_tenant};
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn test_tenant_id_creation() {
239        let id1 = TenantId::new("tenant-123");
240        assert_eq!(id1.as_str(), "tenant-123");
241
242        let id2: TenantId = "tenant-456".into();
243        assert_eq!(id2.as_str(), "tenant-456");
244
245        let id3: TenantId = 123_i64.into();
246        assert_eq!(id3.as_str(), "123");
247    }
248
249    #[test]
250    fn test_tenant_context() {
251        let ctx = TenantContext::new("tenant-123");
252        assert_eq!(ctx.id.as_str(), "tenant-123");
253        assert!(!ctx.should_bypass());
254    }
255
256    #[test]
257    fn test_system_context() {
258        let ctx = TenantContext::system();
259        assert!(ctx.should_bypass());
260        assert!(ctx.info.is_system);
261        assert!(ctx.info.is_superuser);
262    }
263
264    #[test]
265    fn test_tenant_info() {
266        let info = TenantInfo::new()
267            .with_name("Acme Corp")
268            .with_schema("tenant_acme")
269            .with_metadata("plan", "enterprise".to_string());
270
271        assert_eq!(info.name, Some("Acme Corp".to_string()));
272        assert_eq!(info.schema, Some("tenant_acme".to_string()));
273        assert_eq!(
274            info.get_metadata::<String>("plan"),
275            Some(&"enterprise".to_string())
276        );
277    }
278}
prax_query/tenant/context.rs

prax_query/tenant/
context.rs
