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