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

sz_orm_query/
hydration_plugin.rs
