dbnexus/database/pool/
session.rs

1// Copyright (c) 2026 Kirky.X
2//
3// Licensed under the MIT License
4// See LICENSE file in the project root for full license information.
5
6//! Session 模块
7//!
8//! 提供数据库会话管理，包括事务、权限检查和读写分离
9
10use std::sync::Arc;
11use std::time::{Duration, Instant};
12
13#[cfg(feature = "permission")]
14use crate::access::permission::{PermissionAction, PermissionContext};
15#[cfg(feature = "sql-parser")]
16use crate::access::security::{DdlGuard, DdlValidationResult};
17#[cfg(all(feature = "sql-parser", feature = "permission"))]
18use crate::access::sql_parser::SqlParser;
19#[cfg(feature = "sql-parser")]
20use crate::access::sql_parser::is_ddl_operation;
21use crate::database::pool::db_pool::{DatabaseConnection, DbConnection, DbPool, DbPoolInner};
22use crate::foundation::error::{DbError, DbResult};
23#[cfg(feature = "metrics")]
24use crate::observability::metrics::MetricsCollector;
25use async_trait::async_trait;
26
27// 导入 Sea-ORM 的事务 trait 和连接 trait
28use sea_orm::{ConnectionTrait, DatabaseTransaction, ExecResult, TransactionTrait};
29use tokio::sync::Mutex;
30
31/// Session 内部可变状态
32///
33/// 使用 Mutex 包装需要内部可变性的字段，支持 `&self` 方法签名
34struct SessionState {
35    /// 事务对象（用于真实的事务管理）
36    ///
37    /// v0.3.0 性能优化：使用 `Arc<DatabaseTransaction>` 而非 `DatabaseTransaction`，
38    /// 因为 sea-orm 的 `DatabaseTransaction` 未实现 `Clone`，使用 `Arc` 包装后
39    /// 可在 `execute_raw` 中短锁 clone 后锁外执行 async DB 操作，避免持锁 await。
40    transaction: Option<Arc<DatabaseTransaction>>,
41
42    /// 最后写操作时间（用于读写分离）
43    last_write: Option<Instant>,
44}
45
46/// Session 结构
47pub struct Session {
48    /// 数据库连接（统一枚举：SeaORM 或 DuckDB）
49    connection: Option<DbConnection>,
50
51    /// 连接池（用于释放连接）
52    pool: Arc<DbPool>,
53
54    /// 连接池内部状态
55    pool_inner: Arc<DbPoolInner>,
56
57    /// 角色
58    role: String,
59
60    /// 权限上下文
61    #[cfg(feature = "permission")]
62    permission_ctx: PermissionContext,
63
64    /// 内部可变状态（事务和写操作时间）
65    state: Mutex<SessionState>,
66
67    /// 指标收集器（可选，用于 metrics 特性）
68    #[cfg(feature = "metrics")]
69    metrics_collector: Option<Arc<MetricsCollector>>,
70}
71
72impl Session {
73    /// 创建新的 Session
74    pub(crate) fn new(connection: DbConnection, pool: Arc<DbPool>, pool_inner: Arc<DbPoolInner>, role: String) -> Self {
75        #[cfg(feature = "permission")]
76        let permission_ctx = PermissionContext::new(role.clone(), pool_inner.policy_cache.clone());
77
78        #[cfg(feature = "metrics")]
79        let metrics = pool_inner.metrics_collector.clone();
80
81        Session {
82            connection: Some(connection),
83            pool,
84            pool_inner,
85            role,
86            #[cfg(feature = "permission")]
87            permission_ctx,
88            state: Mutex::new(SessionState {
89                transaction: None,
90                last_write: None,
91            }),
92            #[cfg(feature = "metrics")]
93            metrics_collector: metrics,
94        }
95    }
96
97    /// 获取角色
98    pub fn role(&self) -> &str {
99        &self.role
100    }
101
102    /// 获取权限上下文
103    #[cfg(feature = "permission")]
104    pub fn permission_ctx(&self) -> &PermissionContext {
105        &self.permission_ctx
106    }
107
108    /// 标记为写操作
109    pub async fn mark_write(&self) {
110        let mut state = self.state.lock().await;
111        state.last_write = Some(Instant::now());
112    }
113
114    /// 检查权限
115    #[cfg(feature = "permission")]
116    pub async fn check_permission(&self, table: &str, operation: &PermissionAction) -> Result<(), DbError> {
117        // Admin 角色绕过权限检查（拥有完全控制权）
118        if self.role == self.pool_inner.admin_role {
119            return Ok(());
120        }
121
122        if self.permission_ctx.check_table_access(table, operation).await {
123            Ok(())
124        } else {
125            Err(permission_denied(operation, table))
126        }
127    }
128
129    /// 是否在事务中
130    pub async fn is_in_transaction(&self) -> bool {
131        let state = self.state.lock().await;
132        state.transaction.is_some()
133    }
134
135    /// 开始事务
136    ///
137    /// v0.3.0 性能优化：短锁模式，避免持锁期间 async DB 调用。
138    /// 流程：短锁检查 → 锁外 begin → 短锁写入（含并发冲突处理）
139    pub async fn begin_transaction(&self) -> Result<(), DbError> {
140        // 短锁：检查是否已在事务中
141        {
142            let state = self.state.lock().await;
143            if state.transaction.is_some() {
144                return Err(DbError::Transaction("Already in transaction".to_string()));
145            }
146        }
147
148        // 锁外：执行 async DB 操作（避免持锁 await 阻塞其他调用者）
149        let conn = self.connection()?;
150        let transaction = conn
151            .begin()
152            .await
153            .map_err(|e| DbError::Transaction(format!("Failed to begin transaction: {}", e)))?;
154
155        // 短锁：写入 transaction（含并发冲突处理）
156        let mut state = self.state.lock().await;
157        if state.transaction.is_some() {
158            // 并发冲突：两次锁之间有其他调用已开始事务，回滚新创建的事务
159            let _ = transaction.rollback().await;
160            return Err(DbError::Transaction(
161                "Already in transaction (concurrent begin detected)".to_string(),
162            ));
163        }
164        state.transaction = Some(Arc::new(transaction));
165        Ok(())
166    }
167
168    /// 提交事务
169    ///
170    /// v0.3.0 性能优化：短锁模式，take transaction 后锁外执行 commit。
171    ///
172    /// # 并发安全
173    ///
174    /// 如果在 commit 时有其他查询正在执行（持有 transaction 的 Arc clone），
175    /// `Arc::try_unwrap` 会失败并返回错误。这是预期行为：用户不应在查询执行中提交事务。
176    pub async fn commit(&self) -> Result<(), DbError> {
177        // 短锁：take transaction
178        let transaction_arc = {
179            let mut state = self.state.lock().await;
180            state
181                .transaction
182                .take()
183                .ok_or_else(|| DbError::Transaction("No active transaction to commit".to_string()))?
184        };
185
186        // 锁外：try_unwrap 解包 Arc（如果有并发查询持有引用，会失败）
187        let transaction = Arc::try_unwrap(transaction_arc).map_err(|_| {
188            DbError::Transaction("Cannot commit: transaction is in use by a concurrent query".to_string())
189        })?;
190
191        // 锁外：执行 async commit（commit 消耗 self）
192        transaction
193            .commit()
194            .await
195            .map_err(|e| DbError::Transaction(e.to_string()))?;
196
197        // 短锁：清除 last_write
198        let mut state = self.state.lock().await;
199        state.last_write = None;
200        Ok(())
201    }
202
203    /// 回滚事务
204    ///
205    /// v0.3.0 性能优化：短锁模式，take transaction 后锁外执行 rollback。
206    ///
207    /// # 并发安全
208    ///
209    /// 如果在 rollback 时有其他查询正在执行（持有 transaction 的 Arc clone），
210    /// `Arc::try_unwrap` 会失败并返回错误。这是预期行为：用户不应在查询执行中回滚事务。
211    pub async fn rollback(&self) -> Result<(), DbError> {
212        // 短锁：take transaction
213        let transaction_arc = {
214            let mut state = self.state.lock().await;
215            if state.transaction.is_none() {
216                return Err(DbError::Transaction("Not in transaction".to_string()));
217            }
218            state
219                .transaction
220                .take()
221                .ok_or_else(|| DbError::Transaction("No active transaction to rollback".to_string()))?
222        };
223
224        // 锁外：try_unwrap 解包 Arc
225        let transaction = Arc::try_unwrap(transaction_arc).map_err(|_| {
226            DbError::Transaction("Cannot rollback: transaction is in use by a concurrent query".to_string())
227        })?;
228
229        // 锁外：执行 async rollback（rollback 消耗 self）
230        transaction
231            .rollback()
232            .await
233            .map_err(|e| DbError::Transaction(format!("Failed to rollback transaction: {}", e)))?;
234
235        Ok(())
236    }
237
238    /// 是否应该使用主库（基于读写分离配置）
239    pub async fn should_use_master(&self) -> bool {
240        let state = self.state.lock().await;
241        // 如果在事务中，必须使用主库
242        if state.transaction.is_some() {
243            return true;
244        }
245
246        // 如果配置了读写分离且有写操作，使用主库
247        state
248            .last_write
249            .map(|t| t.elapsed() < Duration::from_secs(5))
250            .unwrap_or(false)
251    }
252
253    /// 获取 SeaORM 连接引用（仅内部宏和测试使用）
254    ///
255    /// 用户应通过 Entity 的 CRUD 方法进行数据库操作，不应直接调用此方法。
256    /// 此方法从 `DbConnection` 枚举中提取 SeaORM 连接，若为 DuckDB 连接则返回错误。
257    ///
258    /// # 安全性
259    ///
260    /// 此方法确保连接在使用前是可用的。如果连接已被释放（不应发生），
261    /// 将返回错误。Session 的生命周期管理确保连接始终可用。
262    pub fn connection(&self) -> Result<&DatabaseConnection, DbError> {
263        self.connection
264            .as_ref()
265            .ok_or_else(|| DbError::Config("Connection not available - Session may have been invalidated".to_string()))?
266            .as_sea_orm()
267    }
268
269    /// 创建迁移执行器（仅内部使用）
270    ///
271    /// 用于迁移功能，将底层连接包装成 MigrationExecutor
272    #[allow(dead_code)]
273    #[cfg(feature = "migration")]
274    pub fn create_migration_executor(
275        &self,
276        db_type: crate::foundation::config::DatabaseType,
277    ) -> Result<super::MigrationExecutor, DbError> {
278        let conn = self.connection()?.clone();
279        Ok(super::MigrationExecutor::new(conn, db_type))
280    }
281
282    /// 执行原始 SQL（带权限检查）
283    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), fields(db.role = %self.role)))]
284    pub async fn execute_raw(&self, sql: &str) -> DbResult<ExecResult> {
285        #[cfg(feature = "sql-parser")]
286        {
287            // 检查是否为 DDL 操作
288            if is_ddl_operation(sql) {
289                return Err(DbError::Permission(
290                    "DDL operations are not allowed in this context".to_string(),
291                ));
292            }
293        }
294
295        #[cfg(not(feature = "sql-parser"))]
296        {
297            let _ = sql;
298            Err(DbError::Permission(
299                "execute_raw requires the sql-parser feature to be enabled".to_string(),
300            ))
301        }
302
303        #[cfg(feature = "sql-parser")]
304        {
305            #[cfg(all(feature = "sql-parser", feature = "permission"))]
306            {
307                // 解析 SQL 操作类型和表名（使用全局共享单例，避免重复创建 parser + 缓存）
308                let parser = SqlParser::shared().await;
309                match parser.parse_operation_async(sql).await {
310                    Ok(Some((table_name, action))) => {
311                        if table_name.is_empty() || is_invalid_table_name(&table_name) {
312                            return Err(DbError::Permission(
313                                "Failed to extract table name for permission checking".to_string(),
314                            ));
315                        }
316                        // 检查权限
317                        // Admin 角色绕过权限检查
318                        if self.role == self.pool_inner.admin_role {
319                            // admin 有完全权限，跳过检查
320                        } else if !self.permission_ctx.check_table_access(&table_name, &action).await {
321                            return Err(permission_denied(&action, &table_name));
322                        }
323                    }
324                    Ok(None) => {
325                        // 解析成功但是不支持的语句类型（DDL/DCL/Transaction）或没有表名的语句
326                        // 这些情况需要拒绝执行以确保安全
327                        return Err(DbError::Permission(
328                            "SQL statement requires a valid table name for permission checking".to_string(),
329                        ));
330                    }
331                    Err(_) => {
332                        // 解析失败，拒绝执行
333                        return Err(DbError::Permission(
334                            "Failed to parse SQL statement for permission checking".to_string(),
335                        ));
336                    }
337                }
338            }
339
340            // v0.3.0 性能优化：短锁 clone Arc<DatabaseTransaction>，锁外执行 async DB 调用
341            let tx_opt: Option<Arc<DatabaseTransaction>> = {
342                let state = self.state.lock().await;
343                state.transaction.clone()
344            };
345            if let Some(tx) = tx_opt {
346                return tx.execute_unprepared(sql).await.map_err(DbError::Connection);
347            }
348
349            let conn = self.connection()?;
350            conn.execute_unprepared(sql).await.map_err(DbError::Connection)
351        }
352    }
353
354    /// 执行 DDL 操作（允许创建表、删除表等操作）
355    ///
356    /// 此方法专门用于执行 DDL 操作，绕过常规的 DDL 检查。
357    /// 仅用于测试和迁移场景，生产环境应谨慎使用。
358    ///
359    /// # Arguments
360    ///
361    /// * `sql` - 要执行的 DDL SQL 语句
362    ///
363    /// # Returns
364    ///
365    /// 执行结果
366    ///
367    /// # Note
368    ///
369    /// 此方法只允许管理员角色执行，用于测试和迁移场景。
370    pub async fn execute_raw_ddl(&self, sql: &str) -> DbResult<ExecResult> {
371        // 检查角色白名单（只允许管理员角色执行 DDL）
372        if self.role != self.pool_inner.admin_role {
373            return Err(DbError::Permission(format!(
374                "DDL operations are only allowed for admin role. Current role: '{}', Admin role: '{}'",
375                self.role, self.pool_inner.admin_role
376            )));
377        }
378
379        // DDL 安全验证（基于 AST 解析，防止注入绕过）
380        #[cfg(feature = "sql-parser")]
381        {
382            let guard = DdlGuard::new();
383            match guard.validate(sql) {
384                Ok(DdlValidationResult::Allowed) => {
385                    // 通过验证，继续执行
386                }
387                Ok(DdlValidationResult::Forbidden(reason)) => {
388                    return Err(DbError::Permission(format!("DDL operation not allowed: {}", reason)));
389                }
390                Ok(DdlValidationResult::ParseError(error)) => {
391                    return Err(DbError::Config(format!("Failed to parse DDL SQL: {}", error)));
392                }
393                Err(error) => {
394                    return Err(DbError::Config(format!("DDL validation error: {}", error)));
395                }
396            }
397        }
398
399        // 执行 SQL
400        let conn = self.connection()?;
401        conn.execute_unprepared(sql).await.map_err(DbError::Connection)
402    }
403
404    /// 执行 DuckDB 查询（仅 DuckDB 连接可用）
405    ///
406    /// 当 Session 持有 DuckDB 连接时，通过此方法执行 SQL 查询并返回结果行。
407    /// 若持有 SeaORM 连接则返回错误。
408    ///
409    /// # 参数
410    ///
411    /// * `sql` - 要执行的 SQL 查询语句（SELECT）
412    ///
413    /// # 返回
414    ///
415    /// 查询结果行列表
416    #[cfg(feature = "duckdb")]
417    pub async fn execute_duckdb(&self, sql: &str) -> DbResult<Vec<crate::database::pool::DuckDbRow>> {
418        // 安全检查：与 execute_raw 一致的防御链（DDL 拦截 + SQL 注入检测 + 权限校验）
419        #[cfg(feature = "sql-parser")]
420        {
421            if is_ddl_operation(sql) {
422                return Err(DbError::Permission(
423                    "DDL operations are not allowed in DuckDB query context".to_string(),
424                ));
425            }
426        }
427
428        #[cfg(not(feature = "sql-parser"))]
429        {
430            let _ = sql;
431            Err(DbError::Permission(
432                "execute_duckdb requires the sql-parser feature to be enabled for security checks".to_string(),
433            ))
434        }
435
436        #[cfg(feature = "sql-parser")]
437        {
438            #[cfg(all(feature = "sql-parser", feature = "permission"))]
439            {
440                let parser = SqlParser::shared().await;
441                match parser.parse_operation_async(sql).await {
442                    Ok(Some((table_name, action))) => {
443                        if table_name.is_empty() || is_invalid_table_name(&table_name) {
444                            return Err(DbError::Permission(
445                                "Failed to extract table name for permission checking".to_string(),
446                            ));
447                        }
448                        if self.role != self.pool_inner.admin_role
449                            && !self.permission_ctx.check_table_access(&table_name, &action).await
450                        {
451                            return Err(permission_denied(&action, &table_name));
452                        }
453                    }
454                    Ok(None) => {
455                        // admin role 对无法解析的语句直接执行（对齐 execute 的 None 路径），
456                        // 支持 SELECT 1 / SELECT 1 AS health 等无表名健康检查查询；
457                        // 非 admin role 拒绝（安全默认：无法解析则无法做权限检查）。
458                        if self.role != self.pool_inner.admin_role {
459                            return Err(DbError::Permission(
460                                "SQL statement requires a valid table name for permission checking".to_string(),
461                            ));
462                        }
463                    }
464                    Err(_) => {
465                        return Err(DbError::Permission(
466                            "Failed to parse SQL statement for permission checking".to_string(),
467                        ));
468                    }
469                }
470            }
471
472            let conn = self
473                .connection
474                .as_ref()
475                .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
476            let duck_conn = conn.as_duckdb()?;
477            duck_conn.query(sql).await
478        }
479    }
480
481    /// 执行 DuckDB DDL/DML 语句（仅 DuckDB 连接可用）
482    ///
483    /// 当 Session 持有 DuckDB 连接时，通过此方法执行 CREATE/INSERT/UPDATE/DELETE 等语句。
484    /// 若持有 SeaORM 连接则返回错误。
485    ///
486    /// # 参数
487    ///
488    /// * `sql` - 要执行的 SQL 语句（DDL/DML）
489    ///
490    /// # 返回
491    ///
492    /// 受影响的行数信息
493    #[cfg(feature = "duckdb")]
494    pub async fn execute_duckdb_raw(&self, sql: &str) -> DbResult<crate::database::pool::DuckDbExecResult> {
495        // 安全检查：与 execute_raw_ddl 对齐 —— admin role 通过 DdlGuard 验证后允许 DDL，
496        // 非 admin role 拒绝 DDL。DuckDB 是分析型数据库，admin 需要能创建表/视图，
497        // 与 SeaORM 路径的 execute_raw_ddl 行为保持一致。
498        #[cfg(feature = "sql-parser")]
499        {
500            if is_ddl_operation(sql) {
501                if self.role == self.pool_inner.admin_role {
502                    // admin role 通过 DdlGuard AST 验证后直接执行，不再走 parse_operation 权限检查
503                    // （DDL 语句无法被 parse_operation_async 正确解析，会返回 Err）
504                    let guard = DdlGuard::new();
505                    match guard.validate(sql) {
506                        Ok(DdlValidationResult::Allowed) => {
507                            let conn = self
508                                .connection
509                                .as_ref()
510                                .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
511                            let duck_conn = conn.as_duckdb()?;
512                            return duck_conn.execute(sql).await;
513                        }
514                        Ok(DdlValidationResult::Forbidden(reason)) => {
515                            return Err(DbError::Permission(format!("DDL operation not allowed: {}", reason)));
516                        }
517                        Ok(DdlValidationResult::ParseError(error)) => {
518                            return Err(DbError::Config(format!("Failed to parse DDL SQL: {}", error)));
519                        }
520                        Err(error) => {
521                            return Err(DbError::Config(format!("DDL validation error: {}", error)));
522                        }
523                    }
524                } else {
525                    return Err(DbError::Permission(format!(
526                        "DDL operations are only allowed for admin role in DuckDB context. Current role: '{}', Admin role: '{}'",
527                        self.role, self.pool_inner.admin_role
528                    )));
529                }
530            }
531        }
532
533        #[cfg(not(feature = "sql-parser"))]
534        {
535            let _ = sql;
536            Err(DbError::Permission(
537                "execute_duckdb_raw requires the sql-parser feature to be enabled for security checks".to_string(),
538            ))
539        }
540
541        #[cfg(feature = "sql-parser")]
542        {
543            #[cfg(all(feature = "sql-parser", feature = "permission"))]
544            {
545                let parser = SqlParser::shared().await;
546                match parser.parse_operation_async(sql).await {
547                    Ok(Some((table_name, action))) => {
548                        if table_name.is_empty() || is_invalid_table_name(&table_name) {
549                            return Err(DbError::Permission(
550                                "Failed to extract table name for permission checking".to_string(),
551                            ));
552                        }
553                        if self.role != self.pool_inner.admin_role
554                            && !self.permission_ctx.check_table_access(&table_name, &action).await
555                        {
556                            return Err(permission_denied(&action, &table_name));
557                        }
558                    }
559                    Ok(None) => {
560                        // admin role 对无法解析的语句直接执行（对齐 execute 的 None 路径），
561                        // 非 admin role 拒绝（安全默认：无法解析则无法做权限检查）。
562                        if self.role != self.pool_inner.admin_role {
563                            return Err(DbError::Permission(
564                                "SQL statement requires a valid table name for permission checking".to_string(),
565                            ));
566                        }
567                    }
568                    Err(_) => {
569                        return Err(DbError::Permission(
570                            "Failed to parse SQL statement for permission checking".to_string(),
571                        ));
572                    }
573                }
574            }
575
576            let conn = self
577                .connection
578                .as_ref()
579                .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
580            let duck_conn = conn.as_duckdb()?;
581            duck_conn.execute(sql).await
582        }
583    }
584
585    /// 执行 SQL（带权限检查和操作类型）
586    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), fields(db.role = %self.role)))]
587    pub async fn execute(&self, sql: &str) -> DbResult<ExecResult> {
588        // DDL 检查（sql-parser 启用时）
589        #[cfg(feature = "sql-parser")]
590        check_ddl_operation(sql)?;
591
592        #[cfg(feature = "permission")]
593        {
594            let start = Instant::now();
595            // 解析 SQL 操作类型和表名
596            let parsed = parse_sql_for_permission(sql).await?;
597            match parsed {
598                Some((table_name, action)) => {
599                    // 表名有效性检查
600                    if table_name.is_empty() || is_invalid_table_name(&table_name) {
601                        return Err(DbError::Permission(
602                            "Failed to extract table name for permission checking".to_string(),
603                        ));
604                    }
605                    // 权限检查（含 admin 角色绕过）
606                    self.check_permission(&table_name, &action).await?;
607                    // 执行 SQL
608                    let result = self.execute_raw(sql).await?;
609                    // 记录指标并标记写操作
610                    self.record_metrics_and_mark_write(&action, start).await;
611                    Ok(result)
612                }
613                None => {
614                    // 解析失败或不支持的语句类型，直接执行
615                    // （仅 sql-parser 启用时可能出现 None）
616                    let result = self.execute_raw(sql).await?;
617                    Ok(result)
618                }
619            }
620        }
621
622        #[cfg(not(feature = "permission"))]
623        {
624            // 执行 SQL
625            let result = self.execute_raw(sql).await?;
626            Ok(result)
627        }
628    }
629
630    /// 执行 SQL 并指定操作类型
631    #[cfg(feature = "permission")]
632    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, operation), fields(db.role = %self.role)))]
633    pub async fn execute_with_operation(&self, sql: &str, operation: &PermissionAction) -> DbResult<ExecResult> {
634        let start = Instant::now();
635
636        #[cfg(feature = "sql-parser")]
637        {
638            // 检查是否为 DDL 操作
639            if is_ddl_operation(sql) {
640                return Err(DbError::Permission(
641                    "DDL operations are not allowed in this context".to_string(),
642                ));
643            }
644        }
645
646        // 提取表名
647        let table_name = extract_table_name(sql);
648
649        // 检查权限
650        #[cfg(feature = "permission")]
651        {
652            if !table_name.is_empty() && !self.permission_ctx.check_table_access(&table_name, operation).await {
653                return Err(permission_denied(operation, &table_name));
654            }
655        }
656
657        // 执行 SQL
658        let result = self.execute_raw(sql).await?;
659
660        // 记录指标
661        let duration = start.elapsed();
662        self.record_query_metrics(&format!("{:?}", operation), duration, true);
663
664        // 如果是写操作，标记
665        #[cfg(feature = "permission")]
666        {
667            if is_write_action(operation) {
668                self.mark_write().await;
669            }
670        }
671
672        Ok(result)
673    }
674
675    /// 批量执行 SQL
676    ///
677    /// # Arguments
678    ///
679    /// * `sqls` - 要执行的 SQL 语句列表
680    ///
681    /// # Returns
682    ///
683    /// 返回执行结果列表
684    pub async fn batch_execute(&self, sqls: Vec<&str>) -> DbResult<Vec<DbResult<ExecResult>>> {
685        let mut results = Vec::new();
686
687        for sql in sqls {
688            let result = self.execute(sql).await;
689            results.push(result);
690        }
691
692        Ok(results)
693    }
694
695    /// 批量执行（带事务）
696    ///
697    /// 所有操作在一个事务中执行，任一失败则全部回滚
698    ///
699    /// # Arguments
700    ///
701    /// * `sqls` - 要执行的 SQL 语句列表
702    ///
703    /// # Returns
704    ///
705    /// 返回执行结果列表，任一失败则返回错误
706    pub async fn batch_execute_in_transaction(&self, sqls: Vec<&str>) -> DbResult<Vec<ExecResult>> {
707        self.begin_transaction().await?;
708
709        let mut results = Vec::new();
710        let mut last_error = None;
711
712        for sql in sqls {
713            match self.execute_raw(sql).await {
714                Ok(result) => results.push(result),
715                Err(e) => {
716                    last_error = Some(e);
717                    break;
718                }
719            }
720        }
721
722        if let Some(error) = last_error {
723            self.rollback().await?;
724            Err(error)
725        } else {
726            self.commit().await?;
727            Ok(results)
728        }
729    }
730
731    /// 记录查询指标
732    #[cfg(all(feature = "metrics", feature = "permission"))]
733    fn record_query_metrics(&self, query_type: &str, duration: Duration, success: bool) {
734        if let Some(metrics) = &self.metrics_collector {
735            metrics.record_query(query_type, duration, success, None);
736        }
737    }
738
739    /// 记录查询指标（无 metrics 特性）
740    #[cfg(all(not(feature = "metrics"), feature = "permission"))]
741    fn record_query_metrics(&self, _query_type: &str, _duration: Duration, _success: bool) {
742        // No-op when metrics feature is disabled
743    }
744
745    /// 记录查询指标并标记写操作
746    ///
747    /// 统一 execute 流程中 metrics 记录与 mark_write 逻辑，
748    /// 避免在多个 cfg 分支中重复实现。
749    #[cfg(feature = "permission")]
750    async fn record_metrics_and_mark_write(&self, action: &PermissionAction, start: Instant) {
751        let duration = start.elapsed();
752        self.record_query_metrics(&format!("{:?}", action), duration, true);
753        if is_write_action(action) {
754            self.mark_write().await;
755        }
756    }
757
758    /// 检查表级权限
759    ///
760    /// 此方法为 ORM 操作提供权限检查，确保所有实体操作都经过权限验证
761    pub async fn check_table_permission(&self, _table_name: &str, _operation: &str) -> DbResult<()> {
762        #[cfg(feature = "permission")]
763        {
764            let action = match _operation {
765                "INSERT" => PermissionAction::Insert,
766                "SELECT" => PermissionAction::Select,
767                "UPDATE" => PermissionAction::Update,
768                "DELETE" => PermissionAction::Delete,
769                _ => return Err(DbError::Permission(format!("Unknown operation: {}", _operation))),
770            };
771
772            // Admin 角色绕过权限检查
773            if self.role == self.pool_inner.admin_role {
774                // admin 有完全权限，跳过检查
775            } else if !self.permission_ctx.check_table_access(_table_name, &action).await {
776                return Err(permission_denied(_operation, _table_name));
777            }
778        }
779        Ok(())
780    }
781
782    /// 记录指标
783    #[cfg(feature = "metrics")]
784    pub fn record_metric(&self, operation: &str, table_name: &str, success: bool) {
785        if let Some(metrics) = &self.metrics_collector {
786            // 使用表名的哈希值作为 bytes 参数
787            let bytes = Some(table_name.len() as u64);
788            metrics.record_query(operation, std::time::Duration::from_millis(0), success, bytes);
789        }
790    }
791}
792
793#[cfg(feature = "permission")]
794fn is_invalid_table_name(table_name: &str) -> bool {
795    let table_name = table_name.trim();
796    if table_name.is_empty() {
797        return true;
798    }
799
800    for part in table_name.split('.') {
801        let part = part.trim();
802        if part.is_empty() {
803            return true;
804        }
805
806        let unquoted = part
807            .strip_prefix('"')
808            .and_then(|s| s.strip_suffix('"'))
809            .or_else(|| part.strip_prefix('`').and_then(|s| s.strip_suffix('`')))
810            .or_else(|| part.strip_prefix('\'').and_then(|s| s.strip_suffix('\'')))
811            .unwrap_or(part)
812            .trim();
813
814        if unquoted.is_empty() {
815            return true;
816        }
817    }
818
819    false
820}
821
822/// 构造权限拒绝错误
823///
824/// 统一 "Permission denied for {action} on {table}" 错误消息格式，
825/// 避免在多处调用点重复 `DbError::Permission(format!(...))` 模板。
826#[cfg(feature = "permission")]
827fn permission_denied(action: &(impl std::fmt::Display + ?Sized), table: &(impl std::fmt::Display + ?Sized)) -> DbError {
828    DbError::Permission(format!("Permission denied for {} on {}", action, table))
829}
830
831/// 判断是否为写操作（Insert/Update/Delete）
832#[cfg(feature = "permission")]
833fn is_write_action(action: &PermissionAction) -> bool {
834    matches!(
835        action,
836        PermissionAction::Insert | PermissionAction::Update | PermissionAction::Delete
837    )
838}
839
840/// 检查 DDL 操作，如果 SQL 为 DDL 则返回错误
841///
842/// 统一 execute / execute_raw / execute_with_operation 中的 DDL 拒绝逻辑。
843#[cfg(feature = "sql-parser")]
844fn check_ddl_operation(sql: &str) -> DbResult<()> {
845    if is_ddl_operation(sql) {
846        return Err(DbError::Permission(
847            "DDL operations are not allowed in this context".to_string(),
848        ));
849    }
850    Ok(())
851}
852
853impl Drop for Session {
854    fn drop(&mut self) {
855        // 归还连接到池
856        if let Some(conn) = self.connection.take() {
857            self.pool.release_connection(conn);
858        }
859    }
860}
861
862/// 简化的表名提取（用于权限检查）
863#[cfg(feature = "permission")]
864fn extract_table_name(sql: &str) -> String {
865    // 这是一个简化的实现，实际应该使用 sqlparser
866    let sql_upper = sql.to_uppercase();
867
868    if sql_upper.contains("FROM ") {
869        if let Some(start) = sql_upper.find("FROM ") {
870            let rest = &sql[start + 5..];
871            if let Some(end) = rest.find(|c| [' ', ',', ';', '(', ')'].contains(&c)) {
872                return rest[..end].trim().to_string();
873            } else {
874                return rest.trim().to_string();
875            }
876        }
877    }
878
879    if sql_upper.contains("INTO ") {
880        if let Some(start) = sql_upper.find("INTO ") {
881            let rest = &sql[start + 5..];
882            if let Some(end) = rest.find(|c| [' ', '(', ';'].contains(&c)) {
883                return rest[..end].trim().to_string();
884            } else {
885                return rest.trim().to_string();
886            }
887        }
888    }
889
890    if sql_upper.contains("UPDATE ") {
891        if let Some(start) = sql_upper.find("UPDATE ") {
892            let rest = &sql[start + 7..];
893            if let Some(end) = rest.find(|c| [' ', ';'].contains(&c)) {
894                return rest[..end].trim().to_string();
895            } else {
896                return rest.trim().to_string();
897            }
898        }
899    }
900
901    String::new()
902}
903
904#[cfg(all(feature = "permission", not(feature = "sql-parser")))]
905fn parse_table_and_action(sql: &str) -> (String, PermissionAction) {
906    let table_name = extract_table_name(sql);
907    let sql_upper = sql.trim_start().to_uppercase();
908    let action = if sql_upper.starts_with("INSERT") {
909        PermissionAction::Insert
910    } else if sql_upper.starts_with("UPDATE") {
911        PermissionAction::Update
912    } else if sql_upper.starts_with("DELETE") {
913        PermissionAction::Delete
914    } else {
915        PermissionAction::Select
916    };
917
918    (table_name, action)
919}
920
921/// 解析 SQL 操作类型和表名用于权限检查
922///
923/// 统一 execute 流程中的 SQL 解析入口，消除 permission+sql-parser 与
924/// permission+无 sql-parser 两个 cfg 分支的重复结构：
925/// - sql-parser 启用：返回 None 表示不支持的语句或解析失败（execute 会跳过权限检查直接执行）
926/// - sql-parser 未启用：始终返回 Some（使用简化解析器 parse_table_and_action）
927#[cfg(feature = "permission")]
928async fn parse_sql_for_permission(sql: &str) -> DbResult<Option<(String, PermissionAction)>> {
929    #[cfg(feature = "sql-parser")]
930    {
931        let parser = SqlParser::shared().await;
932        match parser.parse_operation_async(sql).await {
933            Ok(Some((table, action))) => Ok(Some((table, action))),
934            Ok(None) => Ok(None),
935            Err(_) => Ok(None),
936        }
937    }
938    #[cfg(not(feature = "sql-parser"))]
939    {
940        Ok(Some(parse_table_and_action(sql)))
941    }
942}
943
944// 实现 DatabaseSession trait
945#[async_trait]
946impl super::DatabaseSession for Session {
947    async fn execute(&self, sql: &str) -> crate::DbResult<ExecResult> {
948        Ok(self.execute(sql).await?)
949    }
950
951    async fn execute_raw(&self, sql: &str) -> crate::DbResult<ExecResult> {
952        Ok(self.execute_raw(sql).await?)
953    }
954
955    async fn execute_raw_ddl(&self, sql: &str) -> crate::DbResult<ExecResult> {
956        Ok(self.execute_raw_ddl(sql).await?)
957    }
958
959    async fn begin_transaction(&self) -> crate::DbResult<()> {
960        Ok(self.begin_transaction().await?)
961    }
962
963    async fn commit(&self) -> crate::DbResult<()> {
964        Ok(self.commit().await?)
965    }
966
967    async fn rollback(&self) -> crate::DbResult<()> {
968        Ok(self.rollback().await?)
969    }
970
971    fn role(&self) -> &str {
972        self.role()
973    }
974
975    async fn is_in_transaction(&self) -> bool {
976        self.is_in_transaction().await
977    }
978}
dbnexus/database/pool/session.rs

dbnexus/database/pool/
session.rs
