dbnexus 0.1.3

// Copyright (c) 2026 Kirky.X
//
// Licensed under the MIT License
// See LICENSE file in the project root for full license information.

//! Session 模块
//!
//! 提供数据库会话管理，包括事务、权限检查和读写分离

use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::error::{DbError, DbResult};
#[cfg(feature = "metrics")]
use crate::metrics::MetricsCollector;
#[cfg(feature = "permission")]
use crate::permission::{PermissionAction, PermissionContext};
use crate::pool::db_pool::{DatabaseConnection, DbPool, DbPoolInner};
#[cfg(feature = "sql-parser")]
use crate::sql_parser::{SqlParser, is_ddl_operation};
use async_trait::async_trait;

#[cfg(not(any(feature = "permission", feature = "sql-parser")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PermissionAction {
    Select,
}

// 导入 Sea-ORM 的事务 trait 和连接 trait
use sea_orm::{ConnectionTrait, DatabaseTransaction, ExecResult, TransactionTrait};

/// Session 结构
pub struct Session {
    /// 数据库连接
    connection: Option<DatabaseConnection>,

    /// 连接池（用于释放连接）
    pool: Arc<DbPool>,

    /// 连接池内部状态
    pool_inner: Arc<DbPoolInner>,

    /// 角色
    role: String,

    /// 最后写操作时间（用于读写分离）
    last_write: Option<Instant>,

    /// 权限上下文
    #[cfg(feature = "permission")]
    permission_ctx: PermissionContext,

    /// 事务对象（用于真实的事务管理）
    transaction: Option<DatabaseTransaction>,

    /// 事务开始时间（用于超时检查）
    transaction_start_time: Option<Instant>,

    /// 事务超时时间
    transaction_timeout: Duration,

    /// 指标收集器（可选，用于 metrics 特性）
    #[cfg(feature = "metrics")]
    metrics_collector: Option<Arc<MetricsCollector>>,
}

impl Session {
    /// 创建新的 Session
    pub(crate) fn new(
        connection: DatabaseConnection,
        pool: Arc<DbPool>,
        pool_inner: Arc<DbPoolInner>,
        role: String,
    ) -> Self {
        #[cfg(feature = "permission")]
        let permission_ctx = {
            // 从 pool 获取 permission provider
            let provider = pool_inner.take_permission_provider();
            crate::permission::PermissionContext::with_provider_or_default(
                role.clone(),
                pool_inner.policy_cache.clone(),
                provider,
            )
        };

        #[cfg(feature = "metrics")]
        let metrics = pool_inner.metrics_collector.clone();

        // 默认事务超时时间为 5 分钟
        let transaction_timeout = Duration::from_secs(300);

        Session {
            connection: Some(connection),
            pool,
            pool_inner,
            role,
            last_write: None,
            #[cfg(feature = "permission")]
            permission_ctx,
            transaction: None,
            transaction_start_time: None,
            transaction_timeout,
            #[cfg(feature = "metrics")]
            metrics_collector: metrics,
        }
    }

    /// 获取角色
    pub fn role(&self) -> &str {
        &self.role
    }

    /// 获取权限上下文
    #[cfg(feature = "permission")]
    pub fn permission_ctx(&self) -> &PermissionContext {
        &self.permission_ctx
    }

    /// 标记为写操作
    pub fn mark_write(&mut self) {
        self.last_write = Some(Instant::now());
    }

    /// 检查权限
    #[cfg(feature = "permission")]
    pub async fn check_permission(&self, table: &str, operation: &PermissionAction) -> Result<(), DbError> {
        if self.permission_ctx.check_table_access(table, operation).await {
            Ok(())
        } else {
            Err(DbError::Permission(format!(
                "Permission denied for {} on {}",
                operation, table
            )))
        }
    }

    /// 是否在事务中
    pub fn is_in_transaction(&self) -> bool {
        self.transaction.is_some()
    }

    /// 开始事务
    pub async fn begin_transaction(&mut self) -> Result<(), DbError> {
        if self.is_in_transaction() {
            return Err(DbError::Transaction("Already in transaction".to_string()));
        }

        let start_time = Instant::now();

        let conn = self
            .connection
            .as_ref()
            .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;

        let transaction = conn.begin().await.map_err(|e| {
            self.record_transaction_metrics("begin", start_time.elapsed(), false);
            DbError::Transaction(format!("Failed to begin transaction: {}", e))
        })?;

        self.transaction = Some(transaction);
        self.transaction_start_time = Some(Instant::now());

        self.record_transaction_metrics("begin", start_time.elapsed(), true);
        Ok(())
    }

    /// Short form of begin_transaction()
    pub async fn begin(&mut self) -> DbResult<()> {
        self.begin_transaction().await
    }

    /// 提交事务
    pub async fn commit(&mut self) -> Result<(), DbError> {
        let start_time = Instant::now();

        // 检查事务是否超时
        if let Some(start_time) = self.transaction_start_time {
            let elapsed = start_time.elapsed();
            if elapsed > self.transaction_timeout {
                tracing::warn!(
                    "Transaction timeout: elapsed {:?} exceeds timeout {:?}",
                    elapsed,
                    self.transaction_timeout
                );
                // 注意：即使超时，我们仍然尝试提交事务，让数据库决定是否接受
            }
        }

        let transaction = self.transaction.take().ok_or_else(|| {
            self.record_transaction_metrics("commit", start_time.elapsed(), false);
            DbError::Transaction("No active transaction to commit".to_string())
        })?;

        transaction.commit().await.map_err(|e| {
            self.record_transaction_metrics("commit", start_time.elapsed(), false);
            DbError::Transaction(e.to_string())
        })?;

        // 只有在提交成功后才清理事务状态
        self.transaction_start_time = None;
        self.last_write = None;

        self.record_transaction_metrics("commit", start_time.elapsed(), true);
        Ok(())
    }

    /// 回滚事务
    pub async fn rollback(&mut self) -> Result<(), DbError> {
        let start_time = Instant::now();

        if !self.is_in_transaction() {
            self.record_transaction_metrics("rollback", start_time.elapsed(), false);
            return Err(DbError::Transaction("Not in transaction".to_string()));
        }

        // 检查事务是否超时
        if let Some(start_time) = self.transaction_start_time {
            let elapsed = start_time.elapsed();
            if elapsed > self.transaction_timeout {
                tracing::warn!(
                    "Transaction timeout during rollback: elapsed {:?} exceeds timeout {:?}",
                    elapsed,
                    self.transaction_timeout
                );
            }
        }

        let transaction = self.transaction.take().ok_or_else(|| {
            self.record_transaction_metrics("rollback", start_time.elapsed(), false);
            DbError::Transaction("No active transaction to rollback".to_string())
        })?;

        transaction.rollback().await.map_err(|e| {
            self.record_transaction_metrics("rollback", start_time.elapsed(), false);
            DbError::Transaction(format!("Failed to rollback transaction: {}", e))
        })?;

        // 只有在回滚成功后才清理事务状态
        self.transaction_start_time = None;

        self.record_transaction_metrics("rollback", start_time.elapsed(), true);
        Ok(())
    }

    /// 是否应该使用主库（基于读写分离配置）
    pub fn should_use_master(&self) -> bool {
        // 如果在事务中，必须使用主库
        if self.is_in_transaction() {
            return true;
        }

        // 如果配置了读写分离且有写操作，使用主库
        self.last_write
            .map(|t| t.elapsed() < Duration::from_secs(5))
            .unwrap_or(false)
    }

    /// 获取连接引用
    pub fn connection(&mut self) -> Result<&mut DatabaseConnection, DbError> {
        self.connection
            .as_mut()
            .ok_or_else(|| DbError::Config("Connection not available".to_string()))
    }

    /// 执行原始 SQL（带权限检查）
    pub async fn execute_raw(&self, sql: &str) -> DbResult<ExecResult> {
        #[cfg(feature = "sql-parser")]
        {
            // 检查是否为 DDL 操作
            if is_ddl_operation(sql) {
                return Err(DbError::Permission(
                    "DDL operations are not allowed in this context".to_string(),
                ));
            }
        }

        #[cfg(not(feature = "sql-parser"))]
        {
            // 当 sql-parser 特性未启用时，执行基本的 SQL 注入防护检查
            let sql_upper = sql.trim().to_uppercase();

            // 检测多语句执行（SQL 注入常用技术）
            if let Some(semi_colon_pos) = sql_upper.find(';') {
                // 检查分号后是否有非空内容
                let after_semicolon = sql_upper[semi_colon_pos + 1..].trim();
                if !after_semicolon.is_empty() {
                    tracing::warn!(
                        "Rejected SQL with multiple statements (potential SQL injection): {}",
                        sql
                    );
                    return Err(DbError::Permission(
                        "Multiple SQL statements not allowed in this context".to_string(),
                    ));
                }
            }

            // 检测常见的 SQL 注入模式
            let dangerous_patterns = [
                // 注释注入
                ("--", "Line comment"),
                ("/*", "Block comment"),
                ("*/", "Block comment end"),
                // UNION 注入
                ("UNION ALL", "UNION injection"),
                ("UNION SELECT", "UNION injection"),
                ("UNION(", "UNION injection"),
                // 数据操作危险操作
                ("DROP DATABASE", "DROP DATABASE"),
                ("TRUNCATE TABLE", "TRUNCATE TABLE"),
                ("DELETE FROM", "DELETE injection"),
                ("DELETE(", "DELETE injection"),
                // 存储过程注入
                ("EXEC(", "EXEC injection"),
                ("EXECUTE(", "EXECUTE injection"),
                (" xp_", "SQL Server extended stored procedure"),
                (" sp_", "SQL Server stored procedure"),
                (" xp_cmdshell", "SQL Server cmdshell"),
                // 系统表访问
                ("INFORMATION_SCHEMA", "System table access"),
                ("SYSOBJECTS", "System table access (SQL Server)"),
                ("SYSCOLUMNS", "System table access"),
                // 十六进制编码尝试
                ("0x", "Hex-encoded string"),
                // 字符串拼接
                ("||", "String concatenation"),
                ("CONCAT(", "String concatenation function"),
                // 时序攻击
                ("BENCHMARK(", "Timing attack"),
                ("SLEEP(", "Timing attack"),
                ("PG_SLEEP", "Timing attack"),
                // 条件注入
                ("' OR '1'='1", "Classic SQL injection"),
                ("' OR 1=1", "Numeric SQL injection"),
                (" OR 1=1", "Numeric SQL injection"),
                // 负载注入
                ("<script>", "Script injection"),
                ("javascript:", "JavaScript injection"),
                ("VARCHAR", "Type conversion injection"),
            ];

            for (pattern, description) in &dangerous_patterns {
                if sql_upper.contains(pattern) {
                    tracing::warn!(
                        "Rejected SQL containing dangerous pattern '{}' ({})",
                        pattern,
                        description
                    );
                    return Err(DbError::Permission(format!(
                        "SQL statement contains forbidden pattern: {} ({})",
                        pattern, description
                    )));
                }
            }

            // 检测可疑的字符串逃逸模式
            let escape_patterns = [("''", "Escaped single quote"), ("\\\\", "Double backslash")];

            for (pattern, _) in &escape_patterns {
                let count = sql_upper.matches(pattern).count();
                if count > 10 {
                    // 过多的逃逸字符可能表示注入尝试
                    tracing::warn!("Suspicious escape pattern count: {} ({} occurrences)", pattern, count);
                    return Err(DbError::Permission(
                        "SQL statement contains suspicious escape patterns".to_string(),
                    ));
                }
            }
        }

        #[cfg(all(feature = "sql-parser", feature = "permission"))]
        {
            // 解析 SQL 操作类型和表名
            let parser = SqlParser::new();
            if let Some((table_name, action)) = parser.parse_operation(sql) {
                if table_name.is_empty() || is_invalid_table_name(&table_name) {
                    return Err(DbError::Permission(
                        "Failed to extract table name for permission checking".to_string(),
                    ));
                }
                // 检查权限
                if !self.permission_ctx.check_table_access(&table_name, &action).await {
                    return Err(DbError::Permission(format!(
                        "Permission denied for {} on {}",
                        action, table_name
                    )));
                }
            } else {
                // 解析失败，拒绝执行
                return Err(DbError::Permission(
                    "Failed to parse SQL statement for permission checking".to_string(),
                ));
            }
        }

        #[cfg(all(feature = "permission", not(feature = "sql-parser")))]
        {
            let (table_name, action) = parse_table_and_action(sql);
            if table_name.is_empty() {
                return Err(DbError::Permission(
                    "Failed to extract table name for permission checking".to_string(),
                ));
            }

            if !self.permission_ctx.check_table_access(&table_name, &action).await {
                return Err(DbError::Permission(format!(
                    "Permission denied for {} on {}",
                    action, table_name
                )));
            }
        }

        if let Some(tx) = self.transaction.as_ref() {
            return tx.execute_unprepared(sql).await.map_err(DbError::Connection);
        }

        let conn = self
            .connection
            .as_ref()
            .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;

        conn.execute_unprepared(sql).await.map_err(DbError::Connection)
    }

    /// 执行 DDL 操作（允许创建表、删除表等操作）
    ///
    /// 此方法专门用于执行 DDL 操作，只允许管理员角色执行，
    /// 并使用 SQL 解析器验证操作类型。
    ///
    /// # Arguments
    ///
    /// * `sql` - 要执行的 DDL SQL 语句
    ///
    /// # Returns
    ///
    /// 执行结果
    ///
    /// # Errors
    ///
    /// - 如果不是管理员角色，返回权限错误
    /// - 如果 SQL 解析失败，返回解析错误
    /// - 如果包含危险操作，返回权限错误
    pub async fn execute_raw_ddl(&self, sql: &str) -> DbResult<ExecResult> {
        // 检查角色白名单（只允许管理员角色执行 DDL）
        if self.role != self.pool_inner.admin_role {
            return Err(DbError::Permission(format!(
                "DDL operations are only allowed for admin role. Current role: '{}', Admin role: '{}'",
                self.role, self.pool_inner.admin_role
            )));
        }

        #[cfg(feature = "sql-parser")]
        {
            // 使用 SQL 解析器验证操作
            let parser = SqlParser::new();
            let parsed = parser
                .parse_single(sql)
                .map_err(|e| DbError::Permission(format!("DDL operation validation failed: {}", e)))?;

            // 检查是否为允许的 DDL 操作
            match parsed.operation_type {
                crate::sql_parser::SqlOperationType::Ddl => {
                    // 允许标准的 DDL 操作
                    tracing::info!("Executing DDL operation on table: {:?}", parsed.table_name);
                }
                crate::sql_parser::SqlOperationType::Dcl => {
                    return Err(DbError::Permission(
                        "DCL operations (GRANT/REVOKE) are not allowed".to_string(),
                    ));
                }
                crate::sql_parser::SqlOperationType::Transaction => {
                    return Err(DbError::Permission(
                        "Transaction control statements (BEGIN/COMMIT/ROLLBACK) should use Session methods".to_string(),
                    ));
                }
                crate::sql_parser::SqlOperationType::Select => {
                    // SELECT 查询不通过此方法执行
                    return Err(DbError::Permission(
                        "SELECT queries should use execute() method, not execute_raw_ddl()".to_string(),
                    ));
                }
                crate::sql_parser::SqlOperationType::Insert
                | crate::sql_parser::SqlOperationType::Update
                | crate::sql_parser::SqlOperationType::Delete => {
                    return Err(DbError::Permission(
                        "DML operations (INSERT/UPDATE/DELETE) should use execute() method, not execute_raw_ddl()"
                            .to_string(),
                    ));
                }
                _ => {
                    return Err(DbError::Permission("Unknown SQL operation type".to_string()));
                }
            }

            // 额外安全检查：禁止危险操作
            let sql_upper = sql.trim().to_uppercase();
            let forbidden_patterns = ["DROP DATABASE", "DROP SCHEMA", "DROP ALL"];

            for pattern in &forbidden_patterns {
                if sql_upper.contains(pattern) {
                    return Err(DbError::Permission(format!(
                        "DDL operation not allowed: contains forbidden pattern '{}'",
                        pattern
                    )));
                }
            }
        }

        #[cfg(not(feature = "sql-parser"))]
        {
            // 没有 sql-parser 特性时的基本验证
            let sql_upper = sql.trim().to_uppercase();

            // 禁止的危险操作
            let forbidden_patterns = ["DROP DATABASE", "TRUNCATE TABLE", "DROP ALL", "DELETE FROM"];

            for pattern in &forbidden_patterns {
                if sql_upper.contains(pattern) {
                    return Err(DbError::Permission(format!(
                        "DDL operation not allowed: contains forbidden pattern '{}'",
                        pattern
                    )));
                }
            }

            // 只允许特定的 DDL 操作
            let allowed_prefixes = [
                "CREATE TABLE",
                "ALTER TABLE",
                "DROP TABLE",
                "CREATE INDEX",
                "DROP INDEX",
                "CREATE VIEW",
                "DROP VIEW",
            ];

            let is_allowed = allowed_prefixes.iter().any(|prefix| sql_upper.starts_with(prefix));

            if !is_allowed {
                return Err(DbError::Permission(format!(
                    "DDL operation not allowed: {}. Allowed operations: CREATE TABLE, ALTER TABLE, DROP TABLE, CREATE INDEX, DROP INDEX, CREATE VIEW, DROP VIEW",
                    sql_upper.split_whitespace().next().unwrap_or("UNKNOWN")
                )));
            }
        }

        // 执行 SQL
        let conn = self
            .connection
            .as_ref()
            .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;

        conn.execute_unprepared(sql).await.map_err(DbError::Connection)
    }

    /// 执行参数化查询（带权限检查）
    ///
    /// 使用参数化查询防止 SQL 注入攻击
    ///
    /// # Arguments
    ///
    /// * `sql` - 要执行的 SQL 语句，使用 `?` 或 `$1` 等占位符
    /// * `params` - 参数值列表
    ///
    /// # Returns
    ///
    /// 执行结果
    ///
    /// # Errors
    ///
    /// - 如果权限检查失败，返回权限错误
    /// - 如果 SQL 解析失败，返回解析错误
    ///
    /// # Example
    ///
    /// ```ignore
    /// let params = vec![1i32.into(), "Alice".into()];
    /// session.execute_paramized("SELECT * FROM users WHERE id = ? AND name = ?", params).await?;
    /// ```
    #[cfg(feature = "sql-parser")]
    pub async fn execute_paramized(&self, sql: &str, params: Vec<sea_orm::Value>) -> DbResult<ExecResult> {
        // 使用 SQL 解析器验证操作类型
        let parser = SqlParser::new();
        let parsed = parser.parse_single(sql).map_err(|e| {
            DbError::Permission(format!("Failed to parse SQL statement for parameterized query: {}", e))
        })?;

        // 检查是否为 DDL 操作
        if matches!(
            parsed.operation_type,
            crate::sql_parser::SqlOperationType::Ddl | crate::sql_parser::SqlOperationType::Dcl
        ) {
            return Err(DbError::Permission(
                "DDL and DCL operations should use execute_raw_ddl() or execute_raw() methods, not execute_paramized()"
                    .to_string(),
            ));
        }

        #[cfg(feature = "permission")]
        let (table_name, action) = {
            let action = match parsed.operation_type {
                crate::sql_parser::SqlOperationType::Select => PermissionAction::Select,
                crate::sql_parser::SqlOperationType::Insert => PermissionAction::Insert,
                crate::sql_parser::SqlOperationType::Update => PermissionAction::Update,
                crate::sql_parser::SqlOperationType::Delete => PermissionAction::Delete,
                _ => PermissionAction::Select,
            };

            (parsed.table_name, action)
        };

        #[cfg(feature = "permission")]
        {
            if let Some(table_name) = &table_name {
                if table_name.is_empty() || is_invalid_table_name(table_name) {
                    return Err(DbError::Permission(
                        "Failed to extract table name for permission checking".to_string(),
                    ));
                }

                if !self.permission_ctx.check_table_access(table_name, &action).await {
                    return Err(DbError::Permission(format!(
                        "Permission denied for {} on {}",
                        action, table_name
                    )));
                }
            }
        }

        // 执行参数化查询
        let conn = self
            .connection
            .as_ref()
            .ok_or_else(|| DbError::Config("Connection not available".to_string()))?;

        // 将参数值转换为字符串并替换占位符
        let mut sql_with_params = sql.to_string();
        for param in params {
            let param_str = match param {
                sea_orm::Value::String(Some(s)) => format!("'{}'", s.replace('\'', "''")),
                sea_orm::Value::String(None) => "NULL".to_string(),
                sea_orm::Value::Int(Some(i)) => i.to_string(),
                sea_orm::Value::Int(None) => "NULL".to_string(),
                sea_orm::Value::BigInt(Some(i)) => i.to_string(),
                sea_orm::Value::BigInt(None) => "NULL".to_string(),
                sea_orm::Value::TinyInt(Some(i)) => i.to_string(),
                sea_orm::Value::TinyInt(None) => "NULL".to_string(),
                sea_orm::Value::SmallInt(Some(i)) => i.to_string(),
                sea_orm::Value::SmallInt(None) => "NULL".to_string(),
                sea_orm::Value::Float(Some(f)) => f.to_string(),
                sea_orm::Value::Float(None) => "NULL".to_string(),
                sea_orm::Value::Double(Some(f)) => f.to_string(),
                sea_orm::Value::Double(None) => "NULL".to_string(),
                sea_orm::Value::Bool(Some(b)) => {
                    if b {
                        "1".to_string()
                    } else {
                        "0".to_string()
                    }
                }
                sea_orm::Value::Bool(None) => "NULL".to_string(),
                sea_orm::Value::Char(Some(c)) => format!("'{}'", c),
                sea_orm::Value::Char(None) => "NULL".to_string(),
                sea_orm::Value::Bytes(Some(b)) => {
                    // 简单的十六进制编码
                    let hex_str: String = b.iter().map(|byte| format!("{:02x}", byte)).collect();
                    format!("X'{}'", hex_str)
                }
                sea_orm::Value::Bytes(None) => "NULL".to_string(),
                _ => {
                    return Err(DbError::Config(format!(
                        "Unsupported parameter type in parameterized query: {:?}",
                        param
                    )));
                }
            };

            // 替换第一个 ? 占位符
            if let Some(pos) = sql_with_params.find('?') {
                sql_with_params.replace_range(pos..pos + 1, &param_str);
            } else {
                break;
            }
        }

        // 如果在事务中，使用事务执行
        if let Some(tx) = self.transaction.as_ref() {
            return tx
                .execute_unprepared(&sql_with_params)
                .await
                .map_err(DbError::Connection);
        }

        // 否则使用连接执行
        conn.execute_unprepared(&sql_with_params)
            .await
            .map_err(DbError::Connection)
    }

    /// 执行 SQL（带权限检查和操作类型）
    pub async fn execute(&mut self, sql: &str) -> DbResult<ExecResult> {
        let start = Instant::now();

        #[cfg(feature = "sql-parser")]
        {
            // 检查是否为 DDL 操作
            if is_ddl_operation(sql) {
                return Err(DbError::Permission(
                    "DDL operations are not allowed in this context".to_string(),
                ));
            }
        }

        #[cfg(all(feature = "permission", feature = "sql-parser"))]
        let (table_name, action) = {
            let parser = SqlParser::new();
            parser.parse_operation(sql).ok_or_else(|| {
                DbError::Permission("Failed to parse SQL statement for permission checking".to_string())
            })?
        };

        #[cfg(all(feature = "permission", not(feature = "sql-parser")))]
        let (table_name, action) = parse_table_and_action(sql);

        #[cfg(all(not(feature = "permission"), feature = "sql-parser"))]
        let action = crate::sql_parser::PermissionAction::Select;

        #[cfg(not(any(feature = "permission", feature = "sql-parser")))]
        let action = PermissionAction::Select;

        #[cfg(feature = "permission")]
        {
            if table_name.is_empty() || is_invalid_table_name(&table_name) {
                return Err(DbError::Permission(
                    "Failed to extract table name for permission checking".to_string(),
                ));
            }

            if !self.permission_ctx.check_table_access(&table_name, &action).await {
                return Err(DbError::Permission(format!(
                    "Permission denied for {} on {}",
                    action, table_name
                )));
            }
        }

        // 执行 SQL
        let result = self.execute_raw(sql).await?;

        // 记录指标
        let duration = start.elapsed();
        self.record_query_metrics(&format!("{:?}", action), duration, true);

        // 如果是写操作，标记
        #[cfg(feature = "permission")]
        {
            if matches!(
                action,
                PermissionAction::Insert | PermissionAction::Update | PermissionAction::Delete
            ) {
                self.mark_write();
            }
        }

        Ok(result)
    }

    /// 执行 SQL 并指定操作类型
    #[cfg(feature = "permission")]
    pub async fn execute_with_operation(&mut self, sql: &str, operation: &PermissionAction) -> DbResult<ExecResult> {
        let start = Instant::now();

        #[cfg(feature = "sql-parser")]
        {
            // 检查是否为 DDL 操作
            if is_ddl_operation(sql) {
                return Err(DbError::Permission(
                    "DDL operations are not allowed in this context".to_string(),
                ));
            }
        }

        // 提取表名
        let table_name = extract_table_name(sql);

        // 检查权限
        #[cfg(feature = "permission")]
        {
            if !table_name.is_empty() && !self.permission_ctx.check_table_access(&table_name, operation).await {
                return Err(DbError::Permission(format!(
                    "Permission denied for {} on {}",
                    operation, table_name
                )));
            }
        }

        // 执行 SQL
        let result = self.execute_raw(sql).await?;

        // 记录指标
        let duration = start.elapsed();
        self.record_query_metrics(&format!("{:?}", operation), duration, true);

        // 如果是写操作，标记
        #[cfg(feature = "permission")]
        {
            if matches!(
                operation,
                PermissionAction::Insert | PermissionAction::Update | PermissionAction::Delete
            ) {
                self.mark_write();
            }
        }

        Ok(result)
    }

    /// 批量执行 SQL
    ///
    /// # Arguments
    ///
    /// * `sqls` - 要执行的 SQL 语句列表
    ///
    /// # Returns
    ///
    /// 返回执行结果列表
    pub async fn batch_execute(&mut self, sqls: Vec<&str>) -> DbResult<Vec<DbResult<ExecResult>>> {
        let mut results = Vec::new();

        for sql in sqls {
            let result = self.execute(sql).await;
            results.push(result);
        }

        Ok(results)
    }

    /// 批量执行（带事务）
    ///
    /// 所有操作在一个事务中执行，任一失败则全部回滚
    ///
    /// # Arguments
    ///
    /// * `sqls` - 要执行的 SQL 语句列表
    ///
    /// # Returns
    ///
    /// 返回执行结果列表，任一失败则返回错误
    pub async fn batch_execute_in_transaction(&mut self, sqls: Vec<&str>) -> DbResult<Vec<ExecResult>> {
        self.begin_transaction().await?;

        let mut results = Vec::new();
        let mut last_error = None;

        for sql in sqls {
            match self.execute_raw(sql).await {
                Ok(result) => results.push(result),
                Err(e) => {
                    last_error = Some(e);
                    break;
                }
            }
        }

        if let Some(error) = last_error {
            // 尝试回滚事务，如果回滚失败，将两个错误信息合并
            if let Err(rollback_error) = self.rollback().await {
                return Err(DbError::Transaction(format!(
                    "Batch execution failed: {}. Rollback also failed: {}",
                    error, rollback_error
                )));
            }
            Err(error)
        } else {
            self.commit().await?;
            Ok(results)
        }
    }

    /// 记录查询指标
    #[cfg(feature = "metrics")]
    fn record_query_metrics(&self, query_type: &str, duration: Duration, success: bool) {
        if let Some(metrics) = &self.metrics_collector {
            metrics.record_query(query_type, duration, success, None);
        }
    }

    /// 记录查询指标（无 metrics 特性）
    #[cfg(not(feature = "metrics"))]
    fn record_query_metrics(&self, _query_type: &str, _duration: Duration, _success: bool) {
        // No-op when metrics feature is disabled
    }

    /// 记录连接错误
    #[cfg(feature = "metrics")]
    fn record_connection_error(&self) {
        if let Some(metrics) = &self.metrics_collector {
            metrics.record_connection_error();
        }
    }

    /// 记录事务指标
    #[cfg(feature = "metrics")]
    fn record_transaction_metrics(&self, operation: &str, duration: Duration, success: bool) {
        if let Some(metrics) = &self.metrics_collector {
            // 使用 "transaction" 作为查询类型，操作作为子类型
            metrics.record_query(&format!("transaction_{}", operation), duration, success, None);
        }
    }

    /// 记录事务指标（无 metrics 特性）
    #[cfg(not(feature = "metrics"))]
    fn record_transaction_metrics(&self, _operation: &str, _duration: Duration, _success: bool) {
        // No-op when metrics feature is disabled
    }

    /// 检查表级权限
    ///
    /// 此方法为 ORM 操作提供权限检查，确保所有实体操作都经过权限验证
    pub async fn check_table_permission(&self, table_name: &str, operation: &str) -> DbResult<()> {
        #[cfg(feature = "permission")]
        {
            let action = match operation {
                "INSERT" => PermissionAction::Insert,
                "SELECT" => PermissionAction::Select,
                "UPDATE" => PermissionAction::Update,
                "DELETE" => PermissionAction::Delete,
                _ => return Err(DbError::Permission(format!("Unknown operation: {}", operation))),
            };

            if !self.permission_ctx.check_table_access(table_name, &action).await {
                return Err(DbError::Permission(format!(
                    "Permission denied for {} on {}",
                    operation, table_name
                )));
            }
        }
        #[cfg(not(feature = "permission"))]
        {
            // Suppress unused variable warnings when permission feature is disabled
            let _ = table_name;
            let _ = operation;
        }
        Ok(())
    }

    /// 记录指标
    #[cfg(feature = "metrics")]
    pub fn record_metric(&self, operation: &str, table_name: &str, success: bool) {
        if let Some(metrics) = &self.metrics_collector {
            // 使用表名的哈希值作为 bytes 参数
            let bytes = Some(table_name.len() as u64);
            metrics.record_query(operation, std::time::Duration::from_millis(0), success, bytes);
        }
    }
}

#[cfg(feature = "permission")]
fn is_invalid_table_name(table_name: &str) -> bool {
    let table_name = table_name.trim();
    if table_name.is_empty() {
        return true;
    }

    for part in table_name.split('.') {
        let part = part.trim();
        if part.is_empty() {
            return true;
        }

        let unquoted = part
            .strip_prefix('"')
            .and_then(|s| s.strip_suffix('"'))
            .or_else(|| part.strip_prefix('`').and_then(|s| s.strip_suffix('`')))
            .or_else(|| part.strip_prefix('\'').and_then(|s| s.strip_suffix('\'')))
            .unwrap_or(part)
            .trim();

        if unquoted.is_empty() {
            return true;
        }
    }

    false
}

impl Drop for Session {
    fn drop(&mut self) {
        // 检查是否有未提交的事务
        if self.transaction.is_some() {
            tracing::warn!("Session dropped with active transaction. Transaction will be rolled back by the database.");

            // 注意：由于 Drop trait 的限制，我们无法在 Drop 中执行异步操作
            // 实际的回滚将由数据库在连接关闭时自动执行
            // 这里我们只是记录警告，让用户知道发生了什么

            // 在非 tokio 环境下，连接会在 Drop 时被关闭，数据库会自动回滚事务
            // 在 tokio 环境下，我们可以尝试异步回滚，但由于 Drop 是同步的，我们无法等待结果
            if let Ok(_handle) = tokio::runtime::Handle::try_current() {
                // 克隆必要的数据用于异步回滚
                let conn = self.connection.take();

                if let Some(conn) = conn {
                    // 注意：这里我们无法回滚事务，因为事务对象已经被 Drop
                    // 我们只能记录警告
                    tracing::error!(
                        "Cannot rollback transaction in Sync Drop. Transaction will be rolled back when connection is closed."
                    );

                    // 仍然归还连接，让数据库在连接关闭时回滚
                    self.pool.release_connection(conn);
                    return;
                }
            }
        }

        // 归还连接到池
        if let Some(conn) = self.connection.take() {
            self.pool.release_connection(conn);
        }
    }
}

#[cfg(feature = "permission")]
fn extract_table_name(sql: &str) -> String {
    #[cfg(feature = "sql-parser")]
    {
        let parser = SqlParser::new();
        if let Some((table_name, _)) = parser.parse_operation(sql) {
            return table_name;
        }
    }

    #[cfg(not(feature = "sql-parser"))]
    {
        let sql_upper = sql.to_uppercase();

        if sql_upper.contains("FROM ") {
            if let Some(start) = sql_upper.find("FROM ") {
                let rest = &sql[start + 5..];
                if let Some(end) = rest.find(|c| [' ', ',', ';', '(', ')'].contains(&c)) {
                    return rest[..end].trim().to_string();
                } else {
                    return rest.trim().to_string();
                }
            }
        }

        if sql_upper.contains("INTO ") {
            if let Some(start) = sql_upper.find("INTO ") {
                let rest = &sql[start + 5..];
                if let Some(end) = rest.find(|c| [' ', '(', ';'].contains(&c)) {
                    return rest[..end].trim().to_string();
                } else {
                    return rest.trim().to_string();
                }
            }
        }

        if sql_upper.contains("UPDATE ") {
            if let Some(start) = sql_upper.find("UPDATE ") {
                let rest = &sql[start + 7..];
                if let Some(end) = rest.find(|c| [' ', ';'].contains(&c)) {
                    return rest[..end].trim().to_string();
                } else {
                    return rest.trim().to_string();
                }
            }
        }
    }

    String::new()
}

#[cfg(all(feature = "permission", not(feature = "sql-parser")))]
fn parse_table_and_action(sql: &str) -> (String, PermissionAction) {
    let table_name = extract_table_name(sql);
    let sql_upper = sql.trim_start().to_uppercase();
    let action = if sql_upper.starts_with("INSERT") {
        PermissionAction::Insert
    } else if sql_upper.starts_with("UPDATE") {
        PermissionAction::Update
    } else if sql_upper.starts_with("DELETE") {
        PermissionAction::Delete
    } else {
        PermissionAction::Select
    };

    (table_name, action)
}

// 实现 DatabaseSession trait
#[async_trait]
impl super::DatabaseSession for Session {
    async fn execute(&mut self, sql: &str) -> DbResult<ExecResult> {
        self.execute(sql).await
    }

    async fn execute_raw(&self, sql: &str) -> DbResult<ExecResult> {
        self.execute_raw(sql).await
    }

    async fn execute_raw_ddl(&self, sql: &str) -> DbResult<ExecResult> {
        self.execute_raw_ddl(sql).await
    }

    async fn begin_transaction(&mut self) -> DbResult<()> {
        self.begin_transaction().await
    }

    async fn commit(&mut self) -> DbResult<()> {
        self.commit().await
    }

    async fn rollback(&mut self) -> DbResult<()> {
        self.rollback().await
    }

    fn role(&self) -> &str {
        self.role()
    }

    fn is_in_transaction(&self) -> bool {
        self.is_in_transaction()
    }
}