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