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

sql_cli/sql/parser/
ast_formatter.rs
