Skip to main content

sqrust_rules/structure/
union_column_alias.rs

1use sqrust_core::{Diagnostic, FileContext, Rule};
2use sqlparser::ast::{Query, SelectItem, SetExpr, SetOperator, Statement};
3
4pub struct UnionColumnAlias;
5
6impl Rule for UnionColumnAlias {
7    fn name(&self) -> &'static str {
8        "Structure/UnionColumnAlias"
9    }
10
11    fn check(&self, ctx: &FileContext) -> Vec<Diagnostic> {
12        if !ctx.parse_errors.is_empty() {
13            return Vec::new();
14        }
15
16        let mut diags = Vec::new();
17
18        for stmt in &ctx.statements {
19            if let Statement::Query(query) = stmt {
20                check_query(query, &mut diags);
21            }
22        }
23
24        diags
25    }
26}
27
28fn check_query(query: &Query, diags: &mut Vec<Diagnostic>) {
29    // Check CTEs.
30    if let Some(with) = &query.with {
31        for cte in &with.cte_tables {
32            check_query(&cte.query, diags);
33        }
34    }
35
36    check_set_expr(&query.body, diags);
37}
38
39fn check_set_expr(expr: &SetExpr, diags: &mut Vec<Diagnostic>) {
40    match expr {
41        SetExpr::Select(_) => {}
42        SetExpr::Query(inner) => {
43            check_query(inner, diags);
44        }
45        SetExpr::SetOperation {
46            op,
47            left,
48            right,
49            ..
50        } => {
51            if *op == SetOperator::Union {
52                // `right` is always a non-first branch in sqlparser-rs UNION chains.
53                collect_aliases_from_non_first_branch(right, diags);
54                // Recurse into `left` — may contain more UNION operations.
55                check_set_expr(left, diags);
56            } else {
57                // INTERSECT / EXCEPT — just recurse without flagging.
58                check_set_expr(left, diags);
59                check_set_expr(right, diags);
60            }
61        }
62        _ => {}
63    }
64}
65
66/// Collect all aliases from a non-first UNION branch's SELECT projection.
67fn collect_aliases_from_non_first_branch(expr: &SetExpr, diags: &mut Vec<Diagnostic>) {
68    match expr {
69        SetExpr::Select(sel) => {
70            for item in &sel.projection {
71                if let SelectItem::ExprWithAlias { .. } = item {
72                    diags.push(Diagnostic {
73                        rule: "Structure/UnionColumnAlias",
74                        message: "Column alias in non-first UNION branch is ignored — aliases only apply to the first SELECT"
75                            .to_string(),
76                        line: 1,
77                        col: 1,
78                    });
79                }
80            }
81        }
82        SetExpr::Query(inner) => {
83            // Wrapped query — check its body.
84            collect_aliases_from_non_first_branch(&inner.body, diags);
85        }
86        // Nested set operations in non-first branch — recurse.
87        SetExpr::SetOperation { op, left, right, .. } => {
88            if *op == SetOperator::Union {
89                collect_aliases_from_non_first_branch(right, diags);
90                collect_aliases_from_non_first_branch(left, diags);
91            } else {
92                collect_aliases_from_non_first_branch(left, diags);
93                collect_aliases_from_non_first_branch(right, diags);
94            }
95        }
96        _ => {}
97    }
98}