polyglot_sql/

ast_transforms.rs

//! AST transform helpers and convenience getters.
//!
//! This module provides functions for common AST mutations (adding WHERE clauses,
//! setting LIMIT/OFFSET, renaming columns/tables) and read-only extraction helpers
//! (getting column names, table names, functions, etc.).
//!
//! Mutation functions take an owned [`Expression`] and return a new [`Expression`].
//! Read-only getters take `&Expression`.

use std::collections::{HashMap, HashSet};

use crate::ast_mutation as engine;
use crate::expressions::*;
use crate::traversal::{is_aggregate, ExpressionWalk};

/// Allocate case-insensitively unique names for synthesized AST scopes.
#[cfg(feature = "generate")]
#[derive(Default)]
pub(crate) struct AstNames(HashSet<String>);

#[cfg(feature = "generate")]
impl AstNames {
    pub(crate) fn collect(&mut self, expression: &Expression) {
        for node in expression.dfs() {
            match node {
                Expression::Column(c) => {
                    self.reserve(&c.name);
                    if let Some(table) = &c.table {
                        self.reserve(table);
                    }
                }
                Expression::Identifier(id) => self.reserve(id),
                Expression::Alias(a) => {
                    self.reserve(&a.alias);
                    for id in &a.column_aliases {
                        self.reserve(id);
                    }
                }
                Expression::Table(t) => {
                    self.reserve(&t.name);
                    if let Some(alias) = &t.alias {
                        self.reserve(alias);
                    }
                }
                Expression::Subquery(s) => {
                    if let Some(alias) = &s.alias {
                        self.reserve(alias);
                    }
                    for id in &s.column_aliases {
                        self.reserve(id);
                    }
                }
                _ => {}
            }
        }
    }

    fn reserve(&mut self, id: &Identifier) {
        self.0.insert(id.name.to_ascii_lowercase());
    }

    pub(crate) fn fresh(&mut self, stem: &str) -> String {
        let mut name = stem.to_string();
        let mut suffix = 0;
        while !self.0.insert(name.to_ascii_lowercase()) {
            suffix += 1;
            name = format!("{stem}_{suffix}");
        }
        name
    }
}

/// Apply a complete bottom-up rewrite, retaining the embedded wrapper callbacks
/// historically exposed by these helpers. Physical children are visited only once.
fn xform<F: Fn(Expression) -> Expression>(expr: Expression, fun: F) -> Expression {
    crate::traversal::transform_all(expr, &|node| {
        transform_embedded_wrappers(node, &fun).map(&fun)
    })
    .unwrap_or_else(|_| Expression::Null(Null))
}

// Embedded TableRef/Join/Ordered/etc. fields are not Expression nodes in the
// generated visitor. Preserve their helper callbacks after their expression
// children have been rewritten, without traversing those children again.
fn transform_embedded_wrappers<F: Fn(Expression) -> Expression>(
    expr: Expression,
    fun: &F,
) -> crate::Result<Expression> {
    macro_rules! boxed_wrapper {
        ($value:expr, $variant:ident) => {
            match fun(Expression::$variant(Box::new($value))) {
                Expression::$variant(value) => *value,
                _ => {
                    return Err(crate::Error::Internal(
                        concat!(
                            stringify!($variant),
                            " transformation returned a different node type"
                        )
                        .to_string(),
                    ))
                }
            }
        };
    }
    let table = |table: TableRef| -> crate::Result<TableRef> {
        // Time-travel wrappers also belong to embedded DML targets.
        let node = transform_embedded_wrappers(Expression::Table(Box::new(table)), fun)?;
        match fun(node) {
            Expression::Table(table) => Ok(*table),
            _ => Err(crate::Error::Internal(
                "Table transformation returned a different node type".into(),
            )),
        }
    };
    let joins = |joins: Vec<Join>| -> crate::Result<Vec<Join>> {
        joins
            .into_iter()
            .map(|join| Ok(boxed_wrapper!(join, Join)))
            .collect()
    };
    let ordered = |items: Vec<Ordered>| -> Vec<Ordered> {
        items
            .into_iter()
            .map(
                |item| match fun(Expression::Ordered(Box::new(item.clone()))) {
                    Expression::Ordered(item) => *item,
                    _ => item,
                },
            )
            .collect()
    };
    let order_by = |mut order: OrderBy| {
        order.expressions = ordered(order.expressions);
        order
    };
    Ok(match expr {
        Expression::Select(mut select) => {
            select.joins = joins(select.joins)?;
            select.order_by = select.order_by.map(order_by);
            if let Some(windows) = &mut select.windows {
                for window in windows {
                    window.spec.order_by = ordered(std::mem::take(&mut window.spec.order_by));
                }
            }
            Expression::Select(select)
        }
        Expression::Update(mut update) => {
            update.table = table(update.table)?;
            update.extra_tables = update
                .extra_tables
                .into_iter()
                .map(table)
                .collect::<crate::Result<_>>()?;
            update.table_joins = joins(update.table_joins)?;
            if let Some(from) = update.from_clause.take() {
                update.from_clause = Some(boxed_wrapper!(from, From));
            }
            update.from_joins = joins(update.from_joins)?;
            update.order_by = update.order_by.map(order_by);
            Expression::Update(update)
        }
        Expression::Delete(mut delete) => {
            delete.table = table(delete.table)?;
            delete.using = delete
                .using
                .into_iter()
                .map(table)
                .collect::<crate::Result<_>>()?;
            delete.tables = delete
                .tables
                .into_iter()
                .map(table)
                .collect::<crate::Result<_>>()?;
            delete.joins = joins(delete.joins)?;
            delete.order_by = delete.order_by.map(order_by);
            Expression::Delete(delete)
        }
        Expression::Upsert(mut upsert) => {
            upsert.table = table(upsert.table)?;
            Expression::Upsert(upsert)
        }
        Expression::JoinedTable(mut joined) => {
            joined.joins = joins(joined.joins)?;
            Expression::JoinedTable(joined)
        }
        Expression::WindowFunction(mut window) => {
            window.over.order_by = ordered(window.over.order_by);
            Expression::WindowFunction(window)
        }
        Expression::WithinGroup(mut within) => {
            within.order_by = ordered(within.order_by);
            Expression::WithinGroup(within)
        }
        Expression::StringAgg(mut agg) => {
            agg.order_by = agg.order_by.map(ordered);
            Expression::StringAgg(agg)
        }
        Expression::Union(mut union) => {
            union.order_by = union.order_by.take().map(order_by);
            Expression::Union(union)
        }
        Expression::Intersect(mut intersect) => {
            intersect.order_by = intersect.order_by.take().map(order_by);
            Expression::Intersect(intersect)
        }
        Expression::Except(mut except) => {
            except.order_by = except.order_by.take().map(order_by);
            Expression::Except(except)
        }
        Expression::Cast(mut cast) => {
            cast.to =
                crate::dialects::transform_data_type_recursive(cast.to, &|node| Ok(fun(node)))?;
            Expression::Cast(cast)
        }
        Expression::TryCast(mut cast) => {
            cast.to =
                crate::dialects::transform_data_type_recursive(cast.to, &|node| Ok(fun(node)))?;
            Expression::TryCast(cast)
        }
        Expression::SafeCast(mut cast) => {
            cast.to =
                crate::dialects::transform_data_type_recursive(cast.to, &|node| Ok(fun(node)))?;
            Expression::SafeCast(cast)
        }
        Expression::Table(mut table) => {
            if let Some(when) = table.when.take() {
                if let Expression::HistoricalData(when) = fun(Expression::HistoricalData(when)) {
                    table.when = Some(when);
                }
            }
            if let Some(changes) = table.changes.take() {
                if let Expression::Changes(changes) = fun(Expression::Changes(changes)) {
                    table.changes = Some(changes);
                }
            }
            Expression::Table(table)
        }
        other => other,
    })
}

// ---------------------------------------------------------------------------
// SELECT clause
// ---------------------------------------------------------------------------

/// Append columns to the SELECT list of a query.
///
/// If `expr` is a `Select`, the given `columns` are appended to its expression list.
/// Non-SELECT expressions are returned unchanged.
pub fn add_select_columns(expr: Expression, columns: Vec<Expression>) -> Expression {
    let mut expression = expr;
    let _ = engine::append_select(&mut expression, columns, true);
    expression
}

/// Remove columns from the SELECT list where `predicate` returns `true`.
pub fn remove_select_columns<F: Fn(&Expression) -> bool>(
    expr: Expression,
    predicate: F,
) -> Expression {
    if let Expression::Select(mut sel) = expr {
        sel.expressions.retain(|e| !predicate(e));
        Expression::Select(sel)
    } else {
        expr
    }
}

/// Set or remove the DISTINCT flag on a SELECT.
pub fn set_distinct(expr: Expression, distinct: bool) -> Expression {
    let mut expression = expr;
    let _ = engine::apply_distinct(&mut expression, distinct);
    expression
}

// ---------------------------------------------------------------------------
// WHERE clause
// ---------------------------------------------------------------------------

/// Add a condition to the WHERE clause.
///
/// If the SELECT already has a WHERE clause, the new condition is combined with the
/// existing one using AND (default) or OR (when `use_or` is `true`).
/// If there is no WHERE clause, one is created.
pub fn add_where(expr: Expression, condition: Expression, use_or: bool) -> Expression {
    if !use_or {
        if !matches!(expr, Expression::Select(_)) {
            return expr;
        }
        let mut expression = expr;
        let _ = engine::apply_where(&mut expression, condition, true);
        return expression;
    }
    if let Expression::Select(mut sel) = expr {
        sel.where_clause = Some(match sel.where_clause.take() {
            Some(existing) => {
                let combined = if use_or {
                    Expression::Or(Box::new(BinaryOp::new(existing.this, condition)))
                } else {
                    Expression::And(Box::new(BinaryOp::new(existing.this, condition)))
                };
                Where { this: combined }
            }
            None => Where { this: condition },
        });
        Expression::Select(sel)
    } else {
        expr
    }
}

/// Remove the WHERE clause from a SELECT.
pub fn remove_where(expr: Expression) -> Expression {
    if let Expression::Select(mut sel) = expr {
        sel.where_clause = None;
        Expression::Select(sel)
    } else {
        expr
    }
}

// ---------------------------------------------------------------------------
// LIMIT / OFFSET / ORDER BY
// ---------------------------------------------------------------------------

/// Set the LIMIT on a SELECT or set operation.
pub fn set_limit(expr: Expression, limit: usize) -> Expression {
    set_limit_expr(expr, Expression::number(limit as i64))
}

/// Set the LIMIT on a SELECT or set operation using an expression.
pub fn set_limit_expr(expr: Expression, limit: Expression) -> Expression {
    let mut expression = expr;
    let _ = engine::apply_limit(&mut expression, limit);
    expression
}

/// Set the OFFSET on a SELECT or set operation.
pub fn set_offset(expr: Expression, offset: usize) -> Expression {
    set_offset_expr(expr, Expression::number(offset as i64))
}

/// Set the OFFSET on a SELECT or set operation using an expression.
pub fn set_offset_expr(expr: Expression, offset: Expression) -> Expression {
    let mut expression = expr;
    let _ = engine::apply_offset(&mut expression, offset);
    expression
}

/// Set the ORDER BY clause on a SELECT or set operation.
///
/// Bare expressions are normalized to ascending order expressions. Existing
/// `Ordered` expressions preserve their direction and null-ordering metadata.
pub fn set_order_by(expr: Expression, expressions: Vec<Expression>) -> Expression {
    let mut expression = expr;
    let values = expressions.into_iter().map(engine::ordered).collect();
    let _ = engine::apply_order_by(&mut expression, values, false);
    expression
}

/// Remove both LIMIT and OFFSET from a SELECT.
pub fn remove_limit_offset(expr: Expression) -> Expression {
    if let Expression::Select(mut sel) = expr {
        sel.limit = None;
        sel.offset = None;
        Expression::Select(sel)
    } else {
        expr
    }
}

// ---------------------------------------------------------------------------
// Renaming
// ---------------------------------------------------------------------------

/// Rename columns throughout the expression tree using the provided mapping.
///
/// Column names present as keys in `mapping` are replaced with their corresponding
/// values. The replacement is case-sensitive.
pub fn rename_columns(expr: Expression, mapping: &HashMap<String, String>) -> Expression {
    xform(expr, |node| match node {
        Expression::Column(mut col) => {
            if let Some(new_name) = mapping.get(&col.name.name) {
                col.name.name = new_name.clone();
            }
            Expression::Column(col)
        }
        other => other,
    })
}

/// Options for table renaming.
#[derive(Debug, Clone)]
pub struct RenameTablesOptions {
    /// Whether renamed table references should receive aliases.
    pub alias_renamed_tables: bool,
    /// Whether existing aliases should be preserved when aliasing renamed tables.
    pub preserve_existing_aliases: bool,
}

impl Default for RenameTablesOptions {
    fn default() -> Self {
        Self {
            alias_renamed_tables: false,
            preserve_existing_aliases: true,
        }
    }
}

impl RenameTablesOptions {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_alias_renamed_tables(mut self, alias: bool) -> Self {
        self.alias_renamed_tables = alias;
        self
    }

    pub fn with_preserve_existing_aliases(mut self, preserve: bool) -> Self {
        self.preserve_existing_aliases = preserve;
        self
    }
}

/// Rename tables throughout the expression tree using the provided mapping.
pub fn rename_tables(expr: Expression, mapping: &HashMap<String, String>) -> Expression {
    rename_tables_with_options(expr, mapping, &RenameTablesOptions::default())
}

/// Rename tables throughout the expression tree using the provided mapping and options.
pub fn rename_tables_with_options(
    expr: Expression,
    mapping: &HashMap<String, String>,
    options: &RenameTablesOptions,
) -> Expression {
    xform(expr, |node| match node {
        Expression::Table(mut tbl) => {
            if let Some(new_name) = mapping.get(&tbl.name.name) {
                tbl.name.name = new_name.clone();
                if options.alias_renamed_tables
                    && (!options.preserve_existing_aliases || tbl.alias.is_none())
                {
                    tbl.alias = Some(Identifier::new(new_name));
                    tbl.alias_explicit_as = true;
                }
            }
            Expression::Table(tbl)
        }
        Expression::Column(mut col) => {
            if let Some(ref mut table_id) = col.table {
                if let Some(new_name) = mapping.get(&table_id.name) {
                    table_id.name = new_name.clone();
                }
            }
            Expression::Column(col)
        }
        other => other,
    })
}

/// Qualify all unqualified column references with the given `table_name`.
///
/// Columns that already have a table qualifier are left unchanged.
pub fn qualify_columns(expr: Expression, table_name: &str) -> Expression {
    let table = table_name.to_string();
    xform(expr, move |node| match node {
        Expression::Column(mut col) => {
            if col.table.is_none() {
                col.table = Some(Identifier::new(&table));
            }
            Expression::Column(col)
        }
        other => other,
    })
}

// ---------------------------------------------------------------------------
// Generic replacement
// ---------------------------------------------------------------------------

/// Replace nodes matching `predicate` with `replacement` (cloned for each match).
pub fn replace_nodes<F: Fn(&Expression) -> bool>(
    expr: Expression,
    predicate: F,
    replacement: Expression,
) -> Expression {
    xform(expr, |node| {
        if predicate(&node) {
            replacement.clone()
        } else {
            node
        }
    })
}

/// Replace nodes matching `predicate` by applying `replacer` to the matched node.
pub fn replace_by_type<F, R>(expr: Expression, predicate: F, replacer: R) -> Expression
where
    F: Fn(&Expression) -> bool,
    R: Fn(Expression) -> Expression,
{
    xform(expr, |node| {
        if predicate(&node) {
            replacer(node)
        } else {
            node
        }
    })
}

/// Remove (replace with a `Null`) all nodes matching `predicate`.
///
/// This is most useful for removing clauses or sub-expressions from a tree.
/// Note that removing structural elements (e.g. the FROM clause) may produce
/// invalid SQL; use with care.
pub fn remove_nodes<F: Fn(&Expression) -> bool>(expr: Expression, predicate: F) -> Expression {
    xform(expr, |node| {
        if predicate(&node) {
            Expression::Null(Null)
        } else {
            node
        }
    })
}

// ---------------------------------------------------------------------------
// Convenience getters
// ---------------------------------------------------------------------------

/// Return a single named projection while preserving identifier quoting.
#[cfg(feature = "generate")]
pub(crate) fn output_identifier(expression: &Expression) -> Option<&Identifier> {
    match expression {
        Expression::Alias(a) if a.column_aliases.is_empty() => Some(&a.alias),
        Expression::Column(c) => Some(&c.name),
        _ => None,
    }
}

/// Borrow the leftmost projection list through query wrappers and set operations.
#[cfg(feature = "generate")]
pub(crate) fn query_projections(mut query: &Expression) -> Option<&[Expression]> {
    loop {
        query = match query {
            Expression::Select(s) => return Some(&s.expressions),
            Expression::Union(s) => &s.left,
            Expression::Intersect(s) => &s.left,
            Expression::Except(s) => &s.left,
            Expression::Subquery(s) => &s.this,
            Expression::Paren(p) => &p.this,
            _ => return None,
        };
    }
}

/// Collect all column names (as `String`) referenced in the expression tree.
pub fn get_column_names(expr: &Expression) -> Vec<String> {
    expr.find_all(|e| matches!(e, Expression::Column(_)))
        .into_iter()
        .filter_map(|e| {
            if let Expression::Column(col) = e {
                Some(col.name.name.clone())
            } else {
                None
            }
        })
        .collect()
}

/// Collect projected output column names from a query expression.
///
/// This follows sqlglot-style query semantics:
/// - For `SELECT`, returns names from the projection list.
/// - For set operations (`UNION`/`INTERSECT`/`EXCEPT`), uses the left-most branch.
/// - For `Subquery`, unwraps and evaluates the inner query.
///
/// Unlike [`get_column_names`], this does not return every referenced column in
/// the AST and is suitable for result-schema style output names.
pub fn get_output_column_names(expr: &Expression) -> Vec<String> {
    output_column_names_from_query(expr)
}

/// Collect projected output column names using dialect-specific set-operation
/// alignment rules.
///
/// This differs from [`get_output_column_names`] only for name-aligned set
/// operations. When an output shape is not statically knowable (for example an
/// unresolved wildcard), it preserves the existing leftmost-branch behavior.
pub fn get_output_column_names_for_dialect(
    expr: &Expression,
    dialect: Option<crate::dialects::DialectType>,
) -> Vec<String> {
    crate::set_operation::query_output_identifiers(expr, dialect)
        .map(|identifiers| {
            identifiers
                .into_iter()
                .map(|identifier| identifier.name)
                .collect()
        })
        .unwrap_or_else(|_| output_column_names_from_query(expr))
}

fn output_column_names_from_query(expr: &Expression) -> Vec<String> {
    match expr {
        Expression::Select(select) => select_output_column_names(select),
        Expression::Union(union) => output_column_names_from_query(&union.left),
        Expression::Intersect(intersect) => output_column_names_from_query(&intersect.left),
        Expression::Except(except) => output_column_names_from_query(&except.left),
        Expression::Subquery(subquery) => output_column_names_from_query(&subquery.this),
        _ => Vec::new(),
    }
}

fn select_output_column_names(select: &Select) -> Vec<String> {
    let mut names = Vec::new();
    for expr in &select.expressions {
        if let Some(name) = expression_output_name(expr) {
            names.push(name);
        }
    }
    names
}

fn expression_output_name(expr: &Expression) -> Option<String> {
    match expr {
        Expression::Alias(alias) => Some(alias.alias.name.clone()),
        Expression::Column(col) => Some(col.name.name.clone()),
        Expression::Star(_) => Some("*".to_string()),
        Expression::Identifier(id) => Some(id.name.clone()),
        Expression::Aliases(aliases) => aliases.expressions.iter().find_map(|e| match e {
            Expression::Identifier(id) => Some(id.name.clone()),
            _ => None,
        }),
        _ => None,
    }
}

/// Collect all table names (as `String`) referenced in the expression tree.
pub fn get_table_names(expr: &Expression) -> Vec<String> {
    fn collect_cte_aliases(with_clause: &With, aliases: &mut HashSet<String>) {
        for cte in &with_clause.ctes {
            aliases.insert(cte.alias.name.clone());
        }
    }

    fn push_table_ref_name(
        table: &TableRef,
        cte_aliases: &HashSet<String>,
        names: &mut Vec<String>,
    ) {
        let name = table.name.name.clone();
        if !name.is_empty() && !cte_aliases.contains(&name) {
            names.push(name);
        }
    }

    let mut cte_aliases: HashSet<String> = HashSet::new();
    for node in expr.dfs() {
        match node {
            Expression::Select(select) => {
                if let Some(with) = &select.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Insert(insert) => {
                if let Some(with) = &insert.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Update(update) => {
                if let Some(with) = &update.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Delete(delete) => {
                if let Some(with) = &delete.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Union(union) => {
                if let Some(with) = &union.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Intersect(intersect) => {
                if let Some(with) = &intersect.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Except(except) => {
                if let Some(with) = &except.with {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::CreateTable(create) => {
                if let Some(with) = &create.with_cte {
                    collect_cte_aliases(with, &mut cte_aliases);
                }
            }
            Expression::Merge(merge) => {
                if let Some(with_) = &merge.with_ {
                    if let Expression::With(with_clause) = with_.as_ref() {
                        collect_cte_aliases(with_clause, &mut cte_aliases);
                    }
                }
            }
            _ => {}
        }
    }

    let mut names = Vec::new();
    for node in expr.dfs() {
        match node {
            Expression::Table(tbl) => {
                let name = tbl.name.name.clone();
                if !name.is_empty() && !cte_aliases.contains(&name) {
                    names.push(name);
                }
            }
            Expression::Upsert(upsert) => {
                push_table_ref_name(&upsert.table, &cte_aliases, &mut names);
            }
            Expression::Insert(insert) => {
                push_table_ref_name(&insert.table, &cte_aliases, &mut names);
            }
            Expression::Update(update) => {
                push_table_ref_name(&update.table, &cte_aliases, &mut names);
                for table in &update.extra_tables {
                    push_table_ref_name(table, &cte_aliases, &mut names);
                }
            }
            Expression::Delete(delete) => {
                push_table_ref_name(&delete.table, &cte_aliases, &mut names);
                for table in &delete.using {
                    push_table_ref_name(table, &cte_aliases, &mut names);
                }
                for table in &delete.tables {
                    push_table_ref_name(table, &cte_aliases, &mut names);
                }
            }
            Expression::CreateTable(create) => {
                push_table_ref_name(&create.name, &cte_aliases, &mut names);
                if let Some(as_select) = &create.as_select {
                    names.extend(get_table_names(as_select));
                }
                if let Some(with) = &create.with_cte {
                    for cte in &with.ctes {
                        names.extend(get_table_names(&cte.this));
                    }
                }
            }
            Expression::Cache(cache) => {
                let name = cache.table.name.clone();
                if !name.is_empty() && !cte_aliases.contains(&name) {
                    names.push(name);
                }
            }
            Expression::Uncache(uncache) => {
                let name = uncache.table.name.clone();
                if !name.is_empty() && !cte_aliases.contains(&name) {
                    names.push(name);
                }
            }
            Expression::CreateSynonym(synonym) => {
                push_table_ref_name(&synonym.name, &cte_aliases, &mut names);
                push_table_ref_name(&synonym.target, &cte_aliases, &mut names);
            }
            _ => {}
        }
    }

    names
}

/// Collect all identifier references in the expression tree.
pub fn get_identifiers(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(|e| matches!(e, Expression::Identifier(_)))
}

/// Collect all function call nodes in the expression tree.
pub fn get_functions(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(crate::traversal::is_function)
}

/// Collect all literal value nodes in the expression tree.
pub fn get_literals(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(|e| {
        matches!(
            e,
            Expression::Literal(_) | Expression::Boolean(_) | Expression::Null(_)
        )
    })
}

/// Collect all subquery nodes in the expression tree.
pub fn get_subqueries(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(|e| matches!(e, Expression::Subquery(_)))
}

/// Collect all aggregate function nodes in the expression tree.
///
/// Includes typed aggregates (`Count`, `Sum`, `Avg`, `Min`, `Max`, etc.)
/// and generic `AggregateFunction` nodes.
pub fn get_aggregate_functions(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(is_aggregate)
}

/// Collect all window function nodes in the expression tree.
pub fn get_window_functions(expr: &Expression) -> Vec<&Expression> {
    expr.find_all(|e| matches!(e, Expression::WindowFunction(_)))
}

/// Count the total number of AST nodes in the expression tree.
pub fn node_count(expr: &Expression) -> usize {
    expr.dfs().count()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::Parser;

    fn parse_one(sql: &str) -> Expression {
        let mut exprs = Parser::parse_sql(sql).unwrap();
        exprs.remove(0)
    }

    #[test]
    fn test_add_where() {
        let expr = parse_one("SELECT a FROM t");
        let cond = Expression::Eq(Box::new(BinaryOp::new(
            Expression::column("b"),
            Expression::number(1),
        )));
        let result = add_where(expr, cond, false);
        let sql = result.sql();
        assert!(sql.contains("WHERE"), "Expected WHERE in: {}", sql);
        assert!(sql.contains("b = 1"), "Expected condition in: {}", sql);
    }

    #[test]
    fn test_add_where_combines_with_and() {
        let expr = parse_one("SELECT a FROM t WHERE x = 1");
        let cond = Expression::Eq(Box::new(BinaryOp::new(
            Expression::column("y"),
            Expression::number(2),
        )));
        let result = add_where(expr, cond, false);
        let sql = result.sql();
        assert!(sql.contains("AND"), "Expected AND in: {}", sql);
    }

    #[test]
    fn test_remove_where() {
        let expr = parse_one("SELECT a FROM t WHERE x = 1");
        let result = remove_where(expr);
        let sql = result.sql();
        assert!(!sql.contains("WHERE"), "Should not contain WHERE: {}", sql);
    }

    #[test]
    fn test_set_limit() {
        let expr = parse_one("SELECT a FROM t");
        let result = set_limit(expr, 10);
        let sql = result.sql();
        assert!(sql.contains("LIMIT 10"), "Expected LIMIT in: {}", sql);
    }

    #[test]
    fn test_set_limit_on_set_operation() {
        let expr = parse_one("SELECT a FROM t UNION ALL SELECT a FROM u");
        let result = set_limit(expr, 10);
        let sql = result.sql();
        assert_eq!(sql, "SELECT a FROM t UNION ALL SELECT a FROM u LIMIT 10");
    }

    #[test]
    fn test_set_offset() {
        let expr = parse_one("SELECT a FROM t");
        let result = set_offset(expr, 5);
        let sql = result.sql();
        assert!(sql.contains("OFFSET 5"), "Expected OFFSET in: {}", sql);
    }

    #[test]
    fn test_set_offset_on_set_operation() {
        let expr = parse_one("SELECT a FROM t UNION ALL SELECT a FROM u");
        let result = set_offset(expr, 5);
        let sql = result.sql();
        assert_eq!(sql, "SELECT a FROM t UNION ALL SELECT a FROM u OFFSET 5");
    }

    #[test]
    fn test_set_order_by_on_set_operation() {
        let expr = parse_one("SELECT a FROM t UNION ALL SELECT a FROM u");
        let result = set_order_by(expr, vec![Expression::column("a")]);
        let sql = result.sql();
        assert_eq!(sql, "SELECT a FROM t UNION ALL SELECT a FROM u ORDER BY a");
    }

    #[test]
    fn test_remove_limit_offset() {
        let expr = parse_one("SELECT a FROM t LIMIT 10 OFFSET 5");
        let result = remove_limit_offset(expr);
        let sql = result.sql();
        assert!(!sql.contains("LIMIT"), "Should not contain LIMIT: {}", sql);
        assert!(
            !sql.contains("OFFSET"),
            "Should not contain OFFSET: {}",
            sql
        );
    }

    #[test]
    fn test_get_column_names() {
        let expr = parse_one("SELECT a, b, c FROM t");
        let names = get_column_names(&expr);
        assert!(names.contains(&"a".to_string()));
        assert!(names.contains(&"b".to_string()));
        assert!(names.contains(&"c".to_string()));
    }

    #[test]
    fn test_get_output_column_names_select() {
        let expr = parse_one("SELECT a, b AS c, 1 FROM t");
        let names = get_output_column_names(&expr);
        assert_eq!(names, vec!["a".to_string(), "c".to_string()]);
    }

    #[test]
    fn test_get_output_column_names_union_left_projection() {
        let expr =
            parse_one("SELECT id, name FROM customers UNION ALL SELECT id, name FROM employees");
        let names = get_output_column_names(&expr);
        assert_eq!(names, vec!["id".to_string(), "name".to_string()]);
    }

    #[test]
    fn test_get_output_column_names_union_uses_left_aliases() {
        let expr = parse_one("SELECT id AS c1, name AS c2 FROM t1 UNION SELECT x, y FROM t2");
        let names = get_output_column_names(&expr);
        assert_eq!(names, vec!["c1".to_string(), "c2".to_string()]);
    }

    #[test]
    fn test_get_output_column_names_uses_dialect_by_name_layout() {
        for dialect in [
            crate::dialects::DialectType::DuckDB,
            crate::dialects::DialectType::Snowflake,
        ] {
            let expr = crate::parse_one(
                "SELECT 1 AS left_value UNION ALL BY NAME SELECT 2 AS right_value",
                dialect,
            )
            .expect("parse");
            assert_eq!(
                get_output_column_names_for_dialect(&expr, Some(dialect)),
                vec!["left_value", "right_value"]
            );
        }

        let expr = crate::parse_one(
            "SELECT 1 AS a, 2 AS b UNION ALL BY NAME SELECT 3 AS b, 4 AS a",
            crate::dialects::DialectType::BigQuery,
        )
        .expect("parse");
        assert_eq!(
            get_output_column_names_for_dialect(
                &expr,
                Some(crate::dialects::DialectType::BigQuery),
            ),
            vec!["a", "b"]
        );
    }

    #[test]
    fn test_get_column_names_union_still_returns_all_references() {
        let expr =
            parse_one("SELECT id, name FROM customers UNION ALL SELECT id, name FROM employees");
        let names = get_column_names(&expr);
        assert_eq!(
            names,
            vec![
                "id".to_string(),
                "name".to_string(),
                "id".to_string(),
                "name".to_string()
            ]
        );
    }

    #[test]
    fn test_get_table_names() {
        let expr = parse_one("SELECT a FROM users");
        let names = get_table_names(&expr);
        assert_eq!(names, vec!["users".to_string()]);
    }

    #[test]
    fn test_get_table_names_excludes_cte_aliases() {
        let expr = parse_one(
            "WITH cte AS (SELECT * FROM users) SELECT * FROM cte JOIN orders o ON cte.id = o.id",
        );
        let names = get_table_names(&expr);
        assert!(names.iter().any(|n| n == "users"));
        assert!(names.iter().any(|n| n == "orders"));
        assert!(!names.iter().any(|n| n == "cte"));
    }

    #[test]
    fn test_get_table_names_includes_dml_targets() {
        let insert_expr = parse_one("INSERT INTO users (id) VALUES (1)");
        let insert_names = get_table_names(&insert_expr);
        assert!(insert_names.iter().any(|n| n == "users"));

        let update_expr =
            parse_one("UPDATE users SET name = 'x' FROM accounts WHERE users.id = accounts.id");
        let update_names = get_table_names(&update_expr);
        assert!(update_names.iter().any(|n| n == "users"));
        assert!(update_names.iter().any(|n| n == "accounts"));

        let delete_expr =
            parse_one("DELETE FROM users USING accounts WHERE users.id = accounts.id");
        let delete_names = get_table_names(&delete_expr);
        assert!(delete_names.iter().any(|n| n == "users"));
        assert!(delete_names.iter().any(|n| n == "accounts"));

        let create_expr = parse_one("CREATE TABLE out_table AS SELECT 1 AS id FROM src");
        let create_names = get_table_names(&create_expr);
        assert!(create_names.iter().any(|n| n == "out_table"));
        assert!(create_names.iter().any(|n| n == "src"));
    }

    #[test]
    fn test_node_count() {
        let expr = parse_one("SELECT a FROM t");
        let count = node_count(&expr);
        assert!(count > 0, "Expected non-zero node count");
    }

    #[test]
    fn test_rename_columns() {
        let expr = parse_one("SELECT old_name FROM t");
        let mut mapping = HashMap::new();
        mapping.insert("old_name".to_string(), "new_name".to_string());
        let result = rename_columns(expr, &mapping);
        let sql = result.sql();
        assert!(sql.contains("new_name"), "Expected new_name in: {}", sql);
        assert!(
            !sql.contains("old_name"),
            "Should not contain old_name: {}",
            sql
        );
    }

    #[test]
    fn test_rename_tables() {
        let expr = parse_one("SELECT a FROM old_table");
        let mut mapping = HashMap::new();
        mapping.insert("old_table".to_string(), "new_table".to_string());
        let result = rename_tables(expr, &mapping);
        let sql = result.sql();
        assert!(sql.contains("new_table"), "Expected new_table in: {}", sql);
    }

    #[test]
    fn test_rename_tables_with_alias_renamed_tables() {
        let expr = parse_one("SELECT a FROM old_table");
        let mut mapping = HashMap::new();
        mapping.insert("old_table".to_string(), "new_table".to_string());
        let options = RenameTablesOptions::new().with_alias_renamed_tables(true);
        let result = rename_tables_with_options(expr, &mapping, &options);
        let sql = result.sql();

        assert_eq!(sql, "SELECT a FROM new_table AS new_table");
    }

    #[test]
    fn test_rename_tables_with_alias_preserves_existing_alias() {
        let expr = parse_one("SELECT a FROM old_table AS t");
        let mut mapping = HashMap::new();
        mapping.insert("old_table".to_string(), "new_table".to_string());
        let options = RenameTablesOptions::new().with_alias_renamed_tables(true);
        let result = rename_tables_with_options(expr, &mapping, &options);
        let sql = result.sql();

        assert_eq!(sql, "SELECT a FROM new_table AS t");
    }

    #[test]
    fn test_set_distinct() {
        let expr = parse_one("SELECT a FROM t");
        let result = set_distinct(expr, true);
        let sql = result.sql();
        assert!(sql.contains("DISTINCT"), "Expected DISTINCT in: {}", sql);
    }

    #[test]
    fn test_add_select_columns() {
        let expr = parse_one("SELECT a FROM t");
        let result = add_select_columns(expr, vec![Expression::column("b")]);
        let sql = result.sql();
        assert!(
            sql.contains("a, b") || sql.contains("a,b"),
            "Expected a, b in: {}",
            sql
        );
    }

    #[test]
    fn test_qualify_columns() {
        let expr = parse_one("SELECT a, b FROM t");
        let result = qualify_columns(expr, "t");
        let sql = result.sql();
        assert!(sql.contains("t.a"), "Expected t.a in: {}", sql);
        assert!(sql.contains("t.b"), "Expected t.b in: {}", sql);
    }

    #[test]
    fn test_get_functions() {
        let expr = parse_one("SELECT COUNT(*), UPPER(name) FROM t");
        let funcs = get_functions(&expr);
        // UPPER is a typed function (Expression::Upper), not Expression::Function
        // COUNT is Expression::Count, not Expression::AggregateFunction
        // So get_functions (which checks Function | AggregateFunction) may return 0
        // That's OK — we have separate get_aggregate_functions for typed aggs
        let _ = funcs.len();
    }

    #[test]
    fn test_get_aggregate_functions() {
        let expr = crate::parse_one(
            "SELECT COUNT_IF(numeric_value > 0), MEDIAN(numeric_value), FIRST(numeric_value) FROM source_table",
            crate::dialects::DialectType::DuckDB,
        )
        .unwrap();
        let aggs = get_aggregate_functions(&expr);
        let aggregate_types: Vec<_> = aggs.iter().map(|agg| agg.variant_name()).collect();

        assert_eq!(aggregate_types, vec!["count_if", "median", "first"]);
    }

    #[test]
    fn test_get_duckdb_null_preserving_arg_extrema() {
        let expr = crate::parse_one(
            "SELECT ARG_MAX_NULL(label, score), ARG_MIN_NULL(label, score) FROM source_table",
            crate::dialects::DialectType::DuckDB,
        )
        .unwrap();
        let aggregates = get_aggregate_functions(&expr);

        assert_eq!(aggregates.len(), 2);
        for (aggregate, expected_name) in
            aggregates.into_iter().zip(["ARG_MAX_NULL", "ARG_MIN_NULL"])
        {
            let Expression::AggregateFunction(function) = aggregate else {
                panic!("expected a generic aggregate node, got {aggregate:?}");
            };
            assert_eq!(function.name, expected_name);
            assert_eq!(function.args.len(), 2);
        }
    }
}