Skip to main content

qail_core/fmt/
mod.rs

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