Skip to main content

polyglot_sql/optimizer/
qualify_columns.rs

1//! Column Qualification Module
2//!
3//! This module provides functionality for qualifying column references in SQL queries,
4//! adding table qualifiers to column names and expanding star expressions.
5//!
6//! Ported from sqlglot's optimizer/qualify_columns.py
7
8use crate::dialects::transform_recursive;
9use crate::dialects::DialectType;
10use crate::expressions::{
11    Alias, BinaryOp, Column, Expression, Identifier, Join, JoinKind, LateralView, Literal, Over,
12    Paren, Select, TableRef, VarArgFunc, With,
13};
14use crate::resolver::{Resolver, ResolverError};
15use crate::schema::{normalize_name, Schema};
16use crate::scope::{build_scope, traverse_scope, Scope};
17use std::cell::RefCell;
18use std::collections::{HashMap, HashSet};
19use thiserror::Error;
20
21/// Errors that can occur during column qualification
22#[derive(Debug, Error, Clone)]
23pub enum QualifyColumnsError {
24    #[error("Unknown table: {0}")]
25    UnknownTable(String),
26
27    #[error("Unknown column: {0}")]
28    UnknownColumn(String),
29
30    #[error("Ambiguous column: {0}")]
31    AmbiguousColumn(String),
32
33    #[error("Cannot automatically join: {0}")]
34    CannotAutoJoin(String),
35
36    #[error("Unknown output column: {0}")]
37    UnknownOutputColumn(String),
38
39    #[error("Column could not be resolved: {column}{for_table}")]
40    ColumnNotResolved { column: String, for_table: String },
41
42    #[error("Resolver error: {0}")]
43    ResolverError(#[from] ResolverError),
44}
45
46/// Result type for column qualification operations
47pub type QualifyColumnsResult<T> = Result<T, QualifyColumnsError>;
48
49/// Options for column qualification
50#[derive(Debug, Clone, Default)]
51pub struct QualifyColumnsOptions {
52    /// Whether to expand references to aliases
53    pub expand_alias_refs: bool,
54    /// Whether to expand star expressions to explicit columns
55    pub expand_stars: bool,
56    /// Whether to infer schema if not provided
57    pub infer_schema: Option<bool>,
58    /// Whether to allow partial qualification
59    pub allow_partial_qualification: bool,
60    /// The dialect for dialect-specific behavior
61    pub dialect: Option<DialectType>,
62}
63
64impl QualifyColumnsOptions {
65    /// Create new options with defaults
66    pub fn new() -> Self {
67        Self {
68            expand_alias_refs: true,
69            expand_stars: true,
70            infer_schema: None,
71            allow_partial_qualification: false,
72            dialect: None,
73        }
74    }
75
76    /// Set whether to expand alias refs
77    pub fn with_expand_alias_refs(mut self, expand: bool) -> Self {
78        self.expand_alias_refs = expand;
79        self
80    }
81
82    /// Set whether to expand stars
83    pub fn with_expand_stars(mut self, expand: bool) -> Self {
84        self.expand_stars = expand;
85        self
86    }
87
88    /// Set the dialect
89    pub fn with_dialect(mut self, dialect: DialectType) -> Self {
90        self.dialect = Some(dialect);
91        self
92    }
93
94    /// Set whether to allow partial qualification
95    pub fn with_allow_partial(mut self, allow: bool) -> Self {
96        self.allow_partial_qualification = allow;
97        self
98    }
99}
100
101/// Rewrite SQL AST to have fully qualified columns.
102///
103/// # Example
104/// ```ignore
105/// // SELECT col FROM tbl => SELECT tbl.col AS col FROM tbl
106/// ```
107///
108/// # Arguments
109/// * `expression` - Expression to qualify
110/// * `schema` - Database schema for column lookup
111/// * `options` - Qualification options
112///
113/// # Returns
114/// The qualified expression
115pub fn qualify_columns(
116    expression: Expression,
117    schema: &dyn Schema,
118    options: &QualifyColumnsOptions,
119) -> QualifyColumnsResult<Expression> {
120    let infer_schema = options.infer_schema.unwrap_or(schema.is_empty());
121    let dialect = options.dialect.or_else(|| schema.dialect());
122    let first_error: RefCell<Option<QualifyColumnsError>> = RefCell::new(None);
123
124    let transformed = transform_recursive(expression, &|node| {
125        if first_error.borrow().is_some() {
126            return Ok(node);
127        }
128
129        match node {
130            Expression::Select(mut select) => {
131                if let Some(with) = &mut select.with {
132                    pushdown_cte_alias_columns_with(with);
133                }
134
135                let scope_expr = Expression::Select(select.clone());
136                let scope = build_scope(&scope_expr);
137                let mut resolver = Resolver::new(&scope, schema, infer_schema);
138
139                // 1. Expand USING → ON before column qualification
140                let column_tables = if first_error.borrow().is_none() {
141                    match expand_using(&mut select, &scope, &mut resolver) {
142                        Ok(ct) => ct,
143                        Err(err) => {
144                            *first_error.borrow_mut() = Some(err);
145                            HashMap::new()
146                        }
147                    }
148                } else {
149                    HashMap::new()
150                };
151
152                // 2. Expand alias references before qualification so same-select
153                // aliases do not get treated as unresolved physical columns.
154                if first_error.borrow().is_none() && options.expand_alias_refs {
155                    if let Err(err) = expand_alias_refs(&mut select, &mut resolver, dialect) {
156                        *first_error.borrow_mut() = Some(err);
157                    }
158                }
159
160                // 3. Qualify columns (add table qualifiers)
161                if first_error.borrow().is_none() {
162                    if let Err(err) = qualify_columns_in_scope(
163                        &mut select,
164                        &scope,
165                        &mut resolver,
166                        options.allow_partial_qualification,
167                    ) {
168                        *first_error.borrow_mut() = Some(err);
169                    }
170                }
171
172                // 4. Expand star expressions (with USING deduplication)
173                if first_error.borrow().is_none() && options.expand_stars {
174                    if let Err(err) =
175                        expand_stars(&mut select, &scope, &mut resolver, &column_tables)
176                    {
177                        *first_error.borrow_mut() = Some(err);
178                    }
179                }
180
181                // 5. Qualify outputs
182                if first_error.borrow().is_none() {
183                    if let Err(err) = qualify_outputs_select(&mut select) {
184                        *first_error.borrow_mut() = Some(err);
185                    }
186                }
187
188                // 6. Expand GROUP BY positional refs
189                if first_error.borrow().is_none() {
190                    if let Err(err) = expand_group_by(&mut select, dialect) {
191                        *first_error.borrow_mut() = Some(err);
192                    }
193                }
194
195                Ok(Expression::Select(select))
196            }
197            _ => Ok(node),
198        }
199    })
200    .map_err(|err| QualifyColumnsError::CannotAutoJoin(err.to_string()))?;
201
202    if let Some(err) = first_error.into_inner() {
203        return Err(err);
204    }
205
206    Ok(transformed)
207}
208
209/// Validate that all columns in an expression are qualified.
210///
211/// # Returns
212/// The expression if valid, or an error if unqualified columns exist.
213pub fn validate_qualify_columns(expression: &Expression) -> QualifyColumnsResult<()> {
214    let mut all_unqualified = Vec::new();
215
216    for scope in traverse_scope(expression) {
217        if let Expression::Select(_) = &scope.expression {
218            // Get unqualified columns from this scope
219            let unqualified = get_unqualified_columns(&scope);
220
221            // Check for external columns that couldn't be resolved
222            let external = get_external_columns(&scope);
223            if !external.is_empty() && !is_correlated_subquery(&scope) {
224                let first = &external[0];
225                let for_table = if first.table.is_some() {
226                    format!(" for table: '{}'", first.table.as_ref().unwrap())
227                } else {
228                    String::new()
229                };
230                return Err(QualifyColumnsError::ColumnNotResolved {
231                    column: first.name.clone(),
232                    for_table,
233                });
234            }
235
236            all_unqualified.extend(unqualified);
237        }
238    }
239
240    if !all_unqualified.is_empty() {
241        let first = &all_unqualified[0];
242        return Err(QualifyColumnsError::AmbiguousColumn(first.name.clone()));
243    }
244
245    Ok(())
246}
247
248/// Get the alias or table name from a table expression in FROM/JOIN context.
249fn get_source_name(expr: &Expression) -> Option<String> {
250    match expr {
251        Expression::Table(t) => Some(
252            t.alias
253                .as_ref()
254                .map(|a| a.name.clone())
255                .unwrap_or_else(|| t.name.name.clone()),
256        ),
257        Expression::Subquery(sq) => sq.alias.as_ref().map(|a| a.name.clone()),
258        Expression::Pivot(pivot) => Some(pivot_source_name(
259            &pivot.this,
260            pivot.alias.as_ref().map(|alias| alias.name.as_str()),
261        )),
262        Expression::Unpivot(unpivot) => Some(pivot_source_name(
263            &unpivot.this,
264            unpivot.alias.as_ref().map(|alias| alias.name.as_str()),
265        )),
266        _ => None,
267    }
268}
269
270fn pivot_source_name(source: &Expression, explicit_alias: Option<&str>) -> String {
271    if let Some(alias) = explicit_alias {
272        return alias.to_string();
273    }
274
275    match source {
276        Expression::Table(table) => table
277            .alias
278            .as_ref()
279            .map(|alias| alias.name.clone())
280            .unwrap_or_else(|| table.name.name.clone()),
281        Expression::Subquery(subquery) => subquery
282            .alias
283            .as_ref()
284            .map(|alias| alias.name.clone())
285            .unwrap_or_else(|| "_0".to_string()),
286        Expression::Paren(paren) => pivot_source_name(&paren.this, explicit_alias),
287        _ => "_0".to_string(),
288    }
289}
290
291/// Get ordered source names from a SELECT's FROM + JOIN clauses.
292/// FROM tables come first, then JOIN tables in declaration order.
293fn get_ordered_source_names(select: &Select) -> Vec<String> {
294    let mut ordered = Vec::new();
295    if let Some(from) = &select.from {
296        for expr in &from.expressions {
297            if let Some(name) = get_source_name(expr) {
298                ordered.push(name);
299            }
300        }
301    }
302    for join in &select.joins {
303        if is_semi_or_anti_join_kind(join.kind) {
304            continue;
305        }
306        if let Some(name) = get_source_name(&join.this) {
307            ordered.push(name);
308        }
309    }
310    ordered
311}
312
313fn is_semi_or_anti_join_kind(kind: JoinKind) -> bool {
314    matches!(
315        kind,
316        JoinKind::Semi
317            | JoinKind::Anti
318            | JoinKind::LeftSemi
319            | JoinKind::LeftAnti
320            | JoinKind::RightSemi
321            | JoinKind::RightAnti
322    )
323}
324
325/// Create a COALESCE expression over qualified columns from the given tables.
326fn make_coalesce(column_name: &str, tables: &[String]) -> Expression {
327    let args: Vec<Expression> = tables
328        .iter()
329        .map(|t| Expression::qualified_column(t.as_str(), column_name))
330        .collect();
331    Expression::Coalesce(Box::new(VarArgFunc {
332        expressions: args,
333        original_name: None,
334        inferred_type: None,
335    }))
336}
337
338/// Expand JOIN USING clauses into ON conditions and track which columns
339/// participate in USING joins for later COALESCE rewriting.
340///
341/// Returns a mapping from column name → ordered list of table names that
342/// participate in USING for that column.
343fn expand_using(
344    select: &mut Select,
345    _scope: &Scope,
346    resolver: &mut Resolver,
347) -> QualifyColumnsResult<HashMap<String, Vec<String>>> {
348    // columns: normalized column name → first source that owns it
349    // (first-seen-wins)
350    let mut columns: HashMap<String, String> = HashMap::new();
351    // Preserve the first-seen column order for deterministic NATURAL JOIN
352    // expansion. HashMap iteration order must not affect the generated USING list.
353    let mut column_order: Vec<String> = Vec::new();
354
355    // column_tables: normalized column name → ordered list of tables that
356    // participate in USING
357    let mut column_tables: HashMap<String, Vec<String>> = HashMap::new();
358
359    // Get non-join source names from FROM clause
360    let join_names: HashSet<String> = select
361        .joins
362        .iter()
363        .filter_map(|j| get_source_name(&j.this))
364        .collect();
365
366    let all_ordered = get_ordered_source_names(select);
367    let mut ordered: Vec<String> = all_ordered
368        .iter()
369        .filter(|name| !join_names.contains(name.as_str()))
370        .cloned()
371        .collect();
372    let mut accumulated_schema_known = true;
373
374    if join_names.is_empty() {
375        return Ok(column_tables);
376    }
377
378    // Helper closure to update columns map from a source
379    fn update_source_columns(
380        source_name: &str,
381        columns: &mut HashMap<String, String>,
382        column_order: &mut Vec<String>,
383        resolver: &mut Resolver,
384    ) -> bool {
385        let Ok(source_cols) = resolver.get_source_columns(source_name) else {
386            return false;
387        };
388        let schema_known = !source_cols.is_empty()
389            && !source_cols
390                .iter()
391                .any(|column| column == "*" || column.is_empty());
392        for col_name in source_cols {
393            let normalized = normalize_column_name(&col_name, resolver.dialect);
394            if let std::collections::hash_map::Entry::Vacant(entry) = columns.entry(normalized) {
395                entry.insert(source_name.to_string());
396                column_order.push(col_name);
397            }
398        }
399        schema_known
400    }
401
402    // Pre-populate columns from FROM (base) sources
403    for source_name in &ordered {
404        accumulated_schema_known &=
405            update_source_columns(source_name, &mut columns, &mut column_order, resolver);
406    }
407
408    for i in 0..select.joins.len() {
409        // Get source_table (most recently seen non-join table)
410        let source_table = ordered.last().cloned().unwrap_or_default();
411        if !source_table.is_empty() {
412            accumulated_schema_known &=
413                update_source_columns(&source_table, &mut columns, &mut column_order, resolver);
414        }
415
416        // Get join_table name and append to ordered
417        let join_table = get_source_name(&select.joins[i].this).unwrap_or_default();
418        ordered.push(join_table.clone());
419
420        let join_columns: Vec<String> =
421            resolver.get_source_columns(&join_table).unwrap_or_default();
422        let star = normalize_column_name("*", resolver.dialect);
423        let right_schema_known = !join_columns.is_empty()
424            && !join_columns
425                .iter()
426                .any(|column| normalize_column_name(column, resolver.dialect) == star);
427
428        // NATURAL JOIN is an implicit USING join over every column common to
429        // the accumulated left side and the current right side. Only expand
430        // when both schemas are known; otherwise preserve NATURAL rather than
431        // guessing join keys.
432        let join_kind = select.joins[i].kind;
433        if select.joins[i].using.is_empty()
434            && matches!(
435                join_kind,
436                JoinKind::Natural
437                    | JoinKind::NaturalLeft
438                    | JoinKind::NaturalRight
439                    | JoinKind::NaturalFull
440            )
441        {
442            let left_schema_known =
443                accumulated_schema_known && !columns.is_empty() && !columns.contains_key(&star);
444
445            if left_schema_known && right_schema_known {
446                let right_columns: HashSet<String> = join_columns
447                    .iter()
448                    .map(|column| normalize_column_name(column, resolver.dialect))
449                    .collect();
450                let implicit_using: Vec<Identifier> = column_order
451                    .iter()
452                    .filter(|column| {
453                        right_columns.contains(&normalize_column_name(column, resolver.dialect))
454                    })
455                    .map(|column| Identifier::new(column))
456                    .collect();
457
458                if !implicit_using.is_empty() {
459                    select.joins[i].using = implicit_using;
460                    select.joins[i].kind = match join_kind {
461                        JoinKind::Natural => JoinKind::Inner,
462                        JoinKind::NaturalLeft => JoinKind::Left,
463                        JoinKind::NaturalRight => JoinKind::Right,
464                        JoinKind::NaturalFull => JoinKind::Full,
465                        _ => unreachable!("checked NATURAL join kind above"),
466                    };
467                }
468            }
469        }
470
471        // Preserve NATURAL joins with unknown schemas or no common columns,
472        // and skip ordinary joins without a USING clause.
473        if select.joins[i].using.is_empty() {
474            accumulated_schema_known &= right_schema_known;
475            continue;
476        }
477
478        let using_identifiers: Vec<String> = select.joins[i]
479            .using
480            .iter()
481            .map(|id| id.name.clone())
482            .collect();
483
484        let using_count = using_identifiers.len();
485        let is_semi_or_anti = matches!(
486            select.joins[i].kind,
487            crate::expressions::JoinKind::Semi
488                | crate::expressions::JoinKind::Anti
489                | crate::expressions::JoinKind::LeftSemi
490                | crate::expressions::JoinKind::LeftAnti
491                | crate::expressions::JoinKind::RightSemi
492                | crate::expressions::JoinKind::RightAnti
493        );
494
495        let mut conditions: Vec<Expression> = Vec::new();
496
497        for identifier in &using_identifiers {
498            let normalized_identifier = normalize_column_name(identifier, resolver.dialect);
499            let table = columns
500                .get(&normalized_identifier)
501                .cloned()
502                .unwrap_or_else(|| source_table.clone());
503
504            // Build LHS of the equality
505            let lhs = if i == 0 || using_count == 1 {
506                // Simple qualified column for first join or single USING column
507                Expression::qualified_column(table.as_str(), identifier.as_str())
508            } else {
509                // For subsequent joins with multiple USING columns,
510                // COALESCE over all previous sources that have this column
511                let coalesce_cols: Vec<String> = ordered[..ordered.len() - 1]
512                    .iter()
513                    .filter(|t| {
514                        resolver
515                            .get_source_columns(t)
516                            .unwrap_or_default()
517                            .iter()
518                            .any(|column| {
519                                normalize_column_name(column, resolver.dialect)
520                                    == normalized_identifier
521                            })
522                    })
523                    .cloned()
524                    .collect();
525
526                if coalesce_cols.len() > 1 {
527                    make_coalesce(identifier, &coalesce_cols)
528                } else {
529                    Expression::qualified_column(table.as_str(), identifier.as_str())
530                }
531            };
532
533            // Build RHS: qualified column from join table
534            let rhs = Expression::qualified_column(join_table.as_str(), identifier.as_str());
535
536            conditions.push(Expression::Eq(Box::new(BinaryOp::new(lhs, rhs))));
537
538            // Track tables for COALESCE rewriting (skip for semi/anti joins)
539            if !is_semi_or_anti {
540                let tables = column_tables
541                    .entry(normalized_identifier)
542                    .or_insert_with(Vec::new);
543                if !tables.contains(&table) {
544                    tables.push(table.clone());
545                }
546                if !tables.contains(&join_table) {
547                    tables.push(join_table.clone());
548                }
549            }
550        }
551
552        // Combine conditions with AND (left fold)
553        let on_condition = conditions
554            .into_iter()
555            .reduce(|acc, cond| Expression::And(Box::new(BinaryOp::new(acc, cond))))
556            .expect("at least one USING column");
557
558        // Set ON condition and clear USING
559        select.joins[i].on = Some(on_condition);
560        select.joins[i].using = vec![];
561        accumulated_schema_known &= right_schema_known;
562    }
563
564    // Phase 2: Rewrite unqualified USING column references to COALESCE
565    if !column_tables.is_empty() {
566        // Rewrite select.expressions (projections)
567        let mut new_expressions = Vec::with_capacity(select.expressions.len());
568        for expr in &select.expressions {
569            match expr {
570                Expression::Column(col) if col.table.is_none() => {
571                    let normalized = normalize_column_name(&col.name.name, resolver.dialect);
572                    let Some(tables) = column_tables.get(&normalized) else {
573                        new_expressions.push(expr.clone());
574                        continue;
575                    };
576                    let coalesce = make_coalesce(&col.name.name, tables);
577                    // Wrap in alias to preserve column name in projections
578                    new_expressions.push(Expression::Alias(Box::new(Alias {
579                        this: coalesce,
580                        alias: col.name.clone(),
581                        column_aliases: vec![],
582                        alias_explicit_as: false,
583                        alias_keyword: None,
584                        pre_alias_comments: vec![],
585                        trailing_comments: vec![],
586                        inferred_type: None,
587                    })));
588                }
589                _ => {
590                    let mut rewritten = expr.clone();
591                    rewrite_using_columns_in_expression(
592                        &mut rewritten,
593                        &column_tables,
594                        resolver.dialect,
595                    );
596                    new_expressions.push(rewritten);
597                }
598            }
599        }
600        select.expressions = new_expressions;
601
602        // Rewrite WHERE
603        if let Some(where_clause) = &mut select.where_clause {
604            rewrite_using_columns_in_expression(
605                &mut where_clause.this,
606                &column_tables,
607                resolver.dialect,
608            );
609        }
610
611        // Rewrite GROUP BY
612        if let Some(group_by) = &mut select.group_by {
613            for expr in &mut group_by.expressions {
614                rewrite_using_columns_in_expression(expr, &column_tables, resolver.dialect);
615            }
616        }
617
618        // Rewrite HAVING
619        if let Some(having) = &mut select.having {
620            rewrite_using_columns_in_expression(&mut having.this, &column_tables, resolver.dialect);
621        }
622
623        // Rewrite QUALIFY
624        if let Some(qualify) = &mut select.qualify {
625            rewrite_using_columns_in_expression(
626                &mut qualify.this,
627                &column_tables,
628                resolver.dialect,
629            );
630        }
631
632        // Rewrite ORDER BY
633        if let Some(order_by) = &mut select.order_by {
634            for ordered in &mut order_by.expressions {
635                rewrite_using_columns_in_expression(
636                    &mut ordered.this,
637                    &column_tables,
638                    resolver.dialect,
639                );
640            }
641        }
642    }
643
644    Ok(column_tables)
645}
646
647/// Recursively replace unqualified USING column references with COALESCE.
648fn rewrite_using_columns_in_expression(
649    expr: &mut Expression,
650    column_tables: &HashMap<String, Vec<String>>,
651    dialect: Option<DialectType>,
652) {
653    let transformed = transform_recursive(expr.clone(), &|node| match node {
654        Expression::Column(col) if col.table.is_none() => {
655            let normalized = normalize_column_name(&col.name.name, dialect);
656            if let Some(tables) = column_tables.get(&normalized) {
657                Ok(make_coalesce(&col.name.name, tables))
658            } else {
659                Ok(Expression::Column(col))
660            }
661        }
662        other => Ok(other),
663    });
664
665    if let Ok(next) = transformed {
666        *expr = next;
667    }
668}
669
670/// Qualify columns in a scope by adding table qualifiers
671fn qualify_columns_in_scope(
672    select: &mut Select,
673    scope: &Scope,
674    resolver: &mut Resolver,
675    allow_partial: bool,
676) -> QualifyColumnsResult<()> {
677    for expr in &mut select.expressions {
678        qualify_columns_in_expression(expr, scope, resolver, allow_partial)?;
679    }
680    if let Some(where_clause) = &mut select.where_clause {
681        qualify_columns_in_expression(&mut where_clause.this, scope, resolver, allow_partial)?;
682    }
683    if let Some(group_by) = &mut select.group_by {
684        for expr in &mut group_by.expressions {
685            qualify_columns_in_expression(expr, scope, resolver, allow_partial)?;
686        }
687    }
688    if let Some(having) = &mut select.having {
689        qualify_columns_in_expression(&mut having.this, scope, resolver, allow_partial)?;
690    }
691    if let Some(qualify) = &mut select.qualify {
692        qualify_columns_in_expression(&mut qualify.this, scope, resolver, allow_partial)?;
693    }
694    if let Some(order_by) = &mut select.order_by {
695        for ordered in &mut order_by.expressions {
696            qualify_columns_in_expression(&mut ordered.this, scope, resolver, allow_partial)?;
697        }
698    }
699    for join in &mut select.joins {
700        qualify_columns_in_expression(&mut join.this, scope, resolver, allow_partial)?;
701        if let Some(on) = &mut join.on {
702            qualify_columns_in_expression(on, scope, resolver, allow_partial)?;
703        }
704    }
705    Ok(())
706}
707
708/// Expand alias references in a scope.
709///
710/// For example:
711/// `SELECT y.foo AS bar, bar * 2 AS baz FROM y`
712/// becomes:
713/// `SELECT y.foo AS bar, y.foo * 2 AS baz FROM y`
714fn expand_alias_refs(
715    select: &mut Select,
716    _resolver: &mut Resolver,
717    _dialect: Option<DialectType>,
718) -> QualifyColumnsResult<()> {
719    let mut alias_to_expression: HashMap<String, (Expression, usize)> = HashMap::new();
720
721    for (i, expr) in select.expressions.iter_mut().enumerate() {
722        replace_alias_refs_in_expression(expr, &alias_to_expression, false);
723        if let Expression::Alias(alias) = expr {
724            alias_to_expression.insert(alias.alias.name.clone(), (alias.this.clone(), i + 1));
725        }
726    }
727
728    if let Some(where_clause) = &mut select.where_clause {
729        replace_alias_refs_in_expression(&mut where_clause.this, &alias_to_expression, false);
730    }
731    if let Some(group_by) = &mut select.group_by {
732        for expr in &mut group_by.expressions {
733            replace_alias_refs_in_expression(expr, &alias_to_expression, true);
734        }
735    }
736    if let Some(having) = &mut select.having {
737        replace_alias_refs_in_expression(&mut having.this, &alias_to_expression, false);
738    }
739    if let Some(qualify) = &mut select.qualify {
740        replace_alias_refs_in_expression(&mut qualify.this, &alias_to_expression, false);
741    }
742    if let Some(order_by) = &mut select.order_by {
743        for ordered in &mut order_by.expressions {
744            replace_alias_refs_in_expression(&mut ordered.this, &alias_to_expression, false);
745        }
746    }
747
748    Ok(())
749}
750
751/// Expand GROUP BY positional references.
752///
753/// For example:
754/// `SELECT a, b FROM t GROUP BY 1, 2`
755/// becomes:
756/// `SELECT a, b FROM t GROUP BY a, b`
757fn expand_group_by(select: &mut Select, _dialect: Option<DialectType>) -> QualifyColumnsResult<()> {
758    let projections = select.expressions.clone();
759
760    if let Some(group_by) = &mut select.group_by {
761        for group_expr in &mut group_by.expressions {
762            if let Some(index) = positional_reference(group_expr) {
763                let replacement = select_expression_at_position(&projections, index)?;
764                *group_expr = replacement;
765            }
766        }
767    }
768    Ok(())
769}
770
771/// Expand star expressions to explicit column lists, with USING deduplication.
772///
773/// For example:
774/// `SELECT * FROM users`
775/// becomes:
776/// `SELECT users.id, users.name, users.email FROM users`
777///
778/// With USING joins, USING columns appear once as COALESCE and are
779/// deduplicated across sources.
780fn expand_stars(
781    select: &mut Select,
782    _scope: &Scope,
783    resolver: &mut Resolver,
784    column_tables: &HashMap<String, Vec<String>>,
785) -> QualifyColumnsResult<()> {
786    let mut new_selections: Vec<Expression> = Vec::new();
787    let mut has_star = false;
788    let mut coalesced_columns: HashSet<String> = HashSet::new();
789
790    // Use ordered source names (not unordered HashMap keys)
791    let ordered_sources = get_ordered_source_names(select);
792
793    for expr in &select.expressions {
794        match expr {
795            Expression::Star(star) => {
796                has_star = true;
797                if let Some(table) = &star.table {
798                    let table_name = &table.name;
799                    if !ordered_sources.contains(table_name) {
800                        return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
801                    }
802                    if let Ok(columns) = resolver.get_source_columns(table_name) {
803                        if columns.contains(&"*".to_string()) || columns.is_empty() {
804                            return Ok(());
805                        }
806                        for col_name in &columns {
807                            let normalized = normalize_column_name(col_name, resolver.dialect);
808                            if coalesced_columns.contains(&normalized) {
809                                continue;
810                            }
811                            if let Some(tables) = column_tables.get(&normalized) {
812                                if tables.contains(table_name) {
813                                    coalesced_columns.insert(normalized);
814                                    let coalesce = make_coalesce(col_name, tables);
815                                    new_selections.push(Expression::Alias(Box::new(Alias {
816                                        this: coalesce,
817                                        alias: Identifier::new(col_name),
818                                        column_aliases: vec![],
819                                        alias_explicit_as: false,
820                                        alias_keyword: None,
821                                        pre_alias_comments: vec![],
822                                        trailing_comments: vec![],
823                                        inferred_type: None,
824                                    })));
825                                    continue;
826                                }
827                            }
828                            new_selections
829                                .push(create_qualified_column(col_name, Some(table_name)));
830                        }
831                    }
832                } else {
833                    for source_name in &ordered_sources {
834                        if let Ok(columns) = resolver.get_source_columns(source_name) {
835                            if columns.contains(&"*".to_string()) || columns.is_empty() {
836                                return Ok(());
837                            }
838                            for col_name in &columns {
839                                let normalized = normalize_column_name(col_name, resolver.dialect);
840                                if coalesced_columns.contains(&normalized) {
841                                    // Already emitted as COALESCE, skip
842                                    continue;
843                                }
844                                if let Some(tables) = column_tables.get(&normalized) {
845                                    if tables.contains(source_name) {
846                                        // Emit COALESCE and mark as coalesced
847                                        coalesced_columns.insert(normalized);
848                                        let coalesce = make_coalesce(col_name, tables);
849                                        new_selections.push(Expression::Alias(Box::new(Alias {
850                                            this: coalesce,
851                                            alias: Identifier::new(col_name),
852                                            column_aliases: vec![],
853                                            alias_explicit_as: false,
854                                            alias_keyword: None,
855                                            pre_alias_comments: vec![],
856                                            trailing_comments: vec![],
857                                            inferred_type: None,
858                                        })));
859                                        continue;
860                                    }
861                                }
862                                new_selections
863                                    .push(create_qualified_column(col_name, Some(source_name)));
864                            }
865                        }
866                    }
867                }
868            }
869            Expression::Column(col) if is_star_column(col) => {
870                has_star = true;
871                if let Some(table) = &col.table {
872                    let table_name = &table.name;
873                    if !ordered_sources.contains(table_name) {
874                        return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
875                    }
876                    if let Ok(columns) = resolver.get_source_columns(table_name) {
877                        if columns.contains(&"*".to_string()) || columns.is_empty() {
878                            return Ok(());
879                        }
880                        for col_name in &columns {
881                            let normalized = normalize_column_name(col_name, resolver.dialect);
882                            if coalesced_columns.contains(&normalized) {
883                                continue;
884                            }
885                            if let Some(tables) = column_tables.get(&normalized) {
886                                if tables.contains(table_name) {
887                                    coalesced_columns.insert(normalized);
888                                    let coalesce = make_coalesce(col_name, tables);
889                                    new_selections.push(Expression::Alias(Box::new(Alias {
890                                        this: coalesce,
891                                        alias: Identifier::new(col_name),
892                                        column_aliases: vec![],
893                                        alias_explicit_as: false,
894                                        alias_keyword: None,
895                                        pre_alias_comments: vec![],
896                                        trailing_comments: vec![],
897                                        inferred_type: None,
898                                    })));
899                                    continue;
900                                }
901                            }
902                            new_selections
903                                .push(create_qualified_column(col_name, Some(table_name)));
904                        }
905                    }
906                }
907            }
908            _ => new_selections.push(expr.clone()),
909        }
910    }
911
912    if has_star {
913        select.expressions = new_selections;
914    }
915
916    Ok(())
917}
918
919/// Ensure all output columns in a SELECT are aliased.
920///
921/// For example:
922/// `SELECT a + b FROM t`
923/// becomes:
924/// `SELECT a + b AS _col_0 FROM t`
925pub fn qualify_outputs(scope: &Scope) -> QualifyColumnsResult<()> {
926    if let Expression::Select(mut select) = scope.expression.clone() {
927        qualify_outputs_select(&mut select)?;
928    }
929    Ok(())
930}
931
932fn qualify_outputs_select(select: &mut Select) -> QualifyColumnsResult<()> {
933    let mut new_selections: Vec<Expression> = Vec::new();
934
935    for (i, expr) in select.expressions.iter().enumerate() {
936        match expr {
937            Expression::Alias(_) => new_selections.push(expr.clone()),
938            Expression::Column(col) => {
939                new_selections.push(create_alias(expr.clone(), &col.name.name));
940            }
941            Expression::Star(_) => new_selections.push(expr.clone()),
942            _ => {
943                let alias_name = get_output_name(expr).unwrap_or_else(|| format!("_col_{}", i));
944                new_selections.push(create_alias(expr.clone(), &alias_name));
945            }
946        }
947    }
948
949    select.expressions = new_selections;
950    Ok(())
951}
952
953fn qualify_columns_in_expression(
954    expr: &mut Expression,
955    scope: &Scope,
956    resolver: &mut Resolver,
957    allow_partial: bool,
958) -> QualifyColumnsResult<()> {
959    let first_error: RefCell<Option<QualifyColumnsError>> = RefCell::new(None);
960    let resolver_cell: RefCell<&mut Resolver> = RefCell::new(resolver);
961
962    let transformed = transform_recursive(expr.clone(), &|node| {
963        if first_error.borrow().is_some() {
964            return Ok(node);
965        }
966
967        match node {
968            Expression::Column(mut col) => {
969                if let Err(err) = qualify_single_column(
970                    &mut col,
971                    scope,
972                    &mut resolver_cell.borrow_mut(),
973                    allow_partial,
974                ) {
975                    *first_error.borrow_mut() = Some(err);
976                }
977                Ok(Expression::Column(col))
978            }
979            _ => Ok(node),
980        }
981    })
982    .map_err(|err| QualifyColumnsError::CannotAutoJoin(err.to_string()))?;
983
984    if let Some(err) = first_error.into_inner() {
985        return Err(err);
986    }
987
988    *expr = transformed;
989    Ok(())
990}
991
992fn qualify_single_column(
993    col: &mut Column,
994    scope: &Scope,
995    resolver: &mut Resolver,
996    allow_partial: bool,
997) -> QualifyColumnsResult<()> {
998    if is_star_column(col) {
999        return Ok(());
1000    }
1001
1002    if let Some(table) = &col.table {
1003        let table_name = &table.name;
1004        if !scope.sources.contains_key(table_name) {
1005            // Allow correlated references: if the table exists in the schema
1006            // but not in the current scope, it may be referencing an outer scope
1007            // (e.g., in a correlated scalar subquery).
1008            if resolver.table_exists_in_schema(table_name) {
1009                return Ok(());
1010            }
1011            return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
1012        }
1013
1014        if let Ok(source_columns) = resolver.get_source_columns(table_name) {
1015            let normalized_column_name = normalize_column_name(&col.name.name, resolver.dialect);
1016            if !allow_partial
1017                && !source_columns.is_empty()
1018                && !source_columns.iter().any(|column| {
1019                    normalize_column_name(column, resolver.dialect) == normalized_column_name
1020                })
1021                && !source_columns.contains(&"*".to_string())
1022            {
1023                return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone()));
1024            }
1025        }
1026        return Ok(());
1027    }
1028
1029    if let Some(table_name) = resolver.get_table(&col.name.name) {
1030        col.table = Some(Identifier::new(table_name));
1031        return Ok(());
1032    }
1033
1034    // Check for correlated reference: column might belong to an outer scope table.
1035    // Search all schema tables not in the current scope for this column.
1036    if let Some(outer_table) = resolver.find_column_in_outer_schema_tables(&col.name.name) {
1037        col.table = Some(Identifier::new(outer_table));
1038        return Ok(());
1039    }
1040
1041    if !allow_partial {
1042        return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone()));
1043    }
1044
1045    Ok(())
1046}
1047
1048fn normalize_column_name(name: &str, dialect: Option<DialectType>) -> String {
1049    normalize_name(name, dialect, false, true)
1050}
1051
1052fn replace_alias_refs_in_expression(
1053    expr: &mut Expression,
1054    alias_to_expression: &HashMap<String, (Expression, usize)>,
1055    literal_index: bool,
1056) {
1057    let transformed = transform_recursive(expr.clone(), &|node| match node {
1058        Expression::Column(col) if col.table.is_none() => {
1059            if let Some((alias_expr, index)) = alias_to_expression.get(&col.name.name) {
1060                if literal_index && matches!(alias_expr, Expression::Literal(_)) {
1061                    return Ok(Expression::number(*index as i64));
1062                }
1063                return Ok(Expression::Paren(Box::new(Paren {
1064                    this: alias_expr.clone(),
1065                    trailing_comments: vec![],
1066                })));
1067            }
1068            Ok(Expression::Column(col))
1069        }
1070        other => Ok(other),
1071    });
1072
1073    if let Ok(next) = transformed {
1074        *expr = next;
1075    }
1076}
1077
1078fn positional_reference(expr: &Expression) -> Option<usize> {
1079    match expr {
1080        Expression::Literal(lit) if matches!(lit.as_ref(), Literal::Number(_)) => {
1081            let Literal::Number(value) = lit.as_ref() else {
1082                unreachable!()
1083            };
1084            value.parse::<usize>().ok()
1085        }
1086        _ => None,
1087    }
1088}
1089
1090fn select_expression_at_position(
1091    projections: &[Expression],
1092    index: usize,
1093) -> QualifyColumnsResult<Expression> {
1094    if index == 0 || index > projections.len() {
1095        return Err(QualifyColumnsError::UnknownOutputColumn(index.to_string()));
1096    }
1097
1098    let projection = projections[index - 1].clone();
1099    Ok(match projection {
1100        Expression::Alias(alias) => alias.this.clone(),
1101        other => other,
1102    })
1103}
1104
1105/// Returns the set of SQL reserved words for a given dialect.
1106/// If no dialect is provided, returns a comprehensive default set.
1107fn get_reserved_words(dialect: Option<DialectType>) -> HashSet<&'static str> {
1108    // Core SQL reserved words that are common across all dialects
1109    let mut words: HashSet<&'static str> = [
1110        // SQL standard reserved words
1111        "ADD",
1112        "ALL",
1113        "ALTER",
1114        "AND",
1115        "ANY",
1116        "AS",
1117        "ASC",
1118        "BETWEEN",
1119        "BY",
1120        "CASE",
1121        "CAST",
1122        "CHECK",
1123        "COLUMN",
1124        "CONSTRAINT",
1125        "CREATE",
1126        "CROSS",
1127        "CURRENT",
1128        "CURRENT_DATE",
1129        "CURRENT_TIME",
1130        "CURRENT_TIMESTAMP",
1131        "CURRENT_USER",
1132        "DATABASE",
1133        "DEFAULT",
1134        "DELETE",
1135        "DESC",
1136        "DISTINCT",
1137        "DROP",
1138        "ELSE",
1139        "END",
1140        "ESCAPE",
1141        "EXCEPT",
1142        "EXISTS",
1143        "FALSE",
1144        "FETCH",
1145        "FOR",
1146        "FOREIGN",
1147        "FROM",
1148        "FULL",
1149        "GRANT",
1150        "GROUP",
1151        "HAVING",
1152        "IF",
1153        "IN",
1154        "INDEX",
1155        "INNER",
1156        "INSERT",
1157        "INTERSECT",
1158        "INTO",
1159        "IS",
1160        "JOIN",
1161        "KEY",
1162        "LEFT",
1163        "LIKE",
1164        "LIMIT",
1165        "NATURAL",
1166        "NOT",
1167        "NULL",
1168        "OFFSET",
1169        "ON",
1170        "OR",
1171        "ORDER",
1172        "OUTER",
1173        "PRIMARY",
1174        "REFERENCES",
1175        "REPLACE",
1176        "RETURNING",
1177        "RIGHT",
1178        "ROLLBACK",
1179        "ROW",
1180        "ROWS",
1181        "SELECT",
1182        "SESSION_USER",
1183        "SET",
1184        "SOME",
1185        "TABLE",
1186        "THEN",
1187        "TO",
1188        "TRUE",
1189        "TRUNCATE",
1190        "UNION",
1191        "UNIQUE",
1192        "UPDATE",
1193        "USING",
1194        "VALUES",
1195        "VIEW",
1196        "WHEN",
1197        "WHERE",
1198        "WINDOW",
1199        "WITH",
1200    ]
1201    .iter()
1202    .copied()
1203    .collect();
1204
1205    // Add dialect-specific reserved words
1206    match dialect {
1207        Some(DialectType::MySQL) => {
1208            words.extend(
1209                [
1210                    "ANALYZE",
1211                    "BOTH",
1212                    "CHANGE",
1213                    "CONDITION",
1214                    "DATABASES",
1215                    "DAY_HOUR",
1216                    "DAY_MICROSECOND",
1217                    "DAY_MINUTE",
1218                    "DAY_SECOND",
1219                    "DELAYED",
1220                    "DETERMINISTIC",
1221                    "DIV",
1222                    "DUAL",
1223                    "EACH",
1224                    "ELSEIF",
1225                    "ENCLOSED",
1226                    "EXPLAIN",
1227                    "FLOAT4",
1228                    "FLOAT8",
1229                    "FORCE",
1230                    "HOUR_MICROSECOND",
1231                    "HOUR_MINUTE",
1232                    "HOUR_SECOND",
1233                    "IGNORE",
1234                    "INFILE",
1235                    "INT1",
1236                    "INT2",
1237                    "INT3",
1238                    "INT4",
1239                    "INT8",
1240                    "ITERATE",
1241                    "KEYS",
1242                    "KILL",
1243                    "LEADING",
1244                    "LEAVE",
1245                    "LINES",
1246                    "LOAD",
1247                    "LOCK",
1248                    "LONG",
1249                    "LONGBLOB",
1250                    "LONGTEXT",
1251                    "LOOP",
1252                    "LOW_PRIORITY",
1253                    "MATCH",
1254                    "MEDIUMBLOB",
1255                    "MEDIUMINT",
1256                    "MEDIUMTEXT",
1257                    "MINUTE_MICROSECOND",
1258                    "MINUTE_SECOND",
1259                    "MOD",
1260                    "MODIFIES",
1261                    "NO_WRITE_TO_BINLOG",
1262                    "OPTIMIZE",
1263                    "OPTIONALLY",
1264                    "OUT",
1265                    "OUTFILE",
1266                    "PURGE",
1267                    "READS",
1268                    "REGEXP",
1269                    "RELEASE",
1270                    "RENAME",
1271                    "REPEAT",
1272                    "REQUIRE",
1273                    "RESIGNAL",
1274                    "RETURN",
1275                    "REVOKE",
1276                    "RLIKE",
1277                    "SCHEMA",
1278                    "SCHEMAS",
1279                    "SECOND_MICROSECOND",
1280                    "SENSITIVE",
1281                    "SEPARATOR",
1282                    "SHOW",
1283                    "SIGNAL",
1284                    "SPATIAL",
1285                    "SQL",
1286                    "SQLEXCEPTION",
1287                    "SQLSTATE",
1288                    "SQLWARNING",
1289                    "SQL_BIG_RESULT",
1290                    "SQL_CALC_FOUND_ROWS",
1291                    "SQL_SMALL_RESULT",
1292                    "SSL",
1293                    "STARTING",
1294                    "STRAIGHT_JOIN",
1295                    "TERMINATED",
1296                    "TINYBLOB",
1297                    "TINYINT",
1298                    "TINYTEXT",
1299                    "TRAILING",
1300                    "TRIGGER",
1301                    "UNDO",
1302                    "UNLOCK",
1303                    "UNSIGNED",
1304                    "USAGE",
1305                    "UTC_DATE",
1306                    "UTC_TIME",
1307                    "UTC_TIMESTAMP",
1308                    "VARBINARY",
1309                    "VARCHARACTER",
1310                    "WHILE",
1311                    "WRITE",
1312                    "XOR",
1313                    "YEAR_MONTH",
1314                    "ZEROFILL",
1315                ]
1316                .iter()
1317                .copied(),
1318            );
1319        }
1320        Some(DialectType::PostgreSQL) | Some(DialectType::CockroachDB) => {
1321            words.extend(
1322                [
1323                    "ANALYSE",
1324                    "ANALYZE",
1325                    "ARRAY",
1326                    "AUTHORIZATION",
1327                    "BINARY",
1328                    "BOTH",
1329                    "COLLATE",
1330                    "CONCURRENTLY",
1331                    "DO",
1332                    "FREEZE",
1333                    "ILIKE",
1334                    "INITIALLY",
1335                    "ISNULL",
1336                    "LATERAL",
1337                    "LEADING",
1338                    "LOCALTIME",
1339                    "LOCALTIMESTAMP",
1340                    "NOTNULL",
1341                    "ONLY",
1342                    "OVERLAPS",
1343                    "PLACING",
1344                    "SIMILAR",
1345                    "SYMMETRIC",
1346                    "TABLESAMPLE",
1347                    "TRAILING",
1348                    "VARIADIC",
1349                    "VERBOSE",
1350                ]
1351                .iter()
1352                .copied(),
1353            );
1354        }
1355        Some(DialectType::BigQuery) => {
1356            words.extend(
1357                [
1358                    "ASSERT_ROWS_MODIFIED",
1359                    "COLLATE",
1360                    "CONTAINS",
1361                    "CUBE",
1362                    "DEFINE",
1363                    "ENUM",
1364                    "EXTRACT",
1365                    "FOLLOWING",
1366                    "GROUPING",
1367                    "GROUPS",
1368                    "HASH",
1369                    "IGNORE",
1370                    "LATERAL",
1371                    "LOOKUP",
1372                    "MERGE",
1373                    "NEW",
1374                    "NO",
1375                    "NULLS",
1376                    "OF",
1377                    "OVER",
1378                    "PARTITION",
1379                    "PRECEDING",
1380                    "PROTO",
1381                    "RANGE",
1382                    "RECURSIVE",
1383                    "RESPECT",
1384                    "ROLLUP",
1385                    "STRUCT",
1386                    "TABLESAMPLE",
1387                    "TREAT",
1388                    "UNBOUNDED",
1389                    "WITHIN",
1390                ]
1391                .iter()
1392                .copied(),
1393            );
1394        }
1395        Some(DialectType::Snowflake) => {
1396            words.extend(
1397                [
1398                    "ACCOUNT",
1399                    "BOTH",
1400                    "CONNECT",
1401                    "FOLLOWING",
1402                    "ILIKE",
1403                    "INCREMENT",
1404                    "ISSUE",
1405                    "LATERAL",
1406                    "LEADING",
1407                    "LOCALTIME",
1408                    "LOCALTIMESTAMP",
1409                    "MINUS",
1410                    "QUALIFY",
1411                    "REGEXP",
1412                    "RLIKE",
1413                    "SOME",
1414                    "START",
1415                    "TABLESAMPLE",
1416                    "TOP",
1417                    "TRAILING",
1418                    "TRY_CAST",
1419                ]
1420                .iter()
1421                .copied(),
1422            );
1423        }
1424        Some(DialectType::TSQL) | Some(DialectType::Fabric) => {
1425            words.extend(
1426                [
1427                    "BACKUP",
1428                    "BREAK",
1429                    "BROWSE",
1430                    "BULK",
1431                    "CASCADE",
1432                    "CHECKPOINT",
1433                    "CLOSE",
1434                    "CLUSTERED",
1435                    "COALESCE",
1436                    "COMPUTE",
1437                    "CONTAINS",
1438                    "CONTAINSTABLE",
1439                    "CONTINUE",
1440                    "CONVERT",
1441                    "DBCC",
1442                    "DEALLOCATE",
1443                    "DENY",
1444                    "DISK",
1445                    "DISTRIBUTED",
1446                    "DUMP",
1447                    "ERRLVL",
1448                    "EXEC",
1449                    "EXECUTE",
1450                    "EXIT",
1451                    "EXTERNAL",
1452                    "FILE",
1453                    "FILLFACTOR",
1454                    "FREETEXT",
1455                    "FREETEXTTABLE",
1456                    "FUNCTION",
1457                    "GOTO",
1458                    "HOLDLOCK",
1459                    "IDENTITY",
1460                    "IDENTITYCOL",
1461                    "IDENTITY_INSERT",
1462                    "KILL",
1463                    "LINENO",
1464                    "MERGE",
1465                    "NONCLUSTERED",
1466                    "NULLIF",
1467                    "OF",
1468                    "OFF",
1469                    "OFFSETS",
1470                    "OPEN",
1471                    "OPENDATASOURCE",
1472                    "OPENQUERY",
1473                    "OPENROWSET",
1474                    "OPENXML",
1475                    "OVER",
1476                    "PERCENT",
1477                    "PIVOT",
1478                    "PLAN",
1479                    "PRINT",
1480                    "PROC",
1481                    "PROCEDURE",
1482                    "PUBLIC",
1483                    "RAISERROR",
1484                    "READ",
1485                    "READTEXT",
1486                    "RECONFIGURE",
1487                    "REPLICATION",
1488                    "RESTORE",
1489                    "RESTRICT",
1490                    "REVERT",
1491                    "ROWCOUNT",
1492                    "ROWGUIDCOL",
1493                    "RULE",
1494                    "SAVE",
1495                    "SECURITYAUDIT",
1496                    "SEMANTICKEYPHRASETABLE",
1497                    "SEMANTICSIMILARITYDETAILSTABLE",
1498                    "SEMANTICSIMILARITYTABLE",
1499                    "SETUSER",
1500                    "SHUTDOWN",
1501                    "STATISTICS",
1502                    "SYSTEM_USER",
1503                    "TEXTSIZE",
1504                    "TOP",
1505                    "TRAN",
1506                    "TRANSACTION",
1507                    "TRIGGER",
1508                    "TSEQUAL",
1509                    "UNPIVOT",
1510                    "UPDATETEXT",
1511                    "WAITFOR",
1512                    "WRITETEXT",
1513                ]
1514                .iter()
1515                .copied(),
1516            );
1517        }
1518        Some(DialectType::ClickHouse) => {
1519            words.extend(
1520                [
1521                    "ANTI",
1522                    "ARRAY",
1523                    "ASOF",
1524                    "FINAL",
1525                    "FORMAT",
1526                    "GLOBAL",
1527                    "INF",
1528                    "KILL",
1529                    "MATERIALIZED",
1530                    "NAN",
1531                    "PREWHERE",
1532                    "SAMPLE",
1533                    "SEMI",
1534                    "SETTINGS",
1535                    "TOP",
1536                ]
1537                .iter()
1538                .copied(),
1539            );
1540        }
1541        Some(DialectType::DuckDB) => {
1542            words.extend(
1543                [
1544                    "ANALYSE",
1545                    "ANALYZE",
1546                    "ARRAY",
1547                    "BOTH",
1548                    "LATERAL",
1549                    "LEADING",
1550                    "LOCALTIME",
1551                    "LOCALTIMESTAMP",
1552                    "PLACING",
1553                    "QUALIFY",
1554                    "SIMILAR",
1555                    "TABLESAMPLE",
1556                    "TRAILING",
1557                ]
1558                .iter()
1559                .copied(),
1560            );
1561        }
1562        Some(DialectType::Hive) | Some(DialectType::Spark) | Some(DialectType::Databricks) => {
1563            words.extend(
1564                [
1565                    "BOTH",
1566                    "CLUSTER",
1567                    "DISTRIBUTE",
1568                    "EXCHANGE",
1569                    "EXTENDED",
1570                    "FUNCTION",
1571                    "LATERAL",
1572                    "LEADING",
1573                    "MACRO",
1574                    "OVER",
1575                    "PARTITION",
1576                    "PERCENT",
1577                    "RANGE",
1578                    "READS",
1579                    "REDUCE",
1580                    "REGEXP",
1581                    "REVOKE",
1582                    "RLIKE",
1583                    "ROLLUP",
1584                    "SEMI",
1585                    "SORT",
1586                    "TABLESAMPLE",
1587                    "TRAILING",
1588                    "TRANSFORM",
1589                    "UNBOUNDED",
1590                    "UNIQUEJOIN",
1591                ]
1592                .iter()
1593                .copied(),
1594            );
1595        }
1596        Some(DialectType::Trino) | Some(DialectType::Presto) | Some(DialectType::Athena) => {
1597            words.extend(
1598                [
1599                    "CUBE",
1600                    "DEALLOCATE",
1601                    "DESCRIBE",
1602                    "EXECUTE",
1603                    "EXTRACT",
1604                    "GROUPING",
1605                    "LATERAL",
1606                    "LOCALTIME",
1607                    "LOCALTIMESTAMP",
1608                    "NORMALIZE",
1609                    "PREPARE",
1610                    "ROLLUP",
1611                    "SOME",
1612                    "TABLESAMPLE",
1613                    "UESCAPE",
1614                    "UNNEST",
1615                ]
1616                .iter()
1617                .copied(),
1618            );
1619        }
1620        Some(DialectType::Oracle) => {
1621            words.extend(
1622                [
1623                    "ACCESS",
1624                    "AUDIT",
1625                    "CLUSTER",
1626                    "COMMENT",
1627                    "COMPRESS",
1628                    "CONNECT",
1629                    "EXCLUSIVE",
1630                    "FILE",
1631                    "IDENTIFIED",
1632                    "IMMEDIATE",
1633                    "INCREMENT",
1634                    "INITIAL",
1635                    "LEVEL",
1636                    "LOCK",
1637                    "LONG",
1638                    "MAXEXTENTS",
1639                    "MINUS",
1640                    "MODE",
1641                    "NOAUDIT",
1642                    "NOCOMPRESS",
1643                    "NOWAIT",
1644                    "NUMBER",
1645                    "OF",
1646                    "OFFLINE",
1647                    "ONLINE",
1648                    "PCTFREE",
1649                    "PRIOR",
1650                    "RAW",
1651                    "RENAME",
1652                    "RESOURCE",
1653                    "REVOKE",
1654                    "SHARE",
1655                    "SIZE",
1656                    "START",
1657                    "SUCCESSFUL",
1658                    "SYNONYM",
1659                    "SYSDATE",
1660                    "TRIGGER",
1661                    "UID",
1662                    "VALIDATE",
1663                    "VARCHAR2",
1664                    "WHENEVER",
1665                ]
1666                .iter()
1667                .copied(),
1668            );
1669        }
1670        Some(DialectType::Redshift) => {
1671            words.extend(
1672                [
1673                    "AZ64",
1674                    "BZIP2",
1675                    "DELTA",
1676                    "DELTA32K",
1677                    "DISTSTYLE",
1678                    "ENCODE",
1679                    "GZIP",
1680                    "ILIKE",
1681                    "LIMIT",
1682                    "LUNS",
1683                    "LZO",
1684                    "LZOP",
1685                    "MOSTLY13",
1686                    "MOSTLY32",
1687                    "MOSTLY8",
1688                    "RAW",
1689                    "SIMILAR",
1690                    "SNAPSHOT",
1691                    "SORTKEY",
1692                    "SYSDATE",
1693                    "TOP",
1694                    "ZSTD",
1695                ]
1696                .iter()
1697                .copied(),
1698            );
1699        }
1700        _ => {
1701            // For Generic or unknown dialects, add a broad set of commonly reserved words
1702            words.extend(
1703                [
1704                    "ANALYZE",
1705                    "ARRAY",
1706                    "BOTH",
1707                    "CUBE",
1708                    "GROUPING",
1709                    "LATERAL",
1710                    "LEADING",
1711                    "LOCALTIME",
1712                    "LOCALTIMESTAMP",
1713                    "OVER",
1714                    "PARTITION",
1715                    "QUALIFY",
1716                    "RANGE",
1717                    "ROLLUP",
1718                    "SIMILAR",
1719                    "SOME",
1720                    "TABLESAMPLE",
1721                    "TRAILING",
1722                ]
1723                .iter()
1724                .copied(),
1725            );
1726        }
1727    }
1728
1729    words
1730}
1731
1732/// Check whether an identifier name needs quoting.
1733///
1734/// An identifier needs quoting if:
1735/// - It is empty
1736/// - It starts with a digit
1737/// - It contains characters other than `[a-zA-Z0-9_]`
1738/// - It is a SQL reserved word (case-insensitive)
1739fn needs_quoting(name: &str, reserved_words: &HashSet<&str>) -> bool {
1740    if name.is_empty() {
1741        return false;
1742    }
1743
1744    // Starts with a digit
1745    if name.as_bytes()[0].is_ascii_digit() {
1746        return true;
1747    }
1748
1749    // Contains non-identifier characters
1750    if !name.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_') {
1751        return true;
1752    }
1753
1754    // Is a reserved word (case-insensitive check)
1755    let upper = name.to_uppercase();
1756    reserved_words.contains(upper.as_str())
1757}
1758
1759/// Conditionally set `quoted = true` on an identifier if it needs quoting.
1760fn maybe_quote(id: &mut Identifier, reserved_words: &HashSet<&str>) {
1761    // Don't re-quote something already quoted, and don't quote empty identifiers
1762    // or wildcard identifiers
1763    if id.quoted || id.name.is_empty() || id.name == "*" {
1764        return;
1765    }
1766    if needs_quoting(&id.name, reserved_words) {
1767        id.quoted = true;
1768    }
1769}
1770
1771/// Recursively walk an expression and quote identifiers that need quoting.
1772fn quote_identifiers_recursive(expr: &mut Expression, reserved_words: &HashSet<&str>) {
1773    match expr {
1774        // ── Leaf nodes with Identifier ────────────────────────────
1775        Expression::Identifier(id) => {
1776            maybe_quote(id, reserved_words);
1777        }
1778
1779        Expression::Column(col) => {
1780            maybe_quote(&mut col.name, reserved_words);
1781            if let Some(ref mut table) = col.table {
1782                maybe_quote(table, reserved_words);
1783            }
1784        }
1785
1786        Expression::Table(table_ref) => {
1787            maybe_quote(&mut table_ref.name, reserved_words);
1788            if let Some(ref mut schema) = table_ref.schema {
1789                maybe_quote(schema, reserved_words);
1790            }
1791            if let Some(ref mut catalog) = table_ref.catalog {
1792                maybe_quote(catalog, reserved_words);
1793            }
1794            if let Some(ref mut alias) = table_ref.alias {
1795                maybe_quote(alias, reserved_words);
1796            }
1797            for ca in &mut table_ref.column_aliases {
1798                maybe_quote(ca, reserved_words);
1799            }
1800            for p in &mut table_ref.partitions {
1801                maybe_quote(p, reserved_words);
1802            }
1803            // Recurse into hints and other child expressions
1804            for h in &mut table_ref.hints {
1805                quote_identifiers_recursive(h, reserved_words);
1806            }
1807            if let Some(ref mut ver) = table_ref.version {
1808                quote_identifiers_recursive(&mut ver.this, reserved_words);
1809                if let Some(ref mut e) = ver.expression {
1810                    quote_identifiers_recursive(e, reserved_words);
1811                }
1812            }
1813        }
1814
1815        Expression::Star(star) => {
1816            if let Some(ref mut table) = star.table {
1817                maybe_quote(table, reserved_words);
1818            }
1819            if let Some(ref mut except_ids) = star.except {
1820                for id in except_ids {
1821                    maybe_quote(id, reserved_words);
1822                }
1823            }
1824            if let Some(ref mut replace_aliases) = star.replace {
1825                for alias in replace_aliases {
1826                    maybe_quote(&mut alias.alias, reserved_words);
1827                    quote_identifiers_recursive(&mut alias.this, reserved_words);
1828                }
1829            }
1830            if let Some(ref mut rename_pairs) = star.rename {
1831                for (from, to) in rename_pairs {
1832                    maybe_quote(from, reserved_words);
1833                    maybe_quote(to, reserved_words);
1834                }
1835            }
1836        }
1837
1838        // ── Alias ─────────────────────────────────────────────────
1839        Expression::Alias(alias) => {
1840            maybe_quote(&mut alias.alias, reserved_words);
1841            for ca in &mut alias.column_aliases {
1842                maybe_quote(ca, reserved_words);
1843            }
1844            quote_identifiers_recursive(&mut alias.this, reserved_words);
1845        }
1846
1847        // ── SELECT ────────────────────────────────────────────────
1848        Expression::Select(select) => {
1849            for e in &mut select.expressions {
1850                quote_identifiers_recursive(e, reserved_words);
1851            }
1852            if let Some(ref mut from) = select.from {
1853                for e in &mut from.expressions {
1854                    quote_identifiers_recursive(e, reserved_words);
1855                }
1856            }
1857            for join in &mut select.joins {
1858                quote_join(join, reserved_words);
1859            }
1860            for lv in &mut select.lateral_views {
1861                quote_lateral_view(lv, reserved_words);
1862            }
1863            if let Some(ref mut prewhere) = select.prewhere {
1864                quote_identifiers_recursive(prewhere, reserved_words);
1865            }
1866            if let Some(ref mut wh) = select.where_clause {
1867                quote_identifiers_recursive(&mut wh.this, reserved_words);
1868            }
1869            if let Some(ref mut gb) = select.group_by {
1870                for e in &mut gb.expressions {
1871                    quote_identifiers_recursive(e, reserved_words);
1872                }
1873            }
1874            if let Some(ref mut hv) = select.having {
1875                quote_identifiers_recursive(&mut hv.this, reserved_words);
1876            }
1877            if let Some(ref mut q) = select.qualify {
1878                quote_identifiers_recursive(&mut q.this, reserved_words);
1879            }
1880            if let Some(ref mut ob) = select.order_by {
1881                for o in &mut ob.expressions {
1882                    quote_identifiers_recursive(&mut o.this, reserved_words);
1883                }
1884            }
1885            if let Some(ref mut lim) = select.limit {
1886                quote_identifiers_recursive(&mut lim.this, reserved_words);
1887            }
1888            if let Some(ref mut off) = select.offset {
1889                quote_identifiers_recursive(&mut off.this, reserved_words);
1890            }
1891            if let Some(ref mut with) = select.with {
1892                quote_with(with, reserved_words);
1893            }
1894            if let Some(ref mut windows) = select.windows {
1895                for nw in windows {
1896                    maybe_quote(&mut nw.name, reserved_words);
1897                    quote_over(&mut nw.spec, reserved_words);
1898                }
1899            }
1900            if let Some(ref mut distinct_on) = select.distinct_on {
1901                for e in distinct_on {
1902                    quote_identifiers_recursive(e, reserved_words);
1903                }
1904            }
1905            if let Some(ref mut limit_by) = select.limit_by {
1906                for e in limit_by {
1907                    quote_identifiers_recursive(e, reserved_words);
1908                }
1909            }
1910            if let Some(ref mut settings) = select.settings {
1911                for e in settings {
1912                    quote_identifiers_recursive(e, reserved_words);
1913                }
1914            }
1915            if let Some(ref mut format) = select.format {
1916                quote_identifiers_recursive(format, reserved_words);
1917            }
1918        }
1919
1920        // ── Set operations ────────────────────────────────────────
1921        Expression::Union(u) => {
1922            quote_identifiers_recursive(&mut u.left, reserved_words);
1923            quote_identifiers_recursive(&mut u.right, reserved_words);
1924            if let Some(ref mut with) = u.with {
1925                quote_with(with, reserved_words);
1926            }
1927            if let Some(ref mut ob) = u.order_by {
1928                for o in &mut ob.expressions {
1929                    quote_identifiers_recursive(&mut o.this, reserved_words);
1930                }
1931            }
1932            if let Some(ref mut lim) = u.limit {
1933                quote_identifiers_recursive(lim, reserved_words);
1934            }
1935            if let Some(ref mut off) = u.offset {
1936                quote_identifiers_recursive(off, reserved_words);
1937            }
1938        }
1939        Expression::Intersect(i) => {
1940            quote_identifiers_recursive(&mut i.left, reserved_words);
1941            quote_identifiers_recursive(&mut i.right, reserved_words);
1942            if let Some(ref mut with) = i.with {
1943                quote_with(with, reserved_words);
1944            }
1945            if let Some(ref mut ob) = i.order_by {
1946                for o in &mut ob.expressions {
1947                    quote_identifiers_recursive(&mut o.this, reserved_words);
1948                }
1949            }
1950        }
1951        Expression::Except(e) => {
1952            quote_identifiers_recursive(&mut e.left, reserved_words);
1953            quote_identifiers_recursive(&mut e.right, reserved_words);
1954            if let Some(ref mut with) = e.with {
1955                quote_with(with, reserved_words);
1956            }
1957            if let Some(ref mut ob) = e.order_by {
1958                for o in &mut ob.expressions {
1959                    quote_identifiers_recursive(&mut o.this, reserved_words);
1960                }
1961            }
1962        }
1963
1964        // ── Subquery ──────────────────────────────────────────────
1965        Expression::Subquery(sq) => {
1966            quote_identifiers_recursive(&mut sq.this, reserved_words);
1967            if let Some(ref mut alias) = sq.alias {
1968                maybe_quote(alias, reserved_words);
1969            }
1970            for ca in &mut sq.column_aliases {
1971                maybe_quote(ca, reserved_words);
1972            }
1973            if let Some(ref mut ob) = sq.order_by {
1974                for o in &mut ob.expressions {
1975                    quote_identifiers_recursive(&mut o.this, reserved_words);
1976                }
1977            }
1978        }
1979
1980        // ── DML ───────────────────────────────────────────────────
1981        Expression::Insert(ins) => {
1982            quote_table_ref(&mut ins.table, reserved_words);
1983            for c in &mut ins.columns {
1984                maybe_quote(c, reserved_words);
1985            }
1986            for row in &mut ins.values {
1987                for e in row {
1988                    quote_identifiers_recursive(e, reserved_words);
1989                }
1990            }
1991            if let Some(ref mut q) = ins.query {
1992                quote_identifiers_recursive(q, reserved_words);
1993            }
1994            for (id, val) in &mut ins.partition {
1995                maybe_quote(id, reserved_words);
1996                if let Some(ref mut v) = val {
1997                    quote_identifiers_recursive(v, reserved_words);
1998                }
1999            }
2000            for e in &mut ins.returning {
2001                quote_identifiers_recursive(e, reserved_words);
2002            }
2003            if let Some(ref mut on_conflict) = ins.on_conflict {
2004                quote_identifiers_recursive(on_conflict, reserved_words);
2005            }
2006            if let Some(ref mut with) = ins.with {
2007                quote_with(with, reserved_words);
2008            }
2009            if let Some(ref mut alias) = ins.alias {
2010                maybe_quote(alias, reserved_words);
2011            }
2012            if let Some(ref mut src_alias) = ins.source_alias {
2013                maybe_quote(src_alias, reserved_words);
2014            }
2015        }
2016
2017        Expression::Update(upd) => {
2018            quote_table_ref(&mut upd.table, reserved_words);
2019            for tr in &mut upd.extra_tables {
2020                quote_table_ref(tr, reserved_words);
2021            }
2022            for join in &mut upd.table_joins {
2023                quote_join(join, reserved_words);
2024            }
2025            for (id, val) in &mut upd.set {
2026                maybe_quote(id, reserved_words);
2027                quote_identifiers_recursive(val, reserved_words);
2028            }
2029            if let Some(ref mut from) = upd.from_clause {
2030                for e in &mut from.expressions {
2031                    quote_identifiers_recursive(e, reserved_words);
2032                }
2033            }
2034            for join in &mut upd.from_joins {
2035                quote_join(join, reserved_words);
2036            }
2037            if let Some(ref mut wh) = upd.where_clause {
2038                quote_identifiers_recursive(&mut wh.this, reserved_words);
2039            }
2040            for e in &mut upd.returning {
2041                quote_identifiers_recursive(e, reserved_words);
2042            }
2043            if let Some(ref mut with) = upd.with {
2044                quote_with(with, reserved_words);
2045            }
2046        }
2047
2048        Expression::Delete(del) => {
2049            quote_table_ref(&mut del.table, reserved_words);
2050            if let Some(ref mut alias) = del.alias {
2051                maybe_quote(alias, reserved_words);
2052            }
2053            for tr in &mut del.using {
2054                quote_table_ref(tr, reserved_words);
2055            }
2056            if let Some(ref mut wh) = del.where_clause {
2057                quote_identifiers_recursive(&mut wh.this, reserved_words);
2058            }
2059            if let Some(ref mut with) = del.with {
2060                quote_with(with, reserved_words);
2061            }
2062        }
2063
2064        // ── Binary operations ─────────────────────────────────────
2065        Expression::And(bin)
2066        | Expression::Or(bin)
2067        | Expression::Eq(bin)
2068        | Expression::Neq(bin)
2069        | Expression::Lt(bin)
2070        | Expression::Lte(bin)
2071        | Expression::Gt(bin)
2072        | Expression::Gte(bin)
2073        | Expression::Add(bin)
2074        | Expression::Sub(bin)
2075        | Expression::Mul(bin)
2076        | Expression::Div(bin)
2077        | Expression::Mod(bin)
2078        | Expression::BitwiseAnd(bin)
2079        | Expression::BitwiseOr(bin)
2080        | Expression::BitwiseXor(bin)
2081        | Expression::Concat(bin)
2082        | Expression::Adjacent(bin)
2083        | Expression::TsMatch(bin)
2084        | Expression::PropertyEQ(bin)
2085        | Expression::ArrayContainsAll(bin)
2086        | Expression::ArrayContainedBy(bin)
2087        | Expression::ArrayOverlaps(bin)
2088        | Expression::JSONBContainsAllTopKeys(bin)
2089        | Expression::JSONBContainsAnyTopKeys(bin)
2090        | Expression::JSONBDeleteAtPath(bin)
2091        | Expression::ExtendsLeft(bin)
2092        | Expression::ExtendsRight(bin)
2093        | Expression::Is(bin)
2094        | Expression::NullSafeEq(bin)
2095        | Expression::NullSafeNeq(bin)
2096        | Expression::Glob(bin)
2097        | Expression::Match(bin)
2098        | Expression::MemberOf(bin)
2099        | Expression::BitwiseLeftShift(bin)
2100        | Expression::BitwiseRightShift(bin) => {
2101            quote_identifiers_recursive(&mut bin.left, reserved_words);
2102            quote_identifiers_recursive(&mut bin.right, reserved_words);
2103        }
2104
2105        // ── Like operations ───────────────────────────────────────
2106        Expression::Like(like) | Expression::ILike(like) => {
2107            quote_identifiers_recursive(&mut like.left, reserved_words);
2108            quote_identifiers_recursive(&mut like.right, reserved_words);
2109            if let Some(ref mut esc) = like.escape {
2110                quote_identifiers_recursive(esc, reserved_words);
2111            }
2112        }
2113
2114        // ── Unary operations ──────────────────────────────────────
2115        Expression::Not(un) | Expression::Neg(un) | Expression::BitwiseNot(un) => {
2116            quote_identifiers_recursive(&mut un.this, reserved_words);
2117        }
2118
2119        // ── Predicates ────────────────────────────────────────────
2120        Expression::In(in_expr) => {
2121            quote_identifiers_recursive(&mut in_expr.this, reserved_words);
2122            for e in &mut in_expr.expressions {
2123                quote_identifiers_recursive(e, reserved_words);
2124            }
2125            if let Some(ref mut q) = in_expr.query {
2126                quote_identifiers_recursive(q, reserved_words);
2127            }
2128            if let Some(ref mut un) = in_expr.unnest {
2129                quote_identifiers_recursive(un, reserved_words);
2130            }
2131        }
2132
2133        Expression::Between(bw) => {
2134            quote_identifiers_recursive(&mut bw.this, reserved_words);
2135            quote_identifiers_recursive(&mut bw.low, reserved_words);
2136            quote_identifiers_recursive(&mut bw.high, reserved_words);
2137        }
2138
2139        Expression::IsNull(is_null) => {
2140            quote_identifiers_recursive(&mut is_null.this, reserved_words);
2141        }
2142
2143        Expression::IsTrue(is_tf) | Expression::IsFalse(is_tf) => {
2144            quote_identifiers_recursive(&mut is_tf.this, reserved_words);
2145        }
2146
2147        Expression::Exists(ex) => {
2148            quote_identifiers_recursive(&mut ex.this, reserved_words);
2149        }
2150
2151        // ── Functions ─────────────────────────────────────────────
2152        Expression::Function(func) => {
2153            for arg in &mut func.args {
2154                quote_identifiers_recursive(arg, reserved_words);
2155            }
2156        }
2157
2158        Expression::AggregateFunction(agg) => {
2159            for arg in &mut agg.args {
2160                quote_identifiers_recursive(arg, reserved_words);
2161            }
2162            if let Some(ref mut filter) = agg.filter {
2163                quote_identifiers_recursive(filter, reserved_words);
2164            }
2165            for o in &mut agg.order_by {
2166                quote_identifiers_recursive(&mut o.this, reserved_words);
2167            }
2168        }
2169
2170        Expression::WindowFunction(wf) => {
2171            quote_identifiers_recursive(&mut wf.this, reserved_words);
2172            quote_over(&mut wf.over, reserved_words);
2173        }
2174
2175        // ── CASE ──────────────────────────────────────────────────
2176        Expression::Case(case) => {
2177            if let Some(ref mut operand) = case.operand {
2178                quote_identifiers_recursive(operand, reserved_words);
2179            }
2180            for (when, then) in &mut case.whens {
2181                quote_identifiers_recursive(when, reserved_words);
2182                quote_identifiers_recursive(then, reserved_words);
2183            }
2184            if let Some(ref mut else_) = case.else_ {
2185                quote_identifiers_recursive(else_, reserved_words);
2186            }
2187        }
2188
2189        // ── CAST / TryCast / SafeCast ─────────────────────────────
2190        Expression::Cast(cast) | Expression::TryCast(cast) | Expression::SafeCast(cast) => {
2191            quote_identifiers_recursive(&mut cast.this, reserved_words);
2192            if let Some(ref mut fmt) = cast.format {
2193                quote_identifiers_recursive(fmt, reserved_words);
2194            }
2195        }
2196
2197        // ── Paren / Annotated ─────────────────────────────────────
2198        Expression::Paren(paren) => {
2199            quote_identifiers_recursive(&mut paren.this, reserved_words);
2200        }
2201
2202        Expression::Annotated(ann) => {
2203            quote_identifiers_recursive(&mut ann.this, reserved_words);
2204        }
2205
2206        // ── WITH clause (standalone) ──────────────────────────────
2207        Expression::With(with) => {
2208            quote_with(with, reserved_words);
2209        }
2210
2211        Expression::Cte(cte) => {
2212            maybe_quote(&mut cte.alias, reserved_words);
2213            for c in &mut cte.columns {
2214                maybe_quote(c, reserved_words);
2215            }
2216            quote_identifiers_recursive(&mut cte.this, reserved_words);
2217        }
2218
2219        // ── Clauses (standalone) ──────────────────────────────────
2220        Expression::From(from) => {
2221            for e in &mut from.expressions {
2222                quote_identifiers_recursive(e, reserved_words);
2223            }
2224        }
2225
2226        Expression::Join(join) => {
2227            quote_join(join, reserved_words);
2228        }
2229
2230        Expression::JoinedTable(jt) => {
2231            quote_identifiers_recursive(&mut jt.left, reserved_words);
2232            for join in &mut jt.joins {
2233                quote_join(join, reserved_words);
2234            }
2235            if let Some(ref mut alias) = jt.alias {
2236                maybe_quote(alias, reserved_words);
2237            }
2238        }
2239
2240        Expression::Where(wh) => {
2241            quote_identifiers_recursive(&mut wh.this, reserved_words);
2242        }
2243
2244        Expression::GroupBy(gb) => {
2245            for e in &mut gb.expressions {
2246                quote_identifiers_recursive(e, reserved_words);
2247            }
2248        }
2249
2250        Expression::Having(hv) => {
2251            quote_identifiers_recursive(&mut hv.this, reserved_words);
2252        }
2253
2254        Expression::OrderBy(ob) => {
2255            for o in &mut ob.expressions {
2256                quote_identifiers_recursive(&mut o.this, reserved_words);
2257            }
2258        }
2259
2260        Expression::Ordered(ord) => {
2261            quote_identifiers_recursive(&mut ord.this, reserved_words);
2262        }
2263
2264        Expression::Limit(lim) => {
2265            quote_identifiers_recursive(&mut lim.this, reserved_words);
2266        }
2267
2268        Expression::Offset(off) => {
2269            quote_identifiers_recursive(&mut off.this, reserved_words);
2270        }
2271
2272        Expression::Qualify(q) => {
2273            quote_identifiers_recursive(&mut q.this, reserved_words);
2274        }
2275
2276        Expression::Window(ws) => {
2277            for e in &mut ws.partition_by {
2278                quote_identifiers_recursive(e, reserved_words);
2279            }
2280            for o in &mut ws.order_by {
2281                quote_identifiers_recursive(&mut o.this, reserved_words);
2282            }
2283        }
2284
2285        Expression::Over(over) => {
2286            quote_over(over, reserved_words);
2287        }
2288
2289        Expression::WithinGroup(wg) => {
2290            quote_identifiers_recursive(&mut wg.this, reserved_words);
2291            for o in &mut wg.order_by {
2292                quote_identifiers_recursive(&mut o.this, reserved_words);
2293            }
2294        }
2295
2296        // ── Pivot / Unpivot ───────────────────────────────────────
2297        Expression::Pivot(piv) => {
2298            quote_identifiers_recursive(&mut piv.this, reserved_words);
2299            for e in &mut piv.expressions {
2300                quote_identifiers_recursive(e, reserved_words);
2301            }
2302            for f in &mut piv.fields {
2303                quote_identifiers_recursive(f, reserved_words);
2304            }
2305            if let Some(ref mut alias) = piv.alias {
2306                maybe_quote(alias, reserved_words);
2307            }
2308        }
2309
2310        Expression::Unpivot(unpiv) => {
2311            quote_identifiers_recursive(&mut unpiv.this, reserved_words);
2312            maybe_quote(&mut unpiv.value_column, reserved_words);
2313            maybe_quote(&mut unpiv.name_column, reserved_words);
2314            for e in &mut unpiv.columns {
2315                quote_identifiers_recursive(e, reserved_words);
2316            }
2317            if let Some(ref mut alias) = unpiv.alias {
2318                maybe_quote(alias, reserved_words);
2319            }
2320        }
2321
2322        // ── Values ────────────────────────────────────────────────
2323        Expression::Values(vals) => {
2324            for tuple in &mut vals.expressions {
2325                for e in &mut tuple.expressions {
2326                    quote_identifiers_recursive(e, reserved_words);
2327                }
2328            }
2329            if let Some(ref mut alias) = vals.alias {
2330                maybe_quote(alias, reserved_words);
2331            }
2332            for ca in &mut vals.column_aliases {
2333                maybe_quote(ca, reserved_words);
2334            }
2335        }
2336
2337        // ── Array / Struct / Tuple ────────────────────────────────
2338        Expression::Array(arr) => {
2339            for e in &mut arr.expressions {
2340                quote_identifiers_recursive(e, reserved_words);
2341            }
2342        }
2343
2344        Expression::Struct(st) => {
2345            for (_name, e) in &mut st.fields {
2346                quote_identifiers_recursive(e, reserved_words);
2347            }
2348        }
2349
2350        Expression::Tuple(tup) => {
2351            for e in &mut tup.expressions {
2352                quote_identifiers_recursive(e, reserved_words);
2353            }
2354        }
2355
2356        // ── Subscript / Dot / Method ──────────────────────────────
2357        Expression::Subscript(sub) => {
2358            quote_identifiers_recursive(&mut sub.this, reserved_words);
2359            quote_identifiers_recursive(&mut sub.index, reserved_words);
2360        }
2361
2362        Expression::Dot(dot) => {
2363            quote_identifiers_recursive(&mut dot.this, reserved_words);
2364            maybe_quote(&mut dot.field, reserved_words);
2365        }
2366
2367        Expression::ScopeResolution(sr) => {
2368            if let Some(ref mut this) = sr.this {
2369                quote_identifiers_recursive(this, reserved_words);
2370            }
2371            quote_identifiers_recursive(&mut sr.expression, reserved_words);
2372        }
2373
2374        // ── Lateral ───────────────────────────────────────────────
2375        Expression::Lateral(lat) => {
2376            quote_identifiers_recursive(&mut lat.this, reserved_words);
2377            // lat.alias is Option<String>, not Identifier, so we skip it
2378        }
2379
2380        // ── DPipe (|| concatenation) ──────────────────────────────
2381        Expression::DPipe(dpipe) => {
2382            quote_identifiers_recursive(&mut dpipe.this, reserved_words);
2383            quote_identifiers_recursive(&mut dpipe.expression, reserved_words);
2384        }
2385
2386        // ── Merge ─────────────────────────────────────────────────
2387        Expression::Merge(merge) => {
2388            quote_identifiers_recursive(&mut merge.this, reserved_words);
2389            quote_identifiers_recursive(&mut merge.using, reserved_words);
2390            if let Some(ref mut on) = merge.on {
2391                quote_identifiers_recursive(on, reserved_words);
2392            }
2393            if let Some(ref mut whens) = merge.whens {
2394                quote_identifiers_recursive(whens, reserved_words);
2395            }
2396            if let Some(ref mut with) = merge.with_ {
2397                quote_identifiers_recursive(with, reserved_words);
2398            }
2399            if let Some(ref mut ret) = merge.returning {
2400                quote_identifiers_recursive(ret, reserved_words);
2401            }
2402        }
2403
2404        // ── LateralView (standalone) ──────────────────────────────
2405        Expression::LateralView(lv) => {
2406            quote_lateral_view(lv, reserved_words);
2407        }
2408
2409        // ── Anonymous (generic function) ──────────────────────────
2410        Expression::Anonymous(anon) => {
2411            quote_identifiers_recursive(&mut anon.this, reserved_words);
2412            for e in &mut anon.expressions {
2413                quote_identifiers_recursive(e, reserved_words);
2414            }
2415        }
2416
2417        // ── Filter (e.g., FILTER(WHERE ...)) ──────────────────────
2418        Expression::Filter(filter) => {
2419            quote_identifiers_recursive(&mut filter.this, reserved_words);
2420            quote_identifiers_recursive(&mut filter.expression, reserved_words);
2421        }
2422
2423        // ── Returning ─────────────────────────────────────────────
2424        Expression::Returning(ret) => {
2425            for e in &mut ret.expressions {
2426                quote_identifiers_recursive(e, reserved_words);
2427            }
2428        }
2429
2430        // ── BracedWildcard ────────────────────────────────────────
2431        Expression::BracedWildcard(inner) => {
2432            quote_identifiers_recursive(inner, reserved_words);
2433        }
2434
2435        // ── ReturnStmt ────────────────────────────────────────────
2436        Expression::ReturnStmt(inner) => {
2437            quote_identifiers_recursive(inner, reserved_words);
2438        }
2439
2440        // ── Leaf nodes that never contain identifiers ─────────────
2441        Expression::Literal(_)
2442        | Expression::Boolean(_)
2443        | Expression::Null(_)
2444        | Expression::DataType(_)
2445        | Expression::Raw(_)
2446        | Expression::Placeholder(_)
2447        | Expression::CurrentDate(_)
2448        | Expression::CurrentTime(_)
2449        | Expression::CurrentTimestamp(_)
2450        | Expression::CurrentTimestampLTZ(_)
2451        | Expression::SessionUser(_)
2452        | Expression::RowNumber(_)
2453        | Expression::Rank(_)
2454        | Expression::DenseRank(_)
2455        | Expression::PercentRank(_)
2456        | Expression::CumeDist(_)
2457        | Expression::Random(_)
2458        | Expression::Pi(_)
2459        | Expression::JSONPathRoot(_) => {
2460            // Nothing to do – these are leaves or do not contain identifiers
2461        }
2462
2463        // ── Catch-all: many expression variants follow common patterns.
2464        // Rather than listing every single variant, we leave them unchanged.
2465        // The key identifier-bearing variants are covered above.
2466        _ => {}
2467    }
2468}
2469
2470/// Helper: quote identifiers in a Join.
2471fn quote_join(join: &mut Join, reserved_words: &HashSet<&str>) {
2472    quote_identifiers_recursive(&mut join.this, reserved_words);
2473    if let Some(ref mut on) = join.on {
2474        quote_identifiers_recursive(on, reserved_words);
2475    }
2476    for id in &mut join.using {
2477        maybe_quote(id, reserved_words);
2478    }
2479    if let Some(ref mut mc) = join.match_condition {
2480        quote_identifiers_recursive(mc, reserved_words);
2481    }
2482    for piv in &mut join.pivots {
2483        quote_identifiers_recursive(piv, reserved_words);
2484    }
2485}
2486
2487/// Helper: quote identifiers in a WITH clause.
2488fn quote_with(with: &mut With, reserved_words: &HashSet<&str>) {
2489    for cte in &mut with.ctes {
2490        maybe_quote(&mut cte.alias, reserved_words);
2491        for c in &mut cte.columns {
2492            maybe_quote(c, reserved_words);
2493        }
2494        for k in &mut cte.key_expressions {
2495            maybe_quote(k, reserved_words);
2496        }
2497        quote_identifiers_recursive(&mut cte.this, reserved_words);
2498    }
2499}
2500
2501/// Helper: quote identifiers in an Over clause.
2502fn quote_over(over: &mut Over, reserved_words: &HashSet<&str>) {
2503    if let Some(ref mut wn) = over.window_name {
2504        maybe_quote(wn, reserved_words);
2505    }
2506    for e in &mut over.partition_by {
2507        quote_identifiers_recursive(e, reserved_words);
2508    }
2509    for o in &mut over.order_by {
2510        quote_identifiers_recursive(&mut o.this, reserved_words);
2511    }
2512    if let Some(ref mut alias) = over.alias {
2513        maybe_quote(alias, reserved_words);
2514    }
2515}
2516
2517/// Helper: quote identifiers in a TableRef (used by DML statements).
2518fn quote_table_ref(table_ref: &mut TableRef, reserved_words: &HashSet<&str>) {
2519    maybe_quote(&mut table_ref.name, reserved_words);
2520    if let Some(ref mut schema) = table_ref.schema {
2521        maybe_quote(schema, reserved_words);
2522    }
2523    if let Some(ref mut catalog) = table_ref.catalog {
2524        maybe_quote(catalog, reserved_words);
2525    }
2526    if let Some(ref mut alias) = table_ref.alias {
2527        maybe_quote(alias, reserved_words);
2528    }
2529    for ca in &mut table_ref.column_aliases {
2530        maybe_quote(ca, reserved_words);
2531    }
2532    for p in &mut table_ref.partitions {
2533        maybe_quote(p, reserved_words);
2534    }
2535    for h in &mut table_ref.hints {
2536        quote_identifiers_recursive(h, reserved_words);
2537    }
2538}
2539
2540/// Helper: quote identifiers in a LateralView.
2541fn quote_lateral_view(lv: &mut LateralView, reserved_words: &HashSet<&str>) {
2542    quote_identifiers_recursive(&mut lv.this, reserved_words);
2543    if let Some(ref mut ta) = lv.table_alias {
2544        maybe_quote(ta, reserved_words);
2545    }
2546    for ca in &mut lv.column_aliases {
2547        maybe_quote(ca, reserved_words);
2548    }
2549}
2550
2551/// Quote identifiers that need quoting based on dialect rules.
2552///
2553/// Walks the entire AST recursively and sets `quoted = true` on any
2554/// `Identifier` that:
2555/// - contains special characters (anything not `[a-zA-Z0-9_]`)
2556/// - starts with a digit
2557/// - is a SQL reserved word for the given dialect
2558///
2559/// The function takes ownership of the expression, mutates a clone,
2560/// and returns the modified version.
2561pub fn quote_identifiers(expression: Expression, dialect: Option<DialectType>) -> Expression {
2562    let reserved_words = get_reserved_words(dialect);
2563    let mut result = expression;
2564    quote_identifiers_recursive(&mut result, &reserved_words);
2565    result
2566}
2567
2568/// Pushdown CTE alias columns into the projection.
2569///
2570/// This is useful for dialects like Snowflake where CTE alias columns
2571/// can be referenced in HAVING.
2572pub fn pushdown_cte_alias_columns(_scope: &Scope) {
2573    // Kept for API compatibility. The mutating implementation is applied within
2574    // `qualify_columns` where AST ownership is available.
2575}
2576
2577fn pushdown_cte_alias_columns_with(with: &mut With) {
2578    for cte in &mut with.ctes {
2579        if cte.columns.is_empty() {
2580            continue;
2581        }
2582
2583        if let Expression::Select(select) = &mut cte.this {
2584            let mut next_expressions = Vec::with_capacity(select.expressions.len());
2585
2586            for (i, projection) in select.expressions.iter().enumerate() {
2587                let Some(alias_name) = cte.columns.get(i) else {
2588                    next_expressions.push(projection.clone());
2589                    continue;
2590                };
2591
2592                match projection {
2593                    Expression::Alias(existing) => {
2594                        let mut aliased = existing.clone();
2595                        aliased.alias = alias_name.clone();
2596                        next_expressions.push(Expression::Alias(aliased));
2597                    }
2598                    _ => {
2599                        next_expressions.push(create_alias(projection.clone(), &alias_name.name));
2600                    }
2601                }
2602            }
2603
2604            select.expressions = next_expressions;
2605        }
2606    }
2607}
2608
2609// ============================================================================
2610// Helper functions
2611// ============================================================================
2612
2613/// Get all column references in a scope
2614fn get_scope_columns(scope: &Scope) -> Vec<ColumnRef> {
2615    let mut columns = Vec::new();
2616    collect_columns(&scope.expression, &mut columns);
2617    columns
2618}
2619
2620/// Column reference for tracking
2621#[derive(Debug, Clone)]
2622struct ColumnRef {
2623    table: Option<String>,
2624    name: String,
2625}
2626
2627/// Recursively collect column references from an expression
2628fn collect_columns(expr: &Expression, columns: &mut Vec<ColumnRef>) {
2629    match expr {
2630        Expression::Column(col) => {
2631            columns.push(ColumnRef {
2632                table: col.table.as_ref().map(|t| t.name.clone()),
2633                name: col.name.name.clone(),
2634            });
2635        }
2636        Expression::Select(select) => {
2637            for e in &select.expressions {
2638                collect_columns(e, columns);
2639            }
2640            if let Some(from) = &select.from {
2641                for e in &from.expressions {
2642                    collect_columns(e, columns);
2643                }
2644            }
2645            if let Some(where_clause) = &select.where_clause {
2646                collect_columns(&where_clause.this, columns);
2647            }
2648            if let Some(group_by) = &select.group_by {
2649                for e in &group_by.expressions {
2650                    collect_columns(e, columns);
2651                }
2652            }
2653            if let Some(having) = &select.having {
2654                collect_columns(&having.this, columns);
2655            }
2656            if let Some(order_by) = &select.order_by {
2657                for o in &order_by.expressions {
2658                    collect_columns(&o.this, columns);
2659                }
2660            }
2661            for join in &select.joins {
2662                collect_columns(&join.this, columns);
2663                if let Some(on) = &join.on {
2664                    collect_columns(on, columns);
2665                }
2666            }
2667        }
2668        Expression::Alias(alias) => {
2669            collect_columns(&alias.this, columns);
2670        }
2671        Expression::Function(func) => {
2672            for arg in &func.args {
2673                collect_columns(arg, columns);
2674            }
2675        }
2676        Expression::AggregateFunction(agg) => {
2677            for arg in &agg.args {
2678                collect_columns(arg, columns);
2679            }
2680        }
2681        Expression::And(bin)
2682        | Expression::Or(bin)
2683        | Expression::Eq(bin)
2684        | Expression::Neq(bin)
2685        | Expression::Lt(bin)
2686        | Expression::Lte(bin)
2687        | Expression::Gt(bin)
2688        | Expression::Gte(bin)
2689        | Expression::Add(bin)
2690        | Expression::Sub(bin)
2691        | Expression::Mul(bin)
2692        | Expression::Div(bin) => {
2693            collect_columns(&bin.left, columns);
2694            collect_columns(&bin.right, columns);
2695        }
2696        Expression::Not(unary) | Expression::Neg(unary) => {
2697            collect_columns(&unary.this, columns);
2698        }
2699        Expression::Paren(paren) => {
2700            collect_columns(&paren.this, columns);
2701        }
2702        Expression::Case(case) => {
2703            if let Some(operand) = &case.operand {
2704                collect_columns(operand, columns);
2705            }
2706            for (when, then) in &case.whens {
2707                collect_columns(when, columns);
2708                collect_columns(then, columns);
2709            }
2710            if let Some(else_) = &case.else_ {
2711                collect_columns(else_, columns);
2712            }
2713        }
2714        Expression::Cast(cast) => {
2715            collect_columns(&cast.this, columns);
2716        }
2717        Expression::In(in_expr) => {
2718            collect_columns(&in_expr.this, columns);
2719            for e in &in_expr.expressions {
2720                collect_columns(e, columns);
2721            }
2722            if let Some(query) = &in_expr.query {
2723                collect_columns(query, columns);
2724            }
2725        }
2726        Expression::Between(between) => {
2727            collect_columns(&between.this, columns);
2728            collect_columns(&between.low, columns);
2729            collect_columns(&between.high, columns);
2730        }
2731        Expression::Subquery(subquery) => {
2732            collect_columns(&subquery.this, columns);
2733        }
2734        _ => {}
2735    }
2736}
2737
2738/// Get unqualified columns in a scope
2739fn get_unqualified_columns(scope: &Scope) -> Vec<ColumnRef> {
2740    get_scope_columns(scope)
2741        .into_iter()
2742        .filter(|c| c.table.is_none())
2743        .collect()
2744}
2745
2746/// Get external columns (columns not resolvable in current scope)
2747fn get_external_columns(scope: &Scope) -> Vec<ColumnRef> {
2748    let source_names: HashSet<_> = scope.sources.keys().cloned().collect();
2749
2750    get_scope_columns(scope)
2751        .into_iter()
2752        .filter(|c| {
2753            if let Some(table) = &c.table {
2754                !source_names.contains(table)
2755            } else {
2756                false
2757            }
2758        })
2759        .collect()
2760}
2761
2762/// Check if a scope represents a correlated subquery
2763fn is_correlated_subquery(scope: &Scope) -> bool {
2764    scope.can_be_correlated && !get_external_columns(scope).is_empty()
2765}
2766
2767/// Check if a column represents a star (e.g., table.*)
2768fn is_star_column(col: &Column) -> bool {
2769    col.name.name == "*"
2770}
2771
2772/// Create a qualified column expression
2773fn create_qualified_column(name: &str, table: Option<&str>) -> Expression {
2774    Expression::boxed_column(Column {
2775        name: Identifier::new(name),
2776        table: table.map(Identifier::new),
2777        join_mark: false,
2778        trailing_comments: vec![],
2779        span: None,
2780        inferred_type: None,
2781    })
2782}
2783
2784/// Create an alias expression
2785fn create_alias(expr: Expression, alias_name: &str) -> Expression {
2786    Expression::Alias(Box::new(Alias {
2787        this: expr,
2788        alias: Identifier::new(alias_name),
2789        column_aliases: vec![],
2790        alias_explicit_as: false,
2791        alias_keyword: None,
2792        pre_alias_comments: vec![],
2793        trailing_comments: vec![],
2794        inferred_type: None,
2795    }))
2796}
2797
2798/// Get the output name for an expression
2799fn get_output_name(expr: &Expression) -> Option<String> {
2800    match expr {
2801        Expression::Column(col) => Some(col.name.name.clone()),
2802        Expression::Alias(alias) => Some(alias.alias.name.clone()),
2803        Expression::Identifier(id) => Some(id.name.clone()),
2804        _ => None,
2805    }
2806}
2807
2808#[cfg(test)]
2809mod tests {
2810    use super::*;
2811    use crate::expressions::DataType;
2812    use crate::generator::Generator;
2813    use crate::parser::Parser;
2814    use crate::scope::build_scope;
2815    use crate::{MappingSchema, Schema};
2816
2817    fn gen(expr: &Expression) -> String {
2818        Generator::new().generate(expr).unwrap()
2819    }
2820
2821    fn parse(sql: &str) -> Expression {
2822        Parser::parse_sql(sql).expect("Failed to parse")[0].clone()
2823    }
2824
2825    #[test]
2826    fn test_qualify_columns_options() {
2827        let options = QualifyColumnsOptions::new()
2828            .with_expand_alias_refs(true)
2829            .with_expand_stars(false)
2830            .with_dialect(DialectType::PostgreSQL)
2831            .with_allow_partial(true);
2832
2833        assert!(options.expand_alias_refs);
2834        assert!(!options.expand_stars);
2835        assert_eq!(options.dialect, Some(DialectType::PostgreSQL));
2836        assert!(options.allow_partial_qualification);
2837    }
2838
2839    #[test]
2840    fn test_get_scope_columns() {
2841        let expr = parse("SELECT a, b FROM t WHERE c = 1");
2842        let scope = build_scope(&expr);
2843        let columns = get_scope_columns(&scope);
2844
2845        assert!(columns.iter().any(|c| c.name == "a"));
2846        assert!(columns.iter().any(|c| c.name == "b"));
2847        assert!(columns.iter().any(|c| c.name == "c"));
2848    }
2849
2850    #[test]
2851    fn test_get_unqualified_columns() {
2852        let expr = parse("SELECT t.a, b FROM t");
2853        let scope = build_scope(&expr);
2854        let unqualified = get_unqualified_columns(&scope);
2855
2856        // Only 'b' should be unqualified
2857        assert!(unqualified.iter().any(|c| c.name == "b"));
2858        assert!(!unqualified.iter().any(|c| c.name == "a"));
2859    }
2860
2861    #[test]
2862    fn test_is_star_column() {
2863        let col = Column {
2864            name: Identifier::new("*"),
2865            table: Some(Identifier::new("t")),
2866            join_mark: false,
2867            trailing_comments: vec![],
2868            span: None,
2869            inferred_type: None,
2870        };
2871        assert!(is_star_column(&col));
2872
2873        let col2 = Column {
2874            name: Identifier::new("id"),
2875            table: None,
2876            join_mark: false,
2877            trailing_comments: vec![],
2878            span: None,
2879            inferred_type: None,
2880        };
2881        assert!(!is_star_column(&col2));
2882    }
2883
2884    #[test]
2885    fn test_create_qualified_column() {
2886        let expr = create_qualified_column("id", Some("users"));
2887        let sql = gen(&expr);
2888        assert!(sql.contains("users"));
2889        assert!(sql.contains("id"));
2890    }
2891
2892    #[test]
2893    fn test_create_alias() {
2894        let col = Expression::boxed_column(Column {
2895            name: Identifier::new("value"),
2896            table: None,
2897            join_mark: false,
2898            trailing_comments: vec![],
2899            span: None,
2900            inferred_type: None,
2901        });
2902        let aliased = create_alias(col, "total");
2903        let sql = gen(&aliased);
2904        assert!(sql.contains("AS") || sql.contains("total"));
2905    }
2906
2907    #[test]
2908    fn test_validate_qualify_columns_success() {
2909        // All columns qualified
2910        let expr = parse("SELECT t.a, t.b FROM t");
2911        let result = validate_qualify_columns(&expr);
2912        // This may or may not error depending on scope analysis
2913        // The test verifies the function runs without panic
2914        let _ = result;
2915    }
2916
2917    #[test]
2918    fn test_collect_columns_nested() {
2919        let expr = parse("SELECT a + b, c FROM t WHERE d > 0 GROUP BY e HAVING f = 1");
2920        let mut columns = Vec::new();
2921        collect_columns(&expr, &mut columns);
2922
2923        let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
2924        assert!(names.contains(&"a"));
2925        assert!(names.contains(&"b"));
2926        assert!(names.contains(&"c"));
2927        assert!(names.contains(&"d"));
2928        assert!(names.contains(&"e"));
2929        assert!(names.contains(&"f"));
2930    }
2931
2932    #[test]
2933    fn test_collect_columns_in_case() {
2934        let expr = parse("SELECT CASE WHEN a = 1 THEN b ELSE c END FROM t");
2935        let mut columns = Vec::new();
2936        collect_columns(&expr, &mut columns);
2937
2938        let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
2939        assert!(names.contains(&"a"));
2940        assert!(names.contains(&"b"));
2941        assert!(names.contains(&"c"));
2942    }
2943
2944    #[test]
2945    fn test_collect_columns_in_subquery() {
2946        let expr = parse("SELECT a FROM t WHERE b IN (SELECT c FROM s)");
2947        let mut columns = Vec::new();
2948        collect_columns(&expr, &mut columns);
2949
2950        let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
2951        assert!(names.contains(&"a"));
2952        assert!(names.contains(&"b"));
2953        assert!(names.contains(&"c"));
2954    }
2955
2956    #[test]
2957    fn test_qualify_outputs_basic() {
2958        let expr = parse("SELECT a, b + c FROM t");
2959        let scope = build_scope(&expr);
2960        let result = qualify_outputs(&scope);
2961        assert!(result.is_ok());
2962    }
2963
2964    #[test]
2965    fn test_qualify_columns_expands_star_with_schema() {
2966        let expr = parse("SELECT * FROM users");
2967
2968        let mut schema = MappingSchema::new();
2969        schema
2970            .add_table(
2971                "users",
2972                &[
2973                    (
2974                        "id".to_string(),
2975                        DataType::Int {
2976                            length: None,
2977                            integer_spelling: false,
2978                        },
2979                    ),
2980                    ("name".to_string(), DataType::Text),
2981                    ("email".to_string(), DataType::Text),
2982                ],
2983                None,
2984            )
2985            .expect("schema setup");
2986
2987        let result =
2988            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
2989        let sql = gen(&result);
2990
2991        assert!(!sql.contains("SELECT *"));
2992        assert!(sql.contains("users.id"));
2993        assert!(sql.contains("users.name"));
2994        assert!(sql.contains("users.email"));
2995    }
2996
2997    #[test]
2998    fn test_qualify_columns_expands_group_by_positions() {
2999        let expr = parse("SELECT a, b FROM t GROUP BY 1, 2");
3000
3001        let mut schema = MappingSchema::new();
3002        schema
3003            .add_table(
3004                "t",
3005                &[
3006                    (
3007                        "a".to_string(),
3008                        DataType::Int {
3009                            length: None,
3010                            integer_spelling: false,
3011                        },
3012                    ),
3013                    (
3014                        "b".to_string(),
3015                        DataType::Int {
3016                            length: None,
3017                            integer_spelling: false,
3018                        },
3019                    ),
3020                ],
3021                None,
3022            )
3023            .expect("schema setup");
3024
3025        let result =
3026            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3027        let sql = gen(&result);
3028
3029        assert!(!sql.contains("GROUP BY 1"));
3030        assert!(!sql.contains("GROUP BY 2"));
3031        assert!(sql.contains("GROUP BY"));
3032        assert!(sql.contains("t.a"));
3033        assert!(sql.contains("t.b"));
3034    }
3035
3036    // ======================================================================
3037    // USING expansion tests
3038    // ======================================================================
3039
3040    #[test]
3041    fn test_expand_using_simple() {
3042        // Already-qualified column: USING→ON rewrite but no COALESCE needed
3043        let expr = parse("SELECT x.b FROM x JOIN y USING (b)");
3044
3045        let mut schema = MappingSchema::new();
3046        schema
3047            .add_table(
3048                "x",
3049                &[
3050                    ("a".to_string(), DataType::BigInt { length: None }),
3051                    ("b".to_string(), DataType::BigInt { length: None }),
3052                ],
3053                None,
3054            )
3055            .expect("schema setup");
3056        schema
3057            .add_table(
3058                "y",
3059                &[
3060                    ("b".to_string(), DataType::BigInt { length: None }),
3061                    ("c".to_string(), DataType::BigInt { length: None }),
3062                ],
3063                None,
3064            )
3065            .expect("schema setup");
3066
3067        let result =
3068            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3069        let sql = gen(&result);
3070
3071        // USING should be replaced with ON
3072        assert!(
3073            !sql.contains("USING"),
3074            "USING should be replaced with ON: {sql}"
3075        );
3076        assert!(
3077            sql.contains("ON x.b = y.b"),
3078            "ON condition should be x.b = y.b: {sql}"
3079        );
3080        // x.b in SELECT should remain as-is (already qualified)
3081        assert!(sql.contains("SELECT x.b"), "SELECT should keep x.b: {sql}");
3082    }
3083
3084    #[test]
3085    fn test_expand_using_unqualified_coalesce() {
3086        // Unqualified USING column in SELECT should become COALESCE
3087        let expr = parse("SELECT b FROM x JOIN y USING(b)");
3088
3089        let mut schema = MappingSchema::new();
3090        schema
3091            .add_table(
3092                "x",
3093                &[
3094                    ("a".to_string(), DataType::BigInt { length: None }),
3095                    ("b".to_string(), DataType::BigInt { length: None }),
3096                ],
3097                None,
3098            )
3099            .expect("schema setup");
3100        schema
3101            .add_table(
3102                "y",
3103                &[
3104                    ("b".to_string(), DataType::BigInt { length: None }),
3105                    ("c".to_string(), DataType::BigInt { length: None }),
3106                ],
3107                None,
3108            )
3109            .expect("schema setup");
3110
3111        let result =
3112            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3113        let sql = gen(&result);
3114
3115        assert!(
3116            sql.contains("COALESCE(x.b, y.b)"),
3117            "Unqualified USING column should become COALESCE: {sql}"
3118        );
3119        assert!(
3120            sql.contains("AS b"),
3121            "COALESCE should be aliased as 'b': {sql}"
3122        );
3123        assert!(
3124            sql.contains("ON x.b = y.b"),
3125            "ON condition should be generated: {sql}"
3126        );
3127    }
3128
3129    #[test]
3130    fn test_expand_using_with_where() {
3131        // USING column in WHERE should become COALESCE
3132        let expr = parse("SELECT b FROM x JOIN y USING(b) WHERE b = 1");
3133
3134        let mut schema = MappingSchema::new();
3135        schema
3136            .add_table(
3137                "x",
3138                &[("b".to_string(), DataType::BigInt { length: None })],
3139                None,
3140            )
3141            .expect("schema setup");
3142        schema
3143            .add_table(
3144                "y",
3145                &[("b".to_string(), DataType::BigInt { length: None })],
3146                None,
3147            )
3148            .expect("schema setup");
3149
3150        let result =
3151            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3152        let sql = gen(&result);
3153
3154        assert!(
3155            sql.contains("WHERE COALESCE(x.b, y.b)"),
3156            "WHERE should use COALESCE for USING column: {sql}"
3157        );
3158    }
3159
3160    #[test]
3161    fn test_expand_using_multi_join() {
3162        // Three-way join with same USING column
3163        let expr = parse("SELECT b FROM x JOIN y USING(b) JOIN z USING(b)");
3164
3165        let mut schema = MappingSchema::new();
3166        for table in &["x", "y", "z"] {
3167            schema
3168                .add_table(
3169                    table,
3170                    &[("b".to_string(), DataType::BigInt { length: None })],
3171                    None,
3172                )
3173                .expect("schema setup");
3174        }
3175
3176        let result =
3177            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3178        let sql = gen(&result);
3179
3180        // SELECT should have 3-table COALESCE
3181        assert!(
3182            sql.contains("COALESCE(x.b, y.b, z.b)"),
3183            "Should have 3-table COALESCE: {sql}"
3184        );
3185        // First join: simple ON
3186        assert!(
3187            sql.contains("ON x.b = y.b"),
3188            "First join ON condition: {sql}"
3189        );
3190    }
3191
3192    #[test]
3193    fn test_expand_using_multi_column() {
3194        // Two USING columns
3195        let expr = parse("SELECT b, c FROM y JOIN z USING(b, c)");
3196
3197        let mut schema = MappingSchema::new();
3198        schema
3199            .add_table(
3200                "y",
3201                &[
3202                    ("b".to_string(), DataType::BigInt { length: None }),
3203                    ("c".to_string(), DataType::BigInt { length: None }),
3204                ],
3205                None,
3206            )
3207            .expect("schema setup");
3208        schema
3209            .add_table(
3210                "z",
3211                &[
3212                    ("b".to_string(), DataType::BigInt { length: None }),
3213                    ("c".to_string(), DataType::BigInt { length: None }),
3214                ],
3215                None,
3216            )
3217            .expect("schema setup");
3218
3219        let result =
3220            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3221        let sql = gen(&result);
3222
3223        assert!(
3224            sql.contains("COALESCE(y.b, z.b)"),
3225            "column 'b' should get COALESCE: {sql}"
3226        );
3227        assert!(
3228            sql.contains("COALESCE(y.c, z.c)"),
3229            "column 'c' should get COALESCE: {sql}"
3230        );
3231        // ON should have both conditions ANDed
3232        assert!(
3233            sql.contains("y.b = z.b") && sql.contains("y.c = z.c"),
3234            "ON should have both equality conditions: {sql}"
3235        );
3236    }
3237
3238    #[test]
3239    fn test_expand_using_star() {
3240        // SELECT * should deduplicate USING columns
3241        let expr = parse("SELECT * FROM x JOIN y USING(b)");
3242
3243        let mut schema = MappingSchema::new();
3244        schema
3245            .add_table(
3246                "x",
3247                &[
3248                    ("a".to_string(), DataType::BigInt { length: None }),
3249                    ("b".to_string(), DataType::BigInt { length: None }),
3250                ],
3251                None,
3252            )
3253            .expect("schema setup");
3254        schema
3255            .add_table(
3256                "y",
3257                &[
3258                    ("b".to_string(), DataType::BigInt { length: None }),
3259                    ("c".to_string(), DataType::BigInt { length: None }),
3260                ],
3261                None,
3262            )
3263            .expect("schema setup");
3264
3265        let result =
3266            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3267        let sql = gen(&result);
3268
3269        // b should appear once as COALESCE
3270        assert!(
3271            sql.contains("COALESCE(x.b, y.b) AS b"),
3272            "USING column should be COALESCE in star expansion: {sql}"
3273        );
3274        // a and c should be normal qualified columns
3275        assert!(sql.contains("x.a"), "non-USING column a from x: {sql}");
3276        assert!(sql.contains("y.c"), "non-USING column c from y: {sql}");
3277        // b should only appear once (not duplicated from both tables)
3278        let coalesce_count = sql.matches("COALESCE").count();
3279        assert_eq!(
3280            coalesce_count, 1,
3281            "b should appear only once as COALESCE: {sql}"
3282        );
3283    }
3284
3285    #[test]
3286    fn test_expand_using_table_star() {
3287        // table.* with USING column
3288        let expr = parse("SELECT x.* FROM x JOIN y USING(b)");
3289
3290        let mut schema = MappingSchema::new();
3291        schema
3292            .add_table(
3293                "x",
3294                &[
3295                    ("a".to_string(), DataType::BigInt { length: None }),
3296                    ("b".to_string(), DataType::BigInt { length: None }),
3297                ],
3298                None,
3299            )
3300            .expect("schema setup");
3301        schema
3302            .add_table(
3303                "y",
3304                &[
3305                    ("b".to_string(), DataType::BigInt { length: None }),
3306                    ("c".to_string(), DataType::BigInt { length: None }),
3307                ],
3308                None,
3309            )
3310            .expect("schema setup");
3311
3312        let result =
3313            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3314        let sql = gen(&result);
3315
3316        // b should become COALESCE (since x participates in USING for b)
3317        assert!(
3318            sql.contains("COALESCE(x.b, y.b)"),
3319            "USING column from x.* should become COALESCE: {sql}"
3320        );
3321        assert!(sql.contains("x.a"), "non-USING column a: {sql}");
3322    }
3323
3324    #[test]
3325    fn test_expand_natural_join_with_derived_table() {
3326        let expr = parse(
3327            "SELECT shared_key AS output_key FROM source_table \
3328             NATURAL JOIN (SELECT shared_key FROM source_table) AS derived",
3329        );
3330
3331        let mut schema = MappingSchema::new();
3332        schema
3333            .add_table(
3334                "source_table",
3335                &[(
3336                    "shared_key".to_string(),
3337                    DataType::VarChar {
3338                        length: None,
3339                        parenthesized_length: false,
3340                    },
3341                )],
3342                None,
3343            )
3344            .expect("schema setup");
3345
3346        let result = qualify_columns(
3347            expr,
3348            &schema,
3349            &QualifyColumnsOptions::new().with_dialect(DialectType::DuckDB),
3350        )
3351        .expect("qualify");
3352        let sql = gen(&result);
3353
3354        assert!(!sql.contains("NATURAL"), "NATURAL should expand: {sql}");
3355        assert!(
3356            sql.contains("ON source_table.shared_key = derived.shared_key"),
3357            "common column should become an equality condition: {sql}"
3358        );
3359        assert!(
3360            sql.contains("COALESCE(source_table.shared_key, derived.shared_key) AS output_key"),
3361            "merged projection should preserve both sources: {sql}"
3362        );
3363    }
3364
3365    #[test]
3366    fn test_expand_natural_join_star_and_chained_sources() {
3367        let expr = parse("SELECT * FROM x NATURAL JOIN y NATURAL JOIN z");
3368
3369        let mut schema = MappingSchema::new();
3370        schema
3371            .add_table(
3372                "x",
3373                &[
3374                    ("a".to_string(), DataType::BigInt { length: None }),
3375                    ("shared".to_string(), DataType::BigInt { length: None }),
3376                ],
3377                None,
3378            )
3379            .expect("schema setup");
3380        schema
3381            .add_table(
3382                "y",
3383                &[
3384                    ("shared".to_string(), DataType::BigInt { length: None }),
3385                    ("b".to_string(), DataType::BigInt { length: None }),
3386                ],
3387                None,
3388            )
3389            .expect("schema setup");
3390        schema
3391            .add_table(
3392                "z",
3393                &[
3394                    ("shared".to_string(), DataType::BigInt { length: None }),
3395                    ("c".to_string(), DataType::BigInt { length: None }),
3396                ],
3397                None,
3398            )
3399            .expect("schema setup");
3400
3401        let result =
3402            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3403        let sql = gen(&result);
3404
3405        assert!(!sql.contains("NATURAL"), "both joins should expand: {sql}");
3406        assert!(sql.contains("ON x.shared = y.shared"), "first join: {sql}");
3407        assert!(sql.contains("ON x.shared = z.shared"), "second join: {sql}");
3408        assert_eq!(
3409            sql.matches("COALESCE(x.shared, y.shared, z.shared) AS shared")
3410                .count(),
3411            1,
3412            "merged star column should be emitted once: {sql}"
3413        );
3414        assert!(sql.contains("x.a"), "left-only column should remain: {sql}");
3415        assert!(
3416            sql.contains("y.b"),
3417            "middle-only column should remain: {sql}"
3418        );
3419        assert!(
3420            sql.contains("z.c"),
3421            "right-only column should remain: {sql}"
3422        );
3423    }
3424
3425    #[test]
3426    fn test_expand_natural_outer_join_kinds() {
3427        let mut schema = MappingSchema::new();
3428        for table in ["x", "y"] {
3429            schema
3430                .add_table(
3431                    table,
3432                    &[("shared".to_string(), DataType::BigInt { length: None })],
3433                    None,
3434                )
3435                .expect("schema setup");
3436        }
3437
3438        for (input_kind, output_kind) in [
3439            ("NATURAL LEFT JOIN", "LEFT JOIN"),
3440            ("NATURAL RIGHT JOIN", "RIGHT JOIN"),
3441            ("NATURAL FULL JOIN", "FULL JOIN"),
3442        ] {
3443            let expr = parse(&format!("SELECT shared FROM x {input_kind} y"));
3444            let result =
3445                qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3446            let sql = gen(&result);
3447
3448            assert!(
3449                !sql.contains("NATURAL"),
3450                "{input_kind} should expand: {sql}"
3451            );
3452            assert!(sql.contains(output_kind), "join kind should remain: {sql}");
3453            assert!(
3454                sql.contains("ON x.shared = y.shared"),
3455                "join condition should be derived: {sql}"
3456            );
3457        }
3458    }
3459
3460    #[test]
3461    fn test_preserve_natural_join_without_known_common_columns() {
3462        let mut no_common_schema = MappingSchema::new();
3463        no_common_schema
3464            .add_table(
3465                "x",
3466                &[("a".to_string(), DataType::BigInt { length: None })],
3467                None,
3468            )
3469            .expect("schema setup");
3470        no_common_schema
3471            .add_table(
3472                "y",
3473                &[("b".to_string(), DataType::BigInt { length: None })],
3474                None,
3475            )
3476            .expect("schema setup");
3477
3478        let no_common = qualify_columns(
3479            parse("SELECT * FROM x NATURAL JOIN y"),
3480            &no_common_schema,
3481            &QualifyColumnsOptions::new(),
3482        )
3483        .expect("qualify");
3484        assert!(
3485            gen(&no_common).contains("NATURAL JOIN"),
3486            "join without common columns should remain NATURAL"
3487        );
3488
3489        let mut partial_schema = MappingSchema::new();
3490        partial_schema
3491            .add_table(
3492                "x",
3493                &[("a".to_string(), DataType::BigInt { length: None })],
3494                None,
3495            )
3496            .expect("schema setup");
3497        let unknown_right = qualify_columns(
3498            parse("SELECT * FROM x NATURAL JOIN unknown_table"),
3499            &partial_schema,
3500            &QualifyColumnsOptions::new().with_allow_partial(true),
3501        )
3502        .expect("partial qualification");
3503        assert!(
3504            gen(&unknown_right).contains("NATURAL JOIN"),
3505            "join with an unknown schema should remain NATURAL"
3506        );
3507
3508        let unknown_left = qualify_columns(
3509            parse("SELECT * FROM unknown_table NATURAL JOIN x"),
3510            &partial_schema,
3511            &QualifyColumnsOptions::new().with_allow_partial(true),
3512        )
3513        .expect("partial qualification");
3514        assert!(
3515            gen(&unknown_left).contains("NATURAL JOIN"),
3516            "join with an unknown left schema should remain NATURAL"
3517        );
3518
3519        let partially_known_chain = qualify_columns(
3520            parse("SELECT * FROM x NATURAL JOIN unknown_table NATURAL JOIN x AS x2"),
3521            &partial_schema,
3522            &QualifyColumnsOptions::new().with_allow_partial(true),
3523        )
3524        .expect("partial qualification");
3525        assert_eq!(
3526            gen(&partially_known_chain).matches("NATURAL JOIN").count(),
3527            2,
3528            "an unknown earlier source should prevent later inferred join keys"
3529        );
3530    }
3531
3532    #[test]
3533    fn test_qualify_columns_qualified_table_name() {
3534        let expr = parse("SELECT a FROM raw.t1");
3535
3536        let mut schema = MappingSchema::new();
3537        schema
3538            .add_table(
3539                "raw.t1",
3540                &[("a".to_string(), DataType::BigInt { length: None })],
3541                None,
3542            )
3543            .expect("schema setup");
3544
3545        let result =
3546            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3547        let sql = gen(&result);
3548
3549        assert!(
3550            sql.contains("t1.a"),
3551            "column should be qualified with table name: {sql}"
3552        );
3553
3554        // test that columns in agg functions also get qualified
3555        let expr = parse("SELECT MAX(a) FROM raw.t1");
3556        let result =
3557            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3558        let sql = gen(&result);
3559        assert!(
3560            sql.contains("t1.a"),
3561            "column in function should be qualified with table name: {sql}"
3562        );
3563
3564        // test that columns in scalar functions also get qualified
3565        let expr = parse("SELECT ABS(a) FROM raw.t1");
3566        let result =
3567            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3568        let sql = gen(&result);
3569        assert!(
3570            sql.contains("t1.a"),
3571            "column in function should be qualified with table name: {sql}"
3572        );
3573    }
3574
3575    #[test]
3576    fn test_qualify_columns_count_star() {
3577        // COUNT(*) uses Count { this: None } — verify qualify_columns handles it without panic
3578        let expr = parse("SELECT COUNT(*) FROM t1");
3579
3580        let mut schema = MappingSchema::new();
3581        schema
3582            .add_table(
3583                "t1",
3584                &[("id".to_string(), DataType::BigInt { length: None })],
3585                None,
3586            )
3587            .expect("schema setup");
3588
3589        let result =
3590            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3591        let sql = gen(&result);
3592
3593        assert!(
3594            sql.contains("COUNT(*)"),
3595            "COUNT(*) should be preserved: {sql}"
3596        );
3597    }
3598
3599    #[test]
3600    fn test_qualify_columns_correlated_scalar_subquery() {
3601        let expr =
3602            parse("SELECT id, (SELECT AVG(val) FROM t2 WHERE t2.id = t1.id) AS avg_val FROM t1");
3603
3604        let mut schema = MappingSchema::new();
3605        schema
3606            .add_table(
3607                "t1",
3608                &[("id".to_string(), DataType::BigInt { length: None })],
3609                None,
3610            )
3611            .expect("schema setup");
3612        schema
3613            .add_table(
3614                "t2",
3615                &[
3616                    ("id".to_string(), DataType::BigInt { length: None }),
3617                    ("val".to_string(), DataType::BigInt { length: None }),
3618                ],
3619                None,
3620            )
3621            .expect("schema setup");
3622
3623        let result =
3624            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3625        let sql = gen(&result);
3626
3627        assert!(
3628            sql.contains("t1.id"),
3629            "outer column should be qualified: {sql}"
3630        );
3631        assert!(
3632            sql.contains("t2.id"),
3633            "inner column should be qualified: {sql}"
3634        );
3635    }
3636
3637    #[test]
3638    fn test_qualify_columns_correlated_scalar_subquery_unqualified() {
3639        let expr =
3640            parse("SELECT t1_id, (SELECT AVG(val) FROM t2 WHERE t2_id = t1_id) AS avg_val FROM t1");
3641
3642        let mut schema = MappingSchema::new();
3643        schema
3644            .add_table(
3645                "t1",
3646                &[("t1_id".to_string(), DataType::BigInt { length: None })],
3647                None,
3648            )
3649            .expect("schema setup");
3650        schema
3651            .add_table(
3652                "t2",
3653                &[
3654                    ("t2_id".to_string(), DataType::BigInt { length: None }),
3655                    ("val".to_string(), DataType::BigInt { length: None }),
3656                ],
3657                None,
3658            )
3659            .expect("schema setup");
3660
3661        let result =
3662            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3663        let sql = gen(&result);
3664
3665        assert!(
3666            sql.contains("t1.t1_id"),
3667            "outer column should be qualified: {sql}"
3668        );
3669        assert!(
3670            sql.contains("t2.t2_id"),
3671            "inner column should be qualified: {sql}"
3672        );
3673        // Correlated reference t1_id in inner scope should be qualified as t1.t1_id
3674        assert!(
3675            sql.contains("= t1.t1_id"),
3676            "correlated column should be qualified: {sql}"
3677        );
3678    }
3679
3680    #[test]
3681    fn test_qualify_columns_correlated_exists_subquery() {
3682        let expr = parse(
3683            "SELECT o_orderpriority FROM orders \
3684             WHERE EXISTS (SELECT * FROM lineitem WHERE l_orderkey = o_orderkey)",
3685        );
3686
3687        let mut schema = MappingSchema::new();
3688        schema
3689            .add_table(
3690                "orders",
3691                &[
3692                    ("o_orderpriority".to_string(), DataType::Text),
3693                    ("o_orderkey".to_string(), DataType::BigInt { length: None }),
3694                ],
3695                None,
3696            )
3697            .expect("schema setup");
3698        schema
3699            .add_table(
3700                "lineitem",
3701                &[("l_orderkey".to_string(), DataType::BigInt { length: None })],
3702                None,
3703            )
3704            .expect("schema setup");
3705
3706        let result =
3707            qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3708        let sql = gen(&result);
3709
3710        assert!(
3711            sql.contains("orders.o_orderpriority"),
3712            "outer column should be qualified: {sql}"
3713        );
3714        assert!(
3715            sql.contains("lineitem.l_orderkey"),
3716            "inner column should be qualified: {sql}"
3717        );
3718        assert!(
3719            sql.contains("orders.o_orderkey"),
3720            "correlated outer column should be qualified: {sql}"
3721        );
3722    }
3723
3724    #[test]
3725    fn test_qualify_columns_rejects_unknown_table() {
3726        let expr = parse("SELECT id FROM t1 WHERE nonexistent.col = 1");
3727
3728        let mut schema = MappingSchema::new();
3729        schema
3730            .add_table(
3731                "t1",
3732                &[("id".to_string(), DataType::BigInt { length: None })],
3733                None,
3734            )
3735            .expect("schema setup");
3736
3737        let result = qualify_columns(expr, &schema, &QualifyColumnsOptions::new());
3738        assert!(
3739            result.is_err(),
3740            "should reject reference to table not in scope or schema"
3741        );
3742    }
3743
3744    // ======================================================================
3745    // quote_identifiers tests
3746    // ======================================================================
3747
3748    #[test]
3749    fn test_needs_quoting_reserved_word() {
3750        let reserved = get_reserved_words(None);
3751        assert!(needs_quoting("select", &reserved));
3752        assert!(needs_quoting("SELECT", &reserved));
3753        assert!(needs_quoting("from", &reserved));
3754        assert!(needs_quoting("WHERE", &reserved));
3755        assert!(needs_quoting("join", &reserved));
3756        assert!(needs_quoting("table", &reserved));
3757    }
3758
3759    #[test]
3760    fn test_needs_quoting_normal_identifiers() {
3761        let reserved = get_reserved_words(None);
3762        assert!(!needs_quoting("foo", &reserved));
3763        assert!(!needs_quoting("my_column", &reserved));
3764        assert!(!needs_quoting("col1", &reserved));
3765        assert!(!needs_quoting("A", &reserved));
3766        assert!(!needs_quoting("_hidden", &reserved));
3767    }
3768
3769    #[test]
3770    fn test_needs_quoting_special_characters() {
3771        let reserved = get_reserved_words(None);
3772        assert!(needs_quoting("my column", &reserved)); // space
3773        assert!(needs_quoting("my-column", &reserved)); // hyphen
3774        assert!(needs_quoting("my.column", &reserved)); // dot
3775        assert!(needs_quoting("col@name", &reserved)); // at sign
3776        assert!(needs_quoting("col#name", &reserved)); // hash
3777    }
3778
3779    #[test]
3780    fn test_needs_quoting_starts_with_digit() {
3781        let reserved = get_reserved_words(None);
3782        assert!(needs_quoting("1col", &reserved));
3783        assert!(needs_quoting("123", &reserved));
3784        assert!(needs_quoting("0_start", &reserved));
3785    }
3786
3787    #[test]
3788    fn test_needs_quoting_empty() {
3789        let reserved = get_reserved_words(None);
3790        assert!(!needs_quoting("", &reserved));
3791    }
3792
3793    #[test]
3794    fn test_maybe_quote_sets_quoted_flag() {
3795        let reserved = get_reserved_words(None);
3796        let mut id = Identifier::new("select");
3797        assert!(!id.quoted);
3798        maybe_quote(&mut id, &reserved);
3799        assert!(id.quoted);
3800    }
3801
3802    #[test]
3803    fn test_maybe_quote_skips_already_quoted() {
3804        let reserved = get_reserved_words(None);
3805        let mut id = Identifier::quoted("myname");
3806        assert!(id.quoted);
3807        maybe_quote(&mut id, &reserved);
3808        assert!(id.quoted); // still quoted
3809        assert_eq!(id.name, "myname"); // name unchanged
3810    }
3811
3812    #[test]
3813    fn test_maybe_quote_skips_star() {
3814        let reserved = get_reserved_words(None);
3815        let mut id = Identifier::new("*");
3816        maybe_quote(&mut id, &reserved);
3817        assert!(!id.quoted); // star should not be quoted
3818    }
3819
3820    #[test]
3821    fn test_maybe_quote_skips_normal() {
3822        let reserved = get_reserved_words(None);
3823        let mut id = Identifier::new("normal_col");
3824        maybe_quote(&mut id, &reserved);
3825        assert!(!id.quoted);
3826    }
3827
3828    #[test]
3829    fn test_quote_identifiers_column_with_reserved_name() {
3830        // A column named "select" should be quoted
3831        let expr = Expression::boxed_column(Column {
3832            name: Identifier::new("select"),
3833            table: None,
3834            join_mark: false,
3835            trailing_comments: vec![],
3836            span: None,
3837            inferred_type: None,
3838        });
3839        let result = quote_identifiers(expr, None);
3840        if let Expression::Column(col) = &result {
3841            assert!(col.name.quoted, "Column named 'select' should be quoted");
3842        } else {
3843            panic!("Expected Column expression");
3844        }
3845    }
3846
3847    #[test]
3848    fn test_quote_identifiers_column_with_special_chars() {
3849        let expr = Expression::boxed_column(Column {
3850            name: Identifier::new("my column"),
3851            table: None,
3852            join_mark: false,
3853            trailing_comments: vec![],
3854            span: None,
3855            inferred_type: None,
3856        });
3857        let result = quote_identifiers(expr, None);
3858        if let Expression::Column(col) = &result {
3859            assert!(col.name.quoted, "Column with space should be quoted");
3860        } else {
3861            panic!("Expected Column expression");
3862        }
3863    }
3864
3865    #[test]
3866    fn test_quote_identifiers_preserves_normal_column() {
3867        let expr = Expression::boxed_column(Column {
3868            name: Identifier::new("normal_col"),
3869            table: Some(Identifier::new("my_table")),
3870            join_mark: false,
3871            trailing_comments: vec![],
3872            span: None,
3873            inferred_type: None,
3874        });
3875        let result = quote_identifiers(expr, None);
3876        if let Expression::Column(col) = &result {
3877            assert!(!col.name.quoted, "Normal column should not be quoted");
3878            assert!(
3879                !col.table.as_ref().unwrap().quoted,
3880                "Normal table should not be quoted"
3881            );
3882        } else {
3883            panic!("Expected Column expression");
3884        }
3885    }
3886
3887    #[test]
3888    fn test_quote_identifiers_table_ref_reserved() {
3889        let expr = Expression::Table(Box::new(TableRef::new("select")));
3890        let result = quote_identifiers(expr, None);
3891        if let Expression::Table(tr) = &result {
3892            assert!(tr.name.quoted, "Table named 'select' should be quoted");
3893        } else {
3894            panic!("Expected Table expression");
3895        }
3896    }
3897
3898    #[test]
3899    fn test_quote_identifiers_table_ref_schema_and_alias() {
3900        let mut tr = TableRef::new("my_table");
3901        tr.schema = Some(Identifier::new("from"));
3902        tr.alias = Some(Identifier::new("t"));
3903        let expr = Expression::Table(Box::new(tr));
3904        let result = quote_identifiers(expr, None);
3905        if let Expression::Table(tr) = &result {
3906            assert!(!tr.name.quoted, "Normal table name should not be quoted");
3907            assert!(
3908                tr.schema.as_ref().unwrap().quoted,
3909                "Schema named 'from' should be quoted"
3910            );
3911            assert!(
3912                !tr.alias.as_ref().unwrap().quoted,
3913                "Normal alias should not be quoted"
3914            );
3915        } else {
3916            panic!("Expected Table expression");
3917        }
3918    }
3919
3920    #[test]
3921    fn test_quote_identifiers_identifier_node() {
3922        let expr = Expression::Identifier(Identifier::new("order"));
3923        let result = quote_identifiers(expr, None);
3924        if let Expression::Identifier(id) = &result {
3925            assert!(id.quoted, "Identifier named 'order' should be quoted");
3926        } else {
3927            panic!("Expected Identifier expression");
3928        }
3929    }
3930
3931    #[test]
3932    fn test_quote_identifiers_alias() {
3933        let inner = Expression::boxed_column(Column {
3934            name: Identifier::new("val"),
3935            table: None,
3936            join_mark: false,
3937            trailing_comments: vec![],
3938            span: None,
3939            inferred_type: None,
3940        });
3941        let expr = Expression::Alias(Box::new(Alias {
3942            this: inner,
3943            alias: Identifier::new("select"),
3944            column_aliases: vec![Identifier::new("from")],
3945            alias_explicit_as: false,
3946            alias_keyword: None,
3947            pre_alias_comments: vec![],
3948            trailing_comments: vec![],
3949            inferred_type: None,
3950        }));
3951        let result = quote_identifiers(expr, None);
3952        if let Expression::Alias(alias) = &result {
3953            assert!(alias.alias.quoted, "Alias named 'select' should be quoted");
3954            assert!(
3955                alias.column_aliases[0].quoted,
3956                "Column alias named 'from' should be quoted"
3957            );
3958            // Inner column "val" should not be quoted
3959            if let Expression::Column(col) = &alias.this {
3960                assert!(!col.name.quoted);
3961            }
3962        } else {
3963            panic!("Expected Alias expression");
3964        }
3965    }
3966
3967    #[test]
3968    fn test_quote_identifiers_select_recursive() {
3969        // Parse a query and verify quote_identifiers walks through it
3970        let expr = parse("SELECT a, b FROM t WHERE c = 1");
3971        let result = quote_identifiers(expr, None);
3972        // "a", "b", "c", "t" are all normal identifiers, none should be quoted
3973        let sql = gen(&result);
3974        // The SQL should be unchanged since no reserved words are used
3975        assert!(sql.contains("a"));
3976        assert!(sql.contains("b"));
3977        assert!(sql.contains("t"));
3978    }
3979
3980    #[test]
3981    fn test_quote_identifiers_digit_start() {
3982        let expr = Expression::boxed_column(Column {
3983            name: Identifier::new("1col"),
3984            table: None,
3985            join_mark: false,
3986            trailing_comments: vec![],
3987            span: None,
3988            inferred_type: None,
3989        });
3990        let result = quote_identifiers(expr, None);
3991        if let Expression::Column(col) = &result {
3992            assert!(
3993                col.name.quoted,
3994                "Column starting with digit should be quoted"
3995            );
3996        } else {
3997            panic!("Expected Column expression");
3998        }
3999    }
4000
4001    #[test]
4002    fn test_quote_identifiers_with_mysql_dialect() {
4003        let reserved = get_reserved_words(Some(DialectType::MySQL));
4004        // "KILL" is reserved in MySQL
4005        assert!(needs_quoting("KILL", &reserved));
4006        // "FORCE" is reserved in MySQL
4007        assert!(needs_quoting("FORCE", &reserved));
4008    }
4009
4010    #[test]
4011    fn test_quote_identifiers_with_postgresql_dialect() {
4012        let reserved = get_reserved_words(Some(DialectType::PostgreSQL));
4013        // "ILIKE" is reserved in PostgreSQL
4014        assert!(needs_quoting("ILIKE", &reserved));
4015        // "VERBOSE" is reserved in PostgreSQL
4016        assert!(needs_quoting("VERBOSE", &reserved));
4017    }
4018
4019    #[test]
4020    fn test_quote_identifiers_with_bigquery_dialect() {
4021        let reserved = get_reserved_words(Some(DialectType::BigQuery));
4022        // "STRUCT" is reserved in BigQuery
4023        assert!(needs_quoting("STRUCT", &reserved));
4024        // "PROTO" is reserved in BigQuery
4025        assert!(needs_quoting("PROTO", &reserved));
4026    }
4027
4028    #[test]
4029    fn test_quote_identifiers_case_insensitive_reserved() {
4030        let reserved = get_reserved_words(None);
4031        assert!(needs_quoting("Select", &reserved));
4032        assert!(needs_quoting("sElEcT", &reserved));
4033        assert!(needs_quoting("FROM", &reserved));
4034        assert!(needs_quoting("from", &reserved));
4035    }
4036
4037    #[test]
4038    fn test_quote_identifiers_join_using() {
4039        // Build a join with USING identifiers that include reserved words
4040        let mut join = crate::expressions::Join {
4041            this: Expression::Table(Box::new(TableRef::new("other"))),
4042            on: None,
4043            using: vec![Identifier::new("key"), Identifier::new("value")],
4044            kind: crate::expressions::JoinKind::Inner,
4045            use_inner_keyword: false,
4046            use_outer_keyword: false,
4047            deferred_condition: false,
4048            join_hint: None,
4049            match_condition: None,
4050            pivots: vec![],
4051            comments: vec![],
4052            nesting_group: 0,
4053            directed: false,
4054        };
4055        let reserved = get_reserved_words(None);
4056        quote_join(&mut join, &reserved);
4057        // "key" is reserved, "value" is not
4058        assert!(
4059            join.using[0].quoted,
4060            "USING identifier 'key' should be quoted"
4061        );
4062        assert!(
4063            !join.using[1].quoted,
4064            "USING identifier 'value' should not be quoted"
4065        );
4066    }
4067
4068    #[test]
4069    fn test_quote_identifiers_cte() {
4070        // Build a CTE where alias is a reserved word
4071        let mut cte = crate::expressions::Cte {
4072            alias: Identifier::new("select"),
4073            this: Expression::boxed_column(Column {
4074                name: Identifier::new("x"),
4075                table: None,
4076                join_mark: false,
4077                trailing_comments: vec![],
4078                span: None,
4079                inferred_type: None,
4080            }),
4081            columns: vec![Identifier::new("from"), Identifier::new("normal")],
4082            materialized: None,
4083            key_expressions: vec![],
4084            alias_first: false,
4085            comments: Vec::new(),
4086        };
4087        let reserved = get_reserved_words(None);
4088        maybe_quote(&mut cte.alias, &reserved);
4089        for c in &mut cte.columns {
4090            maybe_quote(c, &reserved);
4091        }
4092        assert!(cte.alias.quoted, "CTE alias 'select' should be quoted");
4093        assert!(cte.columns[0].quoted, "CTE column 'from' should be quoted");
4094        assert!(
4095            !cte.columns[1].quoted,
4096            "CTE column 'normal' should not be quoted"
4097        );
4098    }
4099
4100    #[test]
4101    fn test_quote_identifiers_binary_ops_recurse() {
4102        // a_col + select_col should quote "select_col" but that's actually
4103        // just a regular name. Use actual reserved word as column name.
4104        let expr = Expression::Add(Box::new(crate::expressions::BinaryOp::new(
4105            Expression::boxed_column(Column {
4106                name: Identifier::new("select"),
4107                table: None,
4108                join_mark: false,
4109                trailing_comments: vec![],
4110                span: None,
4111                inferred_type: None,
4112            }),
4113            Expression::boxed_column(Column {
4114                name: Identifier::new("normal"),
4115                table: None,
4116                join_mark: false,
4117                trailing_comments: vec![],
4118                span: None,
4119                inferred_type: None,
4120            }),
4121        )));
4122        let result = quote_identifiers(expr, None);
4123        if let Expression::Add(bin) = &result {
4124            if let Expression::Column(left) = &bin.left {
4125                assert!(
4126                    left.name.quoted,
4127                    "'select' column should be quoted in binary op"
4128                );
4129            }
4130            if let Expression::Column(right) = &bin.right {
4131                assert!(!right.name.quoted, "'normal' column should not be quoted");
4132            }
4133        } else {
4134            panic!("Expected Add expression");
4135        }
4136    }
4137
4138    #[test]
4139    fn test_quote_identifiers_already_quoted_preserved() {
4140        // Already-quoted identifier should stay quoted even if it doesn't need it
4141        let expr = Expression::boxed_column(Column {
4142            name: Identifier::quoted("normal_name"),
4143            table: None,
4144            join_mark: false,
4145            trailing_comments: vec![],
4146            span: None,
4147            inferred_type: None,
4148        });
4149        let result = quote_identifiers(expr, None);
4150        if let Expression::Column(col) = &result {
4151            assert!(
4152                col.name.quoted,
4153                "Already-quoted identifier should remain quoted"
4154            );
4155        } else {
4156            panic!("Expected Column expression");
4157        }
4158    }
4159
4160    #[test]
4161    fn test_quote_identifiers_full_parsed_query() {
4162        // Test with a parsed query that uses reserved words as identifiers
4163        // We build the AST manually since the parser would fail on unquoted reserved words
4164        let mut select = crate::expressions::Select::new();
4165        select.expressions.push(Expression::boxed_column(Column {
4166            name: Identifier::new("order"),
4167            table: Some(Identifier::new("t")),
4168            join_mark: false,
4169            trailing_comments: vec![],
4170            span: None,
4171            inferred_type: None,
4172        }));
4173        select.from = Some(crate::expressions::From {
4174            expressions: vec![Expression::Table(Box::new(TableRef::new("t")))],
4175        });
4176        let expr = Expression::Select(Box::new(select));
4177
4178        let result = quote_identifiers(expr, None);
4179        if let Expression::Select(sel) = &result {
4180            if let Expression::Column(col) = &sel.expressions[0] {
4181                assert!(col.name.quoted, "Column named 'order' should be quoted");
4182                assert!(
4183                    !col.table.as_ref().unwrap().quoted,
4184                    "Table 't' should not be quoted"
4185                );
4186            } else {
4187                panic!("Expected Column in SELECT list");
4188            }
4189        } else {
4190            panic!("Expected Select expression");
4191        }
4192    }
4193
4194    #[test]
4195    fn test_get_reserved_words_all_dialects() {
4196        // Ensure get_reserved_words doesn't panic for any dialect
4197        let dialects = [
4198            None,
4199            Some(DialectType::Generic),
4200            Some(DialectType::MySQL),
4201            Some(DialectType::PostgreSQL),
4202            Some(DialectType::BigQuery),
4203            Some(DialectType::Snowflake),
4204            Some(DialectType::TSQL),
4205            Some(DialectType::ClickHouse),
4206            Some(DialectType::DuckDB),
4207            Some(DialectType::Hive),
4208            Some(DialectType::Spark),
4209            Some(DialectType::Trino),
4210            Some(DialectType::Oracle),
4211            Some(DialectType::Redshift),
4212        ];
4213        for dialect in &dialects {
4214            let words = get_reserved_words(*dialect);
4215            // All dialects should have basic SQL reserved words
4216            assert!(
4217                words.contains("SELECT"),
4218                "All dialects should have SELECT as reserved"
4219            );
4220            assert!(
4221                words.contains("FROM"),
4222                "All dialects should have FROM as reserved"
4223            );
4224        }
4225    }
4226}