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

sz_orm_core/
hydration_plugin.rs
