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;
46use parking_lot::RwLock;
47use std::collections::HashMap;
48use std::time::{Duration, Instant};
49
50// ============================================================================
51// HydrationMode — 填充模式
52// ============================================================================
53
54/// Hydration 填充模式
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
56pub enum HydrationMode {
57    /// 每行 → `HashMap<String, Value>`（默认，对象模式）
58    #[default]
59    Object,
60    /// 每行 → `Vec<Value>`（按列顺序）
61    Array,
62    /// 每行 → Value（取第一列）
63    Scalar,
64    /// 唯一行 + 唯一列 → Value（聚合查询常用）
65    SingleScalar,
66    /// 每行的指定列 → `Vec<Value>`
67    Column,
68}
69
70impl HydrationMode {
71    /// 模式名称
72    pub fn name(&self) -> &'static str {
73        match self {
74            HydrationMode::Object => "object",
75            HydrationMode::Array => "array",
76            HydrationMode::Scalar => "scalar",
77            HydrationMode::SingleScalar => "single_scalar",
78            HydrationMode::Column => "column",
79        }
80    }
81}
82
83// ============================================================================
84// HydrationResult / hydrate 函数
85// ============================================================================
86
87/// Hydration 错误
88#[derive(Debug, Clone, PartialEq)]
89pub enum HydrationError {
90    /// SingleScalar 模式下行数不等于 1
91    SingleScalarRequiresSingleRow { 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();
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();
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();
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().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            .iter()
469            .map(|p| p.name().to_string())
470            .collect()
471    }
472
473    /// 清空插件链
474    pub fn clear(&self) {
475        self.plugins.write().clear();
476    }
477
478    /// 执行插件链
479    ///
480    /// 按 before 插件 → 原操作 → after 插件 的顺序执行。
481    /// 任意插件返回 `Abort` 中止整个链。
482    /// 任意插件返回 `Modified` 会修改 context 后继续。
483    /// 任意插件返回 `Skip` 跳过后续插件（但继续原操作）。
484    pub fn execute(&self, context: &mut PluginContext) -> PluginDecision {
485        let plugins = self.plugins.read();
486        let target_stages = [context.stage];
487
488        for plugin in plugins.iter() {
489            // 仅调用订阅了当前阶段的插件
490            if !plugin.stages().iter().any(|s| target_stages.contains(s)) {
491                continue;
492            }
493            match plugin.intercept(context) {
494                PluginDecision::Continue => continue,
495                PluginDecision::Skip => return PluginDecision::Skip,
496                PluginDecision::Modified { sql, parameters } => {
497                    context.sql = sql;
498                    context.parameters = parameters;
499                    continue;
500                }
501                PluginDecision::Abort(reason) => {
502                    return PluginDecision::Abort(reason);
503                }
504                PluginDecision::Kill => return PluginDecision::Kill,
505            }
506        }
507        PluginDecision::Continue
508    }
509}
510
511impl std::fmt::Debug for PluginChain {
512    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
513        let plugins = self.plugins.read();
514        let names: Vec<&str> = plugins.iter().map(|p| p.name()).collect();
515        f.debug_struct("PluginChain")
516            .field("plugins", &names)
517            .finish()
518    }
519}
520
521// ============================================================================
522// 内置插件：SqlLogPlugin
523// ============================================================================
524
525/// SQL 日志插件
526pub struct SqlLogPlugin {
527    logs: RwLock<Vec<String>>,
528}
529
530impl SqlLogPlugin {
531    pub fn new() -> Self {
532        Self {
533            logs: RwLock::new(Vec::new()),
534        }
535    }
536
537    pub fn logs(&self) -> Vec<String> {
538        self.logs.read().clone()
539    }
540
541    pub fn clear(&self) {
542        self.logs.write().clear();
543    }
544
545    pub fn count(&self) -> usize {
546        self.logs.read().len()
547    }
548}
549
550impl Default for SqlLogPlugin {
551    fn default() -> Self {
552        Self::new()
553    }
554}
555
556impl Plugin for SqlLogPlugin {
557    fn name(&self) -> &str {
558        "sql_log"
559    }
560
561    fn stages(&self) -> Vec<ExecutionStage> {
562        vec![
563            ExecutionStage::BeforeQuery,
564            ExecutionStage::AfterQuery,
565            ExecutionStage::BeforeUpdate,
566            ExecutionStage::AfterUpdate,
567        ]
568    }
569
570    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
571        let mut logs = self.logs.write();
572        let entry = match context.stage {
573            ExecutionStage::BeforeQuery => {
574                format!("[{}] QUERY: {}", context.stage.name(), context.sql)
575            }
576            ExecutionStage::AfterQuery => {
577                let elapsed_ms = context.elapsed.map(|d| d.as_millis()).unwrap_or(0);
578                format!(
579                    "[{}] QUERY ({}ms): {}",
580                    context.stage.name(),
581                    elapsed_ms,
582                    context.sql
583                )
584            }
585            ExecutionStage::BeforeUpdate => {
586                format!("[{}] UPDATE: {}", context.stage.name(), context.sql)
587            }
588            ExecutionStage::AfterUpdate => {
589                let rows = context.affected_rows.unwrap_or(0);
590                format!(
591                    "[{}] UPDATE ({} rows): {}",
592                    context.stage.name(),
593                    rows,
594                    context.sql
595                )
596            }
597            _ => return PluginDecision::Continue,
598        };
599        logs.push(entry);
600        PluginDecision::Continue
601    }
602}
603
604// ============================================================================
605// 内置插件：SlowQueryPlugin
606// ============================================================================
607
608/// 慢查询检测插件
609pub struct SlowQueryPlugin {
610    threshold: Duration,
611    slow_queries: RwLock<Vec<SlowQueryRecord>>,
612    /// 任务6：慢查询自动 kill 阈值（超过此阈值返回 Kill 决策）
613    kill_threshold: Option<Duration>,
614}
615
616/// 慢查询记录
617#[derive(Debug, Clone)]
618pub struct SlowQueryRecord {
619    pub sql: String,
620    pub elapsed: Duration,
621    pub threshold: Duration,
622}
623
624impl SlowQueryPlugin {
625    /// 创建插件，指定阈值
626    pub fn new(threshold: Duration) -> Self {
627        Self {
628            threshold,
629            slow_queries: RwLock::new(Vec::new()),
630            kill_threshold: None,
631        }
632    }
633
634    /// 默认 1 秒阈值
635    pub fn default_threshold() -> Self {
636        Self::new(Duration::from_secs(1))
637    }
638
639    /// 任务6：设置 kill 阈值（超过此阈值时返回 Kill 决策）
640    pub fn with_kill_threshold(mut self, kill_threshold: Duration) -> Self {
641        self.kill_threshold = Some(kill_threshold);
642        self
643    }
644
645    /// 获取所有慢查询记录
646    pub fn slow_queries(&self) -> Vec<SlowQueryRecord> {
647        self.slow_queries.read().clone()
648    }
649
650    /// 慢查询数量
651    pub fn count(&self) -> usize {
652        self.slow_queries.read().len()
653    }
654
655    /// 清空记录
656    pub fn clear(&self) {
657        self.slow_queries.write().clear();
658    }
659
660    /// 阈值
661    pub fn threshold(&self) -> Duration {
662        self.threshold
663    }
664}
665
666/// 脱敏 SQL 中的敏感值（password/passwd/secret/token = '...' → '***'）
667///
668/// 由于 sz-orm-core 未依赖 regex crate，使用手动扫描实现大小写不敏感匹配。
669/// 仅处理 `key='value'` 模式；其他形式（如函数参数）不做脱敏。
670fn mask_sql(sql: &str) -> String {
671    const SENSITIVE_KEYS: &[&str] = &["password", "passwd", "secret", "token"];
672
673    let lower = sql.to_ascii_lowercase();
674    let bytes = sql.as_bytes();
675    let lower_bytes = lower.as_bytes();
676    let mut result = String::with_capacity(sql.len());
677    let mut i = 0;
678
679    while i < bytes.len() {
680        let mut matched = false;
681        for kw in SENSITIVE_KEYS {
682            let kw_b = kw.as_bytes();
683            if i + kw_b.len() <= bytes.len() && &lower_bytes[i..i + kw_b.len()] == kw_b {
684                // 边界检查：前后字符不能是标识符字符（避免匹配 "passworded"）
685                let prev_ok = i == 0 || !is_ident_char(bytes[i - 1]);
686                let next = i + kw_b.len();
687                let next_ok = next >= bytes.len() || !is_ident_char(bytes[next]);
688                if !prev_ok || !next_ok {
689                    continue;
690                }
691                // 复制 keyword 原文（保留大小写）
692                result.push_str(&sql[i..next]);
693                i = next;
694                // 跳过空白
695                while i < bytes.len() && bytes[i].is_ascii_whitespace() {
696                    result.push(bytes[i] as char);
697                    i += 1;
698                }
699                // 期待 '='
700                if i < bytes.len() && bytes[i] == b'=' {
701                    result.push('=');
702                    i += 1;
703                    while i < bytes.len() && bytes[i].is_ascii_whitespace() {
704                        result.push(bytes[i] as char);
705                        i += 1;
706                    }
707                    // 期待单引号字符串 'value'
708                    if i < bytes.len() && bytes[i] == b'\'' {
709                        i += 1; // 跳过开始引号
710                        while i < bytes.len() && bytes[i] != b'\'' {
711                            i += 1;
712                        }
713                        if i < bytes.len() {
714                            i += 1; // 跳过结束引号
715                        }
716                        result.push_str("'***'");
717                    }
718                }
719                matched = true;
720                break;
721            }
722        }
723        if !matched {
724            let ch = sql[i..]
725                .chars()
726                .next()
727                .expect("i < bytes.len() guarantees non-empty slice");
728            result.push(ch);
729            i += ch.len_utf8();
730        }
731    }
732    result
733}
734
735/// ASCII 标识符字符判断（字母/数字/下划线）
736fn is_ident_char(b: u8) -> bool {
737    b.is_ascii_alphanumeric() || b == b'_'
738}
739
740impl Plugin for SlowQueryPlugin {
741    fn name(&self) -> &str {
742        "slow_query"
743    }
744
745    fn stages(&self) -> Vec<ExecutionStage> {
746        // 订阅 BeforeQuery：在执行前记录开始时间，并在超过 kill_threshold 时
747        // 通过 PluginDecision::Kill 中止查询（#9 修复：慢查询自动 Kill 生效）
748        vec![
749            ExecutionStage::BeforeQuery,
750            ExecutionStage::AfterQuery,
751            ExecutionStage::AfterUpdate,
752        ]
753    }
754
755    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
756        match context.stage {
757            ExecutionStage::BeforeQuery => {
758                // 执行前检查：若已记录的开始时间显示查询进行中超过 kill_threshold，
759                // 直接 Kill（理论上 BeforeQuery 阶段 elapsed 为 None，此处保留扩展点）
760                // 主要 Kill 逻辑放在 AfterQuery，因为 BeforeQuery 时 elapsed 尚未计算
761                if let Some(kill_threshold) = self.kill_threshold {
762                    if let Some(elapsed) = context.elapsed {
763                        if elapsed > kill_threshold {
764                            return PluginDecision::Kill;
765                        }
766                    }
767                }
768                PluginDecision::Continue
769            }
770            ExecutionStage::AfterQuery | ExecutionStage::AfterUpdate => {
771                if let Some(elapsed) = context.elapsed {
772                    if elapsed > self.threshold {
773                        // 记录前对 SQL 脱敏，避免敏感值（password/token 等）落入慢查询日志
774                        let masked_sql = mask_sql(&context.sql);
775                        let mut records = self.slow_queries.write();
776                        records.push(SlowQueryRecord {
777                            sql: masked_sql,
778                            elapsed,
779                            threshold: self.threshold,
780                        });
781                    }
782                    // #9 修复：超过 kill_threshold 时返回 Kill 决策
783                    // 这会触发插件链中止后续查询/事务，防止慢查询持续占用连接池资源
784                    if let Some(kill_threshold) = self.kill_threshold {
785                        if elapsed > kill_threshold {
786                            return PluginDecision::Kill;
787                        }
788                    }
789                }
790                PluginDecision::Continue
791            }
792            _ => PluginDecision::Continue,
793        }
794    }
795}
796
797// ============================================================================
798// 内置插件：AuditPlugin
799// ============================================================================
800
801/// 审计插件（记录所有写操作）
802pub struct AuditPlugin {
803    audit_log: RwLock<Vec<AuditRecord>>,
804}
805
806/// 审计记录
807#[derive(Debug, Clone)]
808pub struct AuditRecord {
809    pub stage: ExecutionStage,
810    pub sql: String,
811    pub affected_rows: Option<usize>,
812}
813
814impl AuditPlugin {
815    pub fn new() -> Self {
816        Self {
817            audit_log: RwLock::new(Vec::new()),
818        }
819    }
820
821    pub fn records(&self) -> Vec<AuditRecord> {
822        self.audit_log.read().clone()
823    }
824
825    pub fn count(&self) -> usize {
826        self.audit_log.read().len()
827    }
828
829    pub fn clear(&self) {
830        self.audit_log.write().clear();
831    }
832}
833
834impl Default for AuditPlugin {
835    fn default() -> Self {
836        Self::new()
837    }
838}
839
840impl Plugin for AuditPlugin {
841    fn name(&self) -> &str {
842        "audit"
843    }
844
845    fn stages(&self) -> Vec<ExecutionStage> {
846        vec![ExecutionStage::AfterUpdate]
847    }
848
849    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
850        let mut log = self.audit_log.write();
851        log.push(AuditRecord {
852            stage: context.stage,
853            sql: context.sql.clone(),
854            affected_rows: context.affected_rows,
855        });
856        PluginDecision::Continue
857    }
858}
859
860// ============================================================================
861// 内置插件：SqlRewritePlugin（演示 Modified 决策）
862// ============================================================================
863
864/// SQL 改写插件（演示 Modified 决策）
865///
866/// 在执行前将 `SELECT` 替换为 `SELECT /* hint */`，用于演示 SQL 改写能力。
867pub struct SqlRewritePlugin {
868    pattern: String,
869    replacement: String,
870}
871
872impl SqlRewritePlugin {
873    pub fn new(pattern: impl Into<String>, replacement: impl Into<String>) -> Self {
874        Self {
875            pattern: pattern.into(),
876            replacement: replacement.into(),
877        }
878    }
879}
880
881impl Plugin for SqlRewritePlugin {
882    fn name(&self) -> &str {
883        "sql_rewrite"
884    }
885
886    fn stages(&self) -> Vec<ExecutionStage> {
887        vec![ExecutionStage::BeforeQuery, ExecutionStage::BeforeUpdate]
888    }
889
890    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
891        if context.sql.contains(&self.pattern) {
892            let new_sql = context.sql.replace(&self.pattern, &self.replacement);
893            PluginDecision::Modified {
894                sql: new_sql,
895                parameters: context.parameters.clone(),
896            }
897        } else {
898            PluginDecision::Continue
899        }
900    }
901}
902
903// ============================================================================
904// 内置插件：BlockPlugin（演示 Abort 决策）
905// ============================================================================
906
907/// 阻断插件（演示 Abort 决策）
908///
909/// 拦截包含指定关键字的 SQL（如 DROP TABLE），返回 Abort。
910pub struct BlockPlugin {
911    blocked_keywords: Vec<String>,
912}
913
914impl BlockPlugin {
915    pub fn new(keywords: Vec<String>) -> Self {
916        Self {
917            blocked_keywords: keywords,
918        }
919    }
920
921    /// 默认阻断 DROP / TRUNCATE
922    pub fn default_block_ddl() -> Self {
923        Self::new(vec![
924            "DROP TABLE".to_string(),
925            "TRUNCATE".to_string(),
926            "DROP DATABASE".to_string(),
927        ])
928    }
929}
930
931impl Plugin for BlockPlugin {
932    fn name(&self) -> &str {
933        "block"
934    }
935
936    fn stages(&self) -> Vec<ExecutionStage> {
937        vec![ExecutionStage::BeforeUpdate, ExecutionStage::BeforeQuery]
938    }
939
940    fn intercept(&self, context: &mut PluginContext) -> PluginDecision {
941        let upper_sql = context.sql.to_uppercase();
942        for kw in &self.blocked_keywords {
943            if upper_sql.contains(&kw.to_uppercase()) {
944                return PluginDecision::Abort(format!(
945                    "blocked by BlockPlugin: SQL contains forbidden keyword '{}'",
946                    kw
947                ));
948            }
949        }
950        PluginDecision::Continue
951    }
952}
953
954// ============================================================================
955// 单元测试
956// ============================================================================
957
958#[cfg(test)]
959mod tests {
960    use super::*;
961
962    // ===== HydrationMode =====
963
964    #[test]
965    fn test_hydration_mode_default() {
966        assert_eq!(HydrationMode::default(), HydrationMode::Object);
967    }
968
969    #[test]
970    fn test_hydration_mode_name() {
971        assert_eq!(HydrationMode::Object.name(), "object");
972        assert_eq!(HydrationMode::Array.name(), "array");
973        assert_eq!(HydrationMode::Scalar.name(), "scalar");
974        assert_eq!(HydrationMode::SingleScalar.name(), "single_scalar");
975        assert_eq!(HydrationMode::Column.name(), "column");
976    }
977
978    // ===== hydrate_object =====
979
980    #[test]
981    fn test_hydrate_object_basic() {
982        let mut row = RowData::empty();
983        row.set("id", Value::I64(1));
984        row.set("name", Value::String("Alice".to_string()));
985
986        let result = hydrate_object(&[row]).unwrap();
987        assert_eq!(result.len(), 1);
988        assert_eq!(result[0].get("id"), Some(&Value::I64(1)));
989        assert_eq!(
990            result[0].get("name"),
991            Some(&Value::String("Alice".to_string()))
992        );
993    }
994
995    #[test]
996    fn test_hydrate_object_multiple_rows() {
997        let rows = vec![
998            {
999                let mut r = RowData::empty();
1000                r.set("id", Value::I64(1));
1001                r
1002            },
1003            {
1004                let mut r = RowData::empty();
1005                r.set("id", Value::I64(2));
1006                r
1007            },
1008        ];
1009
1010        let result = hydrate_object(&rows).unwrap();
1011        assert_eq!(result.len(), 2);
1012    }
1013
1014    #[test]
1015    fn test_hydrate_object_empty() {
1016        let result = hydrate_object(&[]).unwrap();
1017        assert!(result.is_empty());
1018    }
1019
1020    // ===== hydrate_array =====
1021
1022    #[test]
1023    fn test_hydrate_array_basic() {
1024        let mut row = RowData::empty();
1025        row.set("id", Value::I64(1));
1026        row.set("name", Value::String("Alice".to_string()));
1027
1028        let result = hydrate_array(&[row]).unwrap();
1029        assert_eq!(result.len(), 1);
1030        assert_eq!(result[0].len(), 2);
1031        // 列名排序后顺序：id, name
1032        assert_eq!(result[0][0], Value::I64(1));
1033        assert_eq!(result[0][1], Value::String("Alice".to_string()));
1034    }
1035
1036    #[test]
1037    fn test_hydrate_array_empty() {
1038        let result = hydrate_array(&[]).unwrap();
1039        assert!(result.is_empty());
1040    }
1041
1042    // ===== hydrate_scalar =====
1043
1044    #[test]
1045    fn test_hydrate_scalar_basic() {
1046        let mut row = RowData::empty();
1047        row.set("count", Value::I64(42));
1048
1049        let result = hydrate_scalar(&[row]).unwrap();
1050        assert_eq!(result.len(), 1);
1051        assert_eq!(result[0], Value::I64(42));
1052    }
1053
1054    #[test]
1055    fn test_hydrate_scalar_multiple_rows() {
1056        let rows = vec![
1057            {
1058                let mut r = RowData::empty();
1059                r.set("id", Value::I64(1));
1060                r
1061            },
1062            {
1063                let mut r = RowData::empty();
1064                r.set("id", Value::I64(2));
1065                r
1066            },
1067        ];
1068
1069        let result = hydrate_scalar(&rows).unwrap();
1070        assert_eq!(result, vec![Value::I64(1), Value::I64(2)]);
1071    }
1072
1073    #[test]
1074    fn test_hydrate_scalar_empty_row_error() {
1075        let row = RowData::empty();
1076        let err = hydrate_scalar(&[row]).unwrap_err();
1077        match err {
1078            HydrationError::EmptyRow => {}
1079            _ => panic!("expected EmptyRow error"),
1080        }
1081    }
1082
1083    // ===== hydrate_single_scalar =====
1084
1085    #[test]
1086    fn test_hydrate_single_scalar_ok() {
1087        let mut row = RowData::empty();
1088        row.set("total", Value::I64(100));
1089
1090        let result = hydrate_single_scalar(&[row]).unwrap();
1091        assert_eq!(result, Value::I64(100));
1092    }
1093
1094    #[test]
1095    fn test_hydrate_single_scalar_no_rows() {
1096        let err = hydrate_single_scalar(&[]).unwrap_err();
1097        match err {
1098            HydrationError::SingleScalarRequiresSingleRow { actual_rows } => {
1099                assert_eq!(actual_rows, 0)
1100            }
1101            _ => panic!("expected SingleScalarRequiresSingleRow"),
1102        }
1103    }
1104
1105    #[test]
1106    fn test_hydrate_single_scalar_too_many_rows() {
1107        let rows = vec![
1108            {
1109                let mut r = RowData::empty();
1110                r.set("id", Value::I64(1));
1111                r
1112            },
1113            {
1114                let mut r = RowData::empty();
1115                r.set("id", Value::I64(2));
1116                r
1117            },
1118        ];
1119        let err = hydrate_single_scalar(&rows).unwrap_err();
1120        match err {
1121            HydrationError::SingleScalarRequiresSingleRow { actual_rows } => {
1122                assert_eq!(actual_rows, 2)
1123            }
1124            _ => panic!("expected SingleScalarRequiresSingleRow"),
1125        }
1126    }
1127
1128    #[test]
1129    fn test_hydrate_single_scalar_empty_row() {
1130        let row = RowData::empty();
1131        let err = hydrate_single_scalar(&[row]).unwrap_err();
1132        match err {
1133            HydrationError::EmptyRow => {}
1134            _ => panic!("expected EmptyRow"),
1135        }
1136    }
1137
1138    // ===== hydrate_column =====
1139
1140    #[test]
1141    fn test_hydrate_column_basic() {
1142        let rows = vec![
1143            {
1144                let mut r = RowData::empty();
1145                r.set("id", Value::I64(1));
1146                r.set("name", Value::String("Alice".to_string()));
1147                r
1148            },
1149            {
1150                let mut r = RowData::empty();
1151                r.set("id", Value::I64(2));
1152                r.set("name", Value::String("Bob".to_string()));
1153                r
1154            },
1155        ];
1156
1157        let result = hydrate_column(&rows, "name").unwrap();
1158        assert_eq!(
1159            result,
1160            vec![
1161                Value::String("Alice".to_string()),
1162                Value::String("Bob".to_string()),
1163            ]
1164        );
1165    }
1166
1167    #[test]
1168    fn test_hydrate_column_missing() {
1169        let rows = vec![{
1170            let mut r = RowData::empty();
1171            r.set("id", Value::I64(1));
1172            r
1173        }];
1174
1175        let err = hydrate_column(&rows, "missing").unwrap_err();
1176        match err {
1177            HydrationError::ColumnNotFound { column } => assert_eq!(column, "missing"),
1178            _ => panic!("expected ColumnNotFound"),
1179        }
1180    }
1181
1182    #[test]
1183    fn test_hydrate_column_empty_rows() {
1184        let result = hydrate_column(&[], "name").unwrap();
1185        assert!(result.is_empty());
1186    }
1187
1188    // ===== 通用 hydrate 函数 =====
1189
1190    #[test]
1191    fn test_hydrate_scalar_mode() {
1192        let mut row = RowData::empty();
1193        row.set("count", Value::I64(42));
1194
1195        let result = hydrate(&[row], HydrationMode::Scalar).unwrap();
1196        assert_eq!(result, vec![Value::I64(42)]);
1197    }
1198
1199    #[test]
1200    fn test_hydrate_single_scalar_mode() {
1201        let mut row = RowData::empty();
1202        row.set("total", Value::I64(100));
1203
1204        let result = hydrate(&[row], HydrationMode::SingleScalar).unwrap();
1205        assert_eq!(result, vec![Value::I64(100)]);
1206    }
1207
1208    #[test]
1209    fn test_hydrate_column_mode() {
1210        let rows = vec![
1211            {
1212                let mut r = RowData::empty();
1213                r.set("id", Value::I64(1));
1214                r
1215            },
1216            {
1217                let mut r = RowData::empty();
1218                r.set("id", Value::I64(2));
1219                r
1220            },
1221        ];
1222
1223        let result = hydrate(&rows, HydrationMode::Column).unwrap();
1224        assert_eq!(result, vec![Value::I64(1), Value::I64(2)]);
1225    }
1226
1227    // ===== ExecutionStage =====
1228
1229    #[test]
1230    fn test_execution_stage_name() {
1231        assert_eq!(ExecutionStage::BeforeQuery.name(), "before_query");
1232        assert_eq!(ExecutionStage::AfterQuery.name(), "after_query");
1233        assert_eq!(ExecutionStage::BeforeUpdate.name(), "before_update");
1234        assert_eq!(ExecutionStage::AfterCommit.name(), "after_commit");
1235    }
1236
1237    #[test]
1238    fn test_execution_stage_is_before() {
1239        assert!(ExecutionStage::BeforeQuery.is_before());
1240        assert!(ExecutionStage::BeforeUpdate.is_before());
1241        assert!(ExecutionStage::BeforeCommit.is_before());
1242        assert!(ExecutionStage::BeforeRollback.is_before());
1243        assert!(!ExecutionStage::AfterQuery.is_before());
1244        assert!(!ExecutionStage::AfterUpdate.is_before());
1245    }
1246
1247    #[test]
1248    fn test_execution_stage_is_after() {
1249        assert!(ExecutionStage::AfterQuery.is_after());
1250        assert!(!ExecutionStage::BeforeQuery.is_after());
1251    }
1252
1253    #[test]
1254    fn test_execution_stage_is_query() {
1255        assert!(ExecutionStage::BeforeQuery.is_query());
1256        assert!(ExecutionStage::AfterQuery.is_query());
1257        assert!(!ExecutionStage::BeforeUpdate.is_query());
1258    }
1259
1260    #[test]
1261    fn test_execution_stage_is_update() {
1262        assert!(ExecutionStage::BeforeUpdate.is_update());
1263        assert!(ExecutionStage::AfterUpdate.is_update());
1264        assert!(!ExecutionStage::BeforeQuery.is_update());
1265    }
1266
1267    #[test]
1268    fn test_execution_stage_is_transaction() {
1269        assert!(ExecutionStage::BeforeCommit.is_transaction());
1270        assert!(ExecutionStage::AfterCommit.is_transaction());
1271        assert!(ExecutionStage::BeforeRollback.is_transaction());
1272        assert!(ExecutionStage::AfterRollback.is_transaction());
1273        assert!(!ExecutionStage::BeforeQuery.is_transaction());
1274    }
1275
1276    // ===== PluginContext =====
1277
1278    #[test]
1279    fn test_plugin_context_new() {
1280        let ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1281        assert_eq!(ctx.stage, ExecutionStage::BeforeQuery);
1282        assert_eq!(ctx.sql, "SELECT 1");
1283        assert!(ctx.parameters.is_empty());
1284        assert!(ctx.started_at.is_none());
1285        assert!(ctx.elapsed.is_none());
1286        assert!(ctx.affected_rows.is_none());
1287    }
1288
1289    #[test]
1290    fn test_plugin_context_with_parameters() {
1291        let ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT ?")
1292            .with_parameters(vec![Value::I64(1)]);
1293        assert_eq!(ctx.parameters.len(), 1);
1294    }
1295
1296    #[test]
1297    fn test_plugin_context_with_elapsed() {
1298        let ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1299            .with_elapsed(Duration::from_millis(50));
1300        assert_eq!(ctx.elapsed.unwrap().as_millis(), 50);
1301    }
1302
1303    #[test]
1304    fn test_plugin_context_with_affected_rows() {
1305        let ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users SET ...")
1306            .with_affected_rows(10);
1307        assert_eq!(ctx.affected_rows.unwrap(), 10);
1308    }
1309
1310    #[test]
1311    fn test_plugin_context_metadata() {
1312        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1313        ctx.set_metadata("user_id", Value::I64(42));
1314        assert_eq!(ctx.get_metadata("user_id"), Some(&Value::I64(42)));
1315        assert_eq!(ctx.get_metadata("missing"), None);
1316    }
1317
1318    // ===== PluginChain =====
1319
1320    #[test]
1321    fn test_plugin_chain_empty() {
1322        let chain = PluginChain::new();
1323        assert!(chain.is_empty());
1324        assert_eq!(chain.len(), 0);
1325    }
1326
1327    #[test]
1328    fn test_plugin_chain_register() {
1329        let chain = PluginChain::new();
1330        chain.register(Box::new(SqlLogPlugin::new()));
1331        assert_eq!(chain.len(), 1);
1332    }
1333
1334    #[test]
1335    fn test_plugin_chain_unregister() {
1336        let chain = PluginChain::new();
1337        chain.register(Box::new(SqlLogPlugin::new()));
1338        assert_eq!(chain.len(), 1);
1339
1340        let removed = chain.unregister("sql_log");
1341        assert!(removed);
1342        assert_eq!(chain.len(), 0);
1343    }
1344
1345    #[test]
1346    fn test_plugin_chain_unregister_missing() {
1347        let chain = PluginChain::new();
1348        let removed = chain.unregister("non_existent");
1349        assert!(!removed);
1350    }
1351
1352    #[test]
1353    fn test_plugin_chain_plugin_names() {
1354        let chain = PluginChain::new();
1355        chain.register(Box::new(SqlLogPlugin::new()));
1356        chain.register(Box::new(AuditPlugin::new()));
1357
1358        let names = chain.plugin_names();
1359        assert_eq!(names, vec!["sql_log", "audit"]);
1360    }
1361
1362    #[test]
1363    fn test_plugin_chain_clear() {
1364        let chain = PluginChain::new();
1365        chain.register(Box::new(SqlLogPlugin::new()));
1366        chain.clear();
1367        assert!(chain.is_empty());
1368    }
1369
1370    #[test]
1371    fn test_plugin_chain_insert_at() {
1372        let chain = PluginChain::new();
1373        chain.register(Box::new(SqlLogPlugin::new()));
1374        chain.insert_at(0, Box::new(AuditPlugin::new()));
1375
1376        let names = chain.plugin_names();
1377        assert_eq!(names, vec!["audit", "sql_log"]);
1378    }
1379
1380    #[test]
1381    fn test_plugin_chain_insert_at_end() {
1382        let chain = PluginChain::new();
1383        chain.register(Box::new(SqlLogPlugin::new()));
1384        chain.insert_at(99, Box::new(AuditPlugin::new()));
1385
1386        let names = chain.plugin_names();
1387        assert_eq!(names, vec!["sql_log", "audit"]);
1388    }
1389
1390    #[test]
1391    fn test_plugin_chain_execute_empty() {
1392        let chain = PluginChain::new();
1393        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1394        let decision = chain.execute(&mut ctx);
1395        assert_eq!(decision, PluginDecision::Continue);
1396    }
1397
1398    #[test]
1399    fn test_plugin_chain_execute_continue() {
1400        let chain = PluginChain::new();
1401        chain.register(Box::new(SqlLogPlugin::new()));
1402
1403        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1404        let decision = chain.execute(&mut ctx);
1405        assert_eq!(decision, PluginDecision::Continue);
1406    }
1407
1408    #[test]
1409    fn test_plugin_chain_execute_skip() {
1410        struct SkipPlugin;
1411        impl Plugin for SkipPlugin {
1412            fn name(&self) -> &str {
1413                "skip"
1414            }
1415            fn stages(&self) -> Vec<ExecutionStage> {
1416                vec![ExecutionStage::BeforeQuery]
1417            }
1418            fn intercept(&self, _ctx: &mut PluginContext) -> PluginDecision {
1419                PluginDecision::Skip
1420            }
1421        }
1422
1423        let chain = PluginChain::new();
1424        chain.register(Box::new(SkipPlugin));
1425
1426        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1427        let decision = chain.execute(&mut ctx);
1428        assert_eq!(decision, PluginDecision::Skip);
1429    }
1430
1431    #[test]
1432    fn test_plugin_chain_execute_modified() {
1433        let chain = PluginChain::new();
1434        chain.register(Box::new(SqlRewritePlugin::new(
1435            "SELECT",
1436            "SELECT /* hint */",
1437        )));
1438
1439        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1440        let decision = chain.execute(&mut ctx);
1441        assert_eq!(decision, PluginDecision::Continue);
1442        assert_eq!(ctx.sql, "SELECT /* hint */ 1");
1443    }
1444
1445    #[test]
1446    fn test_plugin_chain_execute_abort() {
1447        let chain = PluginChain::new();
1448        chain.register(Box::new(BlockPlugin::new(vec!["DROP".to_string()])));
1449
1450        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1451        let decision = chain.execute(&mut ctx);
1452        match decision {
1453            PluginDecision::Abort(reason) => assert!(reason.contains("DROP")),
1454            _ => panic!("expected Abort"),
1455        }
1456    }
1457
1458    #[test]
1459    fn test_plugin_chain_skip_unrelated_stages() {
1460        let chain = PluginChain::new();
1461        // SqlLogPlugin 订阅 Before/AfterQuery/Update，不应响应 BeforeCommit
1462        chain.register(Box::new(SqlLogPlugin::new()));
1463
1464        let mut ctx = PluginContext::new(ExecutionStage::BeforeCommit, "COMMIT");
1465        let decision = chain.execute(&mut ctx);
1466        assert_eq!(decision, PluginDecision::Continue);
1467        // 日志中不应有 BeforeCommit 记录
1468    }
1469
1470    // ===== SqlLogPlugin =====
1471
1472    #[test]
1473    fn test_sql_log_plugin_basic() {
1474        let plugin = SqlLogPlugin::new();
1475        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1476        let decision = plugin.intercept(&mut ctx);
1477        assert_eq!(decision, PluginDecision::Continue);
1478        assert_eq!(plugin.count(), 1);
1479    }
1480
1481    #[test]
1482    fn test_sql_log_plugin_after_query_with_elapsed() {
1483        let plugin = SqlLogPlugin::new();
1484        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1485            .with_elapsed(Duration::from_millis(50));
1486        let _ = plugin.intercept(&mut ctx);
1487
1488        let logs = plugin.logs();
1489        assert!(logs[0].contains("50ms"));
1490    }
1491
1492    #[test]
1493    fn test_sql_log_plugin_after_update_with_rows() {
1494        let plugin = SqlLogPlugin::new();
1495        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users SET ...")
1496            .with_affected_rows(10);
1497        let _ = plugin.intercept(&mut ctx);
1498
1499        let logs = plugin.logs();
1500        assert!(logs[0].contains("10 rows"));
1501    }
1502
1503    #[test]
1504    fn test_sql_log_plugin_clear() {
1505        let plugin = SqlLogPlugin::new();
1506        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT 1");
1507        let _ = plugin.intercept(&mut ctx);
1508        assert_eq!(plugin.count(), 1);
1509
1510        plugin.clear();
1511        assert_eq!(plugin.count(), 0);
1512    }
1513
1514    // ===== SlowQueryPlugin =====
1515
1516    #[test]
1517    fn test_slow_query_plugin_below_threshold() {
1518        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1519        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1520            .with_elapsed(Duration::from_millis(50));
1521        let _ = plugin.intercept(&mut ctx);
1522
1523        assert_eq!(plugin.count(), 0);
1524    }
1525
1526    #[test]
1527    fn test_slow_query_plugin_above_threshold() {
1528        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1529        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT * FROM big_table")
1530            .with_elapsed(Duration::from_millis(500));
1531        let _ = plugin.intercept(&mut ctx);
1532
1533        assert_eq!(plugin.count(), 1);
1534        let records = plugin.slow_queries();
1535        assert!(records[0].elapsed > records[0].threshold);
1536    }
1537
1538    #[test]
1539    fn test_slow_query_plugin_no_elapsed() {
1540        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1541        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1");
1542        let _ = plugin.intercept(&mut ctx);
1543
1544        assert_eq!(plugin.count(), 0);
1545    }
1546
1547    #[test]
1548    fn test_slow_query_plugin_clear() {
1549        let plugin = SlowQueryPlugin::new(Duration::from_millis(100));
1550        let mut ctx = PluginContext::new(ExecutionStage::AfterQuery, "SELECT 1")
1551            .with_elapsed(Duration::from_millis(200));
1552        let _ = plugin.intercept(&mut ctx);
1553        assert_eq!(plugin.count(), 1);
1554
1555        plugin.clear();
1556        assert_eq!(plugin.count(), 0);
1557    }
1558
1559    #[test]
1560    fn test_slow_query_plugin_default_threshold() {
1561        let plugin = SlowQueryPlugin::default_threshold();
1562        assert_eq!(plugin.threshold(), Duration::from_secs(1));
1563    }
1564
1565    // ===== AuditPlugin =====
1566
1567    #[test]
1568    fn test_audit_plugin_basic() {
1569        let plugin = AuditPlugin::new();
1570        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "INSERT INTO users ...")
1571            .with_affected_rows(1);
1572        let _ = plugin.intercept(&mut ctx);
1573
1574        assert_eq!(plugin.count(), 1);
1575        let records = plugin.records();
1576        assert_eq!(records[0].stage, ExecutionStage::AfterUpdate);
1577        assert_eq!(records[0].affected_rows, Some(1));
1578    }
1579
1580    #[test]
1581    fn test_audit_plugin_clear() {
1582        let plugin = AuditPlugin::new();
1583        let mut ctx = PluginContext::new(ExecutionStage::AfterUpdate, "UPDATE users");
1584        let _ = plugin.intercept(&mut ctx);
1585        assert_eq!(plugin.count(), 1);
1586
1587        plugin.clear();
1588        assert_eq!(plugin.count(), 0);
1589    }
1590
1591    // ===== SqlRewritePlugin =====
1592
1593    #[test]
1594    fn test_sql_rewrite_plugin_matches() {
1595        let plugin = SqlRewritePlugin::new("SELECT", "SELECT /* hint */");
1596        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT * FROM users");
1597        let decision = plugin.intercept(&mut ctx);
1598        match decision {
1599            PluginDecision::Modified { sql, .. } => {
1600                assert_eq!(sql, "SELECT /* hint */ * FROM users");
1601            }
1602            _ => panic!("expected Modified"),
1603        }
1604    }
1605
1606    #[test]
1607    fn test_sql_rewrite_plugin_no_match() {
1608        let plugin = SqlRewritePlugin::new("SELECT", "SELECT /* hint */");
1609        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SHOW TABLES");
1610        let decision = plugin.intercept(&mut ctx);
1611        assert_eq!(decision, PluginDecision::Continue);
1612    }
1613
1614    // ===== BlockPlugin =====
1615
1616    #[test]
1617    fn test_block_plugin_blocks_drop() {
1618        let plugin = BlockPlugin::default_block_ddl();
1619        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1620        let decision = plugin.intercept(&mut ctx);
1621        match decision {
1622            PluginDecision::Abort(reason) => {
1623                assert!(reason.contains("DROP TABLE"));
1624            }
1625            _ => panic!("expected Abort"),
1626        }
1627    }
1628
1629    #[test]
1630    fn test_block_plugin_blocks_truncate() {
1631        let plugin = BlockPlugin::default_block_ddl();
1632        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "TRUNCATE TABLE logs");
1633        let decision = plugin.intercept(&mut ctx);
1634        assert!(matches!(decision, PluginDecision::Abort(_)));
1635    }
1636
1637    #[test]
1638    fn test_block_plugin_allows_safe_sql() {
1639        let plugin = BlockPlugin::default_block_ddl();
1640        let mut ctx = PluginContext::new(ExecutionStage::BeforeQuery, "SELECT * FROM users");
1641        let decision = plugin.intercept(&mut ctx);
1642        assert_eq!(decision, PluginDecision::Continue);
1643    }
1644
1645    #[test]
1646    fn test_block_plugin_case_insensitive() {
1647        let plugin = BlockPlugin::default_block_ddl();
1648        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "drop table users");
1649        let decision = plugin.intercept(&mut ctx);
1650        assert!(matches!(decision, PluginDecision::Abort(_)));
1651    }
1652
1653    // ===== 端到端场景 =====
1654
1655    #[test]
1656    fn test_e2e_plugin_chain_workflow() {
1657        let chain = PluginChain::new();
1658        let sql_log = std::sync::Arc::new(SqlLogPlugin::new());
1659        let audit = std::sync::Arc::new(AuditPlugin::new());
1660        let slow = std::sync::Arc::new(SlowQueryPlugin::new(Duration::from_millis(100)));
1661
1662        // 由于 Plugin trait 是 Send + Sync 但需要 'static，我们用 Box::new 克隆实例
1663        // 这里简化：直接注册新的实例
1664        chain.register(Box::new(SqlLogPlugin::new()));
1665        chain.register(Box::new(AuditPlugin::new()));
1666        chain.register(Box::new(SlowQueryPlugin::new(Duration::from_millis(100))));
1667
1668        assert_eq!(chain.len(), 3);
1669
1670        // 1. 模拟查询前
1671        let mut before_ctx = PluginContext::new(
1672            ExecutionStage::BeforeQuery,
1673            "SELECT * FROM users WHERE id = ?",
1674        )
1675        .with_parameters(vec![Value::I64(1)]);
1676        let decision = chain.execute(&mut before_ctx);
1677        assert_eq!(decision, PluginDecision::Continue);
1678        // sql_log 应记录 1 条（before_query 阶段）
1679        // audit 不订阅 before_query，不应记录
1680        // slow_query 不订阅 before_query，不应记录
1681
1682        // 2. 模拟查询后（慢查询）
1683        let mut after_ctx = PluginContext::new(
1684            ExecutionStage::AfterQuery,
1685            "SELECT * FROM users WHERE id = ?",
1686        )
1687        .with_elapsed(Duration::from_millis(500));
1688        let decision = chain.execute(&mut after_ctx);
1689        assert_eq!(decision, PluginDecision::Continue);
1690
1691        // 3. 验证插件链中的插件名
1692        let names = chain.plugin_names();
1693        assert_eq!(names, vec!["sql_log", "audit", "slow_query"]);
1694
1695        let _ = (sql_log, audit, slow); // 避免未使用警告
1696    }
1697
1698    #[test]
1699    fn test_e2e_block_plugin_aborts_chain() {
1700        let chain = PluginChain::new();
1701        // 先注册 block，确保 Abort 中止后续插件
1702        chain.register(Box::new(BlockPlugin::default_block_ddl()));
1703        chain.register(Box::new(SqlLogPlugin::new()));
1704
1705        let mut ctx = PluginContext::new(ExecutionStage::BeforeUpdate, "DROP TABLE users");
1706        let decision = chain.execute(&mut ctx);
1707        assert!(matches!(decision, PluginDecision::Abort(_)));
1708    }
1709
1710    #[test]
1711    fn test_e2e_hydrate_scalar_count_query() {
1712        // 模拟 SELECT COUNT(*) AS cnt FROM users 的结果填充
1713        let mut row = RowData::empty();
1714        row.set("cnt", Value::I64(42));
1715
1716        let result = hydrate_single_scalar(&[row]).unwrap();
1717        assert_eq!(result, Value::I64(42));
1718    }
1719
1720    #[test]
1721    fn test_e2e_hydrate_object_user_query() {
1722        // 模拟 SELECT id, name, email FROM users
1723        let rows = vec![
1724            {
1725                let mut r = RowData::empty();
1726                r.set("id", Value::I64(1));
1727                r.set("name", Value::String("Alice".to_string()));
1728                r.set("email", Value::String("alice@example.com".to_string()));
1729                r
1730            },
1731            {
1732                let mut r = RowData::empty();
1733                r.set("id", Value::I64(2));
1734                r.set("name", Value::String("Bob".to_string()));
1735                r.set("email", Value::String("bob@example.com".to_string()));
1736                r
1737            },
1738        ];
1739
1740        let result = hydrate_object(&rows).unwrap();
1741        assert_eq!(result.len(), 2);
1742        assert_eq!(
1743            result[0].get("name"),
1744            Some(&Value::String("Alice".to_string()))
1745        );
1746    }
1747
1748    #[test]
1749    fn test_e2e_hydrate_array_multi_column() {
1750        let rows = vec![{
1751            let mut r = RowData::empty();
1752            r.set("a", Value::I64(1));
1753            r.set("b", Value::I64(2));
1754            r.set("c", Value::I64(3));
1755            r
1756        }];
1757
1758        let result = hydrate_array(&rows).unwrap();
1759        // 按列名排序：a, b, c
1760        assert_eq!(result[0], vec![Value::I64(1), Value::I64(2), Value::I64(3)]);
1761    }
1762}
sz_orm_core/hydration_plugin.rs

sz_orm_core/
hydration_plugin.rs
