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