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
175                .next_connection_id
176                .lock()
177                .unwrap_or_else(|e| e.into_inner());
178            let id = *next;
179            *next += 1;
180            id
181        };
182        let mut bindings = self
183            .tenant_bindings
184            .lock()
185            .unwrap_or_else(|e| e.into_inner());
186        bindings
187            .entry(tenant_id.to_string())
188            .or_default()
189            .push(conn_id);
190        conn_id
191    }
192
193    /// 查找绑定到指定租户的连接
194    pub fn find_bound_connections(&self, tenant_id: &str) -> Vec<ConnectionId> {
195        let bindings = self
196            .tenant_bindings
197            .lock()
198            .unwrap_or_else(|e| e.into_inner());
199        bindings.get(tenant_id).cloned().unwrap_or_default()
200    }
201
202    /// 解绑连接
203    pub fn unbind_connection(&self, tenant_id: &str, conn_id: ConnectionId) {
204        let mut bindings = self
205            .tenant_bindings
206            .lock()
207            .unwrap_or_else(|e| e.into_inner());
208        if let Some(conns) = bindings.get_mut(tenant_id) {
209            conns.retain(|&id| id != conn_id);
210        }
211    }
212
213    /// 获取绑定数量
214    pub fn binding_count(&self, tenant_id: &str) -> usize {
215        let bindings = self
216            .tenant_bindings
217            .lock()
218            .unwrap_or_else(|e| e.into_inner());
219        bindings.get(tenant_id).map(|v| v.len()).unwrap_or(0)
220    }
221
222    /// 获取所有租户绑定
223    pub fn all_bindings(&self) -> Vec<TenantBinding> {
224        let bindings = self
225            .tenant_bindings
226            .lock()
227            .unwrap_or_else(|e| e.into_inner());
228        let mut result = Vec::new();
229        for (tenant_id, conn_ids) in bindings.iter() {
230            for &conn_id in conn_ids {
231                result.push(TenantBinding {
232                    connection_id: conn_id,
233                    tenant_id: tenant_id.clone(),
234                    bound_at: now_ms(),
235                });
236            }
237        }
238        result
239    }
240
241    /// 获取连接池引用
242    pub fn pool(&self) -> &Arc<Pool> {
243        &self.pool
244    }
245
246    /// 验证租户 ID
247    pub fn validate_tenant_id(&self, tenant_id: &str) -> Result<(), TenantError> {
248        if tenant_id.is_empty() {
249            return Err(TenantError::EmptyTenantId);
250        }
251        Ok(())
252    }
253
254    /// 确定实际隔离机制（处理方言降级）
255    pub fn resolve_isolation(&self) -> ConnectionLevelIsolation {
256        if self.config.isolation == ConnectionLevelIsolation::SetTenantId
257            && !self.supports_set_tenant_id()
258        {
259            ConnectionLevelIsolation::SchemaIsolation
260        } else {
261            self.config.isolation.clone()
262        }
263    }
264}
265
266/// 租户连接守卫（RAII）
267pub struct TenantConnectionGuard {
268    binder: Arc<ConnectionTenantBinder>,
269    tenant_id: String,
270    connection_id: ConnectionId,
271    active: bool,
272}
273
274impl TenantConnectionGuard {
275    /// 创建守卫
276    pub fn new(
277        binder: Arc<ConnectionTenantBinder>,
278        tenant_id: String,
279        connection_id: ConnectionId,
280    ) -> Self {
281        Self {
282            binder,
283            tenant_id,
284            connection_id,
285            active: true,
286        }
287    }
288
289    /// 获取租户 ID
290    pub fn tenant_id(&self) -> &str {
291        &self.tenant_id
292    }
293
294    /// 获取连接 ID
295    pub fn connection_id(&self) -> ConnectionId {
296        self.connection_id
297    }
298
299    /// 是否活跃
300    pub fn is_active(&self) -> bool {
301        self.active
302    }
303
304    /// 获取清理 SQL
305    pub fn clear_tenant_sql(&self) -> String {
306        self.binder.build_clear_tenant_sql()
307    }
308
309    /// 手动释放（提前清理）
310    pub fn release(&mut self) -> Result<(), TenantError> {
311        if !self.active {
312            return Ok(());
313        }
314        self.binder
315            .unbind_connection(&self.tenant_id, self.connection_id);
316        self.active = false;
317        Ok(())
318    }
319}
320
321impl Drop for TenantConnectionGuard {
322    fn drop(&mut self) {
323        if self.active {
324            self.binder
325                .unbind_connection(&self.tenant_id, self.connection_id);
326        }
327    }
328}
329
330#[cfg(test)]
331mod tests {
332    use super::*;
333    use crate::pool::{ConnectionFactory, PoolConfigBuilder};
334
335    struct MockFactory;
336
337    #[async_trait::async_trait]
338    impl ConnectionFactory for MockFactory {
339        async fn create(&self) -> Result<Box<dyn crate::pool::Connection>, crate::DbError> {
340            Err(crate::DbError::PoolError(
341                crate::error::PoolError::InvalidConfig("mock".to_string()),
342            ))
343        }
344    }
345
346    fn make_pool() -> Arc<Pool> {
347        let config = PoolConfigBuilder::new().max_size(1).build().unwrap();
348        let factory: Arc<dyn ConnectionFactory> = Arc::new(MockFactory);
349        Arc::new(Pool::new(config, factory).unwrap())
350    }
351
352    #[test]
353    fn test_config_default() {
354        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
355        assert_eq!(config.isolation, ConnectionLevelIsolation::SetTenantId);
356        assert_eq!(config.affinity_policy, ConnectionAffinityPolicy::Preferred);
357        assert_eq!(config.affinity_timeout_ms, 5_000);
358        assert_eq!(config.db_type, DbType::PostgreSQL);
359    }
360
361    #[test]
362    fn test_config_builder() {
363        let config = ConnectionLevelTenantConfig::new(DbType::MySQL)
364            .with_isolation(ConnectionLevelIsolation::ConnectionBinding)
365            .with_affinity_policy(ConnectionAffinityPolicy::Strict)
366            .with_affinity_timeout_ms(10_000);
367        assert_eq!(
368            config.isolation,
369            ConnectionLevelIsolation::ConnectionBinding
370        );
371        assert_eq!(config.affinity_policy, ConnectionAffinityPolicy::Strict);
372        assert_eq!(config.affinity_timeout_ms, 10_000);
373    }
374
375    #[test]
376    fn test_isolation_serde() {
377        let isolations = vec![
378            ConnectionLevelIsolation::SetTenantId,
379            ConnectionLevelIsolation::SchemaIsolation,
380            ConnectionLevelIsolation::ConnectionBinding,
381        ];
382        for i in &isolations {
383            let json = serde_json::to_string(i).unwrap();
384            let decoded: ConnectionLevelIsolation = serde_json::from_str(&json).unwrap();
385            assert_eq!(*i, decoded);
386        }
387    }
388
389    #[test]
390    fn test_affinity_policy_serde() {
391        let policies = vec![
392            ConnectionAffinityPolicy::Strict,
393            ConnectionAffinityPolicy::Preferred,
394            ConnectionAffinityPolicy::None,
395        ];
396        for p in &policies {
397            let json = serde_json::to_string(p).unwrap();
398            let decoded: ConnectionAffinityPolicy = serde_json::from_str(&json).unwrap();
399            assert_eq!(*p, decoded);
400        }
401    }
402
403    #[test]
404    fn test_supports_set_tenant_id() {
405        let pool = make_pool();
406        let pg_config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
407        let pg_binder = ConnectionTenantBinder::new(pool.clone(), pg_config);
408        assert!(pg_binder.supports_set_tenant_id());
409
410        let mysql_config = ConnectionLevelTenantConfig::new(DbType::MySQL);
411        let mysql_binder = ConnectionTenantBinder::new(pool.clone(), mysql_config);
412        assert!(mysql_binder.supports_set_tenant_id());
413
414        let sqlite_config = ConnectionLevelTenantConfig::new(DbType::Sqlite);
415        let sqlite_binder = ConnectionTenantBinder::new(pool.clone(), sqlite_config);
416        assert!(!sqlite_binder.supports_set_tenant_id());
417
418        let oracle_config = ConnectionLevelTenantConfig::new(DbType::Oracle);
419        let oracle_binder = ConnectionTenantBinder::new(pool, oracle_config);
420        assert!(!oracle_binder.supports_set_tenant_id());
421    }
422
423    #[test]
424    fn test_build_set_tenant_sql() {
425        let pool = make_pool();
426        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
427        let binder = ConnectionTenantBinder::new(pool, config);
428        let sql = binder.build_set_tenant_sql("tenant_123");
429        assert!(sql.contains("SET app.tenant_id"));
430        assert!(sql.contains("tenant_123"));
431    }
432
433    #[test]
434    fn test_build_clear_tenant_sql() {
435        let pool = make_pool();
436        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
437        let binder = ConnectionTenantBinder::new(pool, config);
438        let sql = binder.build_clear_tenant_sql();
439        assert!(sql.contains("NULL"));
440    }
441
442    #[test]
443    fn test_bind_and_find_connection() {
444        let pool = make_pool();
445        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
446        let binder = ConnectionTenantBinder::new(pool, config);
447
448        let conn_id1 = binder.bind_connection("tenant_1");
449        let conn_id2 = binder.bind_connection("tenant_1");
450        let conn_id3 = binder.bind_connection("tenant_2");
451
452        assert_ne!(conn_id1, conn_id2);
453        assert_ne!(conn_id1, conn_id3);
454
455        let tenant1_conns = binder.find_bound_connections("tenant_1");
456        assert_eq!(tenant1_conns.len(), 2);
457        assert!(tenant1_conns.contains(&conn_id1));
458        assert!(tenant1_conns.contains(&conn_id2));
459
460        let tenant2_conns = binder.find_bound_connections("tenant_2");
461        assert_eq!(tenant2_conns.len(), 1);
462        assert!(tenant2_conns.contains(&conn_id3));
463
464        assert_eq!(binder.binding_count("tenant_1"), 2);
465        assert_eq!(binder.binding_count("tenant_2"), 1);
466        assert_eq!(binder.binding_count("tenant_3"), 0);
467    }
468
469    #[test]
470    fn test_unbind_connection() {
471        let pool = make_pool();
472        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
473        let binder = ConnectionTenantBinder::new(pool, config);
474
475        let conn_id1 = binder.bind_connection("tenant_1");
476        let conn_id2 = binder.bind_connection("tenant_1");
477        assert_eq!(binder.binding_count("tenant_1"), 2);
478
479        binder.unbind_connection("tenant_1", conn_id1);
480        assert_eq!(binder.binding_count("tenant_1"), 1);
481
482        let conns = binder.find_bound_connections("tenant_1");
483        assert!(conns.contains(&conn_id2));
484        assert!(!conns.contains(&conn_id1));
485
486        binder.unbind_connection("tenant_1", conn_id2);
487        assert_eq!(binder.binding_count("tenant_1"), 0);
488    }
489
490    #[test]
491    fn test_validate_tenant_id() {
492        let pool = make_pool();
493        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
494        let binder = ConnectionTenantBinder::new(pool, config);
495
496        assert!(binder.validate_tenant_id("tenant_1").is_ok());
497        assert_eq!(
498            binder.validate_tenant_id("").unwrap_err(),
499            TenantError::EmptyTenantId
500        );
501    }
502
503    #[test]
504    fn test_resolve_isolation_pg() {
505        let pool = make_pool();
506        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
507        let binder = ConnectionTenantBinder::new(pool, config);
508        assert_eq!(
509            binder.resolve_isolation(),
510            ConnectionLevelIsolation::SetTenantId
511        );
512    }
513
514    #[test]
515    fn test_resolve_isolation_sqlite_fallback() {
516        let pool = make_pool();
517        let config = ConnectionLevelTenantConfig::new(DbType::Sqlite);
518        let binder = ConnectionTenantBinder::new(pool, config);
519        assert_eq!(
520            binder.resolve_isolation(),
521            ConnectionLevelIsolation::SchemaIsolation
522        );
523    }
524
525    #[test]
526    fn test_resolve_isolation_schema_no_fallback() {
527        let pool = make_pool();
528        let config = ConnectionLevelTenantConfig::new(DbType::Sqlite)
529            .with_isolation(ConnectionLevelIsolation::SchemaIsolation);
530        let binder = ConnectionTenantBinder::new(pool, config);
531        assert_eq!(
532            binder.resolve_isolation(),
533            ConnectionLevelIsolation::SchemaIsolation
534        );
535    }
536
537    #[test]
538    fn test_tenant_error_display() {
539        let err = TenantError::NoBoundConnection;
540        assert!(err.to_string().contains("no connection"));
541        let err = TenantError::EmptyTenantId;
542        assert!(err.to_string().contains("empty"));
543        let err = TenantError::UnsupportedDialect;
544        assert!(err.to_string().contains("unsupported"));
545    }
546
547    #[test]
548    fn test_guard_drop_unbinds() {
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        assert_eq!(binder.binding_count("tenant_1"), 1);
555
556        {
557            let _guard =
558                TenantConnectionGuard::new(binder.clone(), "tenant_1".to_string(), conn_id);
559            assert_eq!(binder.binding_count("tenant_1"), 1);
560        }
561
562        assert_eq!(binder.binding_count("tenant_1"), 0);
563    }
564
565    #[test]
566    fn test_guard_release() {
567        let pool = make_pool();
568        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
569        let binder = Arc::new(ConnectionTenantBinder::new(pool, config));
570
571        let conn_id = binder.bind_connection("tenant_1");
572        let mut guard = TenantConnectionGuard::new(binder, "tenant_1".to_string(), conn_id);
573        assert!(guard.is_active());
574        guard.release().unwrap();
575        assert!(!guard.is_active());
576    }
577
578    #[test]
579    fn test_guard_properties() {
580        let pool = make_pool();
581        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
582        let binder = Arc::new(ConnectionTenantBinder::new(pool, config));
583
584        let guard = TenantConnectionGuard::new(binder, "tenant_42".to_string(), 999);
585        assert_eq!(guard.tenant_id(), "tenant_42");
586        assert_eq!(guard.connection_id(), 999);
587        assert!(guard.is_active());
588        assert!(guard.clear_tenant_sql().contains("NULL"));
589    }
590
591    #[test]
592    fn test_all_bindings() {
593        let pool = make_pool();
594        let config = ConnectionLevelTenantConfig::new(DbType::PostgreSQL);
595        let binder = ConnectionTenantBinder::new(pool, config);
596
597        binder.bind_connection("tenant_1");
598        binder.bind_connection("tenant_1");
599        binder.bind_connection("tenant_2");
600
601        let all = binder.all_bindings();
602        assert_eq!(all.len(), 3);
603    }
604}
sz_orm_core/connection_tenant.rs

sz_orm_core/
connection_tenant.rs
