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