sql_cli/sql/parser/
ast_formatter.rs

1//! AST-based SQL Formatter
2//!
3//! This module provides proper SQL formatting by traversing the parsed AST,
4//! which is more reliable than regex-based formatting and handles complex
5//! features like CTEs, subqueries, and expressions correctly.
6
7use crate::sql::parser::ast::*;
8use std::fmt::Write;
9
10/// Configuration for SQL formatting
11pub struct FormatConfig {
12    /// Indentation string (e.g., "  " for 2 spaces, "\t" for tab)
13    pub indent: String,
14    /// Maximum number of items per line for lists (SELECT columns, etc.)
15    pub items_per_line: usize,
16    /// Whether to uppercase keywords
17    pub uppercase_keywords: bool,
18    /// Whether to add newlines between major clauses
19    pub compact: bool,
20}
21
22impl Default for FormatConfig {
23    fn default() -> Self {
24        Self {
25            indent: "    ".to_string(),
26            items_per_line: 5,
27            uppercase_keywords: true,
28            compact: false,
29        }
30    }
31}
32
33/// Format a SELECT statement into pretty SQL
34pub fn format_select_statement(stmt: &SelectStatement) -> String {
35    format_select_with_config(stmt, &FormatConfig::default())
36}
37
38/// Format a SELECT statement with custom configuration
39pub fn format_select_with_config(stmt: &SelectStatement, config: &FormatConfig) -> String {
40    let formatter = AstFormatter::new(config);
41    formatter.format_select(stmt, 0)
42}
43
44/// Format a single SQL expression into a string
45/// This is useful for extracting and displaying parts of a query
46pub fn format_expression(expr: &SqlExpression) -> String {
47    let config = FormatConfig::default();
48    let formatter = AstFormatter::new(&config);
49    formatter.format_expression(expr)
50}
51
52struct AstFormatter<'a> {
53    config: &'a FormatConfig,
54}
55
56impl<'a> AstFormatter<'a> {
57    fn new(config: &'a FormatConfig) -> Self {
58        Self { config }
59    }
60
61    fn keyword(&self, word: &str) -> String {
62        if self.config.uppercase_keywords {
63            word.to_uppercase()
64        } else {
65            word.to_lowercase()
66        }
67    }
68
69    fn indent(&self, level: usize) -> String {
70        self.config.indent.repeat(level)
71    }
72
73    fn format_select(&self, stmt: &SelectStatement, indent_level: usize) -> String {
74        let mut result = String::new();
75        let indent = self.indent(indent_level);
76
77        // CTEs (WITH clause)
78        if !stmt.ctes.is_empty() {
79            writeln!(&mut result, "{}{}", indent, self.keyword("WITH")).unwrap();
80            for (i, cte) in stmt.ctes.iter().enumerate() {
81                let is_last = i == stmt.ctes.len() - 1;
82                self.format_cte(&mut result, cte, indent_level + 1, is_last);
83            }
84        }
85
86        // SELECT clause
87        write!(&mut result, "{}{}", indent, self.keyword("SELECT")).unwrap();
88        if stmt.distinct {
89            write!(&mut result, " {}", self.keyword("DISTINCT")).unwrap();
90        }
91
92        // Format select items
93        if stmt.select_items.is_empty() && !stmt.columns.is_empty() {
94            // Legacy columns field
95            self.format_column_list(&mut result, &stmt.columns, indent_level);
96        } else {
97            self.format_select_items(&mut result, &stmt.select_items, indent_level);
98        }
99
100        // INTO clause (for SELECT INTO #temp)
101        if let Some(ref into_table) = stmt.into_table {
102            writeln!(&mut result).unwrap();
103            write!(
104                &mut result,
105                "{}{} {}",
106                indent,
107                self.keyword("INTO"),
108                into_table.name
109            )
110            .unwrap();
111        }
112
113        // FROM clause
114        if let Some(ref table) = stmt.from_table {
115            writeln!(&mut result).unwrap();
116            write!(&mut result, "{}{} {}", indent, self.keyword("FROM"), table).unwrap();
117        } else if let Some(ref subquery) = stmt.from_subquery {
118            writeln!(&mut result).unwrap();
119            write!(&mut result, "{}{} (", indent, self.keyword("FROM")).unwrap();
120            writeln!(&mut result).unwrap();
121            let subquery_sql = self.format_select(subquery, indent_level + 1);
122            write!(&mut result, "{}", subquery_sql).unwrap();
123            write!(&mut result, "\n{}", indent).unwrap();
124            write!(&mut result, ")").unwrap();
125            if let Some(ref alias) = stmt.from_alias {
126                write!(&mut result, " {} {}", self.keyword("AS"), alias).unwrap();
127            }
128        } else if let Some(ref func) = stmt.from_function {
129            writeln!(&mut result).unwrap();
130            write!(&mut result, "{}{} ", indent, self.keyword("FROM")).unwrap();
131            self.format_table_function(&mut result, func);
132            if let Some(ref alias) = stmt.from_alias {
133                write!(&mut result, " {} {}", self.keyword("AS"), alias).unwrap();
134            }
135        }
136
137        // JOIN clauses
138        for join in &stmt.joins {
139            writeln!(&mut result).unwrap();
140            self.format_join(&mut result, join, indent_level);
141        }
142
143        // WHERE clause
144        if let Some(ref where_clause) = stmt.where_clause {
145            writeln!(&mut result).unwrap();
146            write!(&mut result, "{}{}", indent, self.keyword("WHERE")).unwrap();
147            self.format_where_clause(&mut result, where_clause, indent_level);
148        }
149
150        // GROUP BY clause
151        if let Some(ref group_by) = stmt.group_by {
152            writeln!(&mut result).unwrap();
153            write!(&mut result, "{}{} ", indent, self.keyword("GROUP BY")).unwrap();
154            for (i, expr) in group_by.iter().enumerate() {
155                if i > 0 {
156                    write!(&mut result, ", ").unwrap();
157                }
158                write!(&mut result, "{}", self.format_expression(expr)).unwrap();
159            }
160        }
161
162        // HAVING clause
163        if let Some(ref having) = stmt.having {
164            writeln!(&mut result).unwrap();
165            write!(
166                &mut result,
167                "{}{} {}",
168                indent,
169                self.keyword("HAVING"),
170                self.format_expression(having)
171            )
172            .unwrap();
173        }
174
175        // ORDER BY clause
176        if let Some(ref order_by) = stmt.order_by {
177            writeln!(&mut result).unwrap();
178            write!(&mut result, "{}{} ", indent, self.keyword("ORDER BY")).unwrap();
179            for (i, col) in order_by.iter().enumerate() {
180                if i > 0 {
181                    write!(&mut result, ", ").unwrap();
182                }
183                write!(&mut result, "{}", col.column).unwrap();
184                match col.direction {
185                    SortDirection::Asc => write!(&mut result, " {}", self.keyword("ASC")).unwrap(),
186                    SortDirection::Desc => {
187                        write!(&mut result, " {}", self.keyword("DESC")).unwrap()
188                    }
189                }
190            }
191        }
192
193        // LIMIT clause
194        if let Some(limit) = stmt.limit {
195            writeln!(&mut result).unwrap();
196            write!(&mut result, "{}{} {}", indent, self.keyword("LIMIT"), limit).unwrap();
197        }
198
199        // OFFSET clause
200        if let Some(offset) = stmt.offset {
201            writeln!(&mut result).unwrap();
202            write!(
203                &mut result,
204                "{}{} {}",
205                indent,
206                self.keyword("OFFSET"),
207                offset
208            )
209            .unwrap();
210        }
211
212        result
213    }
214
215    fn format_cte(&self, result: &mut String, cte: &CTE, indent_level: usize, is_last: bool) {
216        let indent = self.indent(indent_level);
217
218        // Add WEB keyword for Web CTEs
219        let is_web = matches!(&cte.cte_type, crate::sql::parser::ast::CTEType::Web(_));
220        if is_web {
221            write!(result, "{}{} {}", indent, self.keyword("WEB"), cte.name).unwrap();
222        } else {
223            write!(result, "{}{}", indent, cte.name).unwrap();
224        }
225
226        if let Some(ref columns) = cte.column_list {
227            write!(result, "(").unwrap();
228            for (i, col) in columns.iter().enumerate() {
229                if i > 0 {
230                    write!(result, ", ").unwrap();
231                }
232                write!(result, "{}", col).unwrap();
233            }
234            write!(result, ")").unwrap();
235        }
236
237        writeln!(result, " {} (", self.keyword("AS")).unwrap();
238        let cte_sql = match &cte.cte_type {
239            crate::sql::parser::ast::CTEType::Standard(query) => {
240                self.format_select(query, indent_level + 1)
241            }
242            crate::sql::parser::ast::CTEType::Web(web_spec) => {
243                // Format WEB CTE
244                let mut web_str = format!(
245                    "{}{} '{}'",
246                    "    ".repeat(indent_level + 1),
247                    self.keyword("URL"),
248                    web_spec.url
249                );
250
251                // Add METHOD if specified
252                if let Some(method) = &web_spec.method {
253                    web_str.push_str(&format!(
254                        " {} {}",
255                        self.keyword("METHOD"),
256                        match method {
257                            crate::sql::parser::ast::HttpMethod::GET => "GET",
258                            crate::sql::parser::ast::HttpMethod::POST => "POST",
259                            crate::sql::parser::ast::HttpMethod::PUT => "PUT",
260                            crate::sql::parser::ast::HttpMethod::DELETE => "DELETE",
261                            crate::sql::parser::ast::HttpMethod::PATCH => "PATCH",
262                        }
263                    ));
264                }
265
266                // Add BODY if specified
267                if let Some(body) = &web_spec.body {
268                    // Check if the body looks like JSON (starts with { or [)
269                    let trimmed_body = body.trim();
270                    if (trimmed_body.starts_with('{') && trimmed_body.ends_with('}'))
271                        || (trimmed_body.starts_with('[') && trimmed_body.ends_with(']'))
272                    {
273                        // Try to prettify JSON
274                        match serde_json::from_str::<serde_json::Value>(trimmed_body) {
275                            Ok(json_val) => {
276                                // Pretty print JSON with 2-space indentation
277                                match serde_json::to_string_pretty(&json_val) {
278                                    Ok(pretty_json) => {
279                                        // Check if JSON is complex (multiline or has special chars)
280                                        let is_complex = pretty_json.lines().count() > 1
281                                            || pretty_json.contains('"')
282                                            || pretty_json.contains('\\');
283
284                                        if is_complex {
285                                            // Use $JSON$ delimiters for complex JSON
286                                            let base_indent = "    ".repeat(indent_level + 1);
287                                            let json_lines: Vec<String> = pretty_json
288                                                .lines()
289                                                .enumerate()
290                                                .map(|(i, line)| {
291                                                    if i == 0 {
292                                                        line.to_string()
293                                                    } else {
294                                                        format!("{}{}", base_indent, line)
295                                                    }
296                                                })
297                                                .collect();
298                                            let formatted_json = json_lines.join("\n");
299
300                                            web_str.push_str(&format!(
301                                                " {} $JSON${}\n{}$JSON$\n{}",
302                                                self.keyword("BODY"),
303                                                formatted_json,
304                                                base_indent,
305                                                base_indent
306                                            ));
307                                        } else {
308                                            // Simple JSON, use regular single quotes
309                                            web_str.push_str(&format!(
310                                                " {} '{}'",
311                                                self.keyword("BODY"),
312                                                pretty_json
313                                            ));
314                                        }
315                                    }
316                                    Err(_) => {
317                                        // Fall back to original if pretty print fails
318                                        web_str.push_str(&format!(
319                                            " {} '{}'",
320                                            self.keyword("BODY"),
321                                            body
322                                        ));
323                                    }
324                                }
325                            }
326                            Err(_) => {
327                                // Not valid JSON, use as-is
328                                web_str.push_str(&format!(" {} '{}'", self.keyword("BODY"), body));
329                            }
330                        }
331                    } else {
332                        // Not JSON, use as-is
333                        web_str.push_str(&format!(" {} '{}'", self.keyword("BODY"), body));
334                    }
335                }
336
337                // Add FORMAT if specified
338                if let Some(format) = &web_spec.format {
339                    web_str.push_str(&format!(
340                        " {} {}",
341                        self.keyword("FORMAT"),
342                        match format {
343                            crate::sql::parser::ast::DataFormat::CSV => "CSV",
344                            crate::sql::parser::ast::DataFormat::JSON => "JSON",
345                            crate::sql::parser::ast::DataFormat::Auto => "AUTO",
346                        }
347                    ));
348                }
349
350                // Add JSON_PATH if specified
351                if let Some(json_path) = &web_spec.json_path {
352                    web_str.push_str(&format!(" {} '{}'", self.keyword("JSON_PATH"), json_path));
353                }
354
355                // Add CACHE if specified
356                if let Some(cache) = web_spec.cache_seconds {
357                    web_str.push_str(&format!(" {} {}", self.keyword("CACHE"), cache));
358                }
359
360                // Add FORM_FILE entries if specified
361                for (field_name, file_path) in &web_spec.form_files {
362                    web_str.push_str(&format!(
363                        "\n{}{} '{}' '{}'",
364                        "    ".repeat(indent_level + 1),
365                        self.keyword("FORM_FILE"),
366                        field_name,
367                        file_path
368                    ));
369                }
370
371                // Add FORM_FIELD entries if specified
372                for (field_name, value) in &web_spec.form_fields {
373                    // Check if the value looks like JSON (starts with { or [)
374                    let trimmed_value = value.trim();
375                    if (trimmed_value.starts_with('{') && trimmed_value.ends_with('}'))
376                        || (trimmed_value.starts_with('[') && trimmed_value.ends_with(']'))
377                    {
378                        // Try to prettify JSON
379                        match serde_json::from_str::<serde_json::Value>(trimmed_value) {
380                            Ok(json_val) => {
381                                // Pretty print JSON with 2-space indentation
382                                match serde_json::to_string_pretty(&json_val) {
383                                    Ok(pretty_json) => {
384                                        // Check if JSON is complex (multiline or has special chars)
385                                        let is_complex = pretty_json.lines().count() > 1
386                                            || pretty_json.contains('"')
387                                            || pretty_json.contains('\\');
388
389                                        if is_complex {
390                                            // Use $JSON$ delimiters for complex JSON
391                                            let base_indent = "    ".repeat(indent_level + 1);
392                                            let json_lines: Vec<String> = pretty_json
393                                                .lines()
394                                                .enumerate()
395                                                .map(|(i, line)| {
396                                                    if i == 0 {
397                                                        line.to_string()
398                                                    } else {
399                                                        format!("{}{}", base_indent, line)
400                                                    }
401                                                })
402                                                .collect();
403                                            let formatted_json = json_lines.join("\n");
404
405                                            web_str.push_str(&format!(
406                                                "\n{}{} '{}' $JSON${}\n{}$JSON$",
407                                                base_indent,
408                                                self.keyword("FORM_FIELD"),
409                                                field_name,
410                                                formatted_json,
411                                                base_indent
412                                            ));
413                                        } else {
414                                            // Simple JSON, use regular single quotes
415                                            web_str.push_str(&format!(
416                                                "\n{}{} '{}' '{}'",
417                                                "    ".repeat(indent_level + 1),
418                                                self.keyword("FORM_FIELD"),
419                                                field_name,
420                                                pretty_json
421                                            ));
422                                        }
423                                    }
424                                    Err(_) => {
425                                        // Fall back to original if pretty print fails
426                                        web_str.push_str(&format!(
427                                            "\n{}{} '{}' '{}'",
428                                            "    ".repeat(indent_level + 1),
429                                            self.keyword("FORM_FIELD"),
430                                            field_name,
431                                            value
432                                        ));
433                                    }
434                                }
435                            }
436                            Err(_) => {
437                                // Not valid JSON, use as-is
438                                web_str.push_str(&format!(
439                                    "\n{}{} '{}' '{}'",
440                                    "    ".repeat(indent_level + 1),
441                                    self.keyword("FORM_FIELD"),
442                                    field_name,
443                                    value
444                                ));
445                            }
446                        }
447                    } else {
448                        // Not JSON, use as-is
449                        web_str.push_str(&format!(
450                            "\n{}{} '{}' '{}'",
451                            "    ".repeat(indent_level + 1),
452                            self.keyword("FORM_FIELD"),
453                            field_name,
454                            value
455                        ));
456                    }
457                }
458
459                // Add HEADERS if specified
460                if !web_spec.headers.is_empty() {
461                    web_str.push_str(&format!(" {} (", self.keyword("HEADERS")));
462                    for (i, (key, value)) in web_spec.headers.iter().enumerate() {
463                        if i > 0 {
464                            web_str.push_str(", ");
465                        }
466                        web_str.push_str(&format!("'{}': '{}'", key, value));
467                    }
468                    web_str.push(')');
469                }
470
471                web_str
472            }
473        };
474        write!(result, "{}", cte_sql).unwrap();
475        writeln!(result).unwrap();
476        write!(result, "{}", indent).unwrap();
477        if is_last {
478            writeln!(result, ")").unwrap();
479        } else {
480            writeln!(result, "),").unwrap();
481        }
482    }
483
484    fn format_column_list(&self, result: &mut String, columns: &[String], indent_level: usize) {
485        if columns.len() <= self.config.items_per_line {
486            // Single line
487            write!(result, " ").unwrap();
488            for (i, col) in columns.iter().enumerate() {
489                if i > 0 {
490                    write!(result, ", ").unwrap();
491                }
492                write!(result, "{}", col).unwrap();
493            }
494        } else {
495            // Multi-line
496            writeln!(result).unwrap();
497            let indent = self.indent(indent_level + 1);
498            for (i, col) in columns.iter().enumerate() {
499                write!(result, "{}{}", indent, col).unwrap();
500                if i < columns.len() - 1 {
501                    writeln!(result, ",").unwrap();
502                }
503            }
504        }
505    }
506
507    fn format_select_items(&self, result: &mut String, items: &[SelectItem], indent_level: usize) {
508        if items.is_empty() {
509            write!(result, " *").unwrap();
510            return;
511        }
512
513        // Count non-star items for formatting decision
514        let _non_star_count = items
515            .iter()
516            .filter(|i| !matches!(i, SelectItem::Star { .. }))
517            .count();
518
519        // Check if any item is complex (function calls, CASE expressions, etc.)
520        let has_complex_items = items.iter().any(|item| match item {
521            SelectItem::Expression { expr, .. } => self.is_complex_expression(expr),
522            _ => false,
523        });
524
525        // Calculate total approximate length if on single line
526        let single_line_length: usize = items
527            .iter()
528            .map(|item| {
529                match item {
530                    SelectItem::Star { .. } => 1,
531                    SelectItem::Column { column: col, .. } => col.name.len(),
532                    SelectItem::Expression { expr, alias, .. } => {
533                        self.format_expression(expr).len() + 4 + alias.len() // " AS " = 4
534                    }
535                }
536            })
537            .sum::<usize>()
538            + (items.len() - 1) * 2; // ", " between items
539
540        // Use multi-line formatting by default unless:
541        // - It's a single simple column or star
542        // - It's 2-3 simple columns with total length < 40 chars
543        let use_single_line = match items.len() {
544            1 => !has_complex_items, // Single item: only if simple
545            2..=3 => !has_complex_items && single_line_length < 40, // 2-3 items: only if very short
546            _ => false,              // 4+ items: always multi-line
547        };
548
549        if !use_single_line {
550            // Multi-line
551            writeln!(result).unwrap();
552            let indent = self.indent(indent_level + 1);
553            for (i, item) in items.iter().enumerate() {
554                write!(result, "{}", indent).unwrap();
555                self.format_select_item(result, item);
556                if i < items.len() - 1 {
557                    writeln!(result, ",").unwrap();
558                }
559            }
560        } else {
561            // Single line
562            write!(result, " ").unwrap();
563            for (i, item) in items.iter().enumerate() {
564                if i > 0 {
565                    write!(result, ", ").unwrap();
566                }
567                self.format_select_item(result, item);
568            }
569        }
570    }
571
572    fn is_complex_expression(&self, expr: &SqlExpression) -> bool {
573        match expr {
574            SqlExpression::CaseExpression { .. } => true,
575            SqlExpression::FunctionCall { .. } => true,
576            SqlExpression::WindowFunction { .. } => true,
577            SqlExpression::ScalarSubquery { .. } => true,
578            SqlExpression::InSubquery { .. } => true,
579            SqlExpression::NotInSubquery { .. } => true,
580            SqlExpression::BinaryOp { left, right, .. } => {
581                self.is_complex_expression(left) || self.is_complex_expression(right)
582            }
583            _ => false,
584        }
585    }
586
587    fn format_select_item(&self, result: &mut String, item: &SelectItem) {
588        match item {
589            SelectItem::Star { .. } => write!(result, "*").unwrap(),
590            SelectItem::Column { column: col, .. } => write!(result, "{}", col.to_sql()).unwrap(),
591            SelectItem::Expression { expr, alias, .. } => {
592                write!(
593                    result,
594                    "{} {} {}",
595                    self.format_expression(expr),
596                    self.keyword("AS"),
597                    alias
598                )
599                .unwrap();
600            }
601        }
602    }
603
604    fn format_expression(&self, expr: &SqlExpression) -> String {
605        match expr {
606            SqlExpression::Column(column_ref) => column_ref.to_sql(),
607            SqlExpression::StringLiteral(s) => format!("'{}'", s),
608            SqlExpression::NumberLiteral(n) => n.clone(),
609            SqlExpression::BooleanLiteral(b) => b.to_string().to_uppercase(),
610            SqlExpression::Null => self.keyword("NULL"),
611            SqlExpression::BinaryOp { left, op, right } => {
612                // Special handling for IS NULL / IS NOT NULL operators
613                if op == "IS NULL" || op == "IS NOT NULL" {
614                    format!("{} {}", self.format_expression(left), op)
615                } else {
616                    format!(
617                        "{} {} {}",
618                        self.format_expression(left),
619                        op,
620                        self.format_expression(right)
621                    )
622                }
623            }
624            SqlExpression::FunctionCall {
625                name,
626                args,
627                distinct,
628            } => {
629                let mut result = name.clone();
630                result.push('(');
631                if *distinct {
632                    result.push_str(&self.keyword("DISTINCT"));
633                    result.push(' ');
634                }
635                for (i, arg) in args.iter().enumerate() {
636                    if i > 0 {
637                        result.push_str(", ");
638                    }
639                    result.push_str(&self.format_expression(arg));
640                }
641                result.push(')');
642                result
643            }
644            SqlExpression::CaseExpression {
645                when_branches,
646                else_branch,
647            } => {
648                // Format CASE expressions on multiple lines for readability
649                let mut result = String::new();
650                result.push_str(&self.keyword("CASE"));
651                result.push('\n');
652
653                // Format each WHEN branch on its own line with indentation
654                for branch in when_branches {
655                    result.push_str("        "); // 8 spaces for WHEN indent
656                    result.push_str(&format!(
657                        "{} {} {} {}",
658                        self.keyword("WHEN"),
659                        self.format_expression(&branch.condition),
660                        self.keyword("THEN"),
661                        self.format_expression(&branch.result)
662                    ));
663                    result.push('\n');
664                }
665
666                // Format ELSE clause if present
667                if let Some(else_expr) = else_branch {
668                    result.push_str("        "); // 8 spaces for ELSE indent
669                    result.push_str(&format!(
670                        "{} {}",
671                        self.keyword("ELSE"),
672                        self.format_expression(else_expr)
673                    ));
674                    result.push('\n');
675                }
676
677                result.push_str("    "); // 4 spaces for END
678                result.push_str(&self.keyword("END"));
679                result
680            }
681            SqlExpression::SimpleCaseExpression {
682                expr,
683                when_branches,
684                else_branch,
685            } => {
686                // Format simple CASE expressions on multiple lines for readability
687                let mut result = String::new();
688                result.push_str(&format!(
689                    "{} {}",
690                    self.keyword("CASE"),
691                    self.format_expression(expr)
692                ));
693                result.push('\n');
694
695                // Format each WHEN branch on its own line with indentation
696                for branch in when_branches {
697                    result.push_str("        "); // 8 spaces for WHEN indent
698                    result.push_str(&format!(
699                        "{} {} {} {}",
700                        self.keyword("WHEN"),
701                        self.format_expression(&branch.value),
702                        self.keyword("THEN"),
703                        self.format_expression(&branch.result)
704                    ));
705                    result.push('\n');
706                }
707
708                // Format ELSE clause if present
709                if let Some(else_expr) = else_branch {
710                    result.push_str("        "); // 8 spaces for ELSE indent
711                    result.push_str(&format!(
712                        "{} {}",
713                        self.keyword("ELSE"),
714                        self.format_expression(else_expr)
715                    ));
716                    result.push('\n');
717                }
718
719                result.push_str("    "); // 4 spaces for END
720                result.push_str(&self.keyword("END"));
721                result
722            }
723            SqlExpression::Between { expr, lower, upper } => {
724                format!(
725                    "{} {} {} {} {}",
726                    self.format_expression(expr),
727                    self.keyword("BETWEEN"),
728                    self.format_expression(lower),
729                    self.keyword("AND"),
730                    self.format_expression(upper)
731                )
732            }
733            SqlExpression::InList { expr, values } => {
734                let mut result =
735                    format!("{} {} (", self.format_expression(expr), self.keyword("IN"));
736                for (i, val) in values.iter().enumerate() {
737                    if i > 0 {
738                        result.push_str(", ");
739                    }
740                    result.push_str(&self.format_expression(val));
741                }
742                result.push(')');
743                result
744            }
745            SqlExpression::NotInList { expr, values } => {
746                let mut result = format!(
747                    "{} {} {} (",
748                    self.format_expression(expr),
749                    self.keyword("NOT"),
750                    self.keyword("IN")
751                );
752                for (i, val) in values.iter().enumerate() {
753                    if i > 0 {
754                        result.push_str(", ");
755                    }
756                    result.push_str(&self.format_expression(val));
757                }
758                result.push(')');
759                result
760            }
761            SqlExpression::Not { expr } => {
762                format!("{} {}", self.keyword("NOT"), self.format_expression(expr))
763            }
764            SqlExpression::ScalarSubquery { query } => {
765                // Check if subquery is complex enough to warrant multi-line formatting
766                let subquery_str = self.format_select(query, 0);
767                if subquery_str.contains('\n') || subquery_str.len() > 60 {
768                    // Multi-line formatting
769                    format!("(\n{}\n)", self.format_select(query, 1))
770                } else {
771                    // Inline formatting
772                    format!("({})", subquery_str)
773                }
774            }
775            SqlExpression::InSubquery { expr, subquery } => {
776                let subquery_str = self.format_select(subquery, 0);
777                if subquery_str.contains('\n') || subquery_str.len() > 60 {
778                    // Multi-line formatting
779                    format!(
780                        "{} {} (\n{}\n)",
781                        self.format_expression(expr),
782                        self.keyword("IN"),
783                        self.format_select(subquery, 1)
784                    )
785                } else {
786                    // Inline formatting
787                    format!(
788                        "{} {} ({})",
789                        self.format_expression(expr),
790                        self.keyword("IN"),
791                        subquery_str
792                    )
793                }
794            }
795            SqlExpression::NotInSubquery { expr, subquery } => {
796                let subquery_str = self.format_select(subquery, 0);
797                if subquery_str.contains('\n') || subquery_str.len() > 60 {
798                    // Multi-line formatting
799                    format!(
800                        "{} {} {} (\n{}\n)",
801                        self.format_expression(expr),
802                        self.keyword("NOT"),
803                        self.keyword("IN"),
804                        self.format_select(subquery, 1)
805                    )
806                } else {
807                    // Inline formatting
808                    format!(
809                        "{} {} {} ({})",
810                        self.format_expression(expr),
811                        self.keyword("NOT"),
812                        self.keyword("IN"),
813                        subquery_str
814                    )
815                }
816            }
817            SqlExpression::MethodCall {
818                object,
819                method,
820                args,
821            } => {
822                let mut result = format!("{}.{}", object, method);
823                result.push('(');
824                for (i, arg) in args.iter().enumerate() {
825                    if i > 0 {
826                        result.push_str(", ");
827                    }
828                    result.push_str(&self.format_expression(arg));
829                }
830                result.push(')');
831                result
832            }
833            SqlExpression::ChainedMethodCall { base, method, args } => {
834                let mut result = format!("{}.{}", self.format_expression(base), method);
835                result.push('(');
836                for (i, arg) in args.iter().enumerate() {
837                    if i > 0 {
838                        result.push_str(", ");
839                    }
840                    result.push_str(&self.format_expression(arg));
841                }
842                result.push(')');
843                result
844            }
845            SqlExpression::WindowFunction {
846                name,
847                args,
848                window_spec,
849            } => {
850                let mut result = format!("{}(", name);
851
852                // Add function arguments
853                for (i, arg) in args.iter().enumerate() {
854                    if i > 0 {
855                        result.push_str(", ");
856                    }
857                    result.push_str(&self.format_expression(arg));
858                }
859                result.push_str(") ");
860                result.push_str(&self.keyword("OVER"));
861                result.push_str(" (");
862
863                // Add PARTITION BY clause if present
864                if !window_spec.partition_by.is_empty() {
865                    result.push_str(&self.keyword("PARTITION BY"));
866                    result.push(' ');
867                    for (i, col) in window_spec.partition_by.iter().enumerate() {
868                        if i > 0 {
869                            result.push_str(", ");
870                        }
871                        result.push_str(col);
872                    }
873                }
874
875                // Add ORDER BY clause if present
876                if !window_spec.order_by.is_empty() {
877                    if !window_spec.partition_by.is_empty() {
878                        result.push(' ');
879                    }
880                    result.push_str(&self.keyword("ORDER BY"));
881                    result.push(' ');
882                    for (i, col) in window_spec.order_by.iter().enumerate() {
883                        if i > 0 {
884                            result.push_str(", ");
885                        }
886                        result.push_str(&col.column);
887                        match col.direction {
888                            SortDirection::Asc => {
889                                result.push(' ');
890                                result.push_str(&self.keyword("ASC"));
891                            }
892                            SortDirection::Desc => {
893                                result.push(' ');
894                                result.push_str(&self.keyword("DESC"));
895                            }
896                        }
897                    }
898                }
899
900                // Add window frame specification if present
901                if let Some(frame) = &window_spec.frame {
902                    // Add space before frame specification
903                    if !window_spec.partition_by.is_empty() || !window_spec.order_by.is_empty() {
904                        result.push(' ');
905                    }
906
907                    // Format frame unit (ROWS or RANGE)
908                    match frame.unit {
909                        FrameUnit::Rows => result.push_str(&self.keyword("ROWS")),
910                        FrameUnit::Range => result.push_str(&self.keyword("RANGE")),
911                    }
912
913                    result.push(' ');
914
915                    // Format frame bounds
916                    if let Some(end) = &frame.end {
917                        // BETWEEN start AND end
918                        result.push_str(&self.keyword("BETWEEN"));
919                        result.push(' ');
920                        result.push_str(&self.format_frame_bound(&frame.start));
921                        result.push(' ');
922                        result.push_str(&self.keyword("AND"));
923                        result.push(' ');
924                        result.push_str(&self.format_frame_bound(end));
925                    } else {
926                        // Just a single bound (uncommon but valid)
927                        result.push_str(&self.format_frame_bound(&frame.start));
928                    }
929                }
930
931                result.push(')');
932                result
933            }
934            SqlExpression::DateTimeConstructor {
935                year,
936                month,
937                day,
938                hour,
939                minute,
940                second,
941            } => {
942                if let (Some(h), Some(m), Some(s)) = (hour, minute, second) {
943                    format!(
944                        "DateTime({}, {}, {}, {}, {}, {})",
945                        year, month, day, h, m, s
946                    )
947                } else {
948                    format!("DateTime({}, {}, {})", year, month, day)
949                }
950            }
951            SqlExpression::DateTimeToday {
952                hour,
953                minute,
954                second,
955            } => {
956                if let (Some(h), Some(m), Some(s)) = (hour, minute, second) {
957                    format!("Today({}, {}, {})", h, m, s)
958                } else {
959                    "Today()".to_string()
960                }
961            }
962            _ => format!("{:?}", expr), // Fallback for unhandled expression types
963        }
964    }
965
966    fn format_where_clause(
967        &self,
968        result: &mut String,
969        where_clause: &WhereClause,
970        indent_level: usize,
971    ) {
972        let needs_multiline = where_clause.conditions.len() > 1;
973
974        if needs_multiline {
975            writeln!(result).unwrap();
976            let indent = self.indent(indent_level + 1);
977            for (i, condition) in where_clause.conditions.iter().enumerate() {
978                if i > 0 {
979                    if let Some(ref connector) = where_clause.conditions[i - 1].connector {
980                        let connector_str = match connector {
981                            LogicalOp::And => self.keyword("AND"),
982                            LogicalOp::Or => self.keyword("OR"),
983                        };
984                        writeln!(result).unwrap();
985                        write!(result, "{}{} ", indent, connector_str).unwrap();
986                    }
987                } else {
988                    write!(result, "{}", indent).unwrap();
989                }
990                write!(result, "{}", self.format_expression(&condition.expr)).unwrap();
991            }
992        } else if let Some(condition) = where_clause.conditions.first() {
993            write!(result, " {}", self.format_expression(&condition.expr)).unwrap();
994        }
995    }
996
997    fn format_frame_bound(&self, bound: &FrameBound) -> String {
998        match bound {
999            FrameBound::UnboundedPreceding => self.keyword("UNBOUNDED PRECEDING"),
1000            FrameBound::CurrentRow => self.keyword("CURRENT ROW"),
1001            FrameBound::UnboundedFollowing => self.keyword("UNBOUNDED FOLLOWING"),
1002            FrameBound::Preceding(n) => format!("{} {}", n, self.keyword("PRECEDING")),
1003            FrameBound::Following(n) => format!("{} {}", n, self.keyword("FOLLOWING")),
1004        }
1005    }
1006
1007    fn format_join(&self, result: &mut String, join: &JoinClause, indent_level: usize) {
1008        let indent = self.indent(indent_level);
1009        let join_type = match join.join_type {
1010            JoinType::Inner => self.keyword("INNER JOIN"),
1011            JoinType::Left => self.keyword("LEFT JOIN"),
1012            JoinType::Right => self.keyword("RIGHT JOIN"),
1013            JoinType::Full => self.keyword("FULL JOIN"),
1014            JoinType::Cross => self.keyword("CROSS JOIN"),
1015        };
1016
1017        write!(result, "{}{} ", indent, join_type).unwrap();
1018
1019        match &join.table {
1020            TableSource::Table(name) => write!(result, "{}", name).unwrap(),
1021            TableSource::DerivedTable { query, alias } => {
1022                writeln!(result, "(").unwrap();
1023                let subquery_sql = self.format_select(query, indent_level + 1);
1024                write!(result, "{}", subquery_sql).unwrap();
1025                writeln!(result).unwrap();
1026                write!(result, "{}) {} {}", indent, self.keyword("AS"), alias).unwrap();
1027            }
1028        }
1029
1030        if let Some(ref alias) = join.alias {
1031            write!(result, " {} {}", self.keyword("AS"), alias).unwrap();
1032        }
1033
1034        if !join.condition.conditions.is_empty() {
1035            write!(result, " {}", self.keyword("ON")).unwrap();
1036            for (i, condition) in join.condition.conditions.iter().enumerate() {
1037                if i > 0 {
1038                    write!(result, " {}", self.keyword("AND")).unwrap();
1039                }
1040                write!(
1041                    result,
1042                    " {} {} {}",
1043                    condition.left_column,
1044                    self.format_join_operator(&condition.operator),
1045                    condition.right_column
1046                )
1047                .unwrap();
1048            }
1049        }
1050    }
1051
1052    fn format_join_operator(&self, op: &JoinOperator) -> String {
1053        match op {
1054            JoinOperator::Equal => "=",
1055            JoinOperator::NotEqual => "!=",
1056            JoinOperator::LessThan => "<",
1057            JoinOperator::GreaterThan => ">",
1058            JoinOperator::LessThanOrEqual => "<=",
1059            JoinOperator::GreaterThanOrEqual => ">=",
1060        }
1061        .to_string()
1062    }
1063
1064    fn format_table_function(&self, result: &mut String, func: &TableFunction) {
1065        match func {
1066            TableFunction::Generator { name, args } => {
1067                write!(result, "{}(", self.keyword(&name.to_uppercase())).unwrap();
1068                for (i, arg) in args.iter().enumerate() {
1069                    if i > 0 {
1070                        write!(result, ", ").unwrap();
1071                    }
1072                    write!(result, "{}", self.format_expression(arg)).unwrap();
1073                }
1074                write!(result, ")").unwrap();
1075            }
1076        }
1077    }
1078}
1079
1080/// Parse and format SQL query using the AST
1081pub fn format_sql_ast(query: &str) -> Result<String, String> {
1082    use crate::sql::recursive_parser::Parser;
1083
1084    let mut parser = Parser::new(query);
1085    match parser.parse() {
1086        Ok(stmt) => Ok(format_select_statement(&stmt)),
1087        Err(e) => Err(format!("Parse error: {}", e)),
1088    }
1089}
1090
1091/// Parse and format SQL with custom configuration
1092pub fn format_sql_ast_with_config(query: &str, config: &FormatConfig) -> Result<String, String> {
1093    use crate::sql::recursive_parser::Parser;
1094
1095    let mut parser = Parser::new(query);
1096    match parser.parse() {
1097        Ok(stmt) => Ok(format_select_with_config(&stmt, &config)),
1098        Err(e) => Err(format!("Parse error: {}", e)),
1099    }
1100}
sql_cli/sql/parser/ast_formatter.rs

sql_cli/sql/parser/
ast_formatter.rs
