polyglot_sql/
traversal.rs

1//! Tree traversal utilities for SQL expression ASTs.
2//!
3//! This module provides read-only traversal, search, and transformation utilities
4//! for the [`Expression`] tree produced by the parser. Because Rust's ownership
5//! model does not allow parent pointers inside the AST, parent information is
6//! tracked externally via [`TreeContext`] (built on demand).
7//!
8//! # Traversal
9//!
10//! Two iterator types are provided:
11//! - [`DfsIter`] -- depth-first (pre-order) traversal using a stack. Visits a node
12//!   before its children. Good for top-down analysis and early termination.
13//! - [`BfsIter`] -- breadth-first (level-order) traversal using a queue. Visits all
14//!   nodes at depth N before any node at depth N+1. Good for level-aware analysis.
15//!
16//! Both are available through the [`ExpressionWalk`] trait methods [`dfs`](ExpressionWalk::dfs)
17//! and [`bfs`](ExpressionWalk::bfs).
18//!
19//! # Searching
20//!
21//! The [`ExpressionWalk`] trait also provides convenience methods for finding expressions:
22//! [`find`](ExpressionWalk::find), [`find_all`](ExpressionWalk::find_all),
23//! [`contains`](ExpressionWalk::contains), and [`count`](ExpressionWalk::count).
24//! Common predicates are available as free functions: [`is_column`], [`is_literal`],
25//! [`is_function`], [`is_aggregate`], [`is_window_function`], [`is_subquery`], and
26//! [`is_select`].
27//!
28//! # Transformation
29//!
30//! The [`transform`] and [`transform_map`] functions perform bottom-up (post-order)
31//! tree rewrites, delegating to [`transform_recursive`](crate::dialects::transform_recursive).
32//! The [`ExpressionWalk::transform_owned`] method provides the same capability as
33//! an owned method on `Expression`.
34//!
35//! Based on traversal patterns from `sqlglot/expressions.py`.
36
37#![cfg_attr(
38    not(any(feature = "ast-tools", feature = "generate", feature = "semantic")),
39    allow(dead_code)
40)]
41
42use crate::expressions::{Expression, TableRef};
43use std::collections::{HashMap, VecDeque};
44
45/// Unique identifier for expression nodes during traversal
46pub type NodeId = usize;
47
48/// Information about a node's parent relationship
49#[derive(Debug, Clone)]
50pub struct ParentInfo {
51    /// The NodeId of the parent (None for root)
52    pub parent_id: Option<NodeId>,
53    /// Which argument/field in the parent this node occupies
54    pub arg_key: String,
55    /// Index if the node is part of a list (e.g., expressions in SELECT)
56    pub index: Option<usize>,
57}
58
59/// External parent-tracking context for an expression tree.
60///
61/// Since Rust's ownership model does not allow intrusive parent pointers in the AST,
62/// `TreeContext` provides an on-demand side-table that maps each node (identified by
63/// a [`NodeId`]) to its [`ParentInfo`] (parent node, field name, and list index).
64///
65/// Build a context from any expression root with [`TreeContext::build`], then query
66/// parent relationships with [`get`](TreeContext::get), ancestry chains with
67/// [`ancestors_of`](TreeContext::ancestors_of), or tree depth with
68/// [`depth_of`](TreeContext::depth_of).
69///
70/// This is useful when analysis requires upward navigation (e.g., determining whether
71/// a column reference appears inside a WHERE clause or a JOIN condition).
72#[derive(Debug, Default)]
73pub struct TreeContext {
74    /// Map from NodeId to parent information
75    nodes: HashMap<NodeId, ParentInfo>,
76    /// Counter for generating NodeIds
77    next_id: NodeId,
78    /// Stack for tracking current path during traversal
79    path: Vec<(NodeId, String, Option<usize>)>,
80}
81
82impl TreeContext {
83    /// Create a new empty tree context
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// Build context from an expression tree
89    pub fn build(root: &Expression) -> Self {
90        let mut ctx = Self::new();
91        ctx.visit_expr(root);
92        ctx
93    }
94
95    /// Visit an expression and record parent information
96    fn visit_expr(&mut self, expr: &Expression) -> NodeId {
97        let id = self.next_id;
98        self.next_id += 1;
99
100        // Record parent info based on current path
101        let parent_info = if let Some((parent_id, arg_key, index)) = self.path.last() {
102            ParentInfo {
103                parent_id: Some(*parent_id),
104                arg_key: arg_key.clone(),
105                index: *index,
106            }
107        } else {
108            ParentInfo {
109                parent_id: None,
110                arg_key: String::new(),
111                index: None,
112            }
113        };
114        self.nodes.insert(id, parent_info);
115
116        crate::ast_children::for_each_child(expr, |child_path, child| {
117            let (key, index) = child_location(child_path);
118            self.path.push((id, key, index));
119            self.visit_expr(child);
120            self.path.pop();
121        });
122
123        id
124    }
125
126    /// Get parent info for a node
127    pub fn get(&self, id: NodeId) -> Option<&ParentInfo> {
128        self.nodes.get(&id)
129    }
130
131    /// Get the depth of a node (0 for root)
132    pub fn depth_of(&self, id: NodeId) -> usize {
133        let mut depth = 0;
134        let mut current = id;
135        while let Some(info) = self.nodes.get(&current) {
136            if let Some(parent_id) = info.parent_id {
137                depth += 1;
138                current = parent_id;
139            } else {
140                break;
141            }
142        }
143        depth
144    }
145
146    /// Get ancestors of a node (parent, grandparent, etc.)
147    pub fn ancestors_of(&self, id: NodeId) -> Vec<NodeId> {
148        let mut ancestors = Vec::new();
149        let mut current = id;
150        while let Some(info) = self.nodes.get(&current) {
151            if let Some(parent_id) = info.parent_id {
152                ancestors.push(parent_id);
153                current = parent_id;
154            } else {
155                break;
156            }
157        }
158        ancestors
159    }
160}
161
162fn child_location(path: &[crate::ast_children::ChildPathSegment]) -> (String, Option<usize>) {
163    use crate::ast_children::ChildPathSegment;
164
165    let mut key = String::new();
166    let mut index = None;
167    for (position, segment) in path.iter().enumerate() {
168        match segment {
169            ChildPathSegment::Field(field) => {
170                if !key.is_empty() {
171                    key.push('.');
172                }
173                key.push_str(field);
174            }
175            ChildPathSegment::Index(value) if position + 1 == path.len() => {
176                index = Some(*value);
177            }
178            ChildPathSegment::Index(value) => {
179                key.push('[');
180                key.push_str(&value.to_string());
181                key.push(']');
182            }
183        }
184    }
185    (key, index)
186}
187
188/// Pre-order depth-first iterator over an expression tree.
189///
190/// Visits each node before its children, using a stack-based approach. This means
191/// the root is yielded first, followed by the entire left subtree (recursively),
192/// then the right subtree. For a binary expression `a + b`, the iteration order
193/// is: `Add`, `a`, `b`.
194///
195/// Created via [`ExpressionWalk::dfs`] or [`DfsIter::new`].
196pub struct DfsIter<'a> {
197    stack: Vec<&'a Expression>,
198}
199
200impl<'a> DfsIter<'a> {
201    /// Create a new DFS iterator starting from the given expression
202    pub fn new(root: &'a Expression) -> Self {
203        Self { stack: vec![root] }
204    }
205}
206
207impl<'a> Iterator for DfsIter<'a> {
208    type Item = &'a Expression;
209
210    fn next(&mut self) -> Option<Self::Item> {
211        let expr = self.stack.pop()?;
212
213        let child_start = self.stack.len();
214        crate::ast_children::for_each_child(expr, |_, child| self.stack.push(child));
215        self.stack[child_start..].reverse();
216
217        Some(expr)
218    }
219}
220
221/// Level-order breadth-first iterator over an expression tree.
222///
223/// Visits all nodes at depth N before any node at depth N+1, using a queue-based
224/// approach. For a tree `(a + b) = c`, the iteration order is: `Eq` (depth 0),
225/// `Add`, `c` (depth 1), `a`, `b` (depth 2).
226///
227/// Created via [`ExpressionWalk::bfs`] or [`BfsIter::new`].
228pub struct BfsIter<'a> {
229    queue: VecDeque<&'a Expression>,
230}
231
232impl<'a> BfsIter<'a> {
233    /// Create a new BFS iterator starting from the given expression
234    pub fn new(root: &'a Expression) -> Self {
235        let mut queue = VecDeque::new();
236        queue.push_back(root);
237        Self { queue }
238    }
239}
240
241impl<'a> Iterator for BfsIter<'a> {
242    type Item = &'a Expression;
243
244    fn next(&mut self) -> Option<Self::Item> {
245        let expr = self.queue.pop_front()?;
246
247        crate::ast_children::for_each_child(expr, |_, child| self.queue.push_back(child));
248
249        Some(expr)
250    }
251}
252
253/// Extension trait that adds traversal and search methods to [`Expression`].
254///
255/// This trait is implemented for `Expression` and provides a fluent API for
256/// iterating, searching, measuring, and transforming expression trees without
257/// needing to import the iterator types directly.
258pub trait ExpressionWalk {
259    /// Returns a depth-first (pre-order) iterator over this expression and all descendants.
260    ///
261    /// The root node is yielded first, then its children are visited recursively
262    /// from left to right.
263    fn dfs(&self) -> DfsIter<'_>;
264
265    /// Returns a breadth-first (level-order) iterator over this expression and all descendants.
266    ///
267    /// All nodes at depth N are yielded before any node at depth N+1.
268    fn bfs(&self) -> BfsIter<'_>;
269
270    /// Finds the first expression matching `predicate` in depth-first order.
271    ///
272    /// Returns `None` if no descendant (including this node) matches.
273    fn find<F>(&self, predicate: F) -> Option<&Expression>
274    where
275        F: Fn(&Expression) -> bool;
276
277    /// Collects all expressions matching `predicate` in depth-first order.
278    ///
279    /// Returns an empty vector if no descendants match.
280    fn find_all<F>(&self, predicate: F) -> Vec<&Expression>
281    where
282        F: Fn(&Expression) -> bool;
283
284    /// Returns `true` if this node or any descendant matches `predicate`.
285    fn contains<F>(&self, predicate: F) -> bool
286    where
287        F: Fn(&Expression) -> bool;
288
289    /// Counts how many nodes (including this one) match `predicate`.
290    fn count<F>(&self, predicate: F) -> usize
291    where
292        F: Fn(&Expression) -> bool;
293
294    /// Returns direct child expressions of this node.
295    ///
296    /// Collects all single-child fields and list-child fields into a flat vector
297    /// of references. Leaf nodes return an empty vector.
298    fn children(&self) -> Vec<&Expression>;
299
300    /// Returns the maximum depth of the expression tree rooted at this node.
301    ///
302    /// A leaf node has depth 0, a node whose deepest child is a leaf has depth 1, etc.
303    fn tree_depth(&self) -> usize;
304
305    /// Transforms this expression tree bottom-up using the given function (owned variant).
306    ///
307    /// Children are transformed first, then `fun` is called on the resulting node.
308    /// Return `Ok(None)` from `fun` to replace a node with `NULL`.
309    /// Return `Ok(Some(expr))` to substitute the node with `expr`.
310    #[cfg(any(
311        feature = "transpile",
312        feature = "ast-tools",
313        feature = "generate",
314        feature = "semantic"
315    ))]
316    fn transform_owned<F>(self, fun: F) -> crate::Result<Expression>
317    where
318        F: Fn(Expression) -> crate::Result<Option<Expression>>,
319        Self: Sized;
320}
321
322impl ExpressionWalk for Expression {
323    fn dfs(&self) -> DfsIter<'_> {
324        DfsIter::new(self)
325    }
326
327    fn bfs(&self) -> BfsIter<'_> {
328        BfsIter::new(self)
329    }
330
331    fn find<F>(&self, predicate: F) -> Option<&Expression>
332    where
333        F: Fn(&Expression) -> bool,
334    {
335        self.dfs().find(|e| predicate(e))
336    }
337
338    fn find_all<F>(&self, predicate: F) -> Vec<&Expression>
339    where
340        F: Fn(&Expression) -> bool,
341    {
342        self.dfs().filter(|e| predicate(e)).collect()
343    }
344
345    fn contains<F>(&self, predicate: F) -> bool
346    where
347        F: Fn(&Expression) -> bool,
348    {
349        self.dfs().any(|e| predicate(e))
350    }
351
352    fn count<F>(&self, predicate: F) -> usize
353    where
354        F: Fn(&Expression) -> bool,
355    {
356        self.dfs().filter(|e| predicate(e)).count()
357    }
358
359    fn children(&self) -> Vec<&Expression> {
360        let mut result: Vec<&Expression> = Vec::new();
361        crate::ast_children::for_each_child(self, |_, child| result.push(child));
362        result
363    }
364
365    fn tree_depth(&self) -> usize {
366        let mut max_depth = 0usize;
367        let mut stack = vec![(self, 0usize)];
368        while let Some((node, depth)) = stack.pop() {
369            max_depth = max_depth.max(depth);
370            crate::ast_children::for_each_child(node, |_, child| {
371                stack.push((child, depth + 1));
372            });
373        }
374        max_depth
375    }
376
377    #[cfg(any(
378        feature = "transpile",
379        feature = "ast-tools",
380        feature = "generate",
381        feature = "semantic"
382    ))]
383    fn transform_owned<F>(self, fun: F) -> crate::Result<Expression>
384    where
385        F: Fn(Expression) -> crate::Result<Option<Expression>>,
386    {
387        transform(self, &fun)
388    }
389}
390
391/// Transforms an expression tree bottom-up, with optional node removal.
392///
393/// Recursively transforms all children first, then applies `fun` to the resulting node.
394/// If `fun` returns `Ok(None)`, the node is replaced with an `Expression::Null`.
395/// If `fun` returns `Ok(Some(expr))`, the node is replaced with `expr`.
396///
397/// This is the primary transformation entry point when callers need the ability to
398/// "delete" nodes by returning `None`.
399///
400/// # Example
401///
402/// ```rust,ignore
403/// use polyglot_sql::traversal::transform;
404///
405/// // Remove all Paren wrapper nodes from a tree
406/// let result = transform(expr, &|e| match e {
407///     Expression::Paren(p) => Ok(Some(p.this)),
408///     other => Ok(Some(other)),
409/// })?;
410/// ```
411#[cfg(any(
412    feature = "transpile",
413    feature = "ast-tools",
414    feature = "generate",
415    feature = "semantic"
416))]
417pub fn transform<F>(expr: Expression, fun: &F) -> crate::Result<Expression>
418where
419    F: Fn(Expression) -> crate::Result<Option<Expression>>,
420{
421    crate::dialects::transform_recursive(expr, &|e| match fun(e)? {
422        Some(transformed) => Ok(transformed),
423        None => Ok(Expression::Null(crate::expressions::Null)),
424    })
425}
426
427/// Transforms an expression tree bottom-up without node removal.
428///
429/// Like [`transform`], but `fun` returns an `Expression` directly rather than
430/// `Option<Expression>`, so nodes cannot be deleted. This is a convenience wrapper
431/// for the common case where every node is mapped to exactly one output node.
432///
433/// # Example
434///
435/// ```rust,ignore
436/// use polyglot_sql::traversal::transform_map;
437///
438/// // Uppercase all column names in a tree
439/// let result = transform_map(expr, &|e| match e {
440///     Expression::Column(mut c) => {
441///         c.name.name = c.name.name.to_uppercase();
442///         Ok(Expression::Column(c))
443///     }
444///     other => Ok(other),
445/// })?;
446/// ```
447#[cfg(any(
448    feature = "transpile",
449    feature = "ast-tools",
450    feature = "generate",
451    feature = "semantic"
452))]
453pub fn transform_map<F>(expr: Expression, fun: &F) -> crate::Result<Expression>
454where
455    F: Fn(Expression) -> crate::Result<Expression>,
456{
457    crate::dialects::transform_recursive(expr, fun)
458}
459
460// ---------------------------------------------------------------------------
461// Common expression predicates
462// ---------------------------------------------------------------------------
463// These free functions are intended for use with the search methods on
464// `ExpressionWalk` (e.g., `expr.find(is_column)`, `expr.contains(is_aggregate)`).
465
466/// Returns `true` if `expr` is a column reference ([`Expression::Column`]).
467pub fn is_column(expr: &Expression) -> bool {
468    matches!(expr, Expression::Column(_))
469}
470
471/// Returns `true` if `expr` is a literal value (number, string, boolean, or NULL).
472pub fn is_literal(expr: &Expression) -> bool {
473    matches!(
474        expr,
475        Expression::Literal(_) | Expression::Boolean(_) | Expression::Null(_)
476    )
477}
478
479/// Returns `true` if `expr` is a function call (regular or aggregate).
480pub fn is_function(expr: &Expression) -> bool {
481    matches!(
482        expr,
483        Expression::Function(_) | Expression::AggregateFunction(_)
484    )
485}
486
487/// Returns `true` if `expr` is a subquery ([`Expression::Subquery`]).
488pub fn is_subquery(expr: &Expression) -> bool {
489    matches!(expr, Expression::Subquery(_))
490}
491
492/// Returns `true` if `expr` is a SELECT statement ([`Expression::Select`]).
493pub fn is_select(expr: &Expression) -> bool {
494    matches!(expr, Expression::Select(_))
495}
496
497/// Returns `true` if `expr` is an aggregate function.
498///
499/// This is the canonical aggregate classifier used throughout the crate. Window-only
500/// functions and scalar helper functions with aggregate-like names are intentionally
501/// excluded.
502pub fn is_aggregate(expr: &Expression) -> bool {
503    matches!(
504        expr,
505        Expression::AggregateFunction(_)
506            | Expression::Count(_)
507            | Expression::Sum(_)
508            | Expression::Avg(_)
509            | Expression::Min(_)
510            | Expression::Max(_)
511            | Expression::GroupConcat(_)
512            | Expression::StringAgg(_)
513            | Expression::ListAgg(_)
514            | Expression::ArrayAgg(_)
515            | Expression::CountIf(_)
516            | Expression::SumIf(_)
517            | Expression::Stddev(_)
518            | Expression::StddevPop(_)
519            | Expression::StddevSamp(_)
520            | Expression::Variance(_)
521            | Expression::VarPop(_)
522            | Expression::VarSamp(_)
523            | Expression::Median(_)
524            | Expression::Mode(_)
525            | Expression::First(_)
526            | Expression::Last(_)
527            | Expression::AnyValue(_)
528            | Expression::ApproxDistinct(_)
529            | Expression::ApproxCountDistinct(_)
530            | Expression::ApproxPercentile(_)
531            | Expression::Percentile(_)
532            | Expression::PercentileCont(_)
533            | Expression::PercentileDisc(_)
534            | Expression::LogicalAnd(_)
535            | Expression::LogicalOr(_)
536            | Expression::Skewness(_)
537            | Expression::ArrayConcatAgg(_)
538            | Expression::ArrayUniqueAgg(_)
539            | Expression::BoolXorAgg(_)
540            | Expression::BitwiseAndAgg(_)
541            | Expression::BitwiseOrAgg(_)
542            | Expression::BitwiseXorAgg(_)
543            | Expression::JsonArrayAgg(_)
544            | Expression::JsonObjectAgg(_)
545            | Expression::JSONArrayAgg(_)
546            | Expression::JSONObjectAgg(_)
547            | Expression::JSONBObjectAgg(_)
548            | Expression::ParameterizedAgg(_)
549            | Expression::ArgMax(_)
550            | Expression::ArgMin(_)
551            | Expression::ApproxTopK(_)
552            | Expression::ApproxTopKAccumulate(_)
553            | Expression::ApproxTopKCombine(_)
554            | Expression::ApproxTopSum(_)
555            | Expression::ApproxQuantiles(_)
556            | Expression::Grouping(_)
557            | Expression::GroupingId(_)
558            | Expression::AnonymousAggFunc(_)
559            | Expression::CombinedAggFunc(_)
560            | Expression::CombinedParameterizedAgg(_)
561            | Expression::HashAgg(_)
562            | Expression::Hll(_)
563            | Expression::Minhash(_)
564            | Expression::ObjectAgg(_)
565            | Expression::AIAgg(_)
566            | Expression::Quantile(_)
567            | Expression::ApproxQuantile(_)
568            | Expression::Corr(_)
569            | Expression::CovarPop(_)
570            | Expression::CovarSamp(_)
571            | Expression::RegrValx(_)
572            | Expression::RegrValy(_)
573            | Expression::RegrAvgx(_)
574            | Expression::RegrAvgy(_)
575            | Expression::RegrCount(_)
576            | Expression::RegrIntercept(_)
577            | Expression::RegrR2(_)
578            | Expression::RegrSxx(_)
579            | Expression::RegrSxy(_)
580            | Expression::RegrSyy(_)
581            | Expression::RegrSlope(_)
582    )
583}
584
585/// Returns `true` if `expr` is a window function ([`Expression::WindowFunction`]).
586pub fn is_window_function(expr: &Expression) -> bool {
587    matches!(expr, Expression::WindowFunction(_))
588}
589
590/// Collects all column references ([`Expression::Column`]) from the expression tree.
591///
592/// Performs a depth-first search and returns references to every column node found.
593pub fn get_columns(expr: &Expression) -> Vec<&Expression> {
594    expr.find_all(is_column)
595}
596
597/// Collects all table references ([`Expression::Table`]) from the expression tree.
598///
599/// Performs a depth-first search and returns references to every table node found.
600///
601/// Note: DML target tables (`Insert.table`, `Update.table`, `Delete.table`) are
602/// stored as `TableRef` struct fields, not as `Expression::Table` nodes, so they
603/// are not reachable via tree traversal. Use [`get_all_tables`] to include those.
604pub fn get_tables(expr: &Expression) -> Vec<&Expression> {
605    expr.find_all(|e| matches!(e, Expression::Table(_)))
606}
607
608/// Collects **all** referenced tables from the expression tree, including DML
609/// target tables that are stored as `TableRef` struct fields and are therefore
610/// not reachable through normal tree traversal.
611///
612/// Returns owned `Expression::Table` values. This is the comprehensive version
613/// of [`get_tables`] — use it when you need to discover every table referenced
614/// in a statement, including inside CTE bodies containing INSERT/UPDATE/DELETE.
615pub fn get_all_tables(expr: &Expression) -> Vec<Expression> {
616    use std::collections::HashSet;
617
618    let mut seen = HashSet::new();
619    let mut result = Vec::new();
620
621    // First: collect all Expression::Table nodes found via DFS.
622    for node in expr.dfs() {
623        if let Expression::Table(t) = node {
624            let qname = table_ref_qualified_name(t);
625            if seen.insert(qname) {
626                result.push(node.clone());
627            }
628        }
629
630        // Also extract DML target TableRef fields not reachable via iter_children.
631        let refs: Vec<&TableRef> = match node {
632            Expression::Insert(ins) => vec![&ins.table],
633            Expression::Update(upd) => {
634                let mut v = vec![&upd.table];
635                v.extend(upd.extra_tables.iter());
636                v
637            }
638            Expression::Delete(del) => {
639                let mut v = vec![&del.table];
640                v.extend(del.using.iter());
641                v
642            }
643            _ => continue,
644        };
645        for tref in refs {
646            if tref.name.name.is_empty() {
647                continue;
648            }
649            let qname = table_ref_qualified_name(tref);
650            if seen.insert(qname) {
651                result.push(Expression::Table(Box::new(tref.clone())));
652            }
653        }
654    }
655
656    result
657}
658
659/// Build a qualified name string from a TableRef for deduplication purposes.
660fn table_ref_qualified_name(t: &TableRef) -> String {
661    let mut name = String::new();
662    if let Some(ref cat) = t.catalog {
663        name.push_str(&cat.name);
664        name.push('.');
665    }
666    if let Some(ref schema) = t.schema {
667        name.push_str(&schema.name);
668        name.push('.');
669    }
670    name.push_str(&t.name.name);
671    name
672}
673
674/// Extracts the underlying [`Expression::Table`] from a MERGE field that may
675/// be a bare `Table`, an `Alias` wrapping a `Table`, or an `Identifier`.
676/// Returns `None` if the expression doesn't contain a recognisable table.
677fn unwrap_merge_table(expr: &Expression) -> Option<&Expression> {
678    match expr {
679        Expression::Table(_) => Some(expr),
680        Expression::Alias(alias) => match &alias.this {
681            Expression::Table(_) => Some(&alias.this),
682            _ => None,
683        },
684        _ => None,
685    }
686}
687
688/// Returns the target table of a MERGE statement (the `Merge.this` field),
689/// unwrapping any alias wrapper to yield the underlying [`Expression::Table`].
690///
691/// Returns `None` if `expr` is not a `Merge` or the target isn't a recognisable table.
692pub fn get_merge_target(expr: &Expression) -> Option<&Expression> {
693    match expr {
694        Expression::Merge(m) => unwrap_merge_table(&m.this),
695        _ => None,
696    }
697}
698
699/// Returns the source table of a MERGE statement (the `Merge.using` field),
700/// unwrapping any alias wrapper to yield the underlying [`Expression::Table`].
701///
702/// Returns `None` if `expr` is not a `Merge`, the source isn't a recognisable
703/// table (e.g. it's a subquery), or the source is otherwise unresolvable.
704pub fn get_merge_source(expr: &Expression) -> Option<&Expression> {
705    match expr {
706        Expression::Merge(m) => unwrap_merge_table(&m.using),
707        _ => None,
708    }
709}
710
711/// Returns `true` if the expression tree contains any aggregate function calls.
712pub fn contains_aggregate(expr: &Expression) -> bool {
713    expr.contains(is_aggregate)
714}
715
716/// Returns `true` if the expression tree contains any window function calls.
717pub fn contains_window_function(expr: &Expression) -> bool {
718    expr.contains(is_window_function)
719}
720
721/// Returns `true` if the expression tree contains any subquery nodes.
722pub fn contains_subquery(expr: &Expression) -> bool {
723    expr.contains(is_subquery)
724}
725
726// ---------------------------------------------------------------------------
727// Extended type predicates
728// ---------------------------------------------------------------------------
729
730/// Macro for generating simple type-predicate functions.
731macro_rules! is_type {
732    ($name:ident, $($variant:pat),+ $(,)?) => {
733        /// Returns `true` if `expr` matches the expected AST variant(s).
734        pub fn $name(expr: &Expression) -> bool {
735            matches!(expr, $($variant)|+)
736        }
737    };
738}
739
740// Query
741is_type!(is_insert, Expression::Insert(_));
742is_type!(is_update, Expression::Update(_));
743is_type!(is_delete, Expression::Delete(_));
744is_type!(is_merge, Expression::Merge(_));
745is_type!(is_union, Expression::Union(_));
746is_type!(is_intersect, Expression::Intersect(_));
747is_type!(is_except, Expression::Except(_));
748
749// Identifiers & literals
750is_type!(is_boolean, Expression::Boolean(_));
751is_type!(is_null_literal, Expression::Null(_));
752is_type!(is_star, Expression::Star(_));
753is_type!(is_identifier, Expression::Identifier(_));
754is_type!(is_table, Expression::Table(_));
755
756// Comparison
757is_type!(is_eq, Expression::Eq(_));
758is_type!(is_neq, Expression::Neq(_));
759is_type!(is_lt, Expression::Lt(_));
760is_type!(is_lte, Expression::Lte(_));
761is_type!(is_gt, Expression::Gt(_));
762is_type!(is_gte, Expression::Gte(_));
763is_type!(is_like, Expression::Like(_));
764is_type!(is_ilike, Expression::ILike(_));
765
766// Arithmetic
767is_type!(is_add, Expression::Add(_));
768is_type!(is_sub, Expression::Sub(_));
769is_type!(is_mul, Expression::Mul(_));
770is_type!(is_div, Expression::Div(_));
771is_type!(is_mod, Expression::Mod(_));
772is_type!(is_concat, Expression::Concat(_));
773
774// Logical
775is_type!(is_and, Expression::And(_));
776is_type!(is_or, Expression::Or(_));
777is_type!(is_not, Expression::Not(_));
778
779// Predicates
780is_type!(is_in, Expression::In(_));
781is_type!(is_between, Expression::Between(_));
782is_type!(is_is_null, Expression::IsNull(_));
783is_type!(is_exists, Expression::Exists(_));
784
785// Functions
786is_type!(is_count, Expression::Count(_));
787is_type!(is_sum, Expression::Sum(_));
788is_type!(is_avg, Expression::Avg(_));
789is_type!(is_min_func, Expression::Min(_));
790is_type!(is_max_func, Expression::Max(_));
791is_type!(is_coalesce, Expression::Coalesce(_));
792is_type!(is_null_if, Expression::NullIf(_));
793is_type!(is_cast, Expression::Cast(_));
794is_type!(is_try_cast, Expression::TryCast(_));
795is_type!(is_safe_cast, Expression::SafeCast(_));
796is_type!(is_case, Expression::Case(_));
797
798// Clauses
799is_type!(is_from, Expression::From(_));
800is_type!(is_join, Expression::Join(_));
801is_type!(is_where, Expression::Where(_));
802is_type!(is_group_by, Expression::GroupBy(_));
803is_type!(is_having, Expression::Having(_));
804is_type!(is_order_by, Expression::OrderBy(_));
805is_type!(is_limit, Expression::Limit(_));
806is_type!(is_offset, Expression::Offset(_));
807is_type!(is_with, Expression::With(_));
808is_type!(is_cte, Expression::Cte(_));
809is_type!(is_alias, Expression::Alias(_));
810is_type!(is_paren, Expression::Paren(_));
811is_type!(is_ordered, Expression::Ordered(_));
812
813// DDL
814is_type!(is_create_table, Expression::CreateTable(_));
815is_type!(is_drop_table, Expression::DropTable(_));
816is_type!(is_alter_table, Expression::AlterTable(_));
817is_type!(is_create_index, Expression::CreateIndex(_));
818is_type!(is_drop_index, Expression::DropIndex(_));
819is_type!(is_create_view, Expression::CreateView(_));
820is_type!(is_drop_view, Expression::DropView(_));
821
822// ---------------------------------------------------------------------------
823// Composite predicates
824// ---------------------------------------------------------------------------
825
826/// Returns `true` if `expr` is a query statement (SELECT, INSERT, UPDATE, DELETE, or MERGE).
827pub fn is_query(expr: &Expression) -> bool {
828    matches!(
829        expr,
830        Expression::Select(_)
831            | Expression::Insert(_)
832            | Expression::Update(_)
833            | Expression::Delete(_)
834            | Expression::Merge(_)
835    )
836}
837
838/// Returns `true` if `expr` is a set operation (UNION, INTERSECT, or EXCEPT).
839pub fn is_set_operation(expr: &Expression) -> bool {
840    matches!(
841        expr,
842        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
843    )
844}
845
846/// Returns `true` if `expr` is a comparison operator.
847pub fn is_comparison(expr: &Expression) -> bool {
848    matches!(
849        expr,
850        Expression::Eq(_)
851            | Expression::Neq(_)
852            | Expression::Lt(_)
853            | Expression::Lte(_)
854            | Expression::Gt(_)
855            | Expression::Gte(_)
856            | Expression::Like(_)
857            | Expression::ILike(_)
858    )
859}
860
861/// Returns `true` if `expr` is an arithmetic operator.
862pub fn is_arithmetic(expr: &Expression) -> bool {
863    matches!(
864        expr,
865        Expression::Add(_)
866            | Expression::Sub(_)
867            | Expression::Mul(_)
868            | Expression::Div(_)
869            | Expression::Mod(_)
870    )
871}
872
873/// Returns `true` if `expr` is a logical operator (AND, OR, NOT).
874pub fn is_logical(expr: &Expression) -> bool {
875    matches!(
876        expr,
877        Expression::And(_) | Expression::Or(_) | Expression::Not(_)
878    )
879}
880
881/// Returns `true` if `expr` is a DDL statement.
882pub fn is_ddl(expr: &Expression) -> bool {
883    matches!(
884        expr,
885        Expression::CreateTable(_)
886            | Expression::DropTable(_)
887            | Expression::Undrop(_)
888            | Expression::AlterTable(_)
889            | Expression::CreateIndex(_)
890            | Expression::DropIndex(_)
891            | Expression::CreateView(_)
892            | Expression::DropView(_)
893            | Expression::AlterView(_)
894            | Expression::CreateSchema(_)
895            | Expression::DropSchema(_)
896            | Expression::CreateDatabase(_)
897            | Expression::DropDatabase(_)
898            | Expression::CreateFunction(_)
899            | Expression::DropFunction(_)
900            | Expression::CreateProcedure(_)
901            | Expression::DropProcedure(_)
902            | Expression::CreateSequence(_)
903            | Expression::CreateSynonym(_)
904            | Expression::DropSequence(_)
905            | Expression::AlterSequence(_)
906            | Expression::CreateTrigger(_)
907            | Expression::DropTrigger(_)
908            | Expression::CreateType(_)
909            | Expression::DropType(_)
910    )
911}
912
913/// Find the parent of `target` within the tree rooted at `root`.
914///
915/// Uses pointer identity ([`std::ptr::eq`]) — `target` must be a reference
916/// obtained from the same tree (e.g., via [`ExpressionWalk::find`] or DFS iteration).
917///
918/// Returns `None` if `target` is the root itself or is not found in the tree.
919pub fn find_parent<'a>(root: &'a Expression, target: &Expression) -> Option<&'a Expression> {
920    fn search<'a>(node: &'a Expression, target: *const Expression) -> Option<&'a Expression> {
921        let mut result = None;
922        crate::ast_children::for_each_child(node, |_, child| {
923            if result.is_none() {
924                if std::ptr::eq(child, target) {
925                    result = Some(node);
926                } else {
927                    result = search(child, target);
928                }
929            }
930        });
931        result
932    }
933
934    search(root, target as *const Expression)
935}
936
937/// Find the first ancestor of `target` matching `predicate`, walking from
938/// parent toward root.
939///
940/// Uses pointer identity for target lookup. Returns `None` if no ancestor
941/// matches or `target` is not found in the tree.
942pub fn find_ancestor<'a, F>(
943    root: &'a Expression,
944    target: &Expression,
945    predicate: F,
946) -> Option<&'a Expression>
947where
948    F: Fn(&Expression) -> bool,
949{
950    // Build path from root to target
951    fn build_path<'a>(
952        node: &'a Expression,
953        target: *const Expression,
954        path: &mut Vec<&'a Expression>,
955    ) -> bool {
956        if std::ptr::eq(node, target) {
957            return true;
958        }
959        path.push(node);
960        let mut found = false;
961        crate::ast_children::for_each_child(node, |_, child| {
962            if !found {
963                found = build_path(child, target, path);
964            }
965        });
966        if found {
967            return true;
968        }
969        path.pop();
970        false
971    }
972
973    let mut path = Vec::new();
974    if !build_path(root, target as *const Expression, &mut path) {
975        return None;
976    }
977
978    // Walk path in reverse (parent first, then grandparent, etc.)
979    for ancestor in path.iter().rev() {
980        if predicate(ancestor) {
981            return Some(ancestor);
982        }
983    }
984    None
985}
986
987#[cfg(test)]
988mod tests {
989    use super::*;
990    use crate::expressions::{BinaryOp, Column, Identifier, LikeOp, Literal, TableRef};
991
992    fn make_column(name: &str) -> Expression {
993        Expression::boxed_column(Column {
994            name: Identifier {
995                name: name.to_string(),
996                quoted: false,
997                trailing_comments: vec![],
998                span: None,
999            },
1000            table: None,
1001            join_mark: false,
1002            trailing_comments: vec![],
1003            span: None,
1004            inferred_type: None,
1005        })
1006    }
1007
1008    fn make_literal(value: i64) -> Expression {
1009        Expression::Literal(Box::new(Literal::Number(value.to_string())))
1010    }
1011
1012    #[test]
1013    fn test_dfs_simple() {
1014        let left = make_column("a");
1015        let right = make_literal(1);
1016        let expr = Expression::Eq(Box::new(BinaryOp {
1017            left,
1018            right,
1019            left_comments: vec![],
1020            operator_comments: vec![],
1021            trailing_comments: vec![],
1022            inferred_type: None,
1023        }));
1024
1025        let nodes: Vec<_> = expr.dfs().collect();
1026        assert_eq!(nodes.len(), 3); // Eq, Column, Literal
1027        assert!(matches!(nodes[0], Expression::Eq(_)));
1028        assert!(matches!(nodes[1], Expression::Column(_)));
1029        assert!(matches!(nodes[2], Expression::Literal(_)));
1030    }
1031
1032    #[test]
1033    fn test_find() {
1034        let left = make_column("a");
1035        let right = make_literal(1);
1036        let expr = Expression::Eq(Box::new(BinaryOp {
1037            left,
1038            right,
1039            left_comments: vec![],
1040            operator_comments: vec![],
1041            trailing_comments: vec![],
1042            inferred_type: None,
1043        }));
1044
1045        let column = expr.find(is_column);
1046        assert!(column.is_some());
1047        assert!(matches!(column.unwrap(), Expression::Column(_)));
1048
1049        let literal = expr.find(is_literal);
1050        assert!(literal.is_some());
1051        assert!(matches!(literal.unwrap(), Expression::Literal(_)));
1052    }
1053
1054    #[test]
1055    fn test_find_all() {
1056        let col1 = make_column("a");
1057        let col2 = make_column("b");
1058        let expr = Expression::And(Box::new(BinaryOp {
1059            left: col1,
1060            right: col2,
1061            left_comments: vec![],
1062            operator_comments: vec![],
1063            trailing_comments: vec![],
1064            inferred_type: None,
1065        }));
1066
1067        let columns = expr.find_all(is_column);
1068        assert_eq!(columns.len(), 2);
1069    }
1070
1071    #[test]
1072    fn test_contains() {
1073        let col = make_column("a");
1074        let lit = make_literal(1);
1075        let expr = Expression::Eq(Box::new(BinaryOp {
1076            left: col,
1077            right: lit,
1078            left_comments: vec![],
1079            operator_comments: vec![],
1080            trailing_comments: vec![],
1081            inferred_type: None,
1082        }));
1083
1084        assert!(expr.contains(is_column));
1085        assert!(expr.contains(is_literal));
1086        assert!(!expr.contains(is_subquery));
1087    }
1088
1089    #[test]
1090    fn test_count() {
1091        let col1 = make_column("a");
1092        let col2 = make_column("b");
1093        let lit = make_literal(1);
1094
1095        let inner = Expression::Add(Box::new(BinaryOp {
1096            left: col2,
1097            right: lit,
1098            left_comments: vec![],
1099            operator_comments: vec![],
1100            trailing_comments: vec![],
1101            inferred_type: None,
1102        }));
1103
1104        let expr = Expression::Eq(Box::new(BinaryOp {
1105            left: col1,
1106            right: inner,
1107            left_comments: vec![],
1108            operator_comments: vec![],
1109            trailing_comments: vec![],
1110            inferred_type: None,
1111        }));
1112
1113        assert_eq!(expr.count(is_column), 2);
1114        assert_eq!(expr.count(is_literal), 1);
1115    }
1116
1117    #[test]
1118    fn test_tree_depth() {
1119        // Single node
1120        let lit = make_literal(1);
1121        assert_eq!(lit.tree_depth(), 0);
1122
1123        // One level
1124        let col = make_column("a");
1125        let expr = Expression::Eq(Box::new(BinaryOp {
1126            left: col,
1127            right: lit.clone(),
1128            left_comments: vec![],
1129            operator_comments: vec![],
1130            trailing_comments: vec![],
1131            inferred_type: None,
1132        }));
1133        assert_eq!(expr.tree_depth(), 1);
1134
1135        // Two levels
1136        let inner = Expression::Add(Box::new(BinaryOp {
1137            left: make_column("b"),
1138            right: lit,
1139            left_comments: vec![],
1140            operator_comments: vec![],
1141            trailing_comments: vec![],
1142            inferred_type: None,
1143        }));
1144        let outer = Expression::Eq(Box::new(BinaryOp {
1145            left: make_column("a"),
1146            right: inner,
1147            left_comments: vec![],
1148            operator_comments: vec![],
1149            trailing_comments: vec![],
1150            inferred_type: None,
1151        }));
1152        assert_eq!(outer.tree_depth(), 2);
1153    }
1154
1155    #[test]
1156    fn test_tree_context() {
1157        let col = make_column("a");
1158        let lit = make_literal(1);
1159        let expr = Expression::Eq(Box::new(BinaryOp {
1160            left: col,
1161            right: lit,
1162            left_comments: vec![],
1163            operator_comments: vec![],
1164            trailing_comments: vec![],
1165            inferred_type: None,
1166        }));
1167
1168        let ctx = TreeContext::build(&expr);
1169
1170        // Root has no parent
1171        let root_info = ctx.get(0).unwrap();
1172        assert!(root_info.parent_id.is_none());
1173
1174        // Children have root as parent
1175        let left_info = ctx.get(1).unwrap();
1176        assert_eq!(left_info.parent_id, Some(0));
1177        assert_eq!(left_info.arg_key, "left");
1178
1179        let right_info = ctx.get(2).unwrap();
1180        assert_eq!(right_info.parent_id, Some(0));
1181        assert_eq!(right_info.arg_key, "right");
1182    }
1183
1184    // -- Step 8: transform / transform_map tests --
1185
1186    #[test]
1187    fn test_transform_rename_columns() {
1188        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1189        let expr = ast[0].clone();
1190        let result = super::transform_map(expr, &|e| {
1191            if let Expression::Column(ref c) = e {
1192                if c.name.name == "a" {
1193                    return Ok(Expression::boxed_column(Column {
1194                        name: Identifier::new("alpha"),
1195                        table: c.table.clone(),
1196                        join_mark: false,
1197                        trailing_comments: vec![],
1198                        span: None,
1199                        inferred_type: None,
1200                    }));
1201                }
1202            }
1203            Ok(e)
1204        })
1205        .unwrap();
1206        let sql = crate::generator::Generator::sql(&result).unwrap();
1207        assert!(sql.contains("alpha"), "Expected 'alpha' in: {}", sql);
1208        assert!(sql.contains("b"), "Expected 'b' in: {}", sql);
1209    }
1210
1211    #[test]
1212    fn test_transform_noop() {
1213        let ast = crate::parser::Parser::parse_sql("SELECT 1 + 2").unwrap();
1214        let expr = ast[0].clone();
1215        let result = super::transform_map(expr.clone(), &|e| Ok(e)).unwrap();
1216        let sql1 = crate::generator::Generator::sql(&expr).unwrap();
1217        let sql2 = crate::generator::Generator::sql(&result).unwrap();
1218        assert_eq!(sql1, sql2);
1219    }
1220
1221    #[test]
1222    fn test_transform_nested() {
1223        let ast = crate::parser::Parser::parse_sql("SELECT a + b FROM t").unwrap();
1224        let expr = ast[0].clone();
1225        let result = super::transform_map(expr, &|e| {
1226            if let Expression::Column(ref c) = e {
1227                return Ok(Expression::Literal(Box::new(Literal::Number(
1228                    if c.name.name == "a" { "1" } else { "2" }.to_string(),
1229                ))));
1230            }
1231            Ok(e)
1232        })
1233        .unwrap();
1234        let sql = crate::generator::Generator::sql(&result).unwrap();
1235        assert_eq!(sql, "SELECT 1 + 2 FROM t");
1236    }
1237
1238    #[test]
1239    fn test_transform_error() {
1240        let ast = crate::parser::Parser::parse_sql("SELECT a FROM t").unwrap();
1241        let expr = ast[0].clone();
1242        let result = super::transform_map(expr, &|e| {
1243            if let Expression::Column(ref c) = e {
1244                if c.name.name == "a" {
1245                    return Err(crate::error::Error::parse("test error", 0, 0, 0, 0));
1246                }
1247            }
1248            Ok(e)
1249        });
1250        assert!(result.is_err());
1251    }
1252
1253    #[test]
1254    fn test_transform_owned_trait() {
1255        let ast = crate::parser::Parser::parse_sql("SELECT x FROM t").unwrap();
1256        let expr = ast[0].clone();
1257        let result = expr.transform_owned(|e| Ok(Some(e))).unwrap();
1258        let sql = crate::generator::Generator::sql(&result).unwrap();
1259        assert_eq!(sql, "SELECT x FROM t");
1260    }
1261
1262    // -- children() tests --
1263
1264    #[test]
1265    fn test_children_leaf() {
1266        let lit = make_literal(1);
1267        assert_eq!(lit.children().len(), 0);
1268    }
1269
1270    #[test]
1271    fn test_children_binary_op() {
1272        let left = make_column("a");
1273        let right = make_literal(1);
1274        let expr = Expression::Eq(Box::new(BinaryOp {
1275            left,
1276            right,
1277            left_comments: vec![],
1278            operator_comments: vec![],
1279            trailing_comments: vec![],
1280            inferred_type: None,
1281        }));
1282        let children = expr.children();
1283        assert_eq!(children.len(), 2);
1284        assert!(matches!(children[0], Expression::Column(_)));
1285        assert!(matches!(children[1], Expression::Literal(_)));
1286    }
1287
1288    #[test]
1289    fn test_children_select() {
1290        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1291        let expr = &ast[0];
1292        let children = expr.children();
1293        // Should include select list items (a, b)
1294        assert!(children.len() >= 2);
1295    }
1296
1297    #[test]
1298    fn test_children_follow_ast_field_order() {
1299        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1300        let children = ast[0].children();
1301
1302        assert!(matches!(children[0], Expression::Column(column) if column.name.name == "a"));
1303        assert!(matches!(children[1], Expression::Column(column) if column.name.name == "b"));
1304        assert!(matches!(children[2], Expression::Table(table) if table.name.name == "t"));
1305    }
1306
1307    #[test]
1308    fn test_traversal_covers_previously_omitted_expression_fields() {
1309        let like = Expression::Like(Box::new(LikeOp {
1310            left: make_column("name"),
1311            right: Expression::Literal(Box::new(Literal::String("x%".to_string()))),
1312            escape: Some(Expression::Literal(Box::new(Literal::String(
1313                "!".to_string(),
1314            )))),
1315            quantifier: None,
1316            inferred_type: None,
1317        }));
1318        let nodes: Vec<_> = like.dfs().collect();
1319        assert_eq!(nodes.len(), 4);
1320        assert!(matches!(
1321            nodes[3],
1322            Expression::Literal(literal)
1323                if matches!(literal.as_ref(), Literal::String(value) if value == "!")
1324        ));
1325
1326        let mut table = TableRef::new("events");
1327        table.hints.push(make_column("table_hint"));
1328        table.identifier_func = Some(Box::new(Expression::identifier("dynamic_table")));
1329        let table = Expression::Table(Box::new(table));
1330        let children = table.children();
1331        assert_eq!(children.len(), 2);
1332        assert!(
1333            matches!(children[0], Expression::Column(column) if column.name.name == "table_hint")
1334        );
1335        assert!(
1336            matches!(children[1], Expression::Identifier(identifier) if identifier.name == "dynamic_table")
1337        );
1338    }
1339
1340    #[test]
1341    fn test_children_select_includes_from_and_join_sources() {
1342        let ast = crate::parser::Parser::parse_sql(
1343            "SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id",
1344        )
1345        .unwrap();
1346        let expr = &ast[0];
1347        let children = expr.children();
1348
1349        let table_names: Vec<&str> = children
1350            .iter()
1351            .filter_map(|e| match e {
1352                Expression::Table(t) => Some(t.name.name.as_str()),
1353                _ => None,
1354            })
1355            .collect();
1356
1357        assert!(table_names.contains(&"users"));
1358        assert!(table_names.contains(&"orders"));
1359    }
1360
1361    #[test]
1362    fn test_get_tables_includes_insert_query_sources() {
1363        let ast = crate::parser::Parser::parse_sql(
1364            "INSERT INTO dst (id) SELECT s.id FROM src s JOIN dim d ON s.id = d.id",
1365        )
1366        .unwrap();
1367        let expr = &ast[0];
1368        let tables = get_tables(expr);
1369        let names: Vec<&str> = tables
1370            .iter()
1371            .filter_map(|e| match e {
1372                Expression::Table(t) => Some(t.name.name.as_str()),
1373                _ => None,
1374            })
1375            .collect();
1376
1377        assert!(names.contains(&"src"));
1378        assert!(names.contains(&"dim"));
1379    }
1380
1381    // -- find_parent() tests --
1382
1383    #[test]
1384    fn test_find_parent_binary() {
1385        let left = make_column("a");
1386        let right = make_literal(1);
1387        let expr = Expression::Eq(Box::new(BinaryOp {
1388            left,
1389            right,
1390            left_comments: vec![],
1391            operator_comments: vec![],
1392            trailing_comments: vec![],
1393            inferred_type: None,
1394        }));
1395
1396        // Find the column child and get its parent
1397        let col = expr.find(is_column).unwrap();
1398        let parent = super::find_parent(&expr, col);
1399        assert!(parent.is_some());
1400        assert!(matches!(parent.unwrap(), Expression::Eq(_)));
1401    }
1402
1403    #[test]
1404    fn test_find_parent_root_has_none() {
1405        let lit = make_literal(1);
1406        let parent = super::find_parent(&lit, &lit);
1407        assert!(parent.is_none());
1408    }
1409
1410    // -- find_ancestor() tests --
1411
1412    #[test]
1413    fn test_find_ancestor_select() {
1414        let ast = crate::parser::Parser::parse_sql("SELECT a FROM t WHERE a > 1").unwrap();
1415        let expr = &ast[0];
1416
1417        // Find a column inside the WHERE clause
1418        let where_col = expr.dfs().find(|e| {
1419            if let Expression::Column(c) = e {
1420                c.name.name == "a"
1421            } else {
1422                false
1423            }
1424        });
1425        assert!(where_col.is_some());
1426
1427        // Find Select ancestor of that column
1428        let ancestor = super::find_ancestor(expr, where_col.unwrap(), is_select);
1429        assert!(ancestor.is_some());
1430        assert!(matches!(ancestor.unwrap(), Expression::Select(_)));
1431    }
1432
1433    #[test]
1434    fn test_find_ancestor_no_match() {
1435        let left = make_column("a");
1436        let right = make_literal(1);
1437        let expr = Expression::Eq(Box::new(BinaryOp {
1438            left,
1439            right,
1440            left_comments: vec![],
1441            operator_comments: vec![],
1442            trailing_comments: vec![],
1443            inferred_type: None,
1444        }));
1445
1446        let col = expr.find(is_column).unwrap();
1447        let ancestor = super::find_ancestor(&expr, col, is_select);
1448        assert!(ancestor.is_none());
1449    }
1450
1451    #[test]
1452    fn test_ancestors() {
1453        let col = make_column("a");
1454        let lit = make_literal(1);
1455        let inner = Expression::Add(Box::new(BinaryOp {
1456            left: col,
1457            right: lit,
1458            left_comments: vec![],
1459            operator_comments: vec![],
1460            trailing_comments: vec![],
1461            inferred_type: None,
1462        }));
1463        let outer = Expression::Eq(Box::new(BinaryOp {
1464            left: make_column("b"),
1465            right: inner,
1466            left_comments: vec![],
1467            operator_comments: vec![],
1468            trailing_comments: vec![],
1469            inferred_type: None,
1470        }));
1471
1472        let ctx = TreeContext::build(&outer);
1473
1474        // The inner Add's left child (column "a") should have ancestors
1475        // Node 0: Eq
1476        // Node 1: Column "b" (left of Eq)
1477        // Node 2: Add (right of Eq)
1478        // Node 3: Column "a" (left of Add)
1479        // Node 4: Literal (right of Add)
1480
1481        let ancestors = ctx.ancestors_of(3);
1482        assert_eq!(ancestors, vec![2, 0]); // Add, then Eq
1483    }
1484
1485    #[test]
1486    fn test_get_merge_target_and_source() {
1487        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1488
1489        // MERGE with aliased target and source tables
1490        let sql = "MERGE INTO orders o USING customers c ON o.customer_id = c.id WHEN MATCHED THEN UPDATE SET amount = amount + 100";
1491        let exprs = dialect.parse(sql).unwrap();
1492        let expr = &exprs[0];
1493
1494        assert!(is_merge(expr));
1495        assert!(is_query(expr));
1496
1497        let target = get_merge_target(expr).expect("should find target table");
1498        assert!(matches!(target, Expression::Table(_)));
1499        if let Expression::Table(t) = target {
1500            assert_eq!(t.name.name, "orders");
1501        }
1502
1503        let source = get_merge_source(expr).expect("should find source table");
1504        assert!(matches!(source, Expression::Table(_)));
1505        if let Expression::Table(t) = source {
1506            assert_eq!(t.name.name, "customers");
1507        }
1508    }
1509
1510    #[test]
1511    fn test_get_merge_source_subquery_returns_none() {
1512        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1513
1514        // MERGE with subquery source — get_merge_source should return None
1515        let sql = "MERGE INTO orders o USING (SELECT * FROM customers) c ON o.customer_id = c.id WHEN MATCHED THEN DELETE";
1516        let exprs = dialect.parse(sql).unwrap();
1517        let expr = &exprs[0];
1518
1519        assert!(get_merge_target(expr).is_some());
1520        assert!(get_merge_source(expr).is_none());
1521    }
1522
1523    #[test]
1524    fn test_get_merge_on_non_merge_returns_none() {
1525        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1526        let exprs = dialect.parse("SELECT 1").unwrap();
1527        assert!(get_merge_target(&exprs[0]).is_none());
1528        assert!(get_merge_source(&exprs[0]).is_none());
1529    }
1530
1531    #[test]
1532    fn test_get_tables_finds_tables_inside_in_subquery() {
1533        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1534        let sql = "SELECT id, name FROM customers WHERE id IN (SELECT customer_id FROM orders WHERE amount > 1000)";
1535        let exprs = dialect.parse(sql).unwrap();
1536        let tables = get_tables(&exprs[0]);
1537        let names: Vec<&str> = tables
1538            .iter()
1539            .filter_map(|e| {
1540                if let Expression::Table(t) = e {
1541                    Some(t.name.name.as_str())
1542                } else {
1543                    None
1544                }
1545            })
1546            .collect();
1547        assert!(names.contains(&"customers"), "should find outer table");
1548        assert!(names.contains(&"orders"), "should find subquery table");
1549    }
1550
1551    #[test]
1552    fn test_get_tables_finds_tables_inside_exists_subquery() {
1553        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1554        let sql = "SELECT * FROM customers c WHERE EXISTS (SELECT 1 FROM orders o WHERE o.customer_id = c.id)";
1555        let exprs = dialect.parse(sql).unwrap();
1556        let tables = get_tables(&exprs[0]);
1557        let names: Vec<&str> = tables
1558            .iter()
1559            .filter_map(|e| {
1560                if let Expression::Table(t) = e {
1561                    Some(t.name.name.as_str())
1562                } else {
1563                    None
1564                }
1565            })
1566            .collect();
1567        assert!(names.contains(&"customers"), "should find outer table");
1568        assert!(
1569            names.contains(&"orders"),
1570            "should find EXISTS subquery table"
1571        );
1572    }
1573
1574    #[test]
1575    fn test_get_tables_finds_tables_in_correlated_subquery() {
1576        let dialect = crate::Dialect::get(crate::dialects::DialectType::TSQL);
1577        let sql = "SELECT id, name FROM customers WHERE id IN (SELECT customer_id FROM orders WHERE amount > 1000)";
1578        let exprs = dialect.parse(sql).unwrap();
1579        let tables = get_tables(&exprs[0]);
1580        let names: Vec<&str> = tables
1581            .iter()
1582            .filter_map(|e| {
1583                if let Expression::Table(t) = e {
1584                    Some(t.name.name.as_str())
1585                } else {
1586                    None
1587                }
1588            })
1589            .collect();
1590        assert!(
1591            names.contains(&"customers"),
1592            "TSQL: should find outer table"
1593        );
1594        assert!(
1595            names.contains(&"orders"),
1596            "TSQL: should find subquery table"
1597        );
1598    }
1599}
polyglot_sql/traversal.rs

polyglot_sql/
traversal.rs
