1use async_trait::async_trait;
10use thiserror::Error;
11
12use crate::safety;
13
14#[cfg(feature = "real")]
15use serde::{Deserialize, Serialize};
16
17#[derive(Debug, Clone)]
21pub struct SqlQuery {
22 pub sql: String,
24 pub explanation: String,
26 pub confidence: f32,
28}
29
30#[derive(Debug, Clone, Default)]
32pub struct SchemaContext {
33 pub tables: Vec<TableInfo>,
34}
35
36#[derive(Debug, Clone)]
38pub struct TableInfo {
39 pub name: String,
40 pub columns: Vec<ColumnInfo>,
41}
42
43#[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#[derive(Debug, Error)]
56pub enum Nl2SqlError {
57 #[error("Invalid query: {0}")]
59 InvalidQuery(String),
60 #[error("Schema error: {0}")]
62 SchemaError(String),
63 #[error("Safety error: {0}")]
65 SafetyError(String),
66 #[error("Generation error: {0}")]
68 GenerationError(String),
69 #[error("API error (status {0}): {1}")]
71 ApiError(u16, String),
72 #[error("Network error: {0}")]
74 NetworkError(String),
75 #[error("Config error: {0}")]
77 ConfigError(String),
78}
79
80#[async_trait]
86pub trait Nl2SqlEngine: Send + Sync {
87 async fn generate(
97 &self,
98 nl_query: &str,
99 schema: &SchemaContext,
100 ) -> Result<SqlQuery, Nl2SqlError>;
101
102 async fn validate(&self, query: &SqlQuery) -> Result<bool, Nl2SqlError>;
108}
109
110pub struct SimpleNl2SqlEngine {
149 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 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 fn find_table<'a>(&self, query: &str, schema: &'a SchemaContext) -> Option<&'a TableInfo> {
175 let lower = query.to_lowercase();
176
177 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 for (alias, target_table_name) in &self.aliases {
197 if let Some(pos) = lower.find(alias) {
198 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 fn extract_columns_from_nl(query: &str, table: &TableInfo) -> Vec<String> {
221 let lower = query.to_lowercase();
222 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 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 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 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 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 if col_lower.len() >= 2
286 && (lower.contains(&col_lower) || lower.contains(&format!(" {} ", col_lower)))
287 {
288 if !columns.contains(&col.name) {
290 columns.push(col.name.clone());
291 }
292 }
293 }
294
295 columns
296 }
297
298 fn find_join_columns(t1: &TableInfo, t2: &TableInfo) -> Option<(String, String)> {
300 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 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 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 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 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 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 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 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 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 let columns = Self::extract_columns_from_nl(nl_query, table);
465
466 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 let after = nl_query[pos + 5..].trim();
474 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 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 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 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 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 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 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 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 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 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 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 let where_clause = if conditions.is_empty() {
694 String::new()
695 } else {
696 format!(" WHERE {}", conditions.join(" AND "))
697 };
698
699 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 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 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 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#[cfg(feature = "real")]
816pub struct OpenAINl2SqlEngine {
817 api_base: String,
819 api_key: String,
821 model: String,
823 http_client: reqwest::Client,
825}
826
827#[cfg(feature = "real")]
828impl OpenAINl2SqlEngine {
829 const DEFAULT_API_BASE: &'static str = "https://api.openai.com/v1";
831 const DEFAULT_MODEL: &'static str = "gpt-4o-mini";
833
834 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 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 pub fn with_model(mut self, model: impl Into<String>) -> Self {
852 self.model = model.into();
853 self
854 }
855
856 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 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 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 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#[cfg(feature = "real")]
1045fn clean_llm_sql_output(raw: &str) -> String {
1046 let trimmed = raw.trim();
1047 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1068pub enum HintSeverity {
1069 Info,
1071 Warning,
1073 Critical,
1075}
1076
1077impl HintSeverity {
1078 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#[derive(Debug, Clone)]
1090pub struct QueryOptimizationHint {
1091 pub title: String,
1093 pub description: String,
1095 pub severity: HintSeverity,
1097 pub suggested_sql: Option<String>,
1099}
1100
1101impl QueryOptimizationHint {
1102 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 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 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 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#[derive(Debug, Clone)]
1141pub struct QueryAnalysis {
1142 pub original_sql: String,
1144 pub hints: Vec<QueryOptimizationHint>,
1146 pub complexity_score: u32,
1148 pub detected_tables: Vec<String>,
1150 pub has_where: bool,
1152 pub has_limit: bool,
1154 pub has_join: bool,
1156 pub has_subquery: bool,
1158 pub uses_select_star: bool,
1160}
1161
1162impl QueryAnalysis {
1163 pub fn critical_count(&self) -> usize {
1165 self.hints
1166 .iter()
1167 .filter(|h| h.severity == HintSeverity::Critical)
1168 .count()
1169 }
1170
1171 pub fn warning_count(&self) -> usize {
1173 self.hints
1174 .iter()
1175 .filter(|h| h.severity == HintSeverity::Warning)
1176 .count()
1177 }
1178
1179 pub fn has_hints(&self) -> bool {
1181 !self.hints.is_empty()
1182 }
1183}
1184
1185pub struct QueryOptimizer {
1190 pub check_select_star: bool,
1192 pub check_missing_limit: bool,
1194 pub check_missing_where: bool,
1196 pub default_limit: usize,
1198 pub join_weight: u32,
1200 pub subquery_weight: u32,
1202 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 pub fn new() -> Self {
1223 Self::default()
1224 }
1225
1226 pub fn with_default_limit(mut self, limit: usize) -> Self {
1228 self.default_limit = limit;
1229 self
1230 }
1231
1232 pub fn disable_select_star_check(mut self) -> Self {
1234 self.check_select_star = false;
1235 self
1236 }
1237
1238 pub fn disable_missing_limit_check(mut self) -> Self {
1240 self.check_missing_limit = false;
1241 self
1242 }
1243
1244 pub fn disable_missing_where_check(mut self) -> Self {
1246 self.check_missing_where = false;
1247 self
1248 }
1249
1250 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 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 if self.check_missing_where && !has_where {
1283 hints.push(QueryOptimizationHint::critical(
1284 "缺少 WHERE 子句",
1285 "查询没有 WHERE 条件,将扫描全表。对于大表这会导致严重的性能问题。",
1286 ));
1287 }
1288
1289 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 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 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 if Self::has_leading_wildcard_like(&lower) {
1339 hints.push(QueryOptimizationHint::warning(
1340 "LIKE 使用前缀通配符",
1341 "LIKE '%keyword' 无法使用索引,会导致全表扫描。如可能,使用 LIKE 'keyword%' 或全文索引。",
1342 ));
1343 }
1344
1345 if Self::count_or_conditions(&lower) >= 3 {
1347 hints.push(QueryOptimizationHint::info(
1348 "多个 OR 条件",
1349 "多个 OR 条件可能导致无法有效使用索引。考虑使用 UNION ALL 重写。",
1350 ));
1351 }
1352
1353 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 if lower.contains("count(*)") {
1363 hints.push(QueryOptimizationHint::info(
1364 "考虑使用 COUNT(1)",
1365 "某些数据库中 COUNT(1) 比 COUNT(*) 略快(虽然现代优化器通常已优化)。",
1366 ));
1367 }
1368
1369 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 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 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 fn normalize_sql(sql: &str) -> String {
1453 sql.split_whitespace().collect::<Vec<_>>().join(" ")
1454 }
1455
1456 fn detect_select_star(&self, lower: &str) -> bool {
1458 lower.contains("select *") || lower.contains("select *")
1459 }
1460
1461 fn detect_where(&self, lower: &str) -> bool {
1463 lower.contains(" where ")
1464 }
1465
1466 fn detect_limit(&self, lower: &str) -> bool {
1468 lower.contains(" limit ")
1469 }
1470
1471 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 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 fn count_joins(lower: &str) -> usize {
1493 lower.matches(" join ").count()
1494 }
1495
1496 fn count_subqueries(lower: &str) -> usize {
1498 lower.matches("(select").count() + lower.matches("( select").count()
1499 }
1500
1501 fn has_leading_wildcard_like(lower: &str) -> bool {
1503 lower.contains("like '%") || lower.contains("like \"%")
1504 }
1505
1506 fn count_or_conditions(lower: &str) -> usize {
1508 lower.matches(" or ").count()
1509 }
1510
1511 fn has_order_by_without_limit(lower: &str) -> bool {
1513 lower.contains(" order by ") && !lower.contains(" limit ")
1514 }
1515
1516 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 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 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 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 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 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 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; 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 if tables.len() > 1 {
1626 score += (tables.len() as u32 - 1) * 10;
1627 }
1628
1629 if lower.contains(" order by ") {
1631 score += 5;
1632 }
1633
1634 if lower.contains(" group by ") {
1636 score += 10;
1637 }
1638
1639 if lower.contains(" distinct ") {
1641 score += 5;
1642 }
1643
1644 score.min(100)
1645 }
1646}
1647
1648#[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 #[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 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 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 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 #[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 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 assert!(analysis.has_hints());
2169 assert!(analysis.critical_count() >= 1); }
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 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 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 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 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 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 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 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 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 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 assert_eq!(analysis.critical_count(), 0);
2406 }
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 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 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}