sz_orm_core/
hydration_plugin.rs

1//! Hydration Modes + Plugin 拦截器链
2//!
3//! 对应文档 6.8 节改进项 44（Hydration Modes）+ 46（Plugin 拦截器链）。
4//!
5//! # 核心概念
6//!
7//! ## Hydration Modes
8//! - **Object**：每行 → `HashMap<String, Value>`（默认）
9//! - **Array**：每行 → `Vec<Value>`（按列顺序）
10//! - **Scalar**：每行 → Value（取第一列）
11//! - **SingleScalar**：唯一行 + 唯一列 → Value
12//! - **Column**：每行的指定列 → `Vec<Value>`
13//!
14//! ## Plugin 拦截器链
15//! - **Plugin trait**：拦截 Executor 操作
16//! - **PluginContext**：上下文（操作类型、SQL、参数、阶段）
17//! - **PluginDecision**：插件决策（Continue / Skip / Modified）
18//! - **PluginChain**：拦截器链（按注册顺序执行）
19//! - **ExecutionStage**：拦截阶段（BeforeQuery / AfterQuery / BeforeUpdate / AfterUpdate / BeforeCommit / AfterCommit / BeforeRollback / AfterRollback）
20//!
21//! # 设计灵感
22//!
23//! - Doctrine `HYDRATE_*`（OBJECT/ARRAY/SCALAR/SINGLE_SCALAR/COLUMN）
24//! - MyBatis `Interceptor`（拦截 Executor.query/update/commit）
25//! - Hibernate `Interceptor` / `EventListeners`
26//!
27//! # 使用示例
28//!
29//! ```
30//! use sz_orm_core::hydration_plugin::{
31//!     HydrationMode, hydrate, PluginChain, PluginContext, ExecutionStage, PluginDecision,
32//! };
33//! use sz_orm_core::result_map::RowData;
34//! use sz_orm_core::Value;
35//! use std::collections::HashMap;
36//!
37//! // HydrationMode::Scalar
38//! let mut row = RowData::empty();
39//! row.set("count", Value::I64(42));
40//! let result = hydrate(&[row], HydrationMode::Scalar).unwrap();
41//! assert_eq!(result.first(), Some(&Value::I64(42)));
42//! ```
43
44use crate::result_map::RowData;
45use crate::value::Value;
46use parking_lot::RwLock;
47use std::collections::HashMap;
48use std::time::{Duration, Instant};
49
50// ============================================================================
51// HydrationMode — 填充模式
52// ============================================================================
53
54/// Hydration 填充模式
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
56pub enum HydrationMode {
57    /// 每行 → `HashMap<String, Value>`（默认，对象模式）
58    #[default]
59    Object,
60    /// 每行 → `Vec<Value>`（按列顺序）
61    Array,
62    /// 每行 → Value（取第一列）
63    Scalar,
64    /// 唯一行 + 唯一列 → Value（聚合查询常用）
65    SingleScalar,
66    /// 每行的指定列 → `Vec<Value>`
67    Column,
68}
69
70impl HydrationMode {
71    /// 模式名称
72    pub fn name(&self) -> &'static str {
73        match self {
74            HydrationMode::Object => "object",
75            HydrationMode::Array => "array",
76            HydrationMode::Scalar => "scalar",
77            HydrationMode::SingleScalar => "single_scalar",
78            HydrationMode::Column => "column",
79        }
80    }
81}
82
83// ============================================================================
84// HydrationResult / hydrate 函数
85// ============================================================================
86
87/// Hydration 错误
88#[derive(Debug, Clone, PartialEq)]
89pub enum HydrationError {
90    /// SingleScalar 模式下行数不等于 1
91    SingleScalarRequiresSingleRow {
92        /// 实际行数
93        actual_rows: usize,
94    },
95    /// 指定列不存在
96    ColumnNotFound {
97        /// 缺失的列名
98        column: String,
99    },
100    /// 行列数为 0
101    EmptyRow,
102}
103
104impl std::fmt::Display for HydrationError {
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        match self {
107            HydrationError::SingleScalarRequiresSingleRow { actual_rows } => {
108                write!(
109                    f,
110                    "SingleScalar mode requires exactly 1 row, got {}",
111                    actual_rows
112                )
113            }
114            HydrationError::ColumnNotFound { column } => {
115                write!(f, "column '{}' not found", column)
116            }
117            HydrationError::EmptyRow => write!(f, "row has no columns"),
118        }
119    }
120}
121
122impl std::error::Error for HydrationError {}
123
124/// Hydration 结果类型
125pub type HydrationResult<T> = Result<T, HydrationError>;
126
127/// Object 模式：每行 → HashMap<String, Value>
128pub fn hydrate_object(rows: &[RowData]) -> HydrationResult<Vec<HashMap<String, Value>>> {
129    Ok(rows
130        .iter()
131        .map(|r| {
132            let mut map = HashMap::new();
133            for (k, v) in r.iter() {
134                map.insert(k.clone(), v.clone());
135            }
136            map
137        })
138        .collect())
139}
140
141/// Array 模式：每行 → `Vec<Value>`（按列名排序后顺序）
142pub fn hydrate_array(rows: &[RowData]) -> HydrationResult<Vec<Vec<Value>>> {
143    let mut result = Vec::with_capacity(rows.len());
144    for row in rows {
145        // 按列名排序，保证顺序稳定
146        let sorted = row.sorted_columns();
147        let values: Vec<Value> = sorted.iter().map(|(_, v)| (*v).clone()).collect();
148        result.push(values);
149    }
150    Ok(result)
151}
152
153/// Scalar 模式：每行 → Value（取第一列，按列名排序）
154pub fn hydrate_scalar(rows: &[RowData]) -> HydrationResult<Vec<Value>> {
155    let mut result = Vec::with_capacity(rows.len());
156    for row in rows {
157        if row.is_empty() {
158            return Err(HydrationError::EmptyRow);
159        }
160        let sorted = row.sorted_columns();
161        let (_, first_value) = sorted
162            .first()
163            .expect("sorted_columns is non-empty after is_empty check"); // SAFETY: 前置 row.is_empty() 校验已 return，保证 sorted_columns 非空
164        result.push((*first_value).clone());
165    }
166    Ok(result)
167}
168
169/// SingleScalar 模式：唯一行 + 唯一列 → Value
170pub fn hydrate_single_scalar(rows: &[RowData]) -> HydrationResult<Value> {
171    if rows.len() != 1 {
172        return Err(HydrationError::SingleScalarRequiresSingleRow {
173            actual_rows: rows.len(),
174        });
175    }
176    let row = &rows[0];
177    if row.is_empty() {
178        return Err(HydrationError::EmptyRow);
179    }
180    let sorted = row.sorted_columns();
181    let (_, first_value) = sorted
182        .first()
183        .expect("sorted_columns is non-empty after is_empty check"); // SAFETY: 前置 row.is_empty() 校验已 return，保证 sorted_columns 非空
184    Ok((*first_value).clone())
185}
186
187/// Column 模式：每行的指定列 → `Vec<Value>`
188pub fn hydrate_column(rows: &[RowData], column: &str) -> HydrationResult<Vec<Value>> {
189    let mut result = Vec::with_capacity(rows.len());
190    for row in rows {
191        match row.get(column) {
192            Some(v) => result.push(v.clone()),
193            None => {
194                return Err(HydrationError::ColumnNotFound {
195                    column: column.to_string(),
196                })
197            }
198        }
199    }
200    Ok(result)
201}
202
203/// 通用 hydrate 函数：根据 mode 自动选择
204///
205/// 注意：Column 模式需要列名参数，请直接使用 `hydrate_column`。
206/// 此函数对 Column 模式使用第一列。
207pub fn hydrate(rows: &[RowData], mode: HydrationMode) -> HydrationResult<Vec<Value>> {
208    match mode {
209        HydrationMode::Scalar => hydrate_scalar(rows),
210        HydrationMode::SingleScalar => {
211            let v = hydrate_single_scalar(rows)?;
212            Ok(vec![v])
213        }
214        HydrationMode::Column => {
215            if rows.is_empty() {
216                return Ok(Vec::new());
217            }
218            let first_row = &rows[0];
219            if first_row.is_empty() {
220                return Err(HydrationError::EmptyRow);
221            }
222            let sorted = first_row.sorted_columns();
223            let first_col = sorted
224                .first()
225                .expect("sorted_columns is non-empty after is_empty check") // SAFETY: 前置 first_row.is_empty() 校验已 return，保证 sorted_columns 非空
226                .0
227                .as_str();
228            hydrate_column(rows, first_col)
229        }
230        HydrationMode::Object | HydrationMode::Array => {
231            // Object/Array 模式下，结果应是 HashMap/Vec 而非 Value
232            // 这里返回每行的第一列 Value 作为简化
233            // 完整 Object/Array 结果请使用 hydrate_object / hydrate_array
234            hydrate_scalar(rows)
235        }
236    }
237}
238
239// ============================================================================
240// Plugin 拦截器链
241// ============================================================================
242
243/// 执行阶段
244#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
245pub enum ExecutionStage {
246    /// 查询前
247    BeforeQuery,
248    /// 查询后
249    AfterQuery,
250    /// 更新前（INSERT/UPDATE/DELETE）
251    BeforeUpdate,
252    /// 更新后
253    AfterUpdate,
254    /// 提交前
255    BeforeCommit,
256    /// 提交后
257    AfterCommit,
258    /// 回滚前
259    BeforeRollback,
260    /// 回滚后
261    AfterRollback,
262}
263
264impl ExecutionStage {
265    /// 阶段名称
266    pub fn name(&self) -> &'static str {
267        match self {
268            ExecutionStage::BeforeQuery => "before_query",
269            ExecutionStage::AfterQuery => "after_query",
270            ExecutionStage::BeforeUpdate => "before_update",
271            ExecutionStage::AfterUpdate => "after_update",
272            ExecutionStage::BeforeCommit => "before_commit",
273            ExecutionStage::AfterCommit => "after_commit",
274            ExecutionStage::BeforeRollback => "before_rollback",
275            ExecutionStage::AfterRollback => "after_rollback",
276        }
277    }
278
279    /// 是否为 before 阶段
280    pub fn is_before(&self) -> bool {
281        matches!(
282            self,
283            ExecutionStage::BeforeQuery
284                | ExecutionStage::BeforeUpdate
285                | ExecutionStage::BeforeCommit
286                | ExecutionStage::BeforeRollback
287        )
288    }
289
290    /// 是否为 after 阶段
291    pub fn is_after(&self) -> bool {
292        !self.is_before()
293    }
294
295    /// 是否为查询阶段
296    pub fn is_query(&self) -> bool {
297        matches!(
298            self,
299            ExecutionStage::BeforeQuery | ExecutionStage::AfterQuery
300        )
301    }
302
303    /// 是否为更新阶段
304    pub fn is_update(&self) -> bool {
305        matches!(
306            self,
307            ExecutionStage::BeforeUpdate | ExecutionStage::AfterUpdate
308        )
309    }
310
311    /// 是否为事务阶段
312    pub fn is_transaction(&self) -> bool {
313        matches!(
314            self,
315            ExecutionStage::BeforeCommit
316                | ExecutionStage::AfterCommit
317                | ExecutionStage::BeforeRollback
318                | ExecutionStage::AfterRollback
319        )
320    }
321}
322
323/// 插件上下文（携带操作信息）
324#[derive(Debug, Clone)]
325pub struct PluginContext {
326    /// 执行阶段
327    pub stage: ExecutionStage,
328    /// SQL 语句（可被插件修改）
329    pub sql: String,
330    /// 绑定参数
331    pub parameters: Vec<Value>,
332    /// 执行开始时间（用于慢查询检测）
333    pub started_at: Option<Instant>,
334    /// 执行耗时（After 阶段才有）
335    pub elapsed: Option<Duration>,
336    /// 影响行数（After 阶段才有）
337    pub affected_rows: Option<usize>,
338    /// 自定义元数据
339    pub metadata: HashMap<String, Value>,
340}
341
342impl PluginContext {
343    /// 创建上下文
344    pub fn new(stage: ExecutionStage, sql: impl Into<String>) -> Self {
345        Self {
346            stage,
347            sql: sql.into(),
348            parameters: Vec::new(),
349            started_at: None,
350            elapsed: None,
351            affected_rows: None,
352            metadata: HashMap::new(),
353        }
354    }
355
356    /// 设置参数
357    pub fn with_parameters(mut self, params: Vec<Value>) -> Self {
358        self.parameters = params;
359        self
360    }
361
362    /// 设置开始时间
363    pub fn with_start_time(mut self, instant: Instant) -> Self {
364        self.started_at = Some(instant);
365        self
366    }
367
368    /// 设置耗时
369    pub fn with_elapsed(mut self, elapsed: Duration) -> Self {
370        self.elapsed = Some(elapsed);
371        self
372    }
373
374    /// 设置影响行数
375    pub fn with_affected_rows(mut self, rows: usize) -> Self {
376        self.affected_rows = Some(rows);
377        self
378    }
379
380    /// 添加元数据
381    pub fn set_metadata(&mut self, key: impl Into<String>, value: Value) {
382        self.metadata.insert(key.into(), value);
383    }
384
385    /// 获取元数据
386    pub fn get_metadata(&self, key: &str) -> Option<&Value> {
387        self.metadata.get(key)
388    }
389}
390
391/// 插件决策
392#[derive(Debug, Clone, PartialEq)]
393pub enum PluginDecision {
394    /// 继续执行（链中下一个插件）
395    Continue,
396    /// 跳过后续插件（但仍执行原 SQL）
397    Skip,
398    /// 修改 SQL/参数后继续
399    Modified {
400        /// 改写后的 SQL
401        sql: String,
402        /// 替换后的参数列表
403        parameters: Vec<Value>,
404    },
405    /// 中止执行（不执行原 SQL，返回错误）
406    Abort(String),
407    /// Kill 当前查询/事务（用于慢查询自动 kill 等场景）
408    Kill,
409}
410
411/// Plugin 拦截器 trait
412pub trait Plugin: Send + Sync {
413    /// 插件名称
414    fn name(&self) -> &str;
415
416    /// 拦截哪些阶段
417    fn stages(&self) -> Vec<ExecutionStage>;
418
419    /// 拦截处理
420    fn intercept(&self, context: &mut PluginContext) -> PluginDecision;
421}
422
423// ============================================================================
424// PluginChain — 插件链
425// ============================================================================
426
427/// 插件链（按注册顺序执行）
428#[derive(Default)]
429pub struct PluginChain {
430    plugins: RwLock<Vec<Box<dyn Plugin>>>,
431}
432
433impl PluginChain {
434    /// 创建空链
435    pub fn new() -> Self {
436        Self {
437            plugins: RwLock::new(Vec::new()),
438        }
439    }
440
441    /// 注册插件（追加到链尾）
442    pub fn register(&self, plugin: Box<dyn Plugin>) {
443        let mut plugins = self.plugins.write();
444        plugins.push(plugin);
445    }
446
447    /// 注册插件到指定位置
448    pub fn insert_at(&self, index: usize, plugin: Box<dyn Plugin>) {
449        let mut plugins = self.plugins.write();
450        let len = plugins.len();
451        plugins.insert(index.min(len), plugin);
452    }
453
454    /// 注销指定名称的插件
455    pub fn unregister(&self, name: &str) -> bool {
456        let mut plugins = self.plugins.write();
457        if let Some(idx) = plugins.iter().position(|p| p.name() == name) {
458            plugins.remove(idx);
459            true
460        } else {
461            false
462        }
463    }
464
465    /// 已注册的插件数量
466    pub fn len(&self) -> usize {
467        self.plugins.read().len()
468    }
469
470    /// 是否为空
471    pub fn is_empty(&self) -> bool {
472        self.len() == 0
473    }
474
475    /// 列出所有插件名
476    pub fn plugin_names(&self) -> Vec<String> {
477        self.plugins
478            .read()
479            .iter()
480            .map(|p| p.name().to_string())
481            .collect()
482    }
483
484    /// 清空插件链
485    pub fn clear(&self) {
486        self.plugins.write().clear();
487    }
488
489    /// 执行插件链
490    ///
491    /// 按 before 插件 → 原操作 → after 插件 的顺序执行。
492    /// 任意插件返回 `Abort` 中止整个链。
493    /// 任意插件返回 `Modified` 会修改 context 后继续。
494    /// 任意插件返回 `Skip` 跳过后续插件（但继续原操作）。
495    pub fn execute(&self, context: &mut PluginContext) -> PluginDecision {
496        let plugins = self.plugins.read();
497        let target_stages = [context.stage];
498
499        for plugin in plugins.iter() {
500            // 仅调用订阅了当前阶段的插件
501            if !plugin.stages().iter().any(|s| target_stages.contains(s)) {
502                continue;
503            }
504            match plugin.intercept(context) {
505                PluginDecision::Continue => continue,
506                PluginDecision::Skip => return PluginDecision::Skip,
507                PluginDecision::Modified { sql, parameters } => {
508                    context.sql = sql;
509                    context.parameters = parameters;
510                    continue;
511                }
512                PluginDecision::Abort(reason) => {
513                    return PluginDecision::Abort(reason);
514                }
515                PluginDecision::Kill => return PluginDecision::Kill,
516            }
517        }
518        PluginDecision::Continue
519    }
520}
521
522impl std::fmt::Debug for PluginChain {
523    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
524        let plugins = self.plugins.read();
525        let names: Vec<&str> = plugins.iter().map(|p| p.name()).collect();
526        f.debug_struct("PluginChain")
527            .field("plugins", &names)
528            .finish()
529    }
530}
531
532// ============================================================================
533// 内置插件：SqlLogPlugin
534// ============================================================================
535
536/// SQL 日志插件
537pub struct SqlLogPlugin {
538    logs: RwLock<Vec<String>>,
539}
540
541impl SqlLogPlugin {
542    /// 创建一个空的 SQL 日志插件
543    pub fn new() -> Self {
544        Self {
545            logs: RwLock::new(Vec::new()),
546        }
547    }
548
549    /// 返回已记录的全部 SQL 日志
550    pub fn logs(&self) -> Vec<String> {
551        self.logs.read().clone()
552    }
553
554    /// 清空已记录的日志
555    pub fn clear(&self) {
556        self.logs.write().clear();
557    }
558
559    /// 返回已记录的日志条数
560    pub fn count(&self) -> usize {
561        self.logs.read().len()
562    }
563}
564
565impl Default for SqlLogPlugin {
566    fn default() -> Self {
567        Self::new()
568    }
569}
570
571impl Plugin for SqlLogPlugin {
572    fn name(&self) -> &str {
573        "sql_log"
574    }
575
576    fn stages(&self) -> Vec<ExecutionStage> {
577        vec![
578            ExecutionStage::BeforeQuery,
579            ExecutionStage::AfterQuery,
580            ExecutionStage::BeforeUpdate,
581            ExecutionStage::AfterUpdate,
582        ]
583    }
584
585    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
586        let mut logs = self.logs.write();
587        let entry = match context.stage {
588            ExecutionStage::BeforeQuery => {
589                format!("[{}] QUERY: {}", context.stage.name(), context.sql)
590            }
591            ExecutionStage::AfterQuery => {
592                let elapsed_ms = context.elapsed.map(|d| d.as_millis()).unwrap_or(0);
593                format!(
594                    "[{}] QUERY ({}ms): {}",
595                    context.stage.name(),
596                    elapsed_ms,
597                    context.sql
598                )
599            }
600            ExecutionStage::BeforeUpdate => {
601                format!("[{}] UPDATE: {}", context.stage.name(), context.sql)
602            }
603            ExecutionStage::AfterUpdate => {
604                let rows = context.affected_rows.unwrap_or(0);
605                format!(
606                    "[{}] UPDATE ({} rows): {}",
607                    context.stage.name(),
608                    rows,
609                    context.sql
610                )
611            }
612            _ => return PluginDecision::Continue,
613        };
614        logs.push(entry);
615        PluginDecision::Continue
616    }
617}
618
619// ============================================================================
620// 内置插件：SlowQueryPlugin
621// ============================================================================
622
623/// 慢查询检测插件
624pub struct SlowQueryPlugin {
625    threshold: Duration,
626    slow_queries: RwLock<Vec<SlowQueryRecord>>,
627    /// 任务6：慢查询自动 kill 阈值（超过此阈值返回 Kill 决策）
628    kill_threshold: Option<Duration>,
629}
630
631/// 慢查询记录
632#[derive(Debug, Clone)]
633pub struct SlowQueryRecord {
634    /// 慢查询 SQL 文本
635    pub sql: String,
636    /// 实际执行耗时
637    pub elapsed: Duration,
638    /// 判定为慢查询的阈值
639    pub threshold: Duration,
640}
641
642impl SlowQueryPlugin {
643    /// 创建插件，指定阈值
644    pub fn new(threshold: Duration) -> Self {
645        Self {
646            threshold,
647            slow_queries: RwLock::new(Vec::new()),
648            kill_threshold: None,
649        }
650    }
651
652    /// 默认 1 秒阈值
653    pub fn default_threshold() -> Self {
654        Self::new(Duration::from_secs(1))
655    }
656
657    /// 任务6：设置 kill 阈值（超过此阈值时返回 Kill 决策）
658    pub fn with_kill_threshold(mut self, kill_threshold: Duration) -> Self {
659        self.kill_threshold = Some(kill_threshold);
660        self
661    }
662
663    /// 获取所有慢查询记录
664    pub fn slow_queries(&self) -> Vec<SlowQueryRecord> {
665        self.slow_queries.read().clone()
666    }
667
668    /// 慢查询数量
669    pub fn count(&self) -> usize {
670        self.slow_queries.read().len()
671    }
672
673    /// 清空记录
674    pub fn clear(&self) {
675        self.slow_queries.write().clear();
676    }
677
678    /// 阈值
679    pub fn threshold(&self) -> Duration {
680        self.threshold
681    }
682}
683
684/// 脱敏 SQL 中的敏感值（password/passwd/secret/token = '...' → '***'）
685///
686/// 由于 sz-orm-core 未依赖 regex crate，使用手动扫描实现大小写不敏感匹配。
687/// 仅处理 `key='value'` 模式；其他形式（如函数参数）不做脱敏。
688fn mask_sql(sql: &str) -> String {
689    const SENSITIVE_KEYS: &[&str] = &["password", "passwd", "secret", "token"];
690
691    let lower = sql.to_ascii_lowercase();
692    let bytes = sql.as_bytes();
693    let lower_bytes = lower.as_bytes();
694    let mut result = String::with_capacity(sql.len());
695    let mut i = 0;
696
697    while i < bytes.len() {
698        let mut matched = false;
699        for kw in SENSITIVE_KEYS {
700            let kw_b = kw.as_bytes();
701            if i + kw_b.len() <= bytes.len() && &lower_bytes[i..i + kw_b.len()] == kw_b {
702                // 边界检查：前后字符不能是标识符字符（避免匹配 "passworded"）
703                let prev_ok = i == 0 || !is_ident_char(bytes[i - 1]);
704                let next = i + kw_b.len();
705                let next_ok = next >= bytes.len() || !is_ident_char(bytes[next]);
706                if !prev_ok || !next_ok {
707                    continue;
708                }
709                // 复制 keyword 原文（保留大小写）
710                result.push_str(&sql[i..next]);
711                i = next;
712                // 跳过空白
713                while i < bytes.len() && bytes[i].is_ascii_whitespace() {
714                    result.push(bytes[i] as char);
715                    i += 1;
716                }
717                // 期待 '='
718                if i < bytes.len() && bytes[i] == b'=' {
719                    result.push('=');
720                    i += 1;
721                    while i < bytes.len() && bytes[i].is_ascii_whitespace() {
722                        result.push(bytes[i] as char);
723                        i += 1;
724                    }
725                    // 期待单引号字符串 'value'
726                    if i < bytes.len() && bytes[i] == b'\'' {
727                        i += 1; // 跳过开始引号
728                        while i < bytes.len() && bytes[i] != b'\'' {
729                            i += 1;
730                        }
731                        if i < bytes.len() {
732                            i += 1; // 跳过结束引号
733                        }
734                        result.push_str("'***'");
735                    }
736                }
737                matched = true;
738                break;
739            }
740        }
741        if !matched {
742            let ch = sql[i..]
743                .chars()
744                .next()
745                .expect("i < bytes.len() guarantees non-empty slice"); // SAFETY: 循环条件 i < bytes.len() 保证 sql[i..] 非空，chars().next() 必返回 Some
746            result.push(ch);
747            i += ch.len_utf8();
748        }
749    }
750    result
751}
752
753/// ASCII 标识符字符判断（字母/数字/下划线）
754fn is_ident_char(b: u8) -> bool {
755    b.is_ascii_alphanumeric() || b == b'_'
756}
757
758impl Plugin for SlowQueryPlugin {
759    fn name(&self) -> &str {
760        "slow_query"
761    }
762
763    fn stages(&self) -> Vec<ExecutionStage> {
764        // 订阅 BeforeQuery：在执行前记录开始时间，并在超过 kill_threshold 时
765        // 通过 PluginDecision::Kill 中止查询（#9 修复：慢查询自动 Kill 生效）
766        vec![
767            ExecutionStage::BeforeQuery,
768            ExecutionStage::AfterQuery,
769            ExecutionStage::AfterUpdate,
770        ]
771    }
772
773    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
774        match context.stage {
775            ExecutionStage::BeforeQuery => {
776                // 执行前检查：若已记录的开始时间显示查询进行中超过 kill_threshold，
777                // 直接 Kill（理论上 BeforeQuery 阶段 elapsed 为 None，此处保留扩展点）
778                // 主要 Kill 逻辑放在 AfterQuery，因为 BeforeQuery 时 elapsed 尚未计算
779                if let Some(kill_threshold) = self.kill_threshold {
780                    if let Some(elapsed) = context.elapsed {
781                        if elapsed > kill_threshold {
782                            return PluginDecision::Kill;
783                        }
784                    }
785                }
786                PluginDecision::Continue
787            }
788            ExecutionStage::AfterQuery | ExecutionStage::AfterUpdate => {
789                if let Some(elapsed) = context.elapsed {
790                    if elapsed > self.threshold {
791                        // 记录前对 SQL 脱敏，避免敏感值（password/token 等）落入慢查询日志
792                        let masked_sql = mask_sql(&context.sql);
793                        let mut records = self.slow_queries.write();
794                        records.push(SlowQueryRecord {
795                            sql: masked_sql,
796                            elapsed,
797                            threshold: self.threshold,
798                        });
799                    }
800                    // #9 修复：超过 kill_threshold 时返回 Kill 决策
801                    // 这会触发插件链中止后续查询/事务，防止慢查询持续占用连接池资源
802                    if let Some(kill_threshold) = self.kill_threshold {
803                        if elapsed > kill_threshold {
804                            return PluginDecision::Kill;
805                        }
806                    }
807                }
808                PluginDecision::Continue
809            }
810            _ => PluginDecision::Continue,
811        }
812    }
813}
814
815// ============================================================================
816// 内置插件：AuditPlugin
817// ============================================================================
818
819/// 审计插件（记录所有写操作）
820pub struct AuditPlugin {
821    audit_log: RwLock<Vec<AuditRecord>>,
822}
823
824/// 审计记录
825#[derive(Debug, Clone)]
826pub struct AuditRecord {
827    /// 审计触发阶段
828    pub stage: ExecutionStage,
829    /// 审计记录的 SQL 文本
830    pub sql: String,
831    /// 受影响的行数（若可用）
832    pub affected_rows: Option<usize>,
833}
834
835impl AuditPlugin {
836    /// 创建一个空的审计插件
837    pub fn new() -> Self {
838        Self {
839            audit_log: RwLock::new(Vec::new()),
840        }
841    }
842
843    /// 返回全部审计记录
844    pub fn records(&self) -> Vec<AuditRecord> {
845        self.audit_log.read().clone()
846    }
847
848    /// 返回审计记录条数
849    pub fn count(&self) -> usize {
850        self.audit_log.read().len()
851    }
852
853    /// 清空审计记录
854    pub fn clear(&self) {
855        self.audit_log.write().clear();
856    }
857}
858
859impl Default for AuditPlugin {
860    fn default() -> Self {
861        Self::new()
862    }
863}
864
865impl Plugin for AuditPlugin {
866    fn name(&self) -> &str {
867        "audit"
868    }
869
870    fn stages(&self) -> Vec<ExecutionStage> {
871        vec![ExecutionStage::AfterUpdate]
872    }
873
874    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
875        let mut log = self.audit_log.write();
876        log.push(AuditRecord {
877            stage: context.stage,
878            sql: context.sql.clone(),
879            affected_rows: context.affected_rows,
880        });
881        PluginDecision::Continue
882    }
883}
884
885// ============================================================================
886// 内置插件：SqlRewritePlugin（演示 Modified 决策）
887// ============================================================================
888
889/// SQL 改写插件（演示 Modified 决策）
890///
891/// 在执行前将 `SELECT` 替换为 `SELECT /* hint */`，用于演示 SQL 改写能力。
892pub struct SqlRewritePlugin {
893    pattern: String,
894    replacement: String,
895}
896
897impl SqlRewritePlugin {
898    /// 创建 SQL 改写插件，指定匹配模式和替换文本
899    pub fn new(pattern: impl Into<String>, replacement: impl Into<String>) -> Self {
900        Self {
901            pattern: pattern.into(),
902            replacement: replacement.into(),
903        }
904    }
905}
906
907impl Plugin for SqlRewritePlugin {
908    fn name(&self) -> &str {
909        "sql_rewrite"
910    }
911
912    fn stages(&self) -> Vec<ExecutionStage> {
913        vec![ExecutionStage::BeforeQuery, ExecutionStage::BeforeUpdate]
914    }
915
916    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
917        if context.sql.contains(&self.pattern) {
918            let new_sql = context.sql.replace(&self.pattern, &self.replacement);
919            PluginDecision::Modified {
920                sql: new_sql,
921                parameters: context.parameters.clone(),
922            }
923        } else {
924            PluginDecision::Continue
925        }
926    }
927}
928
929// ============================================================================
930// 内置插件：BlockPlugin（演示 Abort 决策）
931// ============================================================================
932
933/// 阻断插件（演示 Abort 决策）
934///
935/// 拦截包含指定关键字的 SQL（如 DROP TABLE），返回 Abort。
936pub struct BlockPlugin {
937    blocked_keywords: Vec<String>,
938}
939
940impl BlockPlugin {
941    /// 创建阻断插件，指定需拦截的关键字列表
942    pub fn new(keywords: Vec<String>) -> Self {
943        Self {
944            blocked_keywords: keywords,
945        }
946    }
947
948    /// 默认阻断 DROP / TRUNCATE
949    pub fn default_block_ddl() -> Self {
950        Self::new(vec![
951            "DROP TABLE".to_string(),
952            "TRUNCATE".to_string(),
953            "DROP DATABASE".to_string(),
954        ])
955    }
956}
957
958impl Plugin for BlockPlugin {
959    fn name(&self) -> &str {
960        "block"
961    }
962
963    fn stages(&self) -> Vec<ExecutionStage> {
964        vec![ExecutionStage::BeforeUpdate, ExecutionStage::BeforeQuery]
965    }
966
967    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
968        let upper_sql = context.sql.to_uppercase();
969        for kw in &self.blocked_keywords {
970            if upper_sql.contains(&kw.to_uppercase()) {
971                return PluginDecision::Abort(format!(
972                    "blocked by BlockPlugin: SQL contains forbidden keyword '{}'",
973                    kw
974                ));
975            }
976        }
977        PluginDecision::Continue
978    }
979}
980
981// ============================================================================
982// 单元测试
983// ============================================================================
984
985#[cfg(test)]
986mod tests {
987    use super::*;
988
989    // ===== HydrationMode =====
990
991    #[test]
992    fn test_hydration_mode_default() {
993        assert_eq!(HydrationMode::default(), HydrationMode::Object);
994    }
995
996    #[test]
997    fn test_hydration_mode_name() {
998        assert_eq!(HydrationMode::Object.name(), "object");
999        assert_eq!(HydrationMode::Array.name(), "array");
1000        assert_eq!(HydrationMode::Scalar.name(), "scalar");
1001        assert_eq!(HydrationMode::SingleScalar.name(), "single_scalar");
1002        assert_eq!(HydrationMode::Column.name(), "column");
1003    }
1004
1005    // ===== hydrate_object =====
1006
1007    #[test]
1008    fn test_hydrate_object_basic() {
1009        let mut row = RowData::empty();
1010        row.set("id", Value::I64(1));
1011        row.set("name", Value::String("Alice".to_string()));
1012
1013        let result = hydrate_object(&[row]).unwrap();
1014        assert_eq!(result.len(), 1);
1015        assert_eq!(result[0].get("id"), Some(&Value::I64(1)));
1016        assert_eq!(
1017            result[0].get("name"),
1018            Some(&Value::String("Alice".to_string()))
1019        );
1020    }
1021
1022    #[test]
1023    fn test_hydrate_object_multiple_rows() {
1024        let rows = vec![
1025            {
1026                let mut r = RowData::empty();
1027                r.set("id", Value::I64(1));
1028                r
1029            },
1030            {
1031                let mut r = RowData::empty();
1032                r.set("id", Value::I64(2));
1033                r
1034            },
1035        ];
1036
1037        let result = hydrate_object(&rows).unwrap();
1038        assert_eq!(result.len(), 2);
1039    }
1040
1041    #[test]
1042    fn test_hydrate_object_empty() {
1043        let result = hydrate_object(&[]).unwrap();
1044        assert!(result.is_empty());
1045    }
1046
1047    // ===== hydrate_array =====
1048
1049    #[test]
1050    fn test_hydrate_array_basic() {
1051        let mut row = RowData::empty();
1052        row.set("id", Value::I64(1));
1053        row.set("name", Value::String("Alice".to_string()));
1054
1055        let result = hydrate_array(&[row]).unwrap();
1056        assert_eq!(result.len(), 1);
1057        assert_eq!(result[0].len(), 2);
1058        // 列名排序后顺序：id, name
1059        assert_eq!(result[0][0], Value::I64(1));
1060        assert_eq!(result[0][1], Value::String("Alice".to_string()));
1061    }
1062
1063    #[test]
1064    fn test_hydrate_array_empty() {
1065        let result = hydrate_array(&[]).unwrap();
1066        assert!(result.is_empty());
1067    }
1068
1069    // ===== hydrate_scalar =====
1070
1071    #[test]
1072    fn test_hydrate_scalar_basic() {
1073        let mut row = RowData::empty();
1074        row.set("count", Value::I64(42));
1075
1076        let result = hydrate_scalar(&[row]).unwrap();
1077        assert_eq!(result.len(), 1);
1078        assert_eq!(result[0], Value::I64(42));
1079    }
1080
1081    #[test]
1082    fn test_hydrate_scalar_multiple_rows() {
1083        let rows = vec![
1084            {
1085                let mut r = RowData::empty();
1086                r.set("id", Value::I64(1));
1087                r
1088            },
1089            {
1090                let mut r = RowData::empty();
1091                r.set("id", Value::I64(2));
1092                r
1093            },
1094        ];
1095
1096        let result = hydrate_scalar(&rows).unwrap();
1097        assert_eq!(result, vec![Value::I64(1), Value::I64(2)]);
1098    }
1099
1100    #[test]
1101    fn test_hydrate_scalar_empty_row_error() {
1102        let row = RowData::empty();
1103        let err = hydrate_scalar(&[row]).unwrap_err();
1104        match err {
1105            HydrationError::EmptyRow => {}
1106            _ => panic!("expected EmptyRow error"),
1107        }
1108    }
1109
1110    // ===== hydrate_single_scalar =====
1111
1112    #[test]
1113    fn test_hydrate_single_scalar_ok() {
1114        let mut row = RowData::empty();
1115        row.set("total", Value::I64(100));
1116
1117        let result = hydrate_single_scalar(&[row]).unwrap();
1118        assert_eq!(result, Value::I64(100));
1119    }
1120
1121    #[test]
1122    fn test_hydrate_single_scalar_no_rows() {
1123        let err = hydrate_single_scalar(&[]).unwrap_err();
1124        match err {
1125            HydrationError::SingleScalarRequiresSingleRow { actual_rows } => {
1126                assert_eq!(actual_rows, 0)
1127            }
1128            _ => panic!("expected SingleScalarRequiresSingleRow"),
1129        }
1130    }
1131
1132    #[test]
1133    fn test_hydrate_single_scalar_too_many_rows() {
1134        let rows = vec![
1135            {
1136                let mut r = RowData::empty();
1137                r.set("id", Value::I64(1));
1138                r
1139            },
1140            {
1141                let mut r = RowData::empty();
1142                r.set("id", Value::I64(2));
1143                r
1144            },
1145        ];
1146        let err = hydrate_single_scalar(&rows).unwrap_err();
1147        match err {
1148            HydrationError::SingleScalarRequiresSingleRow { actual_rows } => {
1149                assert_eq!(actual_rows, 2)
1150            }
1151            _ => panic!("expected SingleScalarRequiresSingleRow"),
1152        }
1153    }
1154
1155    #[test]
1156    fn test_hydrate_single_scalar_empty_row() {
1157        let row = RowData::empty();
1158        let err = hydrate_single_scalar(&[row]).unwrap_err();
1159        match err {
1160            HydrationError::EmptyRow => {}
1161            _ => panic!("expected EmptyRow"),
1162        }
1163    }
1164
1165    // ===== hydrate_column =====
1166
1167    #[test]
1168    fn test_hydrate_column_basic() {
1169        let rows = vec![
1170            {
1171                let mut r = RowData::empty();
1172                r.set("id", Value::I64(1));
1173                r.set("name", Value::String("Alice".to_string()));
1174                r
1175            },
1176            {
1177                let mut r = RowData::empty();
1178                r.set("id", Value::I64(2));
1179                r.set("name", Value::String("Bob".to_string()));
1180                r
1181            },
1182        ];
1183
1184        let result = hydrate_column(&rows, "name").unwrap();
1185        assert_eq!(
1186            result,
1187            vec![
1188                Value::String("Alice".to_string()),
1189                Value::String("Bob".to_string()),
1190            ]
1191        );
1192    }
1193
1194    #[test]
1195    fn test_hydrate_column_missing() {
1196        let rows = vec![{
1197            let mut r = RowData::empty();
1198            r.set("id", Value::I64(1));
1199            r
1200        }];
1201
1202        let err = hydrate_column(&rows, "missing").unwrap_err();
1203        match err {
1204            HydrationError::ColumnNotFound { column } => assert_eq!(column, "missing"),
1205            _ => panic!("expected ColumnNotFound"),
1206        }
1207    }
1208
1209    #[test]
1210    fn test_hydrate_column_empty_rows() {
1211        let result = hydrate_column(&[], "name").unwrap();
1212        assert!(result.is_empty());
1213    }
1214
1215    // ===== 通用 hydrate 函数 =====
1216
1217    #[test]
1218    fn test_hydrate_scalar_mode() {
1219        let mut row = RowData::empty();
1220        row.set("count", Value::I64(42));
1221
1222        let result = hydrate(&[row], HydrationMode::Scalar).unwrap();
1223        assert_eq!(result, vec![Value::I64(42)]);
1224    }
1225
1226    #[test]
1227    fn test_hydrate_single_scalar_mode() {
1228        let mut row = RowData::empty();
1229        row.set("total", Value::I64(100));
1230
1231        let result = hydrate(&[row], HydrationMode::SingleScalar).unwrap();
1232        assert_eq!(result, vec![Value::I64(100)]);
1233    }
1234
1235    #[test]
1236    fn test_hydrate_column_mode() {
1237        let rows = vec![
1238            {
1239                let mut r = RowData::empty();
1240                r.set("id", Value::I64(1));
1241                r
1242            },
1243            {
1244                let mut r = RowData::empty();
1245                r.set("id", Value::I64(2));
1246                r
1247            },
1248        ];
1249
1250        let result = hydrate(&rows, HydrationMode::Column).unwrap();
1251        assert_eq!(result, vec![Value::I64(1), Value::I64(2)]);
1252    }
1253
1254    // ===== ExecutionStage =====
1255
1256    #[test]
1257    fn test_execution_stage_name() {
1258        assert_eq!(ExecutionStage::BeforeQuery.name(), "before_query");
1259        assert_eq!(ExecutionStage::AfterQuery.name(), "after_query");
1260        assert_eq!(ExecutionStage::BeforeUpdate.name(), "before_update");
1261        assert_eq!(ExecutionStage::AfterCommit.name(), "after_commit");
1262    }
1263
1264    #[test]
1265    fn test_execution_stage_is_before() {
1266        assert!(ExecutionStage::BeforeQuery.is_before());
1267        assert!(ExecutionStage::BeforeUpdate.is_before());
1268        assert!(ExecutionStage::BeforeCommit.is_before());
1269        assert!(ExecutionStage::BeforeRollback.is_before());
1270        assert!(!ExecutionStage::AfterQuery.is_before());
1271        assert!(!ExecutionStage::AfterUpdate.is_before());
1272    }
1273
1274    #[test]
1275    fn test_execution_stage_is_after() {
1276        assert!(ExecutionStage::AfterQuery.is_after());
1277        assert!(!ExecutionStage::BeforeQuery.is_after());
1278    }
1279
1280    #[test]
1281    fn test_execution_stage_is_query() {
1282        assert!(ExecutionStage::BeforeQuery.is_query());
1283        assert!(ExecutionStage::AfterQuery.is_query());
1284        assert!(!ExecutionStage::BeforeUpdate.is_query());
1285    }
1286
1287    #[test]
1288    fn test_execution_stage_is_update() {
1289        assert!(ExecutionStage::BeforeUpdate.is_update());
1290        assert!(ExecutionStage::AfterUpdate.is_update());
1291        assert!(!ExecutionStage::BeforeQuery.is_update());
1292    }
1293
1294    #[test]
1295    fn test_execution_stage_is_transaction() {
1296        assert!(ExecutionStage::BeforeCommit.is_transaction());
1297        assert!(ExecutionStage::AfterCommit.is_transaction());
1298        assert!(ExecutionStage::BeforeRollback.is_transaction());
1299        assert!(ExecutionStage::AfterRollback.is_transaction());
1300        assert!(!ExecutionStage::BeforeQuery.is_transaction());
1301    }
1302
1303    // ===== PluginContext =====
1304
1305    #[test]
1306    fn test_plugin_context_new() {
1307        let ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1308        assert_eq!(ctx.stage, ExecutionStage::BeforeQuery);
1309        assert_eq!(ctx.sql, "SELECT 1");
1310        assert!(ctx.parameters.is_empty());
1311        assert!(ctx.started_at.is_none());
1312        assert!(ctx.elapsed.is_none());
1313        assert!(ctx.affected_rows.is_none());
1314    }
1315
1316    #[test]
1317    fn test_plugin_context_with_parameters() {
1318        let ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT ?")
1319            .with_parameters(vec![Value::I64(1)]);
1320        assert_eq!(ctx.parameters.len(), 1);
1321    }
1322
1323    #[test]
1324    fn test_plugin_context_with_elapsed() {
1325        let ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1326            .with_elapsed(Duration::from_millis(50));
1327        assert_eq!(ctx.elapsed.unwrap().as_millis(), 50);
1328    }
1329
1330    #[test]
1331    fn test_plugin_context_with_affected_rows() {
1332        let ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users SET ...")
1333            .with_affected_rows(10);
1334        assert_eq!(ctx.affected_rows.unwrap(), 10);
1335    }
1336
1337    #[test]
1338    fn test_plugin_context_metadata() {
1339        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1340        ctx.set_metadata("user_id", Value::I64(42));
1341        assert_eq!(ctx.get_metadata("user_id"), Some(&Value::I64(42)));
1342        assert_eq!(ctx.get_metadata("missing"), None);
1343    }
1344
1345    // ===== PluginChain =====
1346
1347    #[test]
1348    fn test_plugin_chain_empty() {
1349        let chain = PluginChain::new();
1350        assert!(chain.is_empty());
1351        assert_eq!(chain.len(), 0);
1352    }
1353
1354    #[test]
1355    fn test_plugin_chain_register() {
1356        let chain = PluginChain::new();
1357        chain.register(Box::new(SqlLogPlugin::new()));
1358        assert_eq!(chain.len(), 1);
1359    }
1360
1361    #[test]
1362    fn test_plugin_chain_unregister() {
1363        let chain = PluginChain::new();
1364        chain.register(Box::new(SqlLogPlugin::new()));
1365        assert_eq!(chain.len(), 1);
1366
1367        let removed = chain.unregister("sql_log");
1368        assert!(removed);
1369        assert_eq!(chain.len(), 0);
1370    }
1371
1372    #[test]
1373    fn test_plugin_chain_unregister_missing() {
1374        let chain = PluginChain::new();
1375        let removed = chain.unregister("non_existent");
1376        assert!(!removed);
1377    }
1378
1379    #[test]
1380    fn test_plugin_chain_plugin_names() {
1381        let chain = PluginChain::new();
1382        chain.register(Box::new(SqlLogPlugin::new()));
1383        chain.register(Box::new(AuditPlugin::new()));
1384
1385        let names = chain.plugin_names();
1386        assert_eq!(names, vec!["sql_log", "audit"]);
1387    }
1388
1389    #[test]
1390    fn test_plugin_chain_clear() {
1391        let chain = PluginChain::new();
1392        chain.register(Box::new(SqlLogPlugin::new()));
1393        chain.clear();
1394        assert!(chain.is_empty());
1395    }
1396
1397    #[test]
1398    fn test_plugin_chain_insert_at() {
1399        let chain = PluginChain::new();
1400        chain.register(Box::new(SqlLogPlugin::new()));
1401        chain.insert_at(0, Box::new(AuditPlugin::new()));
1402
1403        let names = chain.plugin_names();
1404        assert_eq!(names, vec!["audit", "sql_log"]);
1405    }
1406
1407    #[test]
1408    fn test_plugin_chain_insert_at_end() {
1409        let chain = PluginChain::new();
1410        chain.register(Box::new(SqlLogPlugin::new()));
1411        chain.insert_at(99, Box::new(AuditPlugin::new()));
1412
1413        let names = chain.plugin_names();
1414        assert_eq!(names, vec!["sql_log", "audit"]);
1415    }
1416
1417    #[test]
1418    fn test_plugin_chain_execute_empty() {
1419        let chain = PluginChain::new();
1420        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1421        let decision = chain.execute(&mut ctx);
1422        assert_eq!(decision, PluginDecision::Continue);
1423    }
1424
1425    #[test]
1426    fn test_plugin_chain_execute_continue() {
1427        let chain = PluginChain::new();
1428        chain.register(Box::new(SqlLogPlugin::new()));
1429
1430        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1431        let decision = chain.execute(&mut ctx);
1432        assert_eq!(decision, PluginDecision::Continue);
1433    }
1434
1435    #[test]
1436    fn test_plugin_chain_execute_skip() {
1437        struct SkipPlugin;
1438        impl Plugin for SkipPlugin {
1439            fn name(&self) -> &str {
1440                "skip"
1441            }
1442            fn stages(&self) -> Vec<ExecutionStage> {
1443                vec![ExecutionStage::BeforeQuery]
1444            }
1445            fn intercept(&self, _ctx: &mut PluginContext) -> PluginDecision {
1446                PluginDecision::Skip
1447            }
1448        }
1449
1450        let chain = PluginChain::new();
1451        chain.register(Box::new(SkipPlugin));
1452
1453        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1454        let decision = chain.execute(&mut ctx);
1455        assert_eq!(decision, PluginDecision::Skip);
1456    }
1457
1458    #[test]
1459    fn test_plugin_chain_execute_modified() {
1460        let chain = PluginChain::new();
1461        chain.register(Box::new(SqlRewritePlugin::new(
1462            "SELECT",
1463            "SELECT /* hint */",
1464        )));
1465
1466        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1467        let decision = chain.execute(&mut ctx);
1468        assert_eq!(decision, PluginDecision::Continue);
1469        assert_eq!(ctx.sql, "SELECT /* hint */ 1");
1470    }
1471
1472    #[test]
1473    fn test_plugin_chain_execute_abort() {
1474        let chain = PluginChain::new();
1475        chain.register(Box::new(BlockPlugin::new(vec!["DROP".to_string()])));
1476
1477        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1478        let decision = chain.execute(&mut ctx);
1479        match decision {
1480            PluginDecision::Abort(reason) => assert!(reason.contains("DROP")),
1481            _ => panic!("expected Abort"),
1482        }
1483    }
1484
1485    #[test]
1486    fn test_plugin_chain_skip_unrelated_stages() {
1487        let chain = PluginChain::new();
1488        // SqlLogPlugin 订阅 Before/AfterQuery/Update，不应响应 BeforeCommit
1489        chain.register(Box::new(SqlLogPlugin::new()));
1490
1491        let mut ctx = PluginContext::new(ExecutionStage::BeforeCommit, "COMMIT");
1492        let decision = chain.execute(&mut ctx);
1493        assert_eq!(decision, PluginDecision::Continue);
1494        // 日志中不应有 BeforeCommit 记录
1495    }
1496
1497    // ===== SqlLogPlugin =====
1498
1499    #[test]
1500    fn test_sql_log_plugin_basic() {
1501        let plugin = SqlLogPlugin::new();
1502        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1503        let decision = plugin.intercept(&mut ctx);
1504        assert_eq!(decision, PluginDecision::Continue);
1505        assert_eq!(plugin.count(), 1);
1506    }
1507
1508    #[test]
1509    fn test_sql_log_plugin_after_query_with_elapsed() {
1510        let plugin = SqlLogPlugin::new();
1511        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1512            .with_elapsed(Duration::from_millis(50));
1513        let _ = plugin.intercept(&mut ctx);
1514
1515        let logs = plugin.logs();
1516        assert!(logs[0].contains("50ms"));
1517    }
1518
1519    #[test]
1520    fn test_sql_log_plugin_after_update_with_rows() {
1521        let plugin = SqlLogPlugin::new();
1522        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users SET ...")
1523            .with_affected_rows(10);
1524        let _ = plugin.intercept(&mut ctx);
1525
1526        let logs = plugin.logs();
1527        assert!(logs[0].contains("10 rows"));
1528    }
1529
1530    #[test]
1531    fn test_sql_log_plugin_clear() {
1532        let plugin = SqlLogPlugin::new();
1533        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1534        let _ = plugin.intercept(&mut ctx);
1535        assert_eq!(plugin.count(), 1);
1536
1537        plugin.clear();
1538        assert_eq!(plugin.count(), 0);
1539    }
1540
1541    // ===== SlowQueryPlugin =====
1542
1543    #[test]
1544    fn test_slow_query_plugin_below_threshold() {
1545        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1546        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1547            .with_elapsed(Duration::from_millis(50));
1548        let _ = plugin.intercept(&mut ctx);
1549
1550        assert_eq!(plugin.count(), 0);
1551    }
1552
1553    #[test]
1554    fn test_slow_query_plugin_above_threshold() {
1555        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1556        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT * FROM big_table")
1557            .with_elapsed(Duration::from_millis(500));
1558        let _ = plugin.intercept(&mut ctx);
1559
1560        assert_eq!(plugin.count(), 1);
1561        let records = plugin.slow_queries();
1562        assert!(records[0].elapsed > records[0].threshold);
1563    }
1564
1565    #[test]
1566    fn test_slow_query_plugin_no_elapsed() {
1567        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1568        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1");
1569        let _ = plugin.intercept(&mut ctx);
1570
1571        assert_eq!(plugin.count(), 0);
1572    }
1573
1574    #[test]
1575    fn test_slow_query_plugin_clear() {
1576        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1577        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1578            .with_elapsed(Duration::from_millis(200));
1579        let _ = plugin.intercept(&mut ctx);
1580        assert_eq!(plugin.count(), 1);
1581
1582        plugin.clear();
1583        assert_eq!(plugin.count(), 0);
1584    }
1585
1586    #[test]
1587    fn test_slow_query_plugin_default_threshold() {
1588        let plugin = SlowQueryPlugin::default_threshold();
1589        assert_eq!(plugin.threshold(), Duration::from_secs(1));
1590    }
1591
1592    // ===== AuditPlugin =====
1593
1594    #[test]
1595    fn test_audit_plugin_basic() {
1596        let plugin = AuditPlugin::new();
1597        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "INSERT INTO users ...")
1598            .with_affected_rows(1);
1599        let _ = plugin.intercept(&mut ctx);
1600
1601        assert_eq!(plugin.count(), 1);
1602        let records = plugin.records();
1603        assert_eq!(records[0].stage, ExecutionStage::AfterUpdate);
1604        assert_eq!(records[0].affected_rows, Some(1));
1605    }
1606
1607    #[test]
1608    fn test_audit_plugin_clear() {
1609        let plugin = AuditPlugin::new();
1610        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users");
1611        let _ = plugin.intercept(&mut ctx);
1612        assert_eq!(plugin.count(), 1);
1613
1614        plugin.clear();
1615        assert_eq!(plugin.count(), 0);
1616    }
1617
1618    // ===== SqlRewritePlugin =====
1619
1620    #[test]
1621    fn test_sql_rewrite_plugin_matches() {
1622        let plugin = SqlRewritePlugin::new("SELECT", "SELECT /* hint */");
1623        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT * FROM users");
1624        let decision = plugin.intercept(&mut ctx);
1625        match decision {
1626            PluginDecision::Modified { sql, .. } => {
1627                assert_eq!(sql, "SELECT /* hint */ * FROM users");
1628            }
1629            _ => panic!("expected Modified"),
1630        }
1631    }
1632
1633    #[test]
1634    fn test_sql_rewrite_plugin_no_match() {
1635        let plugin = SqlRewritePlugin::new("SELECT", "SELECT /* hint */");
1636        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SHOW TABLES");
1637        let decision = plugin.intercept(&mut ctx);
1638        assert_eq!(decision, PluginDecision::Continue);
1639    }
1640
1641    // ===== BlockPlugin =====
1642
1643    #[test]
1644    fn test_block_plugin_blocks_drop() {
1645        let plugin = BlockPlugin::default_block_ddl();
1646        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1647        let decision = plugin.intercept(&mut ctx);
1648        match decision {
1649            PluginDecision::Abort(reason) => {
1650                assert!(reason.contains("DROP TABLE"));
1651            }
1652            _ => panic!("expected Abort"),
1653        }
1654    }
1655
1656    #[test]
1657    fn test_block_plugin_blocks_truncate() {
1658        let plugin = BlockPlugin::default_block_ddl();
1659        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "TRUNCATE TABLE logs");
1660        let decision = plugin.intercept(&mut ctx);
1661        assert!(matches!(decision, PluginDecision::Abort(_)));
1662    }
1663
1664    #[test]
1665    fn test_block_plugin_allows_safe_sql() {
1666        let plugin = BlockPlugin::default_block_ddl();
1667        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT * FROM users");
1668        let decision = plugin.intercept(&mut ctx);
1669        assert_eq!(decision, PluginDecision::Continue);
1670    }
1671
1672    #[test]
1673    fn test_block_plugin_case_insensitive() {
1674        let plugin = BlockPlugin::default_block_ddl();
1675        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "drop table users");
1676        let decision = plugin.intercept(&mut ctx);
1677        assert!(matches!(decision, PluginDecision::Abort(_)));
1678    }
1679
1680    // ===== 端到端场景 =====
1681
1682    #[test]
1683    fn test_e2e_plugin_chain_workflow() {
1684        let chain = PluginChain::new();
1685        let sql_log = std::sync::Arc::new(SqlLogPlugin::new());
1686        let audit = std::sync::Arc::new(AuditPlugin::new());
1687        let slow = std::sync::Arc::new(SlowQueryPlugin::new(Duration::from_millis(100)));
1688
1689        // 由于 Plugin trait 是 Send + Sync 但需要 'static，我们用 Box::new 克隆实例
1690        // 这里简化：直接注册新的实例
1691        chain.register(Box::new(SqlLogPlugin::new()));
1692        chain.register(Box::new(AuditPlugin::new()));
1693        chain.register(Box::new(SlowQueryPlugin::new(Duration::from_millis(100))));
1694
1695        assert_eq!(chain.len(), 3);
1696
1697        // 1. 模拟查询前
1698        let mut before_ctx = PluginContext::new(
1699            ExecutionStage::BeforeQuery,
1700            "SELECT * FROM users WHERE id = ?",
1701        )
1702        .with_parameters(vec![Value::I64(1)]);
1703        let decision = chain.execute(&mut before_ctx);
1704        assert_eq!(decision, PluginDecision::Continue);
1705        // sql_log 应记录 1 条（before_query 阶段）
1706        // audit 不订阅 before_query，不应记录
1707        // slow_query 不订阅 before_query，不应记录
1708
1709        // 2. 模拟查询后（慢查询）
1710        let mut after_ctx = PluginContext::new(
1711            ExecutionStage::AfterQuery,
1712            "SELECT * FROM users WHERE id = ?",
1713        )
1714        .with_elapsed(Duration::from_millis(500));
1715        let decision = chain.execute(&mut after_ctx);
1716        assert_eq!(decision, PluginDecision::Continue);
1717
1718        // 3. 验证插件链中的插件名
1719        let names = chain.plugin_names();
1720        assert_eq!(names, vec!["sql_log", "audit", "slow_query"]);
1721
1722        let _ = (sql_log, audit, slow); // 避免未使用警告
1723    }
1724
1725    #[test]
1726    fn test_e2e_block_plugin_aborts_chain() {
1727        let chain = PluginChain::new();
1728        // 先注册 block，确保 Abort 中止后续插件
1729        chain.register(Box::new(BlockPlugin::default_block_ddl()));
1730        chain.register(Box::new(SqlLogPlugin::new()));
1731
1732        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1733        let decision = chain.execute(&mut ctx);
1734        assert!(matches!(decision, PluginDecision::Abort(_)));
1735    }
1736
1737    #[test]
1738    fn test_e2e_hydrate_scalar_count_query() {
1739        // 模拟 SELECT COUNT(*) AS cnt FROM users 的结果填充
1740        let mut row = RowData::empty();
1741        row.set("cnt", Value::I64(42));
1742
1743        let result = hydrate_single_scalar(&[row]).unwrap();
1744        assert_eq!(result, Value::I64(42));
1745    }
1746
1747    #[test]
1748    fn test_e2e_hydrate_object_user_query() {
1749        // 模拟 SELECT id, name, email FROM users
1750        let rows = vec![
1751            {
1752                let mut r = RowData::empty();
1753                r.set("id", Value::I64(1));
1754                r.set("name", Value::String("Alice".to_string()));
1755                r.set("email", Value::String("alice@example.com".to_string()));
1756                r
1757            },
1758            {
1759                let mut r = RowData::empty();
1760                r.set("id", Value::I64(2));
1761                r.set("name", Value::String("Bob".to_string()));
1762                r.set("email", Value::String("bob@example.com".to_string()));
1763                r
1764            },
1765        ];
1766
1767        let result = hydrate_object(&rows).unwrap();
1768        assert_eq!(result.len(), 2);
1769        assert_eq!(
1770            result[0].get("name"),
1771            Some(&Value::String("Alice".to_string()))
1772        );
1773    }
1774
1775    #[test]
1776    fn test_e2e_hydrate_array_multi_column() {
1777        let rows = vec![{
1778            let mut r = RowData::empty();
1779            r.set("a", Value::I64(1));
1780            r.set("b", Value::I64(2));
1781            r.set("c", Value::I64(3));
1782            r
1783        }];
1784
1785        let result = hydrate_array(&rows).unwrap();
1786        // 按列名排序：a, b, c
1787        assert_eq!(result[0], vec![Value::I64(1), Value::I64(2), Value::I64(3)]);
1788    }
1789}
sz_orm_core/hydration_plugin.rs

sz_orm_core/
hydration_plugin.rs
