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

sql_cli/sql/parser/
ast_formatter.rs
