sz_orm_core/
connection_tenant.rs

1//! # 连接级多租户隔离
2//!
3//! 在同一连接池中连接绑定到特定租户（通过 `SET app.tenant_id = ?`），
4//! 避免每租户独立池的资源开销。支持三种连接亲和策略 + RAII 守卫。
5
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8use std::time::{SystemTime, UNIX_EPOCH};
9
10use serde::{Deserialize, Serialize};
11
12use crate::db_type::DbType;
13use crate::pool::Pool;
14
15fn now_ms() -> u64 {
16    SystemTime::now()
17        .duration_since(UNIX_EPOCH)
18        .map(|d| d.as_millis() as u64)
19        .unwrap_or(0)
20}
21
22/// 连接级隔离机制
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub enum ConnectionLevelIsolation {
25    /// 通过 `SET app.tenant_id = ?` 设置租户上下文
26    SetTenantId,
27    /// Schema 隔离，路由到 `tenant_{id}_{table}`
28    SchemaIsolation,
29    /// 连接绑定，连接专属租户
30    ConnectionBinding,
31}
32
33/// 连接亲和策略
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35pub enum ConnectionAffinityPolicy {
36    /// 严格亲和：仅使用绑定到该租户的连接
37    Strict,
38    /// 优先亲和：优先使用绑定连接，无可用时获取任意连接
39    Preferred,
40    /// 无亲和：任意连接，每次设置租户上下文
41    None,
42}
43
44/// 连接级多租户配置
45#[derive(Debug, Clone)]
46pub struct ConnectionLevelTenantConfig {
47    /// 隔离机制
48    pub isolation: ConnectionLevelIsolation,
49    /// 亲和策略
50    pub affinity_policy: ConnectionAffinityPolicy,
51    /// 亲和超时（毫秒）
52    pub affinity_timeout_ms: u64,
53    /// 数据库类型
54    pub db_type: DbType,
55}
56
57impl ConnectionLevelTenantConfig {
58    /// 创建默认配置（SetTenantId, Preferred, 5000ms）
59    pub fn new(db_type: DbType) -> Self {
60        Self {
61            isolation: ConnectionLevelIsolation::SetTenantId,
62            affinity_policy: ConnectionAffinityPolicy::Preferred,
63            affinity_timeout_ms: 5_000,
64            db_type,
65        }
66    }
67
68    /// 设置隔离机制
69    pub fn with_isolation(mut self, isolation: ConnectionLevelIsolation) -> Self {
70        self.isolation = isolation;
71        self
72    }
73
74    /// 设置亲和策略
75    pub fn with_affinity_policy(mut self, policy: ConnectionAffinityPolicy) -> Self {
76        self.affinity_policy = policy;
77        self
78    }
79
80    /// 设置亲和超时（毫秒）
81    pub fn with_affinity_timeout_ms(mut self, ms: u64) -> Self {
82        self.affinity_timeout_ms = ms;
83        self
84    }
85}
86
87/// 租户错误类型
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub enum TenantError {
90    /// 无可用绑定连接
91    NoBoundConnection,
92    /// 篡改被拒绝
93    TamperingRejected,
94    /// 清理失败
95    CleanupFailed,
96    /// 不支持的方言
97    UnsupportedDialect,
98    /// 租户 ID 为空
99    EmptyTenantId,
100}
101
102impl std::fmt::Display for TenantError {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        match self {
105            TenantError::NoBoundConnection => write!(f, "no connection bound to tenant"),
106            TenantError::TamperingRejected => write!(f, "tenant context tampering rejected"),
107            TenantError::CleanupFailed => write!(f, "tenant context cleanup failed"),
108            TenantError::UnsupportedDialect => {
109                write!(f, "unsupported dialect for SET app.tenant_id")
110            }
111            TenantError::EmptyTenantId => write!(f, "tenant_id is empty"),
112        }
113    }
114}
115
116impl std::error::Error for TenantError {}
117
118/// 连接 ID 类型
119pub type ConnectionId = u64;
120
121/// 连接租户绑定记录
122#[derive(Debug, Clone)]
123pub struct TenantBinding {
124    /// 连接 ID
125    pub connection_id: ConnectionId,
126    /// 绑定的租户 ID
127    pub tenant_id: String,
128    /// 绑定时间戳
129    pub bound_at: u64,
130}
131
132/// 连接租户绑定器
133pub struct ConnectionTenantBinder {
134    pool: Arc<Pool>,
135    config: ConnectionLevelTenantConfig,
136    tenant_bindings: Mutex<HashMap<String, Vec<ConnectionId>>>,
137    next_connection_id: Mutex<u64>,
138}
139
140impl ConnectionTenantBinder {
141    /// 创建连接租户绑定器
142    pub fn new(pool: Arc<Pool>, config: ConnectionLevelTenantConfig) -> Self {
143        Self {
144            pool,
145            config,
146            tenant_bindings: Mutex::new(HashMap::new()),
147            next_connection_id: Mutex::new(1),
148        }
149    }
150
151    /// 获取配置
152    pub fn config(&self) -> &ConnectionLevelTenantConfig {
153        &self.config
154    }
155
156    /// 判断是否支持 `SET app.tenant_id`
157    pub fn supports_set_tenant_id(&self) -> bool {
158        matches!(self.config.db_type, DbType::PostgreSQL | DbType::MySQL)
159    }
160
161    /// 生成 `SET app.tenant_id` SQL
162    pub fn build_set_tenant_sql(&self, tenant_id: &str) -> String {
163        format!("SET app.tenant_id = '{}'", tenant_id)
164    }
165
166    /// 生成清理租户上下文 SQL
167    pub fn build_clear_tenant_sql(&self) -> String {
168        "SET app.tenant_id = NULL".to_string()
169    }
170
171    /// 绑定连接到租户
172    pub fn bind_connection(&self, tenant_id: &str) -> ConnectionId {
173        let conn_id = {
174            let mut next = self.next_connection_id.lock().unwrap();
175            let id = *next;
176            *next += 1;
177            id
178        };
179        let mut bindings = self.tenant_bindings.lock().unwrap();
180        bindings
181            .entry(tenant_id.to_string())
182            .or_default()
183            .push(conn_id);
184        conn_id
185    }
186
187    /// 查找绑定到指定租户的连接
188    pub fn find_bound_connections(&self, tenant_id: &str) -> Vec<ConnectionId> {
189        let bindings = self.tenant_bindings.lock().unwrap();
190        bindings.get(tenant_id).cloned().unwrap_or_default()
191    }
192
193    /// 解绑连接
194    pub fn unbind_connection(&self, tenant_id: &str, conn_id: ConnectionId) {
195        let mut bindings = self.tenant_bindings.lock().unwrap();
196        if let Some(conns) = bindings.get_mut(tenant_id) {
197            conns.retain(|&id| id != conn_id);
198        }
199    }
200
201    /// 获取绑定数量
202    pub fn binding_count(&self, tenant_id: &str) -> usize {
203        let bindings = self.tenant_bindings.lock().unwrap();
204        bindings.get(tenant_id).map(|v| v.len()).unwrap_or(0)
205    }
206
207    /// 获取所有租户绑定
208    pub fn all_bindings(&self) -> Vec<TenantBinding> {
209        let bindings = self.tenant_bindings.lock().unwrap();
210        let mut result = Vec::new();
211        for (tenant_id, conn_ids) in bindings.iter() {
212            for &conn_id in conn_ids {
213                result.push(TenantBinding {
214                    connection_id: conn_id,
215                    tenant_id: tenant_id.clone(),
216                    bound_at: now_ms(),
217                });
218            }
219        }
220        result
221    }
222
223    /// 获取连接池引用
224    pub fn pool(&self) -> &Arc<Pool> {
225        &self.pool
226    }
227
228    /// 验证租户 ID
229    pub fn validate_tenant_id(&self, tenant_id: &str) -> Result<(), TenantError> {
230        if tenant_id.is_empty() {
231            return Err(TenantError::EmptyTenantId);
232        }
233        Ok(())
234    }
235
236    /// 确定实际隔离机制（处理方言降级）
237    pub fn resolve_isolation(&self) -> ConnectionLevelIsolation {
238        if self.config.isolation == ConnectionLevelIsolation::SetTenantId
239            && !self.supports_set_tenant_id()
240        {
241            ConnectionLevelIsolation::SchemaIsolation
242        } else {
243            self.config.isolation.clone()
244        }
245    }
246}
247
248/// 租户连接守卫（RAII）
249pub struct TenantConnectionGuard {
250    binder: Arc<ConnectionTenantBinder>,
251    tenant_id: String,
252    connection_id: ConnectionId,
253    active: bool,
254}
255
256impl TenantConnectionGuard {
257    /// 创建守卫
258    pub fn new(
259        binder: Arc<ConnectionTenantBinder>,
260        tenant_id: String,
261        connection_id: ConnectionId,
262    ) -> Self {
263        Self {
264            binder,
265            tenant_id,
266            connection_id,
267            active: true,
268        }
269    }
270
271    /// 获取租户 ID
272    pub fn tenant_id(&self) -> &str {
273        &self.tenant_id
274    }
275
276    /// 获取连接 ID
277    pub fn connection_id(&self) -> ConnectionId {
278        self.connection_id
279    }
280
281    /// 是否活跃
282    pub fn is_active(&self) -> bool {
283        self.active
284    }
285
286    /// 获取清理 SQL
287    pub fn clear_tenant_sql(&self) -> String {
288        self.binder.build_clear_tenant_sql()
289    }
290
291    /// 手动释放（提前清理）
292    pub fn release(&mut self) -> Result<(), TenantError> {
293        if !self.active {
294            return Ok(());
295        }
296        self.binder
297            .unbind_connection(&self.tenant_id, self.connection_id);
298        self.active = false;
299        Ok(())
300    }
301}
302
303impl Drop for TenantConnectionGuard {
304    fn drop(&mut self) {
305        if self.active {
306            self.binder
307                .unbind_connection(&self.tenant_id, self.connection_id);
308        }
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use crate::pool::{ConnectionFactory, PoolConfigBuilder};
316
317    struct MockFactory;
318
319    #[async_trait::async_trait]
320    impl ConnectionFactory for MockFactory {
321        async fn create(&self) -> Result<Box<dyn crate::pool::Connection>, crate::DbError> {
322            Err(crate::DbError::PoolError(
323                crate::error::PoolError::InvalidConfig("mock".to_string()),
324            ))
325        }
326    }
327
328    fn make_pool() -> Arc<Pool> {
329        let config = PoolConfigBuilder::new().max_size(1).build().unwrap();
330        let factory: Arc<dyn ConnectionFactory> = Arc::new(MockFactory);
331        Arc::new(Pool::new(config, factory).unwrap())
332    }
333
334    #[test]
335    fn test_config_default() {
336        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
337        assert_eq!(config.isolation, ConnectionLevelIsolation::SetTenantId);
338        assert_eq!(config.affinity_policy, ConnectionAffinityPolicy::Preferred);
339        assert_eq!(config.affinity_timeout_ms, 5_000);
340        assert_eq!(config.db_type, DbType::PostgreSQL);
341    }
342
343    #[test]
344    fn test_config_builder() {
345        let config = ConnectionLevelTenantConfig::new(DbType::MySQL)
346            .with_isolation(ConnectionLevelIsolation::ConnectionBinding)
347            .with_affinity_policy(ConnectionAffinityPolicy::Strict)
348            .with_affinity_timeout_ms(10_000);
349        assert_eq!(
350            config.isolation,
351            ConnectionLevelIsolation::ConnectionBinding
352        );
353        assert_eq!(config.affinity_policy, ConnectionAffinityPolicy::Strict);
354        assert_eq!(config.affinity_timeout_ms, 10_000);
355    }
356
357    #[test]
358    fn test_isolation_serde() {
359        let isolations = vec![
360            ConnectionLevelIsolation::SetTenantId,
361            ConnectionLevelIsolation::SchemaIsolation,
362            ConnectionLevelIsolation::ConnectionBinding,
363        ];
364        for i in &isolations {
365            let json = serde_json::to_string(i).unwrap();
366            let decoded: ConnectionLevelIsolation = serde_json::from_str(&json).unwrap();
367            assert_eq!(*i, decoded);
368        }
369    }
370
371    #[test]
372    fn test_affinity_policy_serde() {
373        let policies = vec![
374            ConnectionAffinityPolicy::Strict,
375            ConnectionAffinityPolicy::Preferred,
376            ConnectionAffinityPolicy::None,
377        ];
378        for p in &policies {
379            let json = serde_json::to_string(p).unwrap();
380            let decoded: ConnectionAffinityPolicy = serde_json::from_str(&json).unwrap();
381            assert_eq!(*p, decoded);
382        }
383    }
384
385    #[test]
386    fn test_supports_set_tenant_id() {
387        let pool = make_pool();
388        let pg_config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
389        let pg_binder = ConnectionTenantBinder::new(pool.clone(), pg_config);
390        assert!(pg_binder.supports_set_tenant_id());
391
392        let mysql_config = ConnectionLevelTenantConfig::new(DbType::MySQL);
393        let mysql_binder = ConnectionTenantBinder::new(pool.clone(), mysql_config);
394        assert!(mysql_binder.supports_set_tenant_id());
395
396        let sqlite_config = ConnectionLevelTenantConfig::new(DbType::Sqlite);
397        let sqlite_binder = ConnectionTenantBinder::new(pool.clone(), sqlite_config);
398        assert!(!sqlite_binder.supports_set_tenant_id());
399
400        let oracle_config = ConnectionLevelTenantConfig::new(DbType::Oracle);
401        let oracle_binder = ConnectionTenantBinder::new(pool, oracle_config);
402        assert!(!oracle_binder.supports_set_tenant_id());
403    }
404
405    #[test]
406    fn test_build_set_tenant_sql() {
407        let pool = make_pool();
408        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
409        let binder = ConnectionTenantBinder::new(pool, config);
410        let sql = binder.build_set_tenant_sql("tenant_123");
411        assert!(sql.contains("SET app.tenant_id"));
412        assert!(sql.contains("tenant_123"));
413    }
414
415    #[test]
416    fn test_build_clear_tenant_sql() {
417        let pool = make_pool();
418        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
419        let binder = ConnectionTenantBinder::new(pool, config);
420        let sql = binder.build_clear_tenant_sql();
421        assert!(sql.contains("NULL"));
422    }
423
424    #[test]
425    fn test_bind_and_find_connection() {
426        let pool = make_pool();
427        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
428        let binder = ConnectionTenantBinder::new(pool, config);
429
430        let conn_id1 = binder.bind_connection("tenant_1");
431        let conn_id2 = binder.bind_connection("tenant_1");
432        let conn_id3 = binder.bind_connection("tenant_2");
433
434        assert_ne!(conn_id1, conn_id2);
435        assert_ne!(conn_id1, conn_id3);
436
437        let tenant1_conns = binder.find_bound_connections("tenant_1");
438        assert_eq!(tenant1_conns.len(), 2);
439        assert!(tenant1_conns.contains(&conn_id1));
440        assert!(tenant1_conns.contains(&conn_id2));
441
442        let tenant2_conns = binder.find_bound_connections("tenant_2");
443        assert_eq!(tenant2_conns.len(), 1);
444        assert!(tenant2_conns.contains(&conn_id3));
445
446        assert_eq!(binder.binding_count("tenant_1"), 2);
447        assert_eq!(binder.binding_count("tenant_2"), 1);
448        assert_eq!(binder.binding_count("tenant_3"), 0);
449    }
450
451    #[test]
452    fn test_unbind_connection() {
453        let pool = make_pool();
454        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
455        let binder = ConnectionTenantBinder::new(pool, config);
456
457        let conn_id1 = binder.bind_connection("tenant_1");
458        let conn_id2 = binder.bind_connection("tenant_1");
459        assert_eq!(binder.binding_count("tenant_1"), 2);
460
461        binder.unbind_connection("tenant_1", conn_id1);
462        assert_eq!(binder.binding_count("tenant_1"), 1);
463
464        let conns = binder.find_bound_connections("tenant_1");
465        assert!(conns.contains(&conn_id2));
466        assert!(!conns.contains(&conn_id1));
467
468        binder.unbind_connection("tenant_1", conn_id2);
469        assert_eq!(binder.binding_count("tenant_1"), 0);
470    }
471
472    #[test]
473    fn test_validate_tenant_id() {
474        let pool = make_pool();
475        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
476        let binder = ConnectionTenantBinder::new(pool, config);
477
478        assert!(binder.validate_tenant_id("tenant_1").is_ok());
479        assert_eq!(
480            binder.validate_tenant_id("").unwrap_err(),
481            TenantError::EmptyTenantId
482        );
483    }
484
485    #[test]
486    fn test_resolve_isolation_pg() {
487        let pool = make_pool();
488        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
489        let binder = ConnectionTenantBinder::new(pool, config);
490        assert_eq!(
491            binder.resolve_isolation(),
492            ConnectionLevelIsolation::SetTenantId
493        );
494    }
495
496    #[test]
497    fn test_resolve_isolation_sqlite_fallback() {
498        let pool = make_pool();
499        let config = ConnectionLevelTenantConfig::new(DbType::Sqlite);
500        let binder = ConnectionTenantBinder::new(pool, config);
501        assert_eq!(
502            binder.resolve_isolation(),
503            ConnectionLevelIsolation::SchemaIsolation
504        );
505    }
506
507    #[test]
508    fn test_resolve_isolation_schema_no_fallback() {
509        let pool = make_pool();
510        let config = ConnectionLevelTenantConfig::new(DbType::Sqlite)
511            .with_isolation(ConnectionLevelIsolation::SchemaIsolation);
512        let binder = ConnectionTenantBinder::new(pool, config);
513        assert_eq!(
514            binder.resolve_isolation(),
515            ConnectionLevelIsolation::SchemaIsolation
516        );
517    }
518
519    #[test]
520    fn test_tenant_error_display() {
521        let err = TenantError::NoBoundConnection;
522        assert!(err.to_string().contains("no connection"));
523        let err = TenantError::EmptyTenantId;
524        assert!(err.to_string().contains("empty"));
525        let err = TenantError::UnsupportedDialect;
526        assert!(err.to_string().contains("unsupported"));
527    }
528
529    #[test]
530    fn test_guard_drop_unbinds() {
531        let pool = make_pool();
532        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
533        let binder = Arc::new(ConnectionTenantBinder::new(pool, config));
534
535        let conn_id = binder.bind_connection("tenant_1");
536        assert_eq!(binder.binding_count("tenant_1"), 1);
537
538        {
539            let _guard =
540                TenantConnectionGuard::new(binder.clone(), "tenant_1".to_string(), conn_id);
541            assert_eq!(binder.binding_count("tenant_1"), 1);
542        }
543
544        assert_eq!(binder.binding_count("tenant_1"), 0);
545    }
546
547    #[test]
548    fn test_guard_release() {
549        let pool = make_pool();
550        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
551        let binder = Arc::new(ConnectionTenantBinder::new(pool, config));
552
553        let conn_id = binder.bind_connection("tenant_1");
554        let mut guard = TenantConnectionGuard::new(binder, "tenant_1".to_string(), conn_id);
555        assert!(guard.is_active());
556        guard.release().unwrap();
557        assert!(!guard.is_active());
558    }
559
560    #[test]
561    fn test_guard_properties() {
562        let pool = make_pool();
563        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
564        let binder = Arc::new(ConnectionTenantBinder::new(pool, config));
565
566        let guard = TenantConnectionGuard::new(binder, "tenant_42".to_string(), 999);
567        assert_eq!(guard.tenant_id(), "tenant_42");
568        assert_eq!(guard.connection_id(), 999);
569        assert!(guard.is_active());
570        assert!(guard.clear_tenant_sql().contains("NULL"));
571    }
572
573    #[test]
574    fn test_all_bindings() {
575        let pool = make_pool();
576        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
577        let binder = ConnectionTenantBinder::new(pool, config);
578
579        binder.bind_connection("tenant_1");
580        binder.bind_connection("tenant_1");
581        binder.bind_connection("tenant_2");
582
583        let all = binder.all_bindings();
584        assert_eq!(all.len(), 3);
585    }
586}
sz_orm_core/connection_tenant.rs

sz_orm_core/
connection_tenant.rs
