Skip to main content

sz_orm_ai/
nl2sql.rs

1//! NL→SQL(自然语言转 SQL)模块
2//!
3//! 提供将自然语言查询转换为 SQL 语句的能力,支持:
4//! - 内存模拟引擎 [`SimpleNl2SqlEngine`](基于规则匹配,无需外部 API)
5//! - 真实 LLM 引擎 [`OpenAINl2SqlEngine`](调用 OpenAI 兼容 API,需 `real` feature)
6//!
7//! 所有生成的 SQL 均经过安全验证:只允许 SELECT 查询,并检测注入风险。
8
9use async_trait::async_trait;
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12
13use crate::safety;
14
15// ==================== 数据结构 ====================
16
17/// NL→SQL 查询结果
18#[derive(Debug, Clone)]
19pub struct SqlQuery {
20    /// 生成的 SQL(使用 $1, $2 等参数化占位符)
21    pub sql: String,
22    /// 生成过程的自然语言解释
23    pub explanation: String,
24    /// 置信度(0.0 ~ 1.0)
25    pub confidence: f32,
26}
27
28/// Schema 上下文,描述数据库中的表和列信息
29#[derive(Debug, Clone, Default)]
30pub struct SchemaContext {
31    pub tables: Vec<TableInfo>,
32}
33
34/// 表信息
35#[derive(Debug, Clone)]
36pub struct TableInfo {
37    pub name: String,
38    pub columns: Vec<ColumnInfo>,
39}
40
41/// 列信息
42#[derive(Debug, Clone)]
43pub struct ColumnInfo {
44    pub name: String,
45    pub data_type: String,
46    pub nullable: bool,
47    pub is_primary_key: bool,
48}
49
50// ==================== 错误类型 ====================
51
52/// NL→SQL 相关错误
53#[derive(Debug, Error)]
54pub enum Nl2SqlError {
55    /// 无法解析自然语言查询
56    #[error("Invalid query: {0}")]
57    InvalidQuery(String),
58    /// Schema 信息不足(表/列不存在)
59    #[error("Schema error: {0}")]
60    SchemaError(String),
61    /// SQL 安全验证失败
62    #[error("Safety error: {0}")]
63    SafetyError(String),
64    /// SQL 生成失败
65    #[error("Generation error: {0}")]
66    GenerationError(String),
67    /// API 调用错误(仅 real feature)
68    #[error("API error (status {0}): {1}")]
69    ApiError(u16, String),
70    /// 网络错误(仅 real feature)
71    #[error("Network error: {0}")]
72    NetworkError(String),
73    /// 配置错误
74    #[error("Config error: {0}")]
75    ConfigError(String),
76}
77
78// ==================== Trait 定义 ====================
79
80/// NL→SQL 引擎 trait
81///
82/// 所有 NL→SQL 实现必须实现此 trait,以保证一致的接口。
83#[async_trait]
84pub trait Nl2SqlEngine: Send + Sync {
85    /// 将自然语言查询转换为 SQL
86    ///
87    /// # 参数
88    /// - `nl_query`: 自然语言查询(如 "show all users where age > 25")
89    /// - `schema`: 数据库 schema 上下文
90    ///
91    /// # 返回值
92    /// - `Ok(SqlQuery)`: 生成的 SQL 查询
93    /// - `Err(Nl2SqlError)`: 转换失败
94    async fn generate(
95        &self,
96        nl_query: &str,
97        schema: &SchemaContext,
98    ) -> Result<SqlQuery, Nl2SqlError>;
99
100    /// 验证生成的 SQL 是否安全可用
101    ///
102    /// 执行以下检查:
103    /// - 只允许 SELECT 语句
104    /// - 无 SQL 注入风险
105    async fn validate(&self, query: &SqlQuery) -> Result<bool, Nl2SqlError>;
106}
107
108// ==================== SimpleNl2SqlEngine ====================
109
110/// 基于规则匹配的 NL→SQL 引擎(内存模拟,无需 LLM API)
111///
112/// 通过关键词匹配和模式识别,将常见自然语言查询转换为 SQL。
113/// 支持以下查询模式:
114/// - 简单 SELECT(`SELECT *` / `SELECT col1, col2`)
115/// - COUNT / COUNT DISTINCT
116/// - 聚合函数(SUM, AVG, MIN, MAX)
117/// - WHERE 条件(=, >, <, >=, <=, !=, LIKE)
118/// - ORDER BY(ASC / DESC)
119/// - GROUP BY
120/// - JOIN(通过外键名称推断关联)
121/// - LIMIT
122///
123/// 所有生成的 SQL 使用 `$1`, `$2` 等参数化占位符防止注入。
124///
125/// # 示例
126///
127/// ```ignore
128/// use sz_orm_ai::nl2sql::{SimpleNl2SqlEngine, Nl2SqlEngine, SchemaContext, TableInfo, ColumnInfo};
129///
130/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
131/// let engine = SimpleNl2SqlEngine::new();
132/// let schema = SchemaContext {
133///     tables: vec![TableInfo {
134///         name: "users".into(),
135///         columns: vec![
136///             ColumnInfo { name: "id".into(), data_type: "INTEGER".into(), nullable: false, is_primary_key: true },
137///             ColumnInfo { name: "name".into(), data_type: "TEXT".into(), nullable: true, is_primary_key: false },
138///         ],
139///     }],
140/// };
141/// let result = engine.generate("show all users", &schema).await?;
142/// assert_eq!(result.sql, "SELECT * FROM users");
143/// # Ok(())
144/// # }
145/// ```
146pub struct SimpleNl2SqlEngine {
147    /// 表别名映射(如 "user" → "users")
148    aliases: std::collections::HashMap<String, String>,
149}
150
151impl Default for SimpleNl2SqlEngine {
152    fn default() -> Self {
153        Self::new()
154    }
155}
156
157impl SimpleNl2SqlEngine {
158    pub fn new() -> Self {
159        Self {
160            aliases: std::collections::HashMap::new(),
161        }
162    }
163
164    /// 注册表别名(例如 `with_alias("user", "users")`)
165    pub fn with_alias(mut self, alias: &str, table: &str) -> Self {
166        self.aliases
167            .insert(alias.to_lowercase(), table.to_lowercase());
168        self
169    }
170
171    /// 在 schema 中查找表名(匹配原名或别名)
172    fn find_table<'a>(&self, query: &str, schema: &'a SchemaContext) -> Option<&'a TableInfo> {
173        let lower = query.to_lowercase();
174
175        // 按名称在 query 中出现的顺序打分
176        let mut best: Option<&TableInfo> = None;
177        let mut best_pos: Option<usize> = None;
178
179        for table in &schema.tables {
180            let name = table.name.to_lowercase();
181            if let Some(pos) = lower.find(&name) {
182                let is_better = match best_pos {
183                    None => true,
184                    Some(best) => pos < best,
185                };
186                if is_better {
187                    best = Some(table);
188                    best_pos = Some(pos);
189                }
190            }
191        }
192
193        // 检查别名:遍历所有别名,看查询中是否包含别名
194        for (alias, target_table_name) in &self.aliases {
195            if let Some(pos) = lower.find(alias) {
196                // 找到对应表
197                if let Some(table) = schema
198                    .tables
199                    .iter()
200                    .find(|t| t.name.to_lowercase() == *target_table_name)
201                {
202                    let is_better = match best_pos {
203                        None => true,
204                        Some(best) => pos < best,
205                    };
206                    if is_better {
207                        best = Some(table);
208                        best_pos = Some(pos);
209                    }
210                }
211            }
212        }
213
214        best
215    }
216
217    /// 只在 SELECT/NL 部分提取列名(排除 WHERE/ORDER BY/GROUP BY 后的部分)
218    fn extract_columns_from_nl(query: &str, table: &TableInfo) -> Vec<String> {
219        let lower = query.to_lowercase();
220        // 截取到第一个子句关键字之前
221        let select_part = if let Some(pos) = lower.find(" where ") {
222            &query[..pos]
223        } else if let Some(pos) = lower.find(" having ") {
224            &query[..pos]
225        } else if let Some(pos) = lower.find(" order by") {
226            &query[..pos]
227        } else if let Some(pos) = lower.find(" sorted by") {
228            &query[..pos]
229        } else if let Some(pos) = lower.find(" sort by") {
230            &query[..pos]
231        } else if let Some(pos) = lower.find(" ordered by") {
232            &query[..pos]
233        } else if let Some(pos) = lower.find(" group by") {
234            &query[..pos]
235        } else {
236            query
237        };
238
239        Self::extract_columns(select_part, table)
240    }
241
242    /// 确认 query 中提到的表名(不含 FROM 子句的上下文检测)
243    fn find_mentioned_table<'a>(
244        &self,
245        query: &str,
246        schema: &'a SchemaContext,
247    ) -> Option<&'a TableInfo> {
248        self.find_table(query, schema)
249    }
250
251    /// 在 query 中查找提及的所有表名(用于 JOIN 检测)
252    fn find_all_tables<'a>(&self, query: &str, schema: &'a SchemaContext) -> Vec<&'a TableInfo> {
253        let lower = query.to_lowercase();
254        let mut found = Vec::new();
255
256        for table in &schema.tables {
257            let name = table.name.to_lowercase();
258            if lower.contains(&name) {
259                found.push(table);
260            }
261        }
262
263        found
264    }
265
266    /// 检查列名是否存在于 schema 的指定表中
267    fn column_exists(col: &str, table: &TableInfo) -> bool {
268        let col_lower = col.to_lowercase();
269        table
270            .columns
271            .iter()
272            .any(|c| c.name.to_lowercase() == col_lower)
273    }
274
275    /// 从 query 中提取列名列表(匹配 schema 中的列)
276    fn extract_columns(query: &str, table: &TableInfo) -> Vec<String> {
277        let lower = query.to_lowercase();
278        let mut columns = Vec::new();
279
280        for col in &table.columns {
281            let col_lower = col.name.to_lowercase();
282            // 避免匹配短列名导致误匹配(如 "id")
283            if col_lower.len() >= 2
284                && (lower.contains(&col_lower) || lower.contains(&format!(" {} ", col_lower)))
285            {
286                // 确认不是 from 子句后的表名列
287                if !columns.contains(&col.name) {
288                    columns.push(col.name.clone());
289                }
290            }
291        }
292
293        columns
294    }
295
296    /// 尝试寻找两个表之间的外键关联列
297    fn find_join_columns(t1: &TableInfo, t2: &TableInfo) -> Option<(String, String)> {
298        // 模式 1: t2 中有 t1_name + _id → t2.t1_name_id = t1.id
299        for col1 in &t1.columns {
300            if col1.is_primary_key {
301                let expected_fk = format!("{}_{}", t2.name, col1.name).to_lowercase();
302                for col2 in &t2.columns {
303                    if col2.name.to_lowercase() == expected_fk {
304                        return Some((
305                            format!("{}.{}", t1.name, col1.name),
306                            format!("{}.{}", t2.name, col2.name),
307                        ));
308                    }
309                }
310            }
311        }
312        // 模式 2: t1 中有 t2_name + _id → t1.t2_name_id = t2.id
313        for col2 in &t2.columns {
314            if col2.is_primary_key {
315                let expected_fk = format!("{}_{}", t1.name, col2.name).to_lowercase();
316                for col1 in &t1.columns {
317                    if col1.name.to_lowercase() == expected_fk {
318                        return Some((
319                            format!("{}.{}", t1.name, col1.name),
320                            format!("{}.{}", t2.name, col2.name),
321                        ));
322                    }
323                }
324            }
325        }
326        None
327    }
328
329    /// 从 WHERE 条件文本中提取参数化条件
330    fn parse_conditions(where_text: &str, table: &TableInfo) -> (Vec<String>, Vec<String>) {
331        let text = where_text.trim();
332        let mut conditions = Vec::new();
333        let mut params = Vec::new();
334        let param_idx = 1;
335
336        // 处理单个条件表达式
337        let expr = text;
338        if expr.is_empty() {
339            return (conditions, params);
340        }
341
342        let operators = [
343            (">=", ">="),
344            ("<=", "<="),
345            ("!=", "!="),
346            ("=", "="),
347            (">", ">"),
348            ("<", "<"),
349        ];
350
351        let mut resolved = false;
352        for (op_str, sql_op) in &operators {
353            if let Some(pos) = expr.find(op_str) {
354                let field = expr[..pos].trim();
355                let raw_val = expr[pos + op_str.len()..].trim();
356
357                if !field.is_empty() && !raw_val.is_empty() {
358                    // 检查 raw_val 是否为列名
359                    if Self::column_exists(raw_val, table) {
360                        conditions.push(format!("{} {} {}", field, sql_op, raw_val));
361                    } else {
362                        let param = format!("${}", param_idx);
363                        conditions.push(format!("{} {} {}", field, sql_op, param));
364                        params.push(raw_val.to_string());
365                    }
366                    resolved = true;
367                }
368                break;
369            }
370        }
371
372        // 处理 LIKE / contains
373        if !resolved {
374            let lower_expr = expr.to_lowercase();
375            if let Some(pos) = lower_expr.find(" like ") {
376                let field = expr[..pos].trim();
377                let raw_val = expr[pos + 6..].trim();
378                if !field.is_empty() && !raw_val.is_empty() {
379                    let param = "$1".to_string();
380                    conditions.push(format!("{} LIKE {}", field, param));
381                    params.push(raw_val.to_string());
382                    resolved = true;
383                }
384            } else if let Some(pos) = lower_expr.find(" contains ") {
385                let field = expr[..pos].trim();
386                let raw_val = expr[pos + 10..].trim();
387                if !field.is_empty() && !raw_val.is_empty() {
388                    let param = "$1".to_string();
389                    conditions.push(format!("{} LIKE '%' || {} || '%'", field, param));
390                    params.push(raw_val.to_string());
391                    resolved = true;
392                }
393            }
394        }
395
396        // 未解析的条件原样保留
397        if !resolved && !expr.is_empty() {
398            conditions.push(expr.to_string());
399        }
400
401        (conditions, params)
402    }
403}
404
405#[async_trait]
406impl Nl2SqlEngine for SimpleNl2SqlEngine {
407    async fn generate(
408        &self,
409        nl_query: &str,
410        schema: &SchemaContext,
411    ) -> Result<SqlQuery, Nl2SqlError> {
412        if nl_query.trim().is_empty() {
413            return Err(Nl2SqlError::InvalidQuery("自然语言查询不能为空".into()));
414        }
415
416        if schema.tables.is_empty() {
417            return Err(Nl2SqlError::SchemaError("Schema 中未定义任何表".into()));
418        }
419
420        let lower = nl_query.to_lowercase();
421        let table = self.find_mentioned_table(nl_query, schema).ok_or_else(|| {
422            Nl2SqlError::SchemaError(format!(
423                "无法在查询中识别表名,schema 包含: {}",
424                schema
425                    .tables
426                    .iter()
427                    .map(|t| t.name.as_str())
428                    .collect::<Vec<_>>()
429                    .join(", ")
430            ))
431        })?;
432
433        // ============ 识别查询类型 ============
434
435        // 是否是 COUNT 查询
436        let is_count =
437            lower.contains("count ") || lower.contains("how many") || lower.contains("number of");
438        let is_distinct = lower.contains("distinct");
439        let has_aggregation = lower.contains("sum ")
440            || lower.contains("total ")
441            || lower.contains("avg ")
442            || lower.contains("average ")
443            || lower.contains("mean ")
444            || lower.contains("min ")
445            || lower.contains("minimum ")
446            || lower.contains("max ")
447            || lower.contains("maximum ")
448            || lower.contains("highest");
449
450        // 提取 ORDER BY、GROUP BY、LIMIT、JOIN
451        let has_order = lower.contains("order by")
452            || lower.contains("sort by")
453            || lower.contains("ordered by")
454            || lower.contains("sorted by");
455        let has_group = lower.contains("group by");
456        let has_limit =
457            lower.contains("limit") || lower.contains("top ") || lower.contains("first ");
458        let has_join =
459            lower.contains(" join ") || lower.contains("combine with") || lower.contains("with ");
460
461        // ============ 提取列(仅从 SELECT/NL 部分提取,排除子句关键字后的部分) ============
462        let columns = Self::extract_columns_from_nl(nl_query, table);
463
464        // ============ 提取 WHERE 条件 ============
465        let where_text = if let Some(pos) = lower.find(" where ") {
466            Some(nl_query[pos + 7..].trim())
467        } else if let Some(pos) = lower.find(" having ") {
468            Some(nl_query[pos + 8..].trim())
469        } else if let Some(pos) = lower.find(" with ") {
470            // "with" 可能表示条件或 JOIN
471            let after = nl_query[pos + 5..].trim();
472            // 如果 with 后跟列名和操作符,视为条件
473            if table
474                .columns
475                .iter()
476                .any(|c| after.to_lowercase().starts_with(&c.name.to_lowercase()))
477            {
478                Some(after)
479            } else {
480                None
481            }
482        } else {
483            None
484        };
485
486        let (conditions, params) = match where_text {
487            Some(text) => {
488                // 去掉 order by / group by / limit 部分
489                let clean_text = if let Some(pos) = text.to_lowercase().find(" order by") {
490                    &text[..pos]
491                } else if let Some(pos) = text.to_lowercase().find(" group by") {
492                    &text[..pos]
493                } else if let Some(pos) = text.to_lowercase().find(" limit") {
494                    &text[..pos]
495                } else {
496                    text
497                };
498                Self::parse_conditions(clean_text, table)
499            }
500            None => (Vec::new(), Vec::new()),
501        };
502
503        // ============ 提取 ORDER BY ============
504        let order_clause: Option<String> = if has_order {
505            let (after_text, is_desc) = if let Some(pos) = lower.find("order by") {
506                let txt = nl_query[pos + 8..].trim().to_string();
507                let d = txt.to_lowercase().contains("desc")
508                    || txt.to_lowercase().contains("descending");
509                (txt, d)
510            } else if let Some(pos) = lower.find("sort by") {
511                let txt = nl_query[pos + 7..].trim().to_string();
512                let d = txt.to_lowercase().contains("desc")
513                    || txt.to_lowercase().contains("descending");
514                (txt, d)
515            } else if let Some(pos) = lower.find("sorted by") {
516                let txt = nl_query[pos + 9..].trim().to_string();
517                let d = txt.to_lowercase().contains("desc")
518                    || txt.to_lowercase().contains("descending");
519                (txt, d)
520            } else if let Some(pos) = lower.find("ordered by") {
521                let txt = nl_query[pos + 10..].trim().to_string();
522                let d = txt.to_lowercase().contains("desc")
523                    || txt.to_lowercase().contains("descending");
524                (txt, d)
525            } else {
526                (String::new(), false)
527            };
528
529            let field = after_text
530                .split([' ', ','])
531                .next()
532                .unwrap_or("")
533                .trim()
534                .to_string();
535            if field.is_empty() {
536                None
537            } else if is_desc {
538                Some(format!(" ORDER BY {} DESC", field))
539            } else {
540                Some(format!(" ORDER BY {} ASC", field))
541            }
542        } else {
543            None
544        };
545
546        // ============ 提取 GROUP BY ============
547        let group_clause = if has_group {
548            let after = if let Some(pos) = lower.find("group by") {
549                &nl_query[pos + 8..]
550            } else {
551                ""
552            };
553            let group_cols: Vec<String> = after
554                .split(',')
555                .filter_map(|s| {
556                    let first = s.split_whitespace().next()?;
557                    if first.is_empty() {
558                        None
559                    } else {
560                        Some(first.to_string())
561                    }
562                })
563                .collect();
564            if group_cols.is_empty() {
565                None
566            } else {
567                Some(format!(" GROUP BY {}", group_cols.join(", ")))
568            }
569        } else {
570            None
571        };
572
573        // ============ 提取 LIMIT ============
574        let limit_val: Option<usize> = if has_limit {
575            let limit_text = if let Some(pos) = lower.find("limit ") {
576                let after = &nl_query[pos + 6..].trim();
577                after.split_whitespace().next()
578            } else if let Some(pos) = lower.find("top ") {
579                let after = &nl_query[pos + 4..].trim();
580                after.split_whitespace().next()
581            } else if let Some(pos) = lower.find("first ") {
582                let after = &nl_query[pos + 6..].trim();
583                after.split_whitespace().next()
584            } else {
585                None
586            };
587            limit_text.and_then(|s| s.parse::<usize>().ok())
588        } else {
589            None
590        };
591
592        // ============ 提取 JOIN ============
593        let join_tables = if has_join {
594            let all_tables = self.find_all_tables(nl_query, schema);
595            let others: Vec<&&TableInfo> =
596                all_tables.iter().filter(|t| t.name != table.name).collect();
597            others.into_iter().copied().collect()
598        } else {
599            Vec::new()
600        };
601
602        // ============ 构建 SQL ============
603
604        // 构建 SELECT 子句
605        let select_clause = if is_count {
606            if is_distinct && !columns.is_empty() {
607                format!("SELECT COUNT(DISTINCT {})", columns[0])
608            } else {
609                "SELECT COUNT(*)".to_string()
610            }
611        } else if has_aggregation && !columns.is_empty() {
612            // 检测具体聚合函数
613            let agg_func = if lower.contains("sum ") || lower.contains("total ") {
614                "SUM"
615            } else if lower.contains("avg ")
616                || lower.contains("average ")
617                || lower.contains("mean ")
618            {
619                "AVG"
620            } else if lower.contains("min ") || lower.contains("minimum ") {
621                "MIN"
622            } else if lower.contains("max ")
623                || lower.contains("maximum ")
624                || lower.contains("highest")
625            {
626                "MAX"
627            } else {
628                "SUM"
629            };
630            if group_clause.is_some() {
631                // 聚合 + group by
632                format!("SELECT {}, {}({})", columns.join(", "), agg_func, agg_func)
633            } else if !columns.is_empty() {
634                format!("SELECT {}({})", agg_func, columns[0])
635            } else {
636                format!("SELECT {}(*)", agg_func)
637            }
638        } else if !columns.is_empty() {
639            format!("SELECT {}", columns.join(", "))
640        } else {
641            "SELECT *".to_string()
642        };
643
644        // 修正 group by 场景的 select
645        let select_clause = if group_clause.is_some() && has_aggregation {
646            let agg_func = if lower.contains("sum ") || lower.contains("total ") {
647                "SUM"
648            } else if lower.contains("avg ")
649                || lower.contains("average ")
650                || lower.contains("mean ")
651            {
652                "AVG"
653            } else if lower.contains("min ") || lower.contains("minimum ") {
654                "MIN"
655            } else if lower.contains("max ")
656                || lower.contains("maximum ")
657                || lower.contains("highest")
658            {
659                "MAX"
660            } else {
661                "COUNT"
662            };
663
664            if !columns.is_empty() {
665                if columns.len() >= 2 {
666                    format!("SELECT {}, {}({})", columns[1], agg_func, columns[1])
667                } else {
668                    format!("SELECT {}, {}({})", columns[0], agg_func, columns[0])
669                }
670            } else {
671                select_clause
672            }
673        } else {
674            select_clause
675        };
676
677        // 构建 FROM 子句 + JOIN
678        let mut from_clause = format!(" FROM {}", table.name);
679        let mut join_descriptions = Vec::new();
680        for join_table in &join_tables {
681            if let Some((left, right)) = Self::find_join_columns(table, join_table) {
682                from_clause.push_str(&format!(
683                    " JOIN {} ON {} = {}",
684                    join_table.name, left, right
685                ));
686                join_descriptions.push(format!("{} ON {} = {}", join_table.name, left, right));
687            }
688        }
689
690        // 构建 WHERE 子句(参数化)
691        let where_clause = if conditions.is_empty() {
692            String::new()
693        } else {
694            // SAFETY: conditions 来自 NL2SQL 内部解析(field/op/param 由 schema 和 NLP 解析生成),非直接用户输入;AI 生成 SQL 需在执行前经 sql-validator 校验
695            format!(" WHERE {}", conditions.join(" AND "))
696        };
697
698        // 组合 SQL
699        let sql = format!(
700            "{}{}{}{}{}{}",
701            select_clause,
702            from_clause,
703            where_clause,
704            order_clause.as_deref().unwrap_or(""),
705            group_clause.as_deref().unwrap_or(""),
706            limit_val
707                .map(|v| format!(" LIMIT {}", v))
708                .unwrap_or_default(),
709        );
710
711        // 安全验证
712        if !safety::validate_select_only(&sql) {
713            return Err(Nl2SqlError::SafetyError(
714                "生成的 SQL 不是 SELECT 查询".into(),
715            ));
716        }
717        if !safety::validate_no_injection(&sql) {
718            return Err(Nl2SqlError::SafetyError("生成的 SQL 包含注入风险".into()));
719        }
720        let sql = safety::sanitize_sql(&sql);
721
722        // 构建解释
723        let mut explanation_parts = Vec::new();
724        explanation_parts.push(format!("查询 {} 表", table.name));
725        if !columns.is_empty() {
726            explanation_parts.push(format!("列: {}", columns.join(", ")));
727        }
728        if is_count {
729            explanation_parts.push("统计数量".to_string());
730        }
731        if !conditions.is_empty() {
732            let cond_desc: Vec<String> = conditions
733                .iter()
734                .enumerate()
735                .map(|(i, cond)| {
736                    if i < params.len() {
737                        cond.replace(&format!("${}", i + 1), &format!("'{}'", params[i]))
738                    } else {
739                        cond.clone()
740                    }
741                })
742                .collect();
743            explanation_parts.push(format!("条件: {}", cond_desc.join(", ")));
744        }
745        if !params.is_empty() {
746            explanation_parts.push(format!(
747                "参数: [{}]",
748                params
749                    .iter()
750                    .map(|p| format!("'{}'", p))
751                    .collect::<Vec<_>>()
752                    .join(", ")
753            ));
754        }
755        let explanation = explanation_parts.join(";");
756
757        // 计算置信度
758        // 简单的规则:明确的模式匹配给高置信度
759        let mut confidence = 0.7;
760        if !conditions.is_empty() || is_count {
761            confidence = 0.8;
762        }
763        if !columns.is_empty() && !conditions.is_empty() {
764            confidence = 0.9;
765        }
766
767        Ok(SqlQuery {
768            sql,
769            explanation,
770            confidence,
771        })
772    }
773
774    async fn validate(&self, query: &SqlQuery) -> Result<bool, Nl2SqlError> {
775        if !safety::validate_select_only(&query.sql) {
776            return Ok(false);
777        }
778        if !safety::validate_no_injection(&query.sql) {
779            return Ok(false);
780        }
781        if query.confidence < 0.0 || query.confidence > 1.0 {
782            return Err(Nl2SqlError::GenerationError(format!(
783                "confidence 必须在 0.0~1.0 范围内,实际 {}",
784                query.confidence
785            )));
786        }
787        Ok(true)
788    }
789}
790
791// ==================== OpenAINl2SqlEngine ====================
792
793/// OpenAI 兼容的 NL→SQL 引擎(调用 LLM API)
794///
795/// 仅在启用 `real` feature 时编译。
796/// 调用 OpenAI 兼容的 `/v1/chat/completions` 接口生成 SQL。
797///
798/// 生成的 SQL 经过安全验证:只允许 SELECT,并检测注入风险。
799///
800/// # 用法
801///
802/// ```ignore
803/// use sz_orm_ai::nl2sql::{Nl2SqlEngine, OpenAINl2SqlEngine, SchemaContext};
804///
805/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
806/// let engine = OpenAINl2SqlEngine::new("sk-xxxx")
807///     .with_model("gpt-4o-mini");
808/// let schema = SchemaContext::default();
809/// let result = engine.generate("show all users", &schema).await?;
810/// println!("SQL: {}", result.sql);
811/// # Ok(())
812/// # }
813/// ```
814#[cfg(feature = "real")]
815pub struct OpenAINl2SqlEngine {
816    /// API 基础地址(默认 `https://api.openai.com/v1`)
817    api_base: String,
818    /// API Key(Bearer token)
819    api_key: String,
820    /// 模型名称(默认 `gpt-4o-mini`)
821    model: String,
822    /// HTTP 客户端
823    http_client: reqwest::Client,
824}
825
826#[cfg(feature = "real")]
827impl OpenAINl2SqlEngine {
828    /// 默认 API 基础地址
829    const DEFAULT_API_BASE: &'static str = "https://api.openai.com/v1";
830    /// 默认模型
831    const DEFAULT_MODEL: &'static str = "gpt-4o-mini";
832
833    /// 创建客户端实例
834    pub fn new(api_key: impl Into<String>) -> Self {
835        Self {
836            api_base: Self::DEFAULT_API_BASE.to_string(),
837            api_key: api_key.into(),
838            model: Self::DEFAULT_MODEL.to_string(),
839            http_client: reqwest::Client::new(),
840        }
841    }
842
843    /// 设置 API base URL
844    pub fn with_api_base(mut self, api_base: impl Into<String>) -> Self {
845        self.api_base = api_base.into();
846        self
847    }
848
849    /// 设置模型名称
850    pub fn with_model(mut self, model: impl Into<String>) -> Self {
851        self.model = model.into();
852        self
853    }
854
855    /// 校验 API key 非空
856    fn ensure_api_key(&self) -> Result<(), Nl2SqlError> {
857        if self.api_key.is_empty() {
858            return Err(Nl2SqlError::ConfigError(
859                "API key 为空,无法调用 OpenAI API".into(),
860            ));
861        }
862        Ok(())
863    }
864
865    /// 构建 system prompt(包含 schema 信息和 SQL 生成规范)
866    fn build_system_prompt(schema: &SchemaContext) -> String {
867        let mut prompt = String::from(
868            "You are a SQL generator. Given a database schema and a natural language query, ",
869        );
870        prompt.push_str("generate a valid SQL SELECT statement.\n\n");
871        prompt.push_str("Rules:\n");
872        prompt.push_str("- Only generate SELECT statements (no INSERT, UPDATE, DELETE, DROP, ALTER, TRUNCATE)\n");
873        prompt.push_str("- Use parameterized placeholders ($1, $2, ...) for all values to prevent SQL injection\n");
874        prompt.push_str("- If the query is ambiguous, choose the most likely interpretation\n");
875        prompt
876            .push_str("- Return ONLY the SQL statement, no explanation or markdown formatting\n\n");
877
878        prompt.push_str("Database Schema:\n");
879        for table in &schema.tables {
880            prompt.push_str(&format!("CREATE TABLE {} (\n", table.name));
881            for col in &table.columns {
882                prompt.push_str(&format!(
883                    "  {} {} {} {},\n",
884                    col.name,
885                    col.data_type,
886                    if col.is_primary_key {
887                        "PRIMARY KEY"
888                    } else {
889                        ""
890                    },
891                    if col.nullable { "NULL" } else { "NOT NULL" },
892                ));
893            }
894            prompt.push_str(");\n\n");
895        }
896
897        // v3.3.0 M4:ai-nl2sql-enhanced feature 启用时追加关系信息 + 增强指令
898        #[cfg(feature = "ai-nl2sql-enhanced")]
899        {
900            Self::append_relationship_info(&mut prompt, schema);
901            Self::append_enhanced_instructions(&mut prompt);
902        }
903
904        prompt
905    }
906
907    /// 追加表间关系信息(外键推断 + JOIN 关系)
908    #[cfg(feature = "ai-nl2sql-enhanced")]
909    fn append_relationship_info(prompt: &mut String, schema: &SchemaContext) {
910        prompt.push_str("Relationships (inferred from column naming conventions):\n");
911        let mut found_relations = false;
912        for table in &schema.tables {
913            for col in &table.columns {
914                if col.name.ends_with("_id") && !col.is_primary_key {
915                    let ref_table = col.name.trim_end_matches("_id");
916                    let ref_table_plural = if ref_table.ends_with('s') {
917                        ref_table.to_string()
918                    } else {
919                        format!("{}s", ref_table)
920                    };
921                    if schema.tables.iter().any(|t| t.name == ref_table_plural) {
922                        prompt.push_str(&format!(
923                            "  {}.{} -> {}(id) (foreign key)\n",
924                            table.name, col.name, ref_table_plural
925                        ));
926                        found_relations = true;
927                    } else if schema.tables.iter().any(|t| t.name == ref_table) {
928                        prompt.push_str(&format!(
929                            "  {}.{} -> {}(id) (foreign key)\n",
930                            table.name, col.name, ref_table
931                        ));
932                        found_relations = true;
933                    }
934                }
935            }
936        }
937        if !found_relations {
938            prompt.push_str("  (no foreign keys detected from naming conventions)\n");
939        }
940        prompt.push('\n');
941    }
942
943    /// 追加增强 SQL 生成指令(多表 JOIN + 聚合 + 子查询 + 排序 + 分页)
944    #[cfg(feature = "ai-nl2sql-enhanced")]
945    fn append_enhanced_instructions(prompt: &mut String) {
946        prompt.push_str("Enhanced Query Capabilities:\n");
947        prompt.push_str("- Multi-table JOIN: use INNER/LEFT/RIGHT JOIN with ON conditions when query spans multiple tables\n");
948        prompt.push_str("- Aggregation: support COUNT/SUM/AVG/MIN/MAX with GROUP BY and HAVING\n");
949        prompt.push_str("- Subquery: support subqueries in WHERE (IN/EXISTS) and FROM clauses\n");
950        prompt.push_str("- Ordering: support ORDER BY with ASC/DESC for multiple columns\n");
951        prompt.push_str("- Pagination: support LIMIT and OFFSET for result pagination\n");
952        prompt.push_str("- DISTINCT: support SELECT DISTINCT for deduplication\n");
953        prompt.push_str("- Always use table-qualified column names (e.g., users.name) in JOINs to avoid ambiguity\n\n");
954    }
955}
956
957#[cfg(feature = "real")]
958#[derive(Serialize)]
959struct ChatCompletionRequest<'a> {
960    model: &'a str,
961    messages: Vec<Message<'a>>,
962    temperature: f32,
963    max_tokens: u32,
964}
965
966#[cfg(feature = "real")]
967#[derive(Serialize)]
968struct Message<'a> {
969    role: &'a str,
970    content: String,
971}
972
973#[cfg(feature = "real")]
974#[derive(Deserialize)]
975struct ChatCompletionResponse {
976    choices: Vec<Choice>,
977}
978
979#[cfg(feature = "real")]
980#[derive(Deserialize)]
981struct Choice {
982    message: ChoiceMessage,
983    #[allow(dead_code)]
984    finish_reason: Option<String>,
985}
986
987#[cfg(feature = "real")]
988#[derive(Deserialize)]
989struct ChoiceMessage {
990    content: Option<String>,
991}
992
993#[cfg(feature = "real")]
994#[async_trait]
995impl Nl2SqlEngine for OpenAINl2SqlEngine {
996    async fn generate(
997        &self,
998        nl_query: &str,
999        schema: &SchemaContext,
1000    ) -> Result<SqlQuery, Nl2SqlError> {
1001        self.ensure_api_key()?;
1002
1003        if nl_query.trim().is_empty() {
1004            return Err(Nl2SqlError::InvalidQuery("自然语言查询不能为空".into()));
1005        }
1006
1007        let system_prompt = Self::build_system_prompt(schema);
1008        let user_content = format!(
1009            "Given the schema above, generate a SQL query for: {}",
1010            nl_query
1011        );
1012        // v3.3.0 M4:ai-nl2sql-enhanced feature 启用时对 LLM 请求脱敏
1013        let user_message = if cfg!(feature = "ai-nl2sql-enhanced") {
1014            crate::sql_sanitizer::SqlSanitizer::sanitize(&user_content)
1015        } else {
1016            user_content
1017        };
1018
1019        let body = ChatCompletionRequest {
1020            model: &self.model,
1021            messages: vec![
1022                Message {
1023                    role: "system",
1024                    content: system_prompt,
1025                },
1026                Message {
1027                    role: "user",
1028                    content: user_message,
1029                },
1030            ],
1031            temperature: 0.1,
1032            max_tokens: 500,
1033        };
1034
1035        let url = format!("{}/chat/completions", self.api_base);
1036        let resp = self
1037            .http_client
1038            .post(&url)
1039            .bearer_auth(&self.api_key)
1040            .json(&body)
1041            .send()
1042            .await
1043            .map_err(|e| Nl2SqlError::NetworkError(e.to_string()))?;
1044
1045        let status = resp.status().as_u16();
1046        if !resp.status().is_success() {
1047            let message = resp.text().await.unwrap_or_default();
1048            return Err(Nl2SqlError::ApiError(status, message));
1049        }
1050
1051        let parsed: ChatCompletionResponse = resp
1052            .json()
1053            .await
1054            .map_err(|e| Nl2SqlError::NetworkError(e.to_string()))?;
1055
1056        let raw_sql = parsed
1057            .choices
1058            .into_iter()
1059            .next()
1060            .and_then(|c| c.message.content)
1061            .ok_or_else(|| Nl2SqlError::ApiError(status, "API 返回空响应".into()))?;
1062
1063        // 清理 LLM 返回的 SQL(去掉 markdown 代码块标记)
1064        let cleaned_sql = clean_llm_sql_output(&raw_sql);
1065
1066        let query = SqlQuery {
1067            sql: cleaned_sql.clone(),
1068            explanation: format!("由 {} 模型根据自然语言查询生成", self.model),
1069            confidence: 0.8,
1070        };
1071
1072        // 安全验证
1073        if !safety::validate_select_only(&query.sql) {
1074            return Err(Nl2SqlError::SafetyError(
1075                "生成的 SQL 不是 SELECT 查询,已被拦截".into(),
1076            ));
1077        }
1078        if !safety::validate_no_injection(&query.sql) {
1079            return Err(Nl2SqlError::SafetyError(
1080                "生成的 SQL 包含注入风险,已被拦截".into(),
1081            ));
1082        }
1083
1084        Ok(query)
1085    }
1086
1087    async fn validate(&self, query: &SqlQuery) -> Result<bool, Nl2SqlError> {
1088        if !safety::validate_select_only(&query.sql) {
1089            return Ok(false);
1090        }
1091        if !safety::validate_no_injection(&query.sql) {
1092            return Ok(false);
1093        }
1094        if query.confidence < 0.0 || query.confidence > 1.0 {
1095            return Err(Nl2SqlError::GenerationError(format!(
1096                "confidence 必须在 0.0~1.0 范围内,实际 {}",
1097                query.confidence
1098            )));
1099        }
1100        Ok(true)
1101    }
1102}
1103
1104/// 清理 LLM 输出的 SQL(移除 markdown 代码块标记和前后空白)
1105#[cfg(feature = "real")]
1106fn clean_llm_sql_output(raw: &str) -> String {
1107    let trimmed = raw.trim();
1108    // 移除 ```sql ... ``` 或 ``` ... ``` 包装
1109    if trimmed.starts_with("```") {
1110        let content = trimmed.trim_start_matches('`');
1111        if let Some(end) = content.rfind("```") {
1112            let inner = content[..end].trim();
1113            // 移除 "sql" 语言标记(如 ```sql\n...)
1114            return inner
1115                .strip_prefix("sql")
1116                .unwrap_or(inner)
1117                .trim()
1118                .to_string();
1119        }
1120        return content.trim().to_string();
1121    }
1122    trimmed.to_string()
1123}
1124
1125// ==================== 查询优化提示 ====================
1126
1127/// 优化建议的严重级别
1128#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1129pub enum HintSeverity {
1130    /// 信息级:可选的优化建议
1131    Info,
1132    /// 警告级:可能影响性能
1133    Warning,
1134    /// 严重级:强烈建议修改
1135    Critical,
1136}
1137
1138impl HintSeverity {
1139    /// 转换为字符串
1140    pub fn as_str(&self) -> &'static str {
1141        match self {
1142            HintSeverity::Info => "INFO",
1143            HintSeverity::Warning => "WARNING",
1144            HintSeverity::Critical => "CRITICAL",
1145        }
1146    }
1147}
1148
1149/// 单条查询优化建议
1150#[derive(Debug, Clone)]
1151pub struct QueryOptimizationHint {
1152    /// 建议标题
1153    pub title: String,
1154    /// 详细描述
1155    pub description: String,
1156    /// 严重级别
1157    pub severity: HintSeverity,
1158    /// 优化后的 SQL 建议(可选)
1159    pub suggested_sql: Option<String>,
1160}
1161
1162impl QueryOptimizationHint {
1163    /// 创建一条信息级建议
1164    pub fn info(title: impl Into<String>, description: impl Into<String>) -> Self {
1165        Self {
1166            title: title.into(),
1167            description: description.into(),
1168            severity: HintSeverity::Info,
1169            suggested_sql: None,
1170        }
1171    }
1172
1173    /// 创建一条警告级建议
1174    pub fn warning(title: impl Into<String>, description: impl Into<String>) -> Self {
1175        Self {
1176            title: title.into(),
1177            description: description.into(),
1178            severity: HintSeverity::Warning,
1179            suggested_sql: None,
1180        }
1181    }
1182
1183    /// 创建一条严重级建议
1184    pub fn critical(title: impl Into<String>, description: impl Into<String>) -> Self {
1185        Self {
1186            title: title.into(),
1187            description: description.into(),
1188            severity: HintSeverity::Critical,
1189            suggested_sql: None,
1190        }
1191    }
1192
1193    /// 附加优化后的 SQL 建议
1194    pub fn with_suggested_sql(mut self, sql: impl Into<String>) -> Self {
1195        self.suggested_sql = Some(sql.into());
1196        self
1197    }
1198}
1199
1200/// 查询分析结果
1201#[derive(Debug, Clone)]
1202pub struct QueryAnalysis {
1203    /// 原始 SQL
1204    pub original_sql: String,
1205    /// 所有优化建议
1206    pub hints: Vec<QueryOptimizationHint>,
1207    /// 预估的 SQL 复杂度评分(0-100,越高越复杂)
1208    pub complexity_score: u32,
1209    /// 检测到的表名列表
1210    pub detected_tables: Vec<String>,
1211    /// 是否包含 WHERE 子句
1212    pub has_where: bool,
1213    /// 是否包含 LIMIT 子句
1214    pub has_limit: bool,
1215    /// 是否包含 JOIN
1216    pub has_join: bool,
1217    /// 是否包含子查询
1218    pub has_subquery: bool,
1219    /// 是否使用了 SELECT *
1220    pub uses_select_star: bool,
1221}
1222
1223impl QueryAnalysis {
1224    /// 返回严重级别为 Critical 的建议数量
1225    pub fn critical_count(&self) -> usize {
1226        self.hints
1227            .iter()
1228            .filter(|h| h.severity == HintSeverity::Critical)
1229            .count()
1230    }
1231
1232    /// 返回严重级别为 Warning 的建议数量
1233    pub fn warning_count(&self) -> usize {
1234        self.hints
1235            .iter()
1236            .filter(|h| h.severity == HintSeverity::Warning)
1237            .count()
1238    }
1239
1240    /// 是否存在任何建议
1241    pub fn has_hints(&self) -> bool {
1242        !self.hints.is_empty()
1243    }
1244}
1245
1246/// SQL 查询优化分析器
1247///
1248/// 基于规则分析 SQL 查询,生成优化建议。
1249/// 不依赖外部 LLM API,纯规则匹配,适用于离线场景。
1250pub struct QueryOptimizer {
1251    /// 是否检测 SELECT *
1252    pub check_select_star: bool,
1253    /// 是否检测缺失的 LIMIT
1254    pub check_missing_limit: bool,
1255    /// 是否检测缺失的 WHERE
1256    pub check_missing_where: bool,
1257    /// LIMIT 建议的默认行数
1258    pub default_limit: usize,
1259    /// 复杂度评分中 JOIN 的权重
1260    pub join_weight: u32,
1261    /// 复杂度评分中子查询的权重
1262    pub subquery_weight: u32,
1263    /// 复杂度评分中 WHERE 条件的权重
1264    pub where_weight: u32,
1265}
1266
1267impl Default for QueryOptimizer {
1268    fn default() -> Self {
1269        Self {
1270            check_select_star: true,
1271            check_missing_limit: true,
1272            check_missing_where: true,
1273            default_limit: 100,
1274            join_weight: 15,
1275            subquery_weight: 20,
1276            where_weight: 5,
1277        }
1278    }
1279}
1280
1281impl QueryOptimizer {
1282    /// 创建新的查询优化分析器
1283    pub fn new() -> Self {
1284        Self::default()
1285    }
1286
1287    /// 设置默认 LIMIT 行数
1288    pub fn with_default_limit(mut self, limit: usize) -> Self {
1289        self.default_limit = limit;
1290        self
1291    }
1292
1293    /// 禁用 SELECT * 检测
1294    pub fn disable_select_star_check(mut self) -> Self {
1295        self.check_select_star = false;
1296        self
1297    }
1298
1299    /// 禁用缺失 LIMIT 检测
1300    pub fn disable_missing_limit_check(mut self) -> Self {
1301        self.check_missing_limit = false;
1302        self
1303    }
1304
1305    /// 禁用缺失 WHERE 检测
1306    pub fn disable_missing_where_check(mut self) -> Self {
1307        self.check_missing_where = false;
1308        self
1309    }
1310
1311    /// 分析 SQL 查询并生成优化建议
1312    ///
1313    /// # 参数
1314    /// - `sql`: 要分析的 SQL 查询语句
1315    /// - `schema`: 数据库 schema 上下文(用于检测表名和索引)
1316    pub fn analyze(&self, sql: &str, schema: &SchemaContext) -> QueryAnalysis {
1317        let normalized = Self::normalize_sql(sql);
1318        let lower = normalized.to_lowercase();
1319
1320        let uses_select_star = self.detect_select_star(&lower);
1321        let has_where = self.detect_where(&lower);
1322        let has_limit = self.detect_limit(&lower);
1323        let has_join = self.detect_join(&lower);
1324        let has_subquery = self.detect_subquery(&lower);
1325        let detected_tables = self.extract_tables(&lower, schema);
1326
1327        let mut hints = Vec::new();
1328
1329        // 检测 SELECT *
1330        if self.check_select_star && uses_select_star {
1331            let columns_hint = self.suggest_columns(&detected_tables, schema);
1332            let mut hint = QueryOptimizationHint::warning(
1333                "避免使用 SELECT *",
1334                "SELECT * 会返回所有列,可能导致不必要的数据传输和内存消耗。建议显式指定所需列。",
1335            );
1336            if !columns_hint.is_empty() {
1337                hint = hint.with_suggested_sql(columns_hint);
1338            }
1339            hints.push(hint);
1340        }
1341
1342        // 检测缺失的 WHERE
1343        if self.check_missing_where && !has_where {
1344            hints.push(QueryOptimizationHint::critical(
1345                "缺少 WHERE 子句",
1346                "查询没有 WHERE 条件,将扫描全表。对于大表这会导致严重的性能问题。",
1347            ));
1348        }
1349
1350        // 检测缺失的 LIMIT
1351        if self.check_missing_limit && !has_limit {
1352            let suggested = self.add_limit_suggestion(&normalized, self.default_limit);
1353            hints.push(
1354                QueryOptimizationHint::warning(
1355                    "缺少 LIMIT 子句",
1356                    format!(
1357                        "查询没有 LIMIT 限制,可能返回大量数据。建议添加 LIMIT {}。",
1358                        self.default_limit
1359                    ),
1360                )
1361                .with_suggested_sql(suggested),
1362            );
1363        }
1364
1365        // 检测多表 JOIN 无索引建议
1366        if has_join {
1367            let join_count = Self::count_joins(&lower);
1368            if join_count >= 3 {
1369                hints.push(QueryOptimizationHint::critical(
1370                    "JOIN 数量过多",
1371                    format!(
1372                        "查询包含 {} 个 JOIN,可能导致性能下降。建议拆分为多个查询或使用临时表。",
1373                        join_count
1374                    ),
1375                ));
1376            } else if join_count >= 1 {
1377                hints.push(QueryOptimizationHint::info(
1378                    "JOIN 查询建议",
1379                    "确保 JOIN 条件涉及的列已建立索引,避免嵌套循环扫描。",
1380                ));
1381            }
1382        }
1383
1384        // 检测子查询
1385        if has_subquery {
1386            let subquery_count = Self::count_subqueries(&lower);
1387            if subquery_count >= 2 {
1388                hints.push(QueryOptimizationHint::warning(
1389                    "子查询嵌套过深",
1390                    format!(
1391                        "查询包含 {} 个子查询,建议考虑使用 JOIN 重写以提高性能。",
1392                        subquery_count
1393                    ),
1394                ));
1395            }
1396        }
1397
1398        // 检测 LIKE '%...' 前缀通配符
1399        if Self::has_leading_wildcard_like(&lower) {
1400            hints.push(QueryOptimizationHint::warning(
1401                "LIKE 使用前缀通配符",
1402                "LIKE '%keyword' 无法使用索引,会导致全表扫描。如可能,使用 LIKE 'keyword%' 或全文索引。",
1403            ));
1404        }
1405
1406        // 检测 OR 条件(可能导致无法使用索引)
1407        if Self::count_or_conditions(&lower) >= 3 {
1408            hints.push(QueryOptimizationHint::info(
1409                "多个 OR 条件",
1410                "多个 OR 条件可能导致无法有效使用索引。考虑使用 UNION ALL 重写。",
1411            ));
1412        }
1413
1414        // 检测 ORDER BY 无 LIMIT
1415        if Self::has_order_by_without_limit(&lower) {
1416            hints.push(QueryOptimizationHint::warning(
1417                "ORDER BY 无 LIMIT",
1418                "ORDER BY 无 LIMIT 时需要排序全部数据,可能消耗大量内存。建议添加 LIMIT。",
1419            ));
1420        }
1421
1422        // 检测 COUNT(*) 建议使用 COUNT(1)
1423        if lower.contains("count(*)") {
1424            hints.push(QueryOptimizationHint::info(
1425                "考虑使用 COUNT(1)",
1426                "某些数据库中 COUNT(1) 比 COUNT(*) 略快(虽然现代优化器通常已优化)。",
1427            ));
1428        }
1429
1430        // 检测缺失索引建议(基于 WHERE 条件)
1431        if has_where {
1432            let where_columns = self.extract_where_columns(&lower, schema);
1433            for col in &where_columns {
1434                if !self.column_has_index(col, schema) {
1435                    hints.push(QueryOptimizationHint::info(
1436                        format!("建议为列 {} 添加索引", col),
1437                        format!(
1438                            "WHERE 条件中使用了列 {},但该列似乎没有索引。添加索引可提高查询速度。",
1439                            col
1440                        ),
1441                    ));
1442                }
1443            }
1444        }
1445
1446        // 计算复杂度评分
1447        let complexity_score =
1448            self.calculate_complexity(has_join, has_subquery, has_where, &detected_tables, &lower);
1449
1450        QueryAnalysis {
1451            original_sql: sql.to_string(),
1452            hints,
1453            complexity_score,
1454            detected_tables,
1455            has_where,
1456            has_limit,
1457            has_join,
1458            has_subquery,
1459            uses_select_star,
1460        }
1461    }
1462
1463    /// 生成优化报告文本
1464    pub fn format_report(analysis: &QueryAnalysis) -> String {
1465        let mut report = String::new();
1466        report.push_str("=== SQL 查询优化分析报告 ===\n\n");
1467        report.push_str(&format!("原始 SQL: {}\n", analysis.original_sql));
1468        report.push_str(&format!("复杂度评分: {}/100\n", analysis.complexity_score));
1469        report.push_str(&format!(
1470            "检测到的表: {}\n",
1471            if analysis.detected_tables.is_empty() {
1472                "无".to_string()
1473            } else {
1474                analysis.detected_tables.join(", ")
1475            }
1476        ));
1477        report.push_str(&format!("包含 WHERE: {}\n", analysis.has_where)); // SAFETY: 报告文本拼接,非 SQL 执行
1478        report.push_str(&format!("包含 LIMIT: {}\n", analysis.has_limit));
1479        report.push_str(&format!("包含 JOIN: {}\n", analysis.has_join));
1480        report.push_str(&format!("包含子查询: {}\n", analysis.has_subquery));
1481        report.push_str(&format!("使用 SELECT *: {}\n\n", analysis.uses_select_star));
1482
1483        if analysis.hints.is_empty() {
1484            report.push_str("✓ 未发现优化建议,查询看起来良好。\n");
1485        } else {
1486            report.push_str(&format!(
1487                "共 {} 条优化建议({} 严重,{} 警告):\n\n",
1488                analysis.hints.len(),
1489                analysis.critical_count(),
1490                analysis.warning_count()
1491            ));
1492            for (i, hint) in analysis.hints.iter().enumerate() {
1493                report.push_str(&format!(
1494                    "{}. [{}] {}\n   {}\n",
1495                    i + 1,
1496                    hint.severity.as_str(),
1497                    hint.title,
1498                    hint.description
1499                ));
1500                if let Some(ref sql) = hint.suggested_sql {
1501                    report.push_str(&format!("   建议SQL: {}\n", sql));
1502                }
1503                report.push('\n');
1504            }
1505        }
1506
1507        report
1508    }
1509
1510    // ---- 内部辅助方法 ----
1511
1512    /// 规范化 SQL(去除多余空白、换行)
1513    fn normalize_sql(sql: &str) -> String {
1514        sql.split_whitespace().collect::<Vec<_>>().join(" ")
1515    }
1516
1517    /// 检测 SELECT *
1518    fn detect_select_star(&self, lower: &str) -> bool {
1519        lower.contains("select *") || lower.contains("select  *")
1520    }
1521
1522    /// 检测 WHERE 子句
1523    fn detect_where(&self, lower: &str) -> bool {
1524        lower.contains(" where ")
1525    }
1526
1527    /// 检测 LIMIT 子句
1528    fn detect_limit(&self, lower: &str) -> bool {
1529        lower.contains(" limit ")
1530    }
1531
1532    /// 检测 JOIN
1533    fn detect_join(&self, lower: &str) -> bool {
1534        lower.contains(" join ")
1535            || lower.contains(" inner join ")
1536            || lower.contains(" left join ")
1537            || lower.contains(" right join ")
1538            || lower.contains(" full join ")
1539            || lower.contains(" cross join ")
1540    }
1541
1542    /// 检测子查询(括号内的 SELECT)
1543    fn detect_subquery(&self, lower: &str) -> bool {
1544        if let Some(paren_pos) = lower.find('(') {
1545            let after = &lower[paren_pos..];
1546            after.contains("select")
1547        } else {
1548            false
1549        }
1550    }
1551
1552    /// 统计 JOIN 数量
1553    fn count_joins(lower: &str) -> usize {
1554        lower.matches(" join ").count()
1555    }
1556
1557    /// 统计子查询数量
1558    fn count_subqueries(lower: &str) -> usize {
1559        lower.matches("(select").count() + lower.matches("( select").count()
1560    }
1561
1562    /// 检测前缀通配符 LIKE
1563    fn has_leading_wildcard_like(lower: &str) -> bool {
1564        lower.contains("like '%") || lower.contains("like \"%")
1565    }
1566
1567    /// 统计 OR 条件数量
1568    fn count_or_conditions(lower: &str) -> usize {
1569        lower.matches(" or ").count()
1570    }
1571
1572    /// 检测 ORDER BY 无 LIMIT
1573    fn has_order_by_without_limit(lower: &str) -> bool {
1574        lower.contains(" order by ") && !lower.contains(" limit ")
1575    }
1576
1577    /// 从 SQL 中提取表名(基于 schema)
1578    fn extract_tables(&self, lower: &str, schema: &SchemaContext) -> Vec<String> {
1579        let mut tables = Vec::new();
1580        for table in &schema.tables {
1581            let name_lower = table.name.to_lowercase();
1582            if lower.contains(&name_lower) {
1583                tables.push(table.name.clone());
1584            }
1585        }
1586        tables
1587    }
1588
1589    /// 生成列建议 SQL
1590    fn suggest_columns(&self, tables: &[String], schema: &SchemaContext) -> String {
1591        if tables.is_empty() {
1592            return String::new();
1593        }
1594        let mut cols = Vec::new();
1595        for table_name in tables {
1596            if let Some(table) = schema.tables.iter().find(|t| t.name == *table_name) {
1597                for col in &table.columns {
1598                    cols.push(format!("{}.{}", table_name, col.name));
1599                }
1600            }
1601        }
1602        if cols.is_empty() {
1603            String::new()
1604        } else {
1605            format!("SELECT {} FROM {}", cols.join(", "), tables.join(", "))
1606        }
1607    }
1608
1609    /// 为 SQL 添加 LIMIT 建议
1610    fn add_limit_suggestion(&self, sql: &str, limit: usize) -> String {
1611        let trimmed = sql.trim_end_matches(';');
1612        format!("{} LIMIT {}", trimmed, limit)
1613    }
1614
1615    /// 从 WHERE 子句中提取列名
1616    fn extract_where_columns(&self, lower: &str, schema: &SchemaContext) -> Vec<String> {
1617        let mut columns = Vec::new();
1618        if let Some(where_pos) = lower.find(" where ") {
1619            let after_where = &lower[where_pos + 7..];
1620            // 截取到 GROUP BY / ORDER BY / LIMIT 之前
1621            let where_clause = after_where
1622                .split(" group by ")
1623                .next()
1624                .unwrap_or(after_where)
1625                .split(" order by ")
1626                .next()
1627                .unwrap_or(after_where)
1628                .split(" limit ")
1629                .next()
1630                .unwrap_or(after_where);
1631
1632            for table in &schema.tables {
1633                for col in &table.columns {
1634                    let col_lower = col.name.to_lowercase();
1635                    if col_lower.len() >= 2
1636                        && where_clause.contains(&col_lower)
1637                        && !columns.contains(&col.name)
1638                    {
1639                        columns.push(col.name.clone());
1640                    }
1641                }
1642            }
1643        }
1644        columns
1645    }
1646
1647    /// 检查列是否有索引(简化版:主键视为有索引)
1648    fn column_has_index(&self, col: &str, schema: &SchemaContext) -> bool {
1649        for table in &schema.tables {
1650            for c in &table.columns {
1651                if c.name.eq_ignore_ascii_case(col) && c.is_primary_key {
1652                    return true;
1653                }
1654            }
1655        }
1656        false
1657    }
1658
1659    /// 计算 SQL 复杂度评分
1660    fn calculate_complexity(
1661        &self,
1662        has_join: bool,
1663        has_subquery: bool,
1664        has_where: bool,
1665        tables: &[String],
1666        lower: &str,
1667    ) -> u32 {
1668        let mut score: u32 = 10; // 基础分
1669
1670        if has_join {
1671            let join_count = Self::count_joins(lower) as u32;
1672            score += join_count * self.join_weight;
1673        }
1674
1675        if has_subquery {
1676            let sub_count = Self::count_subqueries(lower) as u32;
1677            score += sub_count * self.subquery_weight;
1678        }
1679
1680        if has_where {
1681            let or_count = Self::count_or_conditions(lower) as u32;
1682            score += self.where_weight + or_count * 3;
1683        }
1684
1685        // 表数量影响
1686        if tables.len() > 1 {
1687            score += (tables.len() as u32 - 1) * 10;
1688        }
1689
1690        // ORDER BY 影响
1691        if lower.contains(" order by ") {
1692            score += 5;
1693        }
1694
1695        // GROUP BY 影响
1696        if lower.contains(" group by ") {
1697            score += 10;
1698        }
1699
1700        // DISTINCT 影响
1701        if lower.contains(" distinct ") {
1702            score += 5;
1703        }
1704
1705        score.min(100)
1706    }
1707}
1708
1709// ==================== 单元测试 ====================
1710
1711#[cfg(test)]
1712mod tests {
1713    use super::*;
1714
1715    fn test_schema() -> SchemaContext {
1716        SchemaContext {
1717            tables: vec![
1718                TableInfo {
1719                    name: "users".into(),
1720                    columns: vec![
1721                        ColumnInfo {
1722                            name: "id".into(),
1723                            data_type: "INTEGER".into(),
1724                            nullable: false,
1725                            is_primary_key: true,
1726                        },
1727                        ColumnInfo {
1728                            name: "name".into(),
1729                            data_type: "TEXT".into(),
1730                            nullable: true,
1731                            is_primary_key: false,
1732                        },
1733                        ColumnInfo {
1734                            name: "email".into(),
1735                            data_type: "TEXT".into(),
1736                            nullable: true,
1737                            is_primary_key: false,
1738                        },
1739                        ColumnInfo {
1740                            name: "age".into(),
1741                            data_type: "INTEGER".into(),
1742                            nullable: true,
1743                            is_primary_key: false,
1744                        },
1745                        ColumnInfo {
1746                            name: "city".into(),
1747                            data_type: "TEXT".into(),
1748                            nullable: true,
1749                            is_primary_key: false,
1750                        },
1751                        ColumnInfo {
1752                            name: "score".into(),
1753                            data_type: "REAL".into(),
1754                            nullable: true,
1755                            is_primary_key: false,
1756                        },
1757                    ],
1758                },
1759                TableInfo {
1760                    name: "orders".into(),
1761                    columns: vec![
1762                        ColumnInfo {
1763                            name: "id".into(),
1764                            data_type: "INTEGER".into(),
1765                            nullable: false,
1766                            is_primary_key: true,
1767                        },
1768                        ColumnInfo {
1769                            name: "user_id".into(),
1770                            data_type: "INTEGER".into(),
1771                            nullable: false,
1772                            is_primary_key: false,
1773                        },
1774                        ColumnInfo {
1775                            name: "product".into(),
1776                            data_type: "TEXT".into(),
1777                            nullable: true,
1778                            is_primary_key: false,
1779                        },
1780                        ColumnInfo {
1781                            name: "price".into(),
1782                            data_type: "REAL".into(),
1783                            nullable: true,
1784                            is_primary_key: false,
1785                        },
1786                        ColumnInfo {
1787                            name: "quantity".into(),
1788                            data_type: "INTEGER".into(),
1789                            nullable: true,
1790                            is_primary_key: false,
1791                        },
1792                    ],
1793                },
1794                TableInfo {
1795                    name: "products".into(),
1796                    columns: vec![
1797                        ColumnInfo {
1798                            name: "id".into(),
1799                            data_type: "INTEGER".into(),
1800                            nullable: false,
1801                            is_primary_key: true,
1802                        },
1803                        ColumnInfo {
1804                            name: "name".into(),
1805                            data_type: "TEXT".into(),
1806                            nullable: false,
1807                            is_primary_key: false,
1808                        },
1809                        ColumnInfo {
1810                            name: "price".into(),
1811                            data_type: "REAL".into(),
1812                            nullable: false,
1813                            is_primary_key: false,
1814                        },
1815                        ColumnInfo {
1816                            name: "category".into(),
1817                            data_type: "TEXT".into(),
1818                            nullable: true,
1819                            is_primary_key: false,
1820                        },
1821                    ],
1822                },
1823            ],
1824        }
1825    }
1826
1827    // ============ SimpleNl2SqlEngine ============
1828
1829    #[tokio::test]
1830    async fn test_simple_select_all() {
1831        let engine = SimpleNl2SqlEngine::new();
1832        let schema = test_schema();
1833        let result = engine.generate("show all users", &schema).await.unwrap();
1834        assert_eq!(result.sql, "SELECT * FROM users");
1835        assert!(result.confidence > 0.0);
1836    }
1837
1838    #[tokio::test]
1839    async fn test_simple_select_columns() {
1840        let engine = SimpleNl2SqlEngine::new();
1841        let schema = test_schema();
1842        // "name" and "email" both appear in the query; schema matching picks them up
1843        let result = engine
1844            .generate("show name and email of users", &schema)
1845            .await
1846            .unwrap();
1847        assert_eq!(result.sql, "SELECT name, email FROM users");
1848    }
1849
1850    #[tokio::test]
1851    async fn test_simple_count() {
1852        let engine = SimpleNl2SqlEngine::new();
1853        let schema = test_schema();
1854        let result = engine.generate("how many users", &schema).await.unwrap();
1855        assert_eq!(result.sql, "SELECT COUNT(*) FROM users");
1856    }
1857
1858    #[tokio::test]
1859    async fn test_simple_where_equality() {
1860        let engine = SimpleNl2SqlEngine::new();
1861        let schema = test_schema();
1862        let result = engine
1863            .generate("find users where name = John", &schema)
1864            .await
1865            .unwrap();
1866        assert_eq!(result.sql, "SELECT * FROM users WHERE name = $1");
1867        assert!(result.explanation.contains("John"));
1868    }
1869
1870    #[tokio::test]
1871    async fn test_simple_where_comparison() {
1872        let engine = SimpleNl2SqlEngine::new();
1873        let schema = test_schema();
1874        let result = engine
1875            .generate("find users where age > 25", &schema)
1876            .await
1877            .unwrap();
1878        assert_eq!(result.sql, "SELECT * FROM users WHERE age > $1");
1879    }
1880
1881    #[tokio::test]
1882    async fn test_simple_order_by() {
1883        let engine = SimpleNl2SqlEngine::new();
1884        let schema = test_schema();
1885        let result = engine
1886            .generate("list users ordered by name", &schema)
1887            .await
1888            .unwrap();
1889        assert_eq!(result.sql, "SELECT * FROM users ORDER BY name ASC");
1890    }
1891
1892    #[tokio::test]
1893    async fn test_simple_order_by_desc() {
1894        let engine = SimpleNl2SqlEngine::new();
1895        let schema = test_schema();
1896        let result = engine
1897            .generate("list users sorted by name descending", &schema)
1898            .await
1899            .unwrap();
1900        assert_eq!(result.sql, "SELECT * FROM users ORDER BY name DESC");
1901    }
1902
1903    #[tokio::test]
1904    async fn test_simple_limit() {
1905        let engine = SimpleNl2SqlEngine::new();
1906        let schema = test_schema();
1907        let result = engine.generate("first 10 users", &schema).await.unwrap();
1908        assert_eq!(result.sql, "SELECT * FROM users LIMIT 10");
1909    }
1910
1911    #[tokio::test]
1912    async fn test_simple_order_by_with_limit() {
1913        let engine = SimpleNl2SqlEngine::new();
1914        let schema = test_schema();
1915        let result = engine
1916            .generate("top 5 users by score", &schema)
1917            .await
1918            .unwrap();
1919        // "top" triggers LIMIT; "score" is a known column, and "by score" matched as sort
1920        // This generates: SELECT * FROM users ORDER BY score ASC LIMIT 5
1921        assert!(result.sql.contains("LIMIT 5"));
1922    }
1923
1924    #[tokio::test]
1925    async fn test_simple_aggregation() {
1926        let engine = SimpleNl2SqlEngine::new();
1927        let schema = test_schema();
1928        let result = engine
1929            .generate("total price from orders", &schema)
1930            .await
1931            .unwrap();
1932        assert_eq!(result.sql, "SELECT SUM(price) FROM orders");
1933    }
1934
1935    #[tokio::test]
1936    async fn test_simple_avg() {
1937        let engine = SimpleNl2SqlEngine::new();
1938        let schema = test_schema();
1939        let result = engine
1940            .generate("average age of users", &schema)
1941            .await
1942            .unwrap();
1943        assert_eq!(result.sql, "SELECT AVG(age) FROM users");
1944    }
1945
1946    #[tokio::test]
1947    async fn test_simple_empty_query() {
1948        let engine = SimpleNl2SqlEngine::new();
1949        let schema = test_schema();
1950        let result = engine.generate("", &schema).await;
1951        assert!(result.is_err());
1952        match result {
1953            Err(Nl2SqlError::InvalidQuery(_)) => {}
1954            Err(e) => panic!("期望 InvalidQuery,实际: {:?}", e),
1955            Ok(_) => panic!("期望错误"),
1956        }
1957    }
1958
1959    #[tokio::test]
1960    async fn test_simple_empty_schema() {
1961        let engine = SimpleNl2SqlEngine::new();
1962        let schema = SchemaContext { tables: vec![] };
1963        let result = engine.generate("show users", &schema).await;
1964        assert!(result.is_err());
1965        match result {
1966            Err(Nl2SqlError::SchemaError(_)) => {}
1967            _ => panic!("期望 SchemaError"),
1968        }
1969    }
1970
1971    #[tokio::test]
1972    async fn test_simple_validate_valid() {
1973        let engine = SimpleNl2SqlEngine::new();
1974        let query = SqlQuery {
1975            sql: "SELECT * FROM users".into(),
1976            explanation: "test".into(),
1977            confidence: 0.9,
1978        };
1979        assert!(engine.validate(&query).await.unwrap());
1980    }
1981
1982    #[tokio::test]
1983    async fn test_simple_validate_invalid_confidence() {
1984        let engine = SimpleNl2SqlEngine::new();
1985        let query = SqlQuery {
1986            sql: "SELECT * FROM users".into(),
1987            explanation: "test".into(),
1988            confidence: 1.5,
1989        };
1990        let result = engine.validate(&query).await;
1991        assert!(result.is_err());
1992    }
1993
1994    #[tokio::test]
1995    async fn test_simple_validate_rejects_drop() {
1996        let engine = SimpleNl2SqlEngine::new();
1997        let query = SqlQuery {
1998            sql: "DROP TABLE users".into(),
1999            explanation: "test".into(),
2000            confidence: 0.9,
2001        };
2002        assert!(!engine.validate(&query).await.unwrap());
2003    }
2004
2005    #[tokio::test]
2006    async fn test_simple_alias() {
2007        let engine = SimpleNl2SqlEngine::new().with_alias("person", "users");
2008        let schema = test_schema();
2009        let result = engine.generate("show all persons", &schema).await.unwrap();
2010        assert_eq!(result.sql, "SELECT * FROM users");
2011    }
2012
2013    #[tokio::test]
2014    async fn test_simple_table_not_found() {
2015        let engine = SimpleNl2SqlEngine::new();
2016        let schema = test_schema();
2017        let result = engine
2018            .generate("show something from nonexistent_table", &schema)
2019            .await;
2020        assert!(result.is_err());
2021    }
2022
2023    #[tokio::test]
2024    async fn test_simple_sql_in_schema_passthrough() {
2025        let engine = SimpleNl2SqlEngine::new();
2026        let schema = test_schema();
2027        // "name" and "email" are both schema columns
2028        let result = engine
2029            .generate("select name and email from users", &schema)
2030            .await
2031            .unwrap();
2032        assert!(result.sql.contains("name"));
2033        assert!(result.sql.contains("email"));
2034        assert!(result.sql.contains("users"));
2035    }
2036
2037    // ============ 仅在 real feature 下测试 OpenAINl2SqlEngine ============
2038
2039    #[cfg(feature = "real")]
2040    #[test]
2041    fn test_openai_engine_new_with_defaults() {
2042        let engine = OpenAINl2SqlEngine::new("sk-test-key");
2043        assert_eq!(engine.api_base, "https://api.openai.com/v1");
2044        assert_eq!(engine.api_key, "sk-test-key");
2045        assert_eq!(engine.model, "gpt-4o-mini");
2046    }
2047
2048    #[cfg(feature = "real")]
2049    #[test]
2050    fn test_openai_engine_with_options() {
2051        let engine = OpenAINl2SqlEngine::new("sk-test")
2052            .with_api_base("https://api.deepseek.com/v1")
2053            .with_model("deepseek-chat");
2054        assert_eq!(engine.api_base, "https://api.deepseek.com/v1");
2055        assert_eq!(engine.model, "deepseek-chat");
2056    }
2057
2058    #[cfg(feature = "real")]
2059    #[tokio::test]
2060    async fn test_openai_engine_missing_api_key() {
2061        let engine = OpenAINl2SqlEngine::new("");
2062        let schema = test_schema();
2063        let result = engine.generate("show users", &schema).await;
2064        match result {
2065            Err(Nl2SqlError::ConfigError(_)) => {}
2066            other => panic!("期望 ConfigError,实际: {:?}", other),
2067        }
2068    }
2069
2070    #[cfg(feature = "real")]
2071    #[test]
2072    fn test_clean_llm_sql_output() {
2073        assert_eq!(
2074            clean_llm_sql_output("SELECT * FROM users"),
2075            "SELECT * FROM users"
2076        );
2077        assert_eq!(
2078            clean_llm_sql_output("```sql\nSELECT * FROM users\n```"),
2079            "SELECT * FROM users"
2080        );
2081        assert_eq!(
2082            clean_llm_sql_output("```\nSELECT * FROM users\n```"),
2083            "SELECT * FROM users"
2084        );
2085        assert_eq!(
2086            clean_llm_sql_output("\n  SELECT * FROM users  \n"),
2087            "SELECT * FROM users"
2088        );
2089    }
2090
2091    // ============ 查询优化提示测试 ============
2092
2093    fn optimizer_test_schema() -> SchemaContext {
2094        SchemaContext {
2095            tables: vec![
2096                TableInfo {
2097                    name: "users".into(),
2098                    columns: vec![
2099                        ColumnInfo {
2100                            name: "id".into(),
2101                            data_type: "INTEGER".into(),
2102                            nullable: false,
2103                            is_primary_key: true,
2104                        },
2105                        ColumnInfo {
2106                            name: "name".into(),
2107                            data_type: "TEXT".into(),
2108                            nullable: true,
2109                            is_primary_key: false,
2110                        },
2111                        ColumnInfo {
2112                            name: "email".into(),
2113                            data_type: "TEXT".into(),
2114                            nullable: true,
2115                            is_primary_key: false,
2116                        },
2117                        ColumnInfo {
2118                            name: "age".into(),
2119                            data_type: "INTEGER".into(),
2120                            nullable: true,
2121                            is_primary_key: false,
2122                        },
2123                    ],
2124                },
2125                TableInfo {
2126                    name: "orders".into(),
2127                    columns: vec![
2128                        ColumnInfo {
2129                            name: "id".into(),
2130                            data_type: "INTEGER".into(),
2131                            nullable: false,
2132                            is_primary_key: true,
2133                        },
2134                        ColumnInfo {
2135                            name: "user_id".into(),
2136                            data_type: "INTEGER".into(),
2137                            nullable: false,
2138                            is_primary_key: false,
2139                        },
2140                        ColumnInfo {
2141                            name: "amount".into(),
2142                            data_type: "DECIMAL".into(),
2143                            nullable: true,
2144                            is_primary_key: false,
2145                        },
2146                    ],
2147                },
2148            ],
2149        }
2150    }
2151
2152    #[test]
2153    fn test_hint_severity_as_str() {
2154        assert_eq!(HintSeverity::Info.as_str(), "INFO");
2155        assert_eq!(HintSeverity::Warning.as_str(), "WARNING");
2156        assert_eq!(HintSeverity::Critical.as_str(), "CRITICAL");
2157    }
2158
2159    #[test]
2160    fn test_query_optimization_hint_info() {
2161        let hint = QueryOptimizationHint::info("标题", "描述");
2162        assert_eq!(hint.title, "标题");
2163        assert_eq!(hint.description, "描述");
2164        assert_eq!(hint.severity, HintSeverity::Info);
2165        assert!(hint.suggested_sql.is_none());
2166    }
2167
2168    #[test]
2169    fn test_query_optimization_hint_warning() {
2170        let hint = QueryOptimizationHint::warning("警告", "警告描述");
2171        assert_eq!(hint.severity, HintSeverity::Warning);
2172    }
2173
2174    #[test]
2175    fn test_query_optimization_hint_critical() {
2176        let hint = QueryOptimizationHint::critical("严重", "严重描述");
2177        assert_eq!(hint.severity, HintSeverity::Critical);
2178    }
2179
2180    #[test]
2181    fn test_query_optimization_hint_with_suggested_sql() {
2182        let hint =
2183            QueryOptimizationHint::info("建议", "描述").with_suggested_sql("SELECT id FROM users");
2184        assert_eq!(hint.suggested_sql.as_deref(), Some("SELECT id FROM users"));
2185    }
2186
2187    #[test]
2188    fn test_query_optimizer_default() {
2189        let opt = QueryOptimizer::default();
2190        assert!(opt.check_select_star);
2191        assert!(opt.check_missing_limit);
2192        assert!(opt.check_missing_where);
2193        assert_eq!(opt.default_limit, 100);
2194    }
2195
2196    #[test]
2197    fn test_query_optimizer_new() {
2198        let opt = QueryOptimizer::new();
2199        assert!(opt.check_select_star);
2200    }
2201
2202    #[test]
2203    fn test_query_optimizer_with_default_limit() {
2204        let opt = QueryOptimizer::new().with_default_limit(50);
2205        assert_eq!(opt.default_limit, 50);
2206    }
2207
2208    #[test]
2209    fn test_query_optimizer_disable_checks() {
2210        let opt = QueryOptimizer::new()
2211            .disable_select_star_check()
2212            .disable_missing_limit_check()
2213            .disable_missing_where_check();
2214        assert!(!opt.check_select_star);
2215        assert!(!opt.check_missing_limit);
2216        assert!(!opt.check_missing_where);
2217    }
2218
2219    #[test]
2220    fn test_analyze_select_star() {
2221        let opt = QueryOptimizer::new();
2222        let schema = optimizer_test_schema();
2223        let analysis = opt.analyze("SELECT * FROM users", &schema);
2224
2225        assert!(analysis.uses_select_star);
2226        assert!(!analysis.has_where);
2227        assert!(!analysis.has_limit);
2228        // 应该有 SELECT * 建议、缺失 WHERE 建议、缺失 LIMIT 建议
2229        assert!(analysis.has_hints());
2230        assert!(analysis.critical_count() >= 1); // 缺失 WHERE 是 critical
2231    }
2232
2233    #[test]
2234    fn test_analyze_select_star_with_suggested_columns() {
2235        let opt = QueryOptimizer::new();
2236        let schema = optimizer_test_schema();
2237        let analysis = opt.analyze("SELECT * FROM users", &schema);
2238
2239        // 应包含 SELECT * 警告,且附带建议 SQL
2240        let select_star_hint = analysis
2241            .hints
2242            .iter()
2243            .find(|h| h.title == "避免使用 SELECT *");
2244        assert!(select_star_hint.is_some());
2245        let hint = select_star_hint.unwrap();
2246        assert!(hint.suggested_sql.is_some());
2247        let suggested = hint.suggested_sql.as_ref().unwrap();
2248        assert!(suggested.contains("users.id"));
2249        assert!(suggested.contains("users.name"));
2250    }
2251
2252    #[test]
2253    fn test_analyze_missing_where_critical() {
2254        let opt = QueryOptimizer::new();
2255        let schema = optimizer_test_schema();
2256        let analysis = opt.analyze("SELECT id FROM users LIMIT 10", &schema);
2257
2258        // 没有 WHERE 应产生 critical 建议
2259        assert!(analysis.critical_count() >= 1);
2260        let where_hint = analysis.hints.iter().find(|h| h.title == "缺少 WHERE 子句");
2261        assert!(where_hint.is_some());
2262        assert_eq!(where_hint.unwrap().severity, HintSeverity::Critical);
2263    }
2264
2265    #[test]
2266    fn test_analyze_missing_limit_warning() {
2267        let opt = QueryOptimizer::new();
2268        let schema = optimizer_test_schema();
2269        let analysis = opt.analyze("SELECT id FROM users WHERE age > 18", &schema);
2270
2271        // 没有 LIMIT 应产生 warning 建议
2272        assert!(analysis.warning_count() >= 1);
2273        let limit_hint = analysis.hints.iter().find(|h| h.title == "缺少 LIMIT 子句");
2274        assert!(limit_hint.is_some());
2275        let hint = limit_hint.unwrap();
2276        assert!(hint.suggested_sql.is_some());
2277        assert!(hint.suggested_sql.as_ref().unwrap().contains("LIMIT 100"));
2278    }
2279
2280    #[test]
2281    fn test_analyze_with_limit_no_limit_hint() {
2282        let opt = QueryOptimizer::new();
2283        let schema = optimizer_test_schema();
2284        let analysis = opt.analyze("SELECT id FROM users WHERE age > 18 LIMIT 10", &schema);
2285
2286        let limit_hint = analysis.hints.iter().find(|h| h.title == "缺少 LIMIT 子句");
2287        assert!(limit_hint.is_none());
2288    }
2289
2290    #[test]
2291    fn test_analyze_join_detection() {
2292        let opt = QueryOptimizer::new();
2293        let schema = optimizer_test_schema();
2294        let analysis = opt.analyze(
2295            "SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id WHERE u.age > 18 LIMIT 10",
2296            &schema,
2297        );
2298
2299        assert!(analysis.has_join);
2300        assert!(analysis.detected_tables.contains(&"users".to_string()));
2301        assert!(analysis.detected_tables.contains(&"orders".to_string()));
2302    }
2303
2304    #[test]
2305    fn test_analyze_multiple_joins_critical() {
2306        let opt = QueryOptimizer::new();
2307        let schema = SchemaContext {
2308            tables: vec![
2309                TableInfo {
2310                    name: "t1".into(),
2311                    columns: vec![ColumnInfo {
2312                        name: "id".into(),
2313                        data_type: "INT".into(),
2314                        nullable: false,
2315                        is_primary_key: true,
2316                    }],
2317                },
2318                TableInfo {
2319                    name: "t2".into(),
2320                    columns: vec![ColumnInfo {
2321                        name: "id".into(),
2322                        data_type: "INT".into(),
2323                        nullable: false,
2324                        is_primary_key: true,
2325                    }],
2326                },
2327                TableInfo {
2328                    name: "t3".into(),
2329                    columns: vec![ColumnInfo {
2330                        name: "id".into(),
2331                        data_type: "INT".into(),
2332                        nullable: false,
2333                        is_primary_key: true,
2334                    }],
2335                },
2336                TableInfo {
2337                    name: "t4".into(),
2338                    columns: vec![ColumnInfo {
2339                        name: "id".into(),
2340                        data_type: "INT".into(),
2341                        nullable: false,
2342                        is_primary_key: true,
2343                    }],
2344                },
2345            ],
2346        };
2347        let analysis = opt.analyze(
2348            "SELECT * FROM t1 JOIN t2 ON t1.id = t2.id JOIN t3 ON t2.id = t3.id JOIN t4 ON t3.id = t4.id LIMIT 10",
2349            &schema,
2350        );
2351        // 4 个 JOIN 应触发 critical
2352        let join_hint = analysis.hints.iter().find(|h| h.title == "JOIN 数量过多");
2353        assert!(join_hint.is_some());
2354        assert_eq!(join_hint.unwrap().severity, HintSeverity::Critical);
2355    }
2356
2357    #[test]
2358    fn test_analyze_subquery_detection() {
2359        let opt = QueryOptimizer::new();
2360        let schema = optimizer_test_schema();
2361        let analysis = opt.analyze(
2362            "SELECT id FROM users WHERE id IN (SELECT user_id FROM orders) LIMIT 10",
2363            &schema,
2364        );
2365
2366        assert!(analysis.has_subquery);
2367    }
2368
2369    #[test]
2370    fn test_analyze_leading_wildcard_like() {
2371        let opt = QueryOptimizer::new();
2372        let schema = optimizer_test_schema();
2373        let analysis = opt.analyze(
2374            "SELECT id FROM users WHERE name LIKE '%john' LIMIT 10",
2375            &schema,
2376        );
2377
2378        let like_hint = analysis
2379            .hints
2380            .iter()
2381            .find(|h| h.title == "LIKE 使用前缀通配符");
2382        assert!(like_hint.is_some());
2383    }
2384
2385    #[test]
2386    fn test_analyze_order_by_without_limit() {
2387        let opt = QueryOptimizer::new();
2388        let schema = optimizer_test_schema();
2389        let analysis = opt.analyze("SELECT id FROM users WHERE age > 18 ORDER BY id", &schema);
2390
2391        let order_hint = analysis
2392            .hints
2393            .iter()
2394            .find(|h| h.title == "ORDER BY 无 LIMIT");
2395        assert!(order_hint.is_some());
2396    }
2397
2398    #[test]
2399    fn test_analyze_count_star_hint() {
2400        let opt = QueryOptimizer::new();
2401        let schema = optimizer_test_schema();
2402        let analysis = opt.analyze("SELECT COUNT(*) FROM users LIMIT 1", &schema);
2403
2404        let count_hint = analysis
2405            .hints
2406            .iter()
2407            .find(|h| h.title == "考虑使用 COUNT(1)");
2408        assert!(count_hint.is_some());
2409    }
2410
2411    #[test]
2412    fn test_analyze_missing_index_hint() {
2413        let opt = QueryOptimizer::new();
2414        let schema = optimizer_test_schema();
2415        // age 列不是主键,应建议添加索引
2416        let analysis = opt.analyze("SELECT id FROM users WHERE age > 18 LIMIT 10", &schema);
2417
2418        let index_hint = analysis.hints.iter().find(|h| h.title.contains("添加索引"));
2419        assert!(index_hint.is_some());
2420        assert!(index_hint.unwrap().title.contains("age"));
2421    }
2422
2423    #[test]
2424    fn test_analyze_primary_key_no_index_hint() {
2425        let opt = QueryOptimizer::new();
2426        let schema = optimizer_test_schema();
2427        // id 列是主键,不应建议添加索引
2428        let analysis = opt.analyze("SELECT name FROM users WHERE id = 1 LIMIT 10", &schema);
2429
2430        let index_hint = analysis
2431            .hints
2432            .iter()
2433            .find(|h| h.title.contains("添加索引") && h.title.contains("id"));
2434        // id 是主键,不应有索引建议
2435        assert!(index_hint.is_none());
2436    }
2437
2438    #[test]
2439    fn test_analyze_complexity_score_simple() {
2440        let opt = QueryOptimizer::new();
2441        let schema = optimizer_test_schema();
2442        let analysis = opt.analyze("SELECT id FROM users WHERE id = 1 LIMIT 10", &schema);
2443        // 简单查询应该低分
2444        assert!(analysis.complexity_score < 30);
2445    }
2446
2447    #[test]
2448    fn test_analyze_complexity_score_complex() {
2449        let opt = QueryOptimizer::new();
2450        let schema = optimizer_test_schema();
2451        let analysis = opt.analyze(
2452            "SELECT u.id, o.amount FROM users u JOIN orders o ON u.id = o.user_id WHERE u.age > 18 OR u.name LIKE '%a%' GROUP BY u.id ORDER BY o.amount DESC",
2453            &schema,
2454        );
2455        // 复杂查询应该高分
2456        assert!(analysis.complexity_score > 30);
2457    }
2458
2459    #[test]
2460    fn test_analyze_well_optimized_query() {
2461        let opt = QueryOptimizer::new();
2462        let schema = optimizer_test_schema();
2463        let analysis = opt.analyze("SELECT id, name FROM users WHERE id = 1 LIMIT 10", &schema);
2464
2465        // 这个查询写得很好,不应该有 critical 或 warning 建议
2466        assert_eq!(analysis.critical_count(), 0);
2467        // id 是主键,不应有索引建议
2468        // 有 LIMIT,不应有 LIMIT 建议
2469        // 有 WHERE,不应有 WHERE 建议
2470        // 没有 SELECT *,不应有 SELECT * 建议
2471    }
2472
2473    #[test]
2474    fn test_query_analysis_has_hints() {
2475        let opt = QueryOptimizer::new();
2476        let schema = optimizer_test_schema();
2477        let analysis = opt.analyze("SELECT * FROM users", &schema);
2478        assert!(analysis.has_hints());
2479
2480        let good_analysis = opt.analyze("SELECT id FROM users WHERE id = 1 LIMIT 1", &schema);
2481        // 可能仍有 info 级建议,但不应有 critical
2482        assert_eq!(good_analysis.critical_count(), 0);
2483    }
2484
2485    #[test]
2486    fn test_format_report_contains_key_info() {
2487        let opt = QueryOptimizer::new();
2488        let schema = optimizer_test_schema();
2489        let analysis = opt.analyze("SELECT * FROM users", &schema);
2490        let report = QueryOptimizer::format_report(&analysis);
2491
2492        assert!(report.contains("SQL 查询优化分析报告"));
2493        assert!(report.contains("原始 SQL"));
2494        assert!(report.contains("复杂度评分"));
2495        assert!(report.contains("SELECT *"));
2496    }
2497
2498    #[test]
2499    fn test_format_report_no_hints() {
2500        let opt = QueryOptimizer::new();
2501        let schema = optimizer_test_schema();
2502        let analysis = opt.analyze("SELECT id FROM users WHERE id = 1 LIMIT 1", &schema);
2503        let report = QueryOptimizer::format_report(&analysis);
2504        // 即使没有建议,报告也应包含基本字段
2505        assert!(report.contains("复杂度评分"));
2506    }
2507
2508    #[test]
2509    fn test_analyze_disable_select_star() {
2510        let opt = QueryOptimizer::new().disable_select_star_check();
2511        let schema = optimizer_test_schema();
2512        let analysis = opt.analyze("SELECT * FROM users WHERE id = 1 LIMIT 10", &schema);
2513
2514        let star_hint = analysis
2515            .hints
2516            .iter()
2517            .find(|h| h.title == "避免使用 SELECT *");
2518        assert!(star_hint.is_none());
2519    }
2520
2521    #[test]
2522    fn test_analyze_disable_missing_where() {
2523        let opt = QueryOptimizer::new().disable_missing_where_check();
2524        let schema = optimizer_test_schema();
2525        let analysis = opt.analyze("SELECT id FROM users LIMIT 10", &schema);
2526
2527        let where_hint = analysis.hints.iter().find(|h| h.title == "缺少 WHERE 子句");
2528        assert!(where_hint.is_none());
2529    }
2530
2531    #[test]
2532    fn test_analyze_disable_missing_limit() {
2533        let opt = QueryOptimizer::new().disable_missing_limit_check();
2534        let schema = optimizer_test_schema();
2535        let analysis = opt.analyze("SELECT id FROM users WHERE id = 1", &schema);
2536
2537        let limit_hint = analysis.hints.iter().find(|h| h.title == "缺少 LIMIT 子句");
2538        assert!(limit_hint.is_none());
2539    }
2540
2541    #[test]
2542    fn test_analyze_multiple_or_conditions() {
2543        let opt = QueryOptimizer::new();
2544        let schema = optimizer_test_schema();
2545        let analysis = opt.analyze(
2546            "SELECT id FROM users WHERE age = 1 OR age = 2 OR age = 3 OR age = 4 LIMIT 10",
2547            &schema,
2548        );
2549
2550        let or_hint = analysis.hints.iter().find(|h| h.title == "多个 OR 条件");
2551        assert!(or_hint.is_some());
2552    }
2553
2554    #[test]
2555    fn test_analyze_detected_tables() {
2556        let opt = QueryOptimizer::new();
2557        let schema = optimizer_test_schema();
2558        let analysis = opt.analyze("SELECT id FROM users LIMIT 10", &schema);
2559
2560        assert_eq!(analysis.detected_tables, vec!["users".to_string()]);
2561    }
2562
2563    #[test]
2564    fn test_analyze_no_detected_tables() {
2565        let opt = QueryOptimizer::new();
2566        let schema = optimizer_test_schema();
2567        let analysis = opt.analyze("SELECT 1 LIMIT 10", &schema);
2568
2569        assert!(analysis.detected_tables.is_empty());
2570    }
2571
2572    #[test]
2573    fn test_normalize_sql() {
2574        let normalized = QueryOptimizer::normalize_sql("SELECT  id\nFROM   users");
2575        assert_eq!(normalized, "SELECT id FROM users");
2576    }
2577}