Skip to main content

qail_core/fmt/
mod.rs

1use crate::ast::{
2    Action, Cage, CageKind, Condition, Expr, Join, LogicalOp, Operator, Qail, SortOrder, Value,
3};
4use std::fmt::{Result, Write};
5
6#[cfg(test)]
7mod tests;
8
9/// Pretty-printer for QAIL AST nodes.
10///
11/// Renders a `Qail` command back into human-readable QAIL syntax
12/// with indentation for readability.
13pub struct Formatter {
14    /// Current indentation depth.
15    indent_level: usize,
16    /// Output buffer.
17    buffer: String,
18}
19
20impl Default for Formatter {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl Formatter {
27    /// Create a new formatter with default settings.
28    pub fn new() -> Self {
29        Self {
30            indent_level: 0,
31            buffer: String::new(),
32        }
33    }
34
35    /// Format a query into a readable QAIL string.
36    pub fn format(mut self, cmd: &Qail) -> std::result::Result<String, std::fmt::Error> {
37        self.visit_cmd(cmd)?;
38        Ok(self.buffer)
39    }
40
41    fn indent(&mut self) -> Result {
42        for _ in 0..self.indent_level {
43            write!(self.buffer, "  ")?;
44        }
45        Ok(())
46    }
47
48    fn visit_cmd(&mut self, cmd: &Qail) -> Result {
49        for cte in &cmd.ctes {
50            write!(self.buffer, "with {} = ", cte.name)?;
51            self.indent_level += 1;
52            writeln!(self.buffer)?;
53            self.indent()?;
54            self.visit_cmd(&cte.base_query)?;
55
56            if cte.recursive
57                && let Some(ref recursive_query) = cte.recursive_query
58            {
59                writeln!(self.buffer)?;
60                self.indent()?;
61                writeln!(self.buffer, "union all")?;
62                self.indent()?;
63                self.visit_cmd(recursive_query)?;
64            }
65
66            self.indent_level -= 1;
67            writeln!(self.buffer)?;
68        }
69
70        // Action and Table
71        match cmd.action {
72            Action::Get => write!(self.buffer, "get {}", cmd.table)?,
73            Action::Set => write!(self.buffer, "set {}", cmd.table)?,
74            Action::Del => write!(self.buffer, "del {}", cmd.table)?,
75            Action::Add => write!(self.buffer, "add {}", cmd.table)?,
76            _ => write!(self.buffer, "{} {}", cmd.action, cmd.table)?, // Fallback for others
77        }
78        writeln!(self.buffer)?;
79
80        // self.indent_level += 1; // Removed: Clauses should act at same level as command
81
82        // Cages: Group By (if any "by" equivalent exists? No, "by" is usually implicit in AST or explicit in group_by_mode?)
83        // The proposal example shows "by phone_number".
84        // In AST `cmd.rs`, there isn't a direct "Group By" list, usually inferred or group_by_mode.
85        // Wait, where is `by phone_number` stored in AST?
86        // Checking `ast/cmd.rs`: `group_by_mode: GroupByMode`.
87        // Usually group by is inferred from aggregates or explicit.
88        // If the AST doesn't have explicit group by columns, we might need to derive it or it's in `cages`?
89        // Let's check `cages.rs` again. `CageKind` has `Filter`, `Sort`, `Limit`... no `GroupBy`.
90        // Maybe it's implied by non-aggregated columns in a `Get` with aggregates?
91        // For now, I will skip "by" unless I find it in AST.
92
93        if !cmd.columns.is_empty() {
94            // But proposal says "Canonical".
95            // "get table" implies "get table fields *" usually?
96            // If manual explicit columns:
97            if !(cmd.columns.len() == 1 && matches!(cmd.columns[0], Expr::Star)) {
98                self.indent()?;
99                writeln!(self.buffer, "fields")?;
100                self.indent_level += 1;
101                for (i, col) in cmd.columns.iter().enumerate() {
102                    self.indent()?;
103                    self.format_column(col)?;
104                    if i < cmd.columns.len() - 1 {
105                        writeln!(self.buffer, ",")?;
106                    } else {
107                        writeln!(self.buffer)?;
108                    }
109                }
110                self.indent_level -= 1;
111            }
112        }
113
114        // Joins
115        for join in &cmd.joins {
116            self.indent()?;
117            self.format_join(join)?;
118            writeln!(self.buffer)?;
119        }
120
121        // Where (Filter Cages)
122        let filters: Vec<&Cage> = cmd
123            .cages
124            .iter()
125            .filter(|c| matches!(c.kind, CageKind::Filter))
126            .collect();
127        if !filters.is_empty() {
128            // We need to merge them or print them?
129            // Proposal says: "where rn = 1"
130            self.indent()?;
131            write!(self.buffer, "where ")?;
132            for (i, cage) in filters.iter().enumerate() {
133                if i > 0 {
134                    write!(self.buffer, " and ")?;
135                }
136                let wrap_or_group = cage.logical_op == LogicalOp::Or && cage.conditions.len() > 1;
137                if wrap_or_group {
138                    write!(self.buffer, "(")?;
139                }
140                self.format_conditions(&cage.conditions, cage.logical_op)?;
141                if wrap_or_group {
142                    write!(self.buffer, ")")?;
143                }
144            }
145            writeln!(self.buffer)?;
146        }
147
148        // Order By (Sort Cages)
149        let sorts: Vec<&Cage> = cmd
150            .cages
151            .iter()
152            .filter(|c| matches!(c.kind, CageKind::Sort(_)))
153            .collect();
154        if !sorts.is_empty() {
155            self.indent()?;
156            writeln!(self.buffer, "order by")?;
157            self.indent_level += 1;
158            for (i, cage) in sorts.iter().enumerate() {
159                if let CageKind::Sort(order) = cage.kind {
160                    for (j, cond) in cage.conditions.iter().enumerate() {
161                        self.indent()?;
162                        write!(self.buffer, "{}", cond.left)?;
163                        self.format_sort_order(order)?;
164                        if i < sorts.len() - 1 || j < cage.conditions.len() - 1 {
165                            writeln!(self.buffer, ",")?;
166                        } else {
167                            writeln!(self.buffer)?;
168                        }
169                    }
170                }
171            }
172            self.indent_level -= 1;
173        }
174
175        for cage in &cmd.cages {
176            match cage.kind {
177                CageKind::Limit(n) => {
178                    self.indent()?;
179                    writeln!(self.buffer, "limit {}", n)?;
180                }
181                CageKind::Offset(n) => {
182                    self.indent()?;
183                    writeln!(self.buffer, "offset {}", n)?;
184                }
185                _ => {}
186            }
187        }
188
189        // self.indent_level -= 1; // Removed matching decrement
190        Ok(())
191    }
192
193    fn format_column(&mut self, col: &Expr) -> Result {
194        match col {
195            Expr::Star => write!(self.buffer, "*")?,
196            Expr::Named(name) => write!(self.buffer, "{}", name)?,
197            Expr::Aliased { name, alias } => write!(self.buffer, "{} as {}", name, alias)?,
198            Expr::Aggregate {
199                col,
200                func,
201                distinct,
202                filter,
203                alias,
204            } => {
205                let func_name = match func {
206                    crate::ast::AggregateFunc::Count => "count",
207                    crate::ast::AggregateFunc::Sum => "sum",
208                    crate::ast::AggregateFunc::Avg => "avg",
209                    crate::ast::AggregateFunc::Min => "min",
210                    crate::ast::AggregateFunc::Max => "max",
211                    crate::ast::AggregateFunc::ArrayAgg => "array_agg",
212                    crate::ast::AggregateFunc::StringAgg => "string_agg",
213                    crate::ast::AggregateFunc::JsonAgg => "json_agg",
214                    crate::ast::AggregateFunc::JsonbAgg => "jsonb_agg",
215                    crate::ast::AggregateFunc::BoolAnd => "bool_and",
216                    crate::ast::AggregateFunc::BoolOr => "bool_or",
217                };
218                if *distinct {
219                    write!(self.buffer, "{}(distinct {})", func_name, col)?;
220                } else {
221                    write!(self.buffer, "{}({})", func_name, col)?;
222                }
223                if let Some(conditions) = filter {
224                    write!(
225                        self.buffer,
226                        " filter (where {})",
227                        conditions
228                            .iter()
229                            .map(|c| c.to_string())
230                            .collect::<Vec<_>>()
231                            .join(" and ")
232                    )?;
233                }
234                if let Some(a) = alias {
235                    write!(self.buffer, " as {}", a)?;
236                }
237            }
238            Expr::FunctionCall { name, args, alias } => {
239                let args_str: Vec<String> = args.iter().map(|a| a.to_string()).collect();
240                write!(self.buffer, "{}({})", name, args_str.join(", "))?;
241                if let Some(a) = alias {
242                    write!(self.buffer, " as {}", a)?;
243                }
244            }
245            Expr::Window {
246                name,
247                func,
248                params,
249                partition,
250                ..
251            } => {
252                // Use Window function format: func(params) OVER (PARTITION BY ...)
253                let params_str: Vec<String> = params.iter().map(|p| p.to_string()).collect();
254                write!(self.buffer, "{}({})", func, params_str.join(", "))?;
255                if !partition.is_empty() {
256                    write!(self.buffer, " over (partition by {})", partition.join(", "))?;
257                }
258                write!(self.buffer, " as {}", name)?;
259            }
260            Expr::Case {
261                when_clauses,
262                else_value,
263                alias,
264            } => {
265                write!(self.buffer, "case")?;
266                for (cond, val) in when_clauses {
267                    write!(self.buffer, " when {} then {}", cond.left, val)?;
268                }
269                if let Some(e) = else_value {
270                    write!(self.buffer, " else {}", e)?;
271                }
272                write!(self.buffer, " end")?;
273                if let Some(a) = alias {
274                    write!(self.buffer, " as {}", a)?;
275                }
276            }
277            Expr::JsonAccess {
278                column,
279                path_segments,
280                alias,
281            } => {
282                write!(self.buffer, "{}", column)?;
283                for (path, as_text) in path_segments {
284                    let op = if *as_text { "->>" } else { "->" };
285                    if path.parse::<i64>().is_ok() {
286                        write!(self.buffer, "{}{}", op, path)?;
287                    } else {
288                        write!(self.buffer, "{}'{}'", op, path)?;
289                    }
290                }
291                if let Some(a) = alias {
292                    write!(self.buffer, " as {}", a)?;
293                }
294            }
295            Expr::Cast {
296                expr,
297                target_type,
298                alias,
299            } => {
300                write!(self.buffer, "{}::{}", expr, target_type)?;
301                if let Some(a) = alias {
302                    write!(self.buffer, " as {}", a)?;
303                }
304            }
305            Expr::Binary {
306                left,
307                op,
308                right,
309                alias,
310            } => {
311                write!(self.buffer, "({} {} {})", left, op, right)?;
312                if let Some(a) = alias {
313                    write!(self.buffer, " as {}", a)?;
314                }
315            }
316            Expr::SpecialFunction { name, args, alias } => {
317                write!(self.buffer, "{}(", name)?;
318                for (i, (keyword, expr)) in args.iter().enumerate() {
319                    if i > 0 {
320                        write!(self.buffer, " ")?;
321                    }
322                    if let Some(kw) = keyword {
323                        write!(self.buffer, "{} ", kw)?;
324                    }
325                    write!(self.buffer, "{}", expr)?;
326                }
327                write!(self.buffer, ")")?;
328                if let Some(a) = alias {
329                    write!(self.buffer, " as {}", a)?;
330                }
331            }
332            Expr::Literal(val) => self.format_value(val)?,
333            Expr::Def {
334                name,
335                data_type,
336                constraints,
337            } => {
338                write!(self.buffer, "{}:{}", name, data_type)?;
339                for c in constraints {
340                    write!(self.buffer, "^{}", c)?;
341                }
342            }
343            Expr::Mod { kind, col } => {
344                let prefix = match kind {
345                    crate::ast::ModKind::Add => "+",
346                    crate::ast::ModKind::Drop => "-",
347                };
348                write!(self.buffer, "{}{}", prefix, col)?;
349            }
350            Expr::ArrayConstructor { elements, alias } => {
351                write!(self.buffer, "ARRAY[")?;
352                for (i, elem) in elements.iter().enumerate() {
353                    if i > 0 {
354                        write!(self.buffer, ", ")?;
355                    }
356                    self.format_column(elem)?;
357                }
358                write!(self.buffer, "]")?;
359                if let Some(a) = alias {
360                    write!(self.buffer, " as {}", a)?;
361                }
362            }
363            Expr::RowConstructor { elements, alias } => {
364                write!(self.buffer, "ROW(")?;
365                for (i, elem) in elements.iter().enumerate() {
366                    if i > 0 {
367                        write!(self.buffer, ", ")?;
368                    }
369                    self.format_column(elem)?;
370                }
371                write!(self.buffer, ")")?;
372                if let Some(a) = alias {
373                    write!(self.buffer, " as {}", a)?;
374                }
375            }
376            Expr::Subscript { expr, index, alias } => {
377                self.format_column(expr)?;
378                write!(self.buffer, "[")?;
379                self.format_column(index)?;
380                write!(self.buffer, "]")?;
381                if let Some(a) = alias {
382                    write!(self.buffer, " as {}", a)?;
383                }
384            }
385            Expr::Collate {
386                expr,
387                collation,
388                alias,
389            } => {
390                self.format_column(expr)?;
391                write!(self.buffer, " COLLATE \"{}\"", collation)?;
392                if let Some(a) = alias {
393                    write!(self.buffer, " as {}", a)?;
394                }
395            }
396            Expr::FieldAccess { expr, field, alias } => {
397                write!(self.buffer, "(")?;
398                self.format_column(expr)?;
399                write!(self.buffer, ").{}", field)?;
400                if let Some(a) = alias {
401                    write!(self.buffer, " as {}", a)?;
402                }
403            }
404            Expr::Subquery { query, alias } => {
405                write!(self.buffer, "(")?;
406                self.visit_cmd(query)?;
407                write!(self.buffer, ")")?;
408                if let Some(a) = alias {
409                    write!(self.buffer, " as {}", a)?;
410                }
411            }
412            Expr::Exists {
413                query,
414                negated,
415                alias,
416            } => {
417                if *negated {
418                    write!(self.buffer, "not ")?;
419                }
420                write!(self.buffer, "exists (")?;
421                self.visit_cmd(query)?;
422                write!(self.buffer, ")")?;
423                if let Some(a) = alias {
424                    write!(self.buffer, " as {}", a)?;
425                }
426            }
427        }
428        Ok(())
429    }
430
431    fn format_join(&mut self, join: &Join) -> Result {
432        match join.kind {
433            crate::ast::JoinKind::Inner => write!(self.buffer, "join {}", join.table)?,
434            crate::ast::JoinKind::Left => write!(self.buffer, "left join {}", join.table)?,
435            crate::ast::JoinKind::Right => write!(self.buffer, "right join {}", join.table)?,
436            crate::ast::JoinKind::Full => write!(self.buffer, "full join {}", join.table)?,
437            crate::ast::JoinKind::Cross => write!(self.buffer, "cross join {}", join.table)?,
438            crate::ast::JoinKind::Lateral => write!(self.buffer, "lateral join {}", join.table)?,
439        }
440
441        if let Some(conditions) = &join.on
442            && !conditions.is_empty()
443        {
444            writeln!(self.buffer)?;
445            self.indent_level += 1;
446            self.indent()?;
447            write!(self.buffer, "on ")?;
448            self.format_conditions(conditions, LogicalOp::And)?;
449            self.indent_level -= 1;
450        }
451        Ok(())
452    }
453
454    fn format_conditions(&mut self, conditions: &[Condition], logical_op: LogicalOp) -> Result {
455        for (i, cond) in conditions.iter().enumerate() {
456            if i > 0 {
457                match logical_op {
458                    LogicalOp::And => write!(self.buffer, " and ")?,
459                    LogicalOp::Or => write!(self.buffer, " or ")?,
460                }
461            }
462
463            write!(self.buffer, "{}", cond.left)?;
464
465            match cond.op {
466                Operator::Eq => write!(self.buffer, " = ")?,
467                Operator::Ne => write!(self.buffer, " != ")?,
468                Operator::Gt => write!(self.buffer, " > ")?,
469                Operator::Gte => write!(self.buffer, " >= ")?,
470                Operator::Lt => write!(self.buffer, " < ")?,
471                Operator::Lte => write!(self.buffer, " <= ")?,
472                Operator::Fuzzy => write!(self.buffer, " ~ ")?, // ILIKE
473                Operator::In => write!(self.buffer, " in ")?,
474                Operator::NotIn => write!(self.buffer, " not in ")?,
475                Operator::IsNull => write!(self.buffer, " is null")?,
476                Operator::IsNotNull => write!(self.buffer, " is not null")?,
477                Operator::Contains => write!(self.buffer, " @> ")?,
478                Operator::KeyExists => write!(self.buffer, " ? ")?,
479                _ => write!(self.buffer, " {:?} ", cond.op)?,
480            }
481
482            // Some operators like IsNull don't need a value printed
483            if !matches!(cond.op, Operator::IsNull | Operator::IsNotNull) {
484                self.format_value(&cond.value)?;
485            }
486        }
487        Ok(())
488    }
489
490    fn format_value(&mut self, val: &Value) -> Result {
491        match val {
492            Value::Null => write!(self.buffer, "null")?,
493            Value::Bool(b) => write!(self.buffer, "{}", b)?,
494            Value::Int(n) => write!(self.buffer, "{}", n)?,
495            Value::Float(n) => write!(self.buffer, "{}", n)?,
496            Value::Param(n) => write!(self.buffer, "${}", n)?,
497            Value::Function(f) => write!(self.buffer, "{}", f)?,
498            Value::Column(c) => write!(self.buffer, "{}", c)?,
499            Value::String(s) => write!(self.buffer, "'{}'", s)?, // Simple quoting, might need escaping
500            // Value::Date and Value::Interval are not in AST, likely Strings
501            // Value::Date(d) => write!(self.buffer, "'{}'", d)?,
502            // Value::Interval(i) => write!(self.buffer, "interval '{}'", i)?,
503            Value::Array(arr) => {
504                write!(self.buffer, "[")?;
505                for (i, v) in arr.iter().enumerate() {
506                    if i > 0 {
507                        write!(self.buffer, ", ")?;
508                    }
509                    self.format_value(v)?;
510                }
511                write!(self.buffer, "]")?;
512            }
513            Value::NamedParam(name) => write!(self.buffer, ":{}", name)?,
514            Value::Uuid(u) => write!(self.buffer, "'{}'", u)?,
515            Value::NullUuid => write!(self.buffer, "null")?,
516            Value::Interval { amount, unit } => {
517                write!(self.buffer, "interval '{} {}'", amount, unit)?
518            }
519            Value::Timestamp(ts) => write!(self.buffer, "'{}'", ts)?,
520            Value::Bytes(bytes) => {
521                write!(self.buffer, "'\\x")?;
522                for byte in bytes {
523                    write!(self.buffer, "{:02x}", byte)?;
524                }
525                write!(self.buffer, "'")?;
526            }
527            Value::Subquery(cmd) => {
528                write!(self.buffer, "(")?;
529                self.visit_cmd(cmd)?;
530                write!(self.buffer, ")")?;
531            }
532            Value::Expr(expr) => write!(self.buffer, "{}", expr)?,
533            Value::Vector(v) => {
534                write!(self.buffer, "[")?;
535                for (i, val) in v.iter().enumerate() {
536                    if i > 0 {
537                        write!(self.buffer, ", ")?;
538                    }
539                    write!(self.buffer, "{}", val)?;
540                }
541                write!(self.buffer, "]")?;
542            }
543            Value::Json(json) => write!(self.buffer, "'{}'::jsonb", json.replace('\'', "''"))?,
544        }
545        Ok(())
546    }
547
548    fn format_sort_order(&mut self, order: SortOrder) -> Result {
549        match order {
550            SortOrder::Asc => {}
551            SortOrder::Desc => write!(self.buffer, " desc")?,
552            SortOrder::AscNullsFirst => write!(self.buffer, " nulls first")?,
553            SortOrder::AscNullsLast => write!(self.buffer, " nulls last")?,
554            SortOrder::DescNullsFirst => write!(self.buffer, " desc nulls first")?,
555            SortOrder::DescNullsLast => write!(self.buffer, " desc nulls last")?,
556        }
557        Ok(())
558    }
559}