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.
498pub fn is_aggregate(expr: &Expression) -> bool {
499    matches!(
500        expr,
501        Expression::AggregateFunction(_)
502            | Expression::Count(_)
503            | Expression::Sum(_)
504            | Expression::Avg(_)
505            | Expression::Min(_)
506            | Expression::Max(_)
507            | Expression::GroupConcat(_)
508            | Expression::StringAgg(_)
509            | Expression::ListAgg(_)
510            | Expression::CountIf(_)
511            | Expression::SumIf(_)
512    )
513}
514
515/// Returns `true` if `expr` is a window function ([`Expression::WindowFunction`]).
516pub fn is_window_function(expr: &Expression) -> bool {
517    matches!(expr, Expression::WindowFunction(_))
518}
519
520/// Collects all column references ([`Expression::Column`]) from the expression tree.
521///
522/// Performs a depth-first search and returns references to every column node found.
523pub fn get_columns(expr: &Expression) -> Vec<&Expression> {
524    expr.find_all(is_column)
525}
526
527/// Collects all table references ([`Expression::Table`]) from the expression tree.
528///
529/// Performs a depth-first search and returns references to every table node found.
530///
531/// Note: DML target tables (`Insert.table`, `Update.table`, `Delete.table`) are
532/// stored as `TableRef` struct fields, not as `Expression::Table` nodes, so they
533/// are not reachable via tree traversal. Use [`get_all_tables`] to include those.
534pub fn get_tables(expr: &Expression) -> Vec<&Expression> {
535    expr.find_all(|e| matches!(e, Expression::Table(_)))
536}
537
538/// Collects **all** referenced tables from the expression tree, including DML
539/// target tables that are stored as `TableRef` struct fields and are therefore
540/// not reachable through normal tree traversal.
541///
542/// Returns owned `Expression::Table` values. This is the comprehensive version
543/// of [`get_tables`] — use it when you need to discover every table referenced
544/// in a statement, including inside CTE bodies containing INSERT/UPDATE/DELETE.
545pub fn get_all_tables(expr: &Expression) -> Vec<Expression> {
546    use std::collections::HashSet;
547
548    let mut seen = HashSet::new();
549    let mut result = Vec::new();
550
551    // First: collect all Expression::Table nodes found via DFS.
552    for node in expr.dfs() {
553        if let Expression::Table(t) = node {
554            let qname = table_ref_qualified_name(t);
555            if seen.insert(qname) {
556                result.push(node.clone());
557            }
558        }
559
560        // Also extract DML target TableRef fields not reachable via iter_children.
561        let refs: Vec<&TableRef> = match node {
562            Expression::Insert(ins) => vec![&ins.table],
563            Expression::Update(upd) => {
564                let mut v = vec![&upd.table];
565                v.extend(upd.extra_tables.iter());
566                v
567            }
568            Expression::Delete(del) => {
569                let mut v = vec![&del.table];
570                v.extend(del.using.iter());
571                v
572            }
573            _ => continue,
574        };
575        for tref in refs {
576            if tref.name.name.is_empty() {
577                continue;
578            }
579            let qname = table_ref_qualified_name(tref);
580            if seen.insert(qname) {
581                result.push(Expression::Table(Box::new(tref.clone())));
582            }
583        }
584    }
585
586    result
587}
588
589/// Build a qualified name string from a TableRef for deduplication purposes.
590fn table_ref_qualified_name(t: &TableRef) -> String {
591    let mut name = String::new();
592    if let Some(ref cat) = t.catalog {
593        name.push_str(&cat.name);
594        name.push('.');
595    }
596    if let Some(ref schema) = t.schema {
597        name.push_str(&schema.name);
598        name.push('.');
599    }
600    name.push_str(&t.name.name);
601    name
602}
603
604/// Extracts the underlying [`Expression::Table`] from a MERGE field that may
605/// be a bare `Table`, an `Alias` wrapping a `Table`, or an `Identifier`.
606/// Returns `None` if the expression doesn't contain a recognisable table.
607fn unwrap_merge_table(expr: &Expression) -> Option<&Expression> {
608    match expr {
609        Expression::Table(_) => Some(expr),
610        Expression::Alias(alias) => match &alias.this {
611            Expression::Table(_) => Some(&alias.this),
612            _ => None,
613        },
614        _ => None,
615    }
616}
617
618/// Returns the target table of a MERGE statement (the `Merge.this` field),
619/// unwrapping any alias wrapper to yield the underlying [`Expression::Table`].
620///
621/// Returns `None` if `expr` is not a `Merge` or the target isn't a recognisable table.
622pub fn get_merge_target(expr: &Expression) -> Option<&Expression> {
623    match expr {
624        Expression::Merge(m) => unwrap_merge_table(&m.this),
625        _ => None,
626    }
627}
628
629/// Returns the source table of a MERGE statement (the `Merge.using` field),
630/// unwrapping any alias wrapper to yield the underlying [`Expression::Table`].
631///
632/// Returns `None` if `expr` is not a `Merge`, the source isn't a recognisable
633/// table (e.g. it's a subquery), or the source is otherwise unresolvable.
634pub fn get_merge_source(expr: &Expression) -> Option<&Expression> {
635    match expr {
636        Expression::Merge(m) => unwrap_merge_table(&m.using),
637        _ => None,
638    }
639}
640
641/// Returns `true` if the expression tree contains any aggregate function calls.
642pub fn contains_aggregate(expr: &Expression) -> bool {
643    expr.contains(is_aggregate)
644}
645
646/// Returns `true` if the expression tree contains any window function calls.
647pub fn contains_window_function(expr: &Expression) -> bool {
648    expr.contains(is_window_function)
649}
650
651/// Returns `true` if the expression tree contains any subquery nodes.
652pub fn contains_subquery(expr: &Expression) -> bool {
653    expr.contains(is_subquery)
654}
655
656// ---------------------------------------------------------------------------
657// Extended type predicates
658// ---------------------------------------------------------------------------
659
660/// Macro for generating simple type-predicate functions.
661macro_rules! is_type {
662    ($name:ident, $($variant:pat),+ $(,)?) => {
663        /// Returns `true` if `expr` matches the expected AST variant(s).
664        pub fn $name(expr: &Expression) -> bool {
665            matches!(expr, $($variant)|+)
666        }
667    };
668}
669
670// Query
671is_type!(is_insert, Expression::Insert(_));
672is_type!(is_update, Expression::Update(_));
673is_type!(is_delete, Expression::Delete(_));
674is_type!(is_merge, Expression::Merge(_));
675is_type!(is_union, Expression::Union(_));
676is_type!(is_intersect, Expression::Intersect(_));
677is_type!(is_except, Expression::Except(_));
678
679// Identifiers & literals
680is_type!(is_boolean, Expression::Boolean(_));
681is_type!(is_null_literal, Expression::Null(_));
682is_type!(is_star, Expression::Star(_));
683is_type!(is_identifier, Expression::Identifier(_));
684is_type!(is_table, Expression::Table(_));
685
686// Comparison
687is_type!(is_eq, Expression::Eq(_));
688is_type!(is_neq, Expression::Neq(_));
689is_type!(is_lt, Expression::Lt(_));
690is_type!(is_lte, Expression::Lte(_));
691is_type!(is_gt, Expression::Gt(_));
692is_type!(is_gte, Expression::Gte(_));
693is_type!(is_like, Expression::Like(_));
694is_type!(is_ilike, Expression::ILike(_));
695
696// Arithmetic
697is_type!(is_add, Expression::Add(_));
698is_type!(is_sub, Expression::Sub(_));
699is_type!(is_mul, Expression::Mul(_));
700is_type!(is_div, Expression::Div(_));
701is_type!(is_mod, Expression::Mod(_));
702is_type!(is_concat, Expression::Concat(_));
703
704// Logical
705is_type!(is_and, Expression::And(_));
706is_type!(is_or, Expression::Or(_));
707is_type!(is_not, Expression::Not(_));
708
709// Predicates
710is_type!(is_in, Expression::In(_));
711is_type!(is_between, Expression::Between(_));
712is_type!(is_is_null, Expression::IsNull(_));
713is_type!(is_exists, Expression::Exists(_));
714
715// Functions
716is_type!(is_count, Expression::Count(_));
717is_type!(is_sum, Expression::Sum(_));
718is_type!(is_avg, Expression::Avg(_));
719is_type!(is_min_func, Expression::Min(_));
720is_type!(is_max_func, Expression::Max(_));
721is_type!(is_coalesce, Expression::Coalesce(_));
722is_type!(is_null_if, Expression::NullIf(_));
723is_type!(is_cast, Expression::Cast(_));
724is_type!(is_try_cast, Expression::TryCast(_));
725is_type!(is_safe_cast, Expression::SafeCast(_));
726is_type!(is_case, Expression::Case(_));
727
728// Clauses
729is_type!(is_from, Expression::From(_));
730is_type!(is_join, Expression::Join(_));
731is_type!(is_where, Expression::Where(_));
732is_type!(is_group_by, Expression::GroupBy(_));
733is_type!(is_having, Expression::Having(_));
734is_type!(is_order_by, Expression::OrderBy(_));
735is_type!(is_limit, Expression::Limit(_));
736is_type!(is_offset, Expression::Offset(_));
737is_type!(is_with, Expression::With(_));
738is_type!(is_cte, Expression::Cte(_));
739is_type!(is_alias, Expression::Alias(_));
740is_type!(is_paren, Expression::Paren(_));
741is_type!(is_ordered, Expression::Ordered(_));
742
743// DDL
744is_type!(is_create_table, Expression::CreateTable(_));
745is_type!(is_drop_table, Expression::DropTable(_));
746is_type!(is_alter_table, Expression::AlterTable(_));
747is_type!(is_create_index, Expression::CreateIndex(_));
748is_type!(is_drop_index, Expression::DropIndex(_));
749is_type!(is_create_view, Expression::CreateView(_));
750is_type!(is_drop_view, Expression::DropView(_));
751
752// ---------------------------------------------------------------------------
753// Composite predicates
754// ---------------------------------------------------------------------------
755
756/// Returns `true` if `expr` is a query statement (SELECT, INSERT, UPDATE, DELETE, or MERGE).
757pub fn is_query(expr: &Expression) -> bool {
758    matches!(
759        expr,
760        Expression::Select(_)
761            | Expression::Insert(_)
762            | Expression::Update(_)
763            | Expression::Delete(_)
764            | Expression::Merge(_)
765    )
766}
767
768/// Returns `true` if `expr` is a set operation (UNION, INTERSECT, or EXCEPT).
769pub fn is_set_operation(expr: &Expression) -> bool {
770    matches!(
771        expr,
772        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
773    )
774}
775
776/// Returns `true` if `expr` is a comparison operator.
777pub fn is_comparison(expr: &Expression) -> bool {
778    matches!(
779        expr,
780        Expression::Eq(_)
781            | Expression::Neq(_)
782            | Expression::Lt(_)
783            | Expression::Lte(_)
784            | Expression::Gt(_)
785            | Expression::Gte(_)
786            | Expression::Like(_)
787            | Expression::ILike(_)
788    )
789}
790
791/// Returns `true` if `expr` is an arithmetic operator.
792pub fn is_arithmetic(expr: &Expression) -> bool {
793    matches!(
794        expr,
795        Expression::Add(_)
796            | Expression::Sub(_)
797            | Expression::Mul(_)
798            | Expression::Div(_)
799            | Expression::Mod(_)
800    )
801}
802
803/// Returns `true` if `expr` is a logical operator (AND, OR, NOT).
804pub fn is_logical(expr: &Expression) -> bool {
805    matches!(
806        expr,
807        Expression::And(_) | Expression::Or(_) | Expression::Not(_)
808    )
809}
810
811/// Returns `true` if `expr` is a DDL statement.
812pub fn is_ddl(expr: &Expression) -> bool {
813    matches!(
814        expr,
815        Expression::CreateTable(_)
816            | Expression::DropTable(_)
817            | Expression::Undrop(_)
818            | Expression::AlterTable(_)
819            | Expression::CreateIndex(_)
820            | Expression::DropIndex(_)
821            | Expression::CreateView(_)
822            | Expression::DropView(_)
823            | Expression::AlterView(_)
824            | Expression::CreateSchema(_)
825            | Expression::DropSchema(_)
826            | Expression::CreateDatabase(_)
827            | Expression::DropDatabase(_)
828            | Expression::CreateFunction(_)
829            | Expression::DropFunction(_)
830            | Expression::CreateProcedure(_)
831            | Expression::DropProcedure(_)
832            | Expression::CreateSequence(_)
833            | Expression::CreateSynonym(_)
834            | Expression::DropSequence(_)
835            | Expression::AlterSequence(_)
836            | Expression::CreateTrigger(_)
837            | Expression::DropTrigger(_)
838            | Expression::CreateType(_)
839            | Expression::DropType(_)
840    )
841}
842
843/// Find the parent of `target` within the tree rooted at `root`.
844///
845/// Uses pointer identity ([`std::ptr::eq`]) — `target` must be a reference
846/// obtained from the same tree (e.g., via [`ExpressionWalk::find`] or DFS iteration).
847///
848/// Returns `None` if `target` is the root itself or is not found in the tree.
849pub fn find_parent<'a>(root: &'a Expression, target: &Expression) -> Option<&'a Expression> {
850    fn search<'a>(node: &'a Expression, target: *const Expression) -> Option<&'a Expression> {
851        let mut result = None;
852        crate::ast_children::for_each_child(node, |_, child| {
853            if result.is_none() {
854                if std::ptr::eq(child, target) {
855                    result = Some(node);
856                } else {
857                    result = search(child, target);
858                }
859            }
860        });
861        result
862    }
863
864    search(root, target as *const Expression)
865}
866
867/// Find the first ancestor of `target` matching `predicate`, walking from
868/// parent toward root.
869///
870/// Uses pointer identity for target lookup. Returns `None` if no ancestor
871/// matches or `target` is not found in the tree.
872pub fn find_ancestor<'a, F>(
873    root: &'a Expression,
874    target: &Expression,
875    predicate: F,
876) -> Option<&'a Expression>
877where
878    F: Fn(&Expression) -> bool,
879{
880    // Build path from root to target
881    fn build_path<'a>(
882        node: &'a Expression,
883        target: *const Expression,
884        path: &mut Vec<&'a Expression>,
885    ) -> bool {
886        if std::ptr::eq(node, target) {
887            return true;
888        }
889        path.push(node);
890        let mut found = false;
891        crate::ast_children::for_each_child(node, |_, child| {
892            if !found {
893                found = build_path(child, target, path);
894            }
895        });
896        if found {
897            return true;
898        }
899        path.pop();
900        false
901    }
902
903    let mut path = Vec::new();
904    if !build_path(root, target as *const Expression, &mut path) {
905        return None;
906    }
907
908    // Walk path in reverse (parent first, then grandparent, etc.)
909    for ancestor in path.iter().rev() {
910        if predicate(ancestor) {
911            return Some(ancestor);
912        }
913    }
914    None
915}
916
917#[cfg(test)]
918mod tests {
919    use super::*;
920    use crate::expressions::{BinaryOp, Column, Identifier, LikeOp, Literal, TableRef};
921
922    fn make_column(name: &str) -> Expression {
923        Expression::boxed_column(Column {
924            name: Identifier {
925                name: name.to_string(),
926                quoted: false,
927                trailing_comments: vec![],
928                span: None,
929            },
930            table: None,
931            join_mark: false,
932            trailing_comments: vec![],
933            span: None,
934            inferred_type: None,
935        })
936    }
937
938    fn make_literal(value: i64) -> Expression {
939        Expression::Literal(Box::new(Literal::Number(value.to_string())))
940    }
941
942    #[test]
943    fn test_dfs_simple() {
944        let left = make_column("a");
945        let right = make_literal(1);
946        let expr = Expression::Eq(Box::new(BinaryOp {
947            left,
948            right,
949            left_comments: vec![],
950            operator_comments: vec![],
951            trailing_comments: vec![],
952            inferred_type: None,
953        }));
954
955        let nodes: Vec<_> = expr.dfs().collect();
956        assert_eq!(nodes.len(), 3); // Eq, Column, Literal
957        assert!(matches!(nodes[0], Expression::Eq(_)));
958        assert!(matches!(nodes[1], Expression::Column(_)));
959        assert!(matches!(nodes[2], Expression::Literal(_)));
960    }
961
962    #[test]
963    fn test_find() {
964        let left = make_column("a");
965        let right = make_literal(1);
966        let expr = Expression::Eq(Box::new(BinaryOp {
967            left,
968            right,
969            left_comments: vec![],
970            operator_comments: vec![],
971            trailing_comments: vec![],
972            inferred_type: None,
973        }));
974
975        let column = expr.find(is_column);
976        assert!(column.is_some());
977        assert!(matches!(column.unwrap(), Expression::Column(_)));
978
979        let literal = expr.find(is_literal);
980        assert!(literal.is_some());
981        assert!(matches!(literal.unwrap(), Expression::Literal(_)));
982    }
983
984    #[test]
985    fn test_find_all() {
986        let col1 = make_column("a");
987        let col2 = make_column("b");
988        let expr = Expression::And(Box::new(BinaryOp {
989            left: col1,
990            right: col2,
991            left_comments: vec![],
992            operator_comments: vec![],
993            trailing_comments: vec![],
994            inferred_type: None,
995        }));
996
997        let columns = expr.find_all(is_column);
998        assert_eq!(columns.len(), 2);
999    }
1000
1001    #[test]
1002    fn test_contains() {
1003        let col = make_column("a");
1004        let lit = make_literal(1);
1005        let expr = Expression::Eq(Box::new(BinaryOp {
1006            left: col,
1007            right: lit,
1008            left_comments: vec![],
1009            operator_comments: vec![],
1010            trailing_comments: vec![],
1011            inferred_type: None,
1012        }));
1013
1014        assert!(expr.contains(is_column));
1015        assert!(expr.contains(is_literal));
1016        assert!(!expr.contains(is_subquery));
1017    }
1018
1019    #[test]
1020    fn test_count() {
1021        let col1 = make_column("a");
1022        let col2 = make_column("b");
1023        let lit = make_literal(1);
1024
1025        let inner = Expression::Add(Box::new(BinaryOp {
1026            left: col2,
1027            right: lit,
1028            left_comments: vec![],
1029            operator_comments: vec![],
1030            trailing_comments: vec![],
1031            inferred_type: None,
1032        }));
1033
1034        let expr = Expression::Eq(Box::new(BinaryOp {
1035            left: col1,
1036            right: inner,
1037            left_comments: vec![],
1038            operator_comments: vec![],
1039            trailing_comments: vec![],
1040            inferred_type: None,
1041        }));
1042
1043        assert_eq!(expr.count(is_column), 2);
1044        assert_eq!(expr.count(is_literal), 1);
1045    }
1046
1047    #[test]
1048    fn test_tree_depth() {
1049        // Single node
1050        let lit = make_literal(1);
1051        assert_eq!(lit.tree_depth(), 0);
1052
1053        // One level
1054        let col = make_column("a");
1055        let expr = Expression::Eq(Box::new(BinaryOp {
1056            left: col,
1057            right: lit.clone(),
1058            left_comments: vec![],
1059            operator_comments: vec![],
1060            trailing_comments: vec![],
1061            inferred_type: None,
1062        }));
1063        assert_eq!(expr.tree_depth(), 1);
1064
1065        // Two levels
1066        let inner = Expression::Add(Box::new(BinaryOp {
1067            left: make_column("b"),
1068            right: lit,
1069            left_comments: vec![],
1070            operator_comments: vec![],
1071            trailing_comments: vec![],
1072            inferred_type: None,
1073        }));
1074        let outer = Expression::Eq(Box::new(BinaryOp {
1075            left: make_column("a"),
1076            right: inner,
1077            left_comments: vec![],
1078            operator_comments: vec![],
1079            trailing_comments: vec![],
1080            inferred_type: None,
1081        }));
1082        assert_eq!(outer.tree_depth(), 2);
1083    }
1084
1085    #[test]
1086    fn test_tree_context() {
1087        let col = make_column("a");
1088        let lit = make_literal(1);
1089        let expr = Expression::Eq(Box::new(BinaryOp {
1090            left: col,
1091            right: lit,
1092            left_comments: vec![],
1093            operator_comments: vec![],
1094            trailing_comments: vec![],
1095            inferred_type: None,
1096        }));
1097
1098        let ctx = TreeContext::build(&expr);
1099
1100        // Root has no parent
1101        let root_info = ctx.get(0).unwrap();
1102        assert!(root_info.parent_id.is_none());
1103
1104        // Children have root as parent
1105        let left_info = ctx.get(1).unwrap();
1106        assert_eq!(left_info.parent_id, Some(0));
1107        assert_eq!(left_info.arg_key, "left");
1108
1109        let right_info = ctx.get(2).unwrap();
1110        assert_eq!(right_info.parent_id, Some(0));
1111        assert_eq!(right_info.arg_key, "right");
1112    }
1113
1114    // -- Step 8: transform / transform_map tests --
1115
1116    #[test]
1117    fn test_transform_rename_columns() {
1118        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1119        let expr = ast[0].clone();
1120        let result = super::transform_map(expr, &|e| {
1121            if let Expression::Column(ref c) = e {
1122                if c.name.name == "a" {
1123                    return Ok(Expression::boxed_column(Column {
1124                        name: Identifier::new("alpha"),
1125                        table: c.table.clone(),
1126                        join_mark: false,
1127                        trailing_comments: vec![],
1128                        span: None,
1129                        inferred_type: None,
1130                    }));
1131                }
1132            }
1133            Ok(e)
1134        })
1135        .unwrap();
1136        let sql = crate::generator::Generator::sql(&result).unwrap();
1137        assert!(sql.contains("alpha"), "Expected 'alpha' in: {}", sql);
1138        assert!(sql.contains("b"), "Expected 'b' in: {}", sql);
1139    }
1140
1141    #[test]
1142    fn test_transform_noop() {
1143        let ast = crate::parser::Parser::parse_sql("SELECT 1 + 2").unwrap();
1144        let expr = ast[0].clone();
1145        let result = super::transform_map(expr.clone(), &|e| Ok(e)).unwrap();
1146        let sql1 = crate::generator::Generator::sql(&expr).unwrap();
1147        let sql2 = crate::generator::Generator::sql(&result).unwrap();
1148        assert_eq!(sql1, sql2);
1149    }
1150
1151    #[test]
1152    fn test_transform_nested() {
1153        let ast = crate::parser::Parser::parse_sql("SELECT a + b FROM t").unwrap();
1154        let expr = ast[0].clone();
1155        let result = super::transform_map(expr, &|e| {
1156            if let Expression::Column(ref c) = e {
1157                return Ok(Expression::Literal(Box::new(Literal::Number(
1158                    if c.name.name == "a" { "1" } else { "2" }.to_string(),
1159                ))));
1160            }
1161            Ok(e)
1162        })
1163        .unwrap();
1164        let sql = crate::generator::Generator::sql(&result).unwrap();
1165        assert_eq!(sql, "SELECT 1 + 2 FROM t");
1166    }
1167
1168    #[test]
1169    fn test_transform_error() {
1170        let ast = crate::parser::Parser::parse_sql("SELECT a FROM t").unwrap();
1171        let expr = ast[0].clone();
1172        let result = super::transform_map(expr, &|e| {
1173            if let Expression::Column(ref c) = e {
1174                if c.name.name == "a" {
1175                    return Err(crate::error::Error::parse("test error", 0, 0, 0, 0));
1176                }
1177            }
1178            Ok(e)
1179        });
1180        assert!(result.is_err());
1181    }
1182
1183    #[test]
1184    fn test_transform_owned_trait() {
1185        let ast = crate::parser::Parser::parse_sql("SELECT x FROM t").unwrap();
1186        let expr = ast[0].clone();
1187        let result = expr.transform_owned(|e| Ok(Some(e))).unwrap();
1188        let sql = crate::generator::Generator::sql(&result).unwrap();
1189        assert_eq!(sql, "SELECT x FROM t");
1190    }
1191
1192    // -- children() tests --
1193
1194    #[test]
1195    fn test_children_leaf() {
1196        let lit = make_literal(1);
1197        assert_eq!(lit.children().len(), 0);
1198    }
1199
1200    #[test]
1201    fn test_children_binary_op() {
1202        let left = make_column("a");
1203        let right = make_literal(1);
1204        let expr = Expression::Eq(Box::new(BinaryOp {
1205            left,
1206            right,
1207            left_comments: vec![],
1208            operator_comments: vec![],
1209            trailing_comments: vec![],
1210            inferred_type: None,
1211        }));
1212        let children = expr.children();
1213        assert_eq!(children.len(), 2);
1214        assert!(matches!(children[0], Expression::Column(_)));
1215        assert!(matches!(children[1], Expression::Literal(_)));
1216    }
1217
1218    #[test]
1219    fn test_children_select() {
1220        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1221        let expr = &ast[0];
1222        let children = expr.children();
1223        // Should include select list items (a, b)
1224        assert!(children.len() >= 2);
1225    }
1226
1227    #[test]
1228    fn test_children_follow_ast_field_order() {
1229        let ast = crate::parser::Parser::parse_sql("SELECT a, b FROM t").unwrap();
1230        let children = ast[0].children();
1231
1232        assert!(matches!(children[0], Expression::Column(column) if column.name.name == "a"));
1233        assert!(matches!(children[1], Expression::Column(column) if column.name.name == "b"));
1234        assert!(matches!(children[2], Expression::Table(table) if table.name.name == "t"));
1235    }
1236
1237    #[test]
1238    fn test_traversal_covers_previously_omitted_expression_fields() {
1239        let like = Expression::Like(Box::new(LikeOp {
1240            left: make_column("name"),
1241            right: Expression::Literal(Box::new(Literal::String("x%".to_string()))),
1242            escape: Some(Expression::Literal(Box::new(Literal::String(
1243                "!".to_string(),
1244            )))),
1245            quantifier: None,
1246            inferred_type: None,
1247        }));
1248        let nodes: Vec<_> = like.dfs().collect();
1249        assert_eq!(nodes.len(), 4);
1250        assert!(matches!(
1251            nodes[3],
1252            Expression::Literal(literal)
1253                if matches!(literal.as_ref(), Literal::String(value) if value == "!")
1254        ));
1255
1256        let mut table = TableRef::new("events");
1257        table.hints.push(make_column("table_hint"));
1258        table.identifier_func = Some(Box::new(Expression::identifier("dynamic_table")));
1259        let table = Expression::Table(Box::new(table));
1260        let children = table.children();
1261        assert_eq!(children.len(), 2);
1262        assert!(
1263            matches!(children[0], Expression::Column(column) if column.name.name == "table_hint")
1264        );
1265        assert!(
1266            matches!(children[1], Expression::Identifier(identifier) if identifier.name == "dynamic_table")
1267        );
1268    }
1269
1270    #[test]
1271    fn test_children_select_includes_from_and_join_sources() {
1272        let ast = crate::parser::Parser::parse_sql(
1273            "SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id",
1274        )
1275        .unwrap();
1276        let expr = &ast[0];
1277        let children = expr.children();
1278
1279        let table_names: Vec<&str> = children
1280            .iter()
1281            .filter_map(|e| match e {
1282                Expression::Table(t) => Some(t.name.name.as_str()),
1283                _ => None,
1284            })
1285            .collect();
1286
1287        assert!(table_names.contains(&"users"));
1288        assert!(table_names.contains(&"orders"));
1289    }
1290
1291    #[test]
1292    fn test_get_tables_includes_insert_query_sources() {
1293        let ast = crate::parser::Parser::parse_sql(
1294            "INSERT INTO dst (id) SELECT s.id FROM src s JOIN dim d ON s.id = d.id",
1295        )
1296        .unwrap();
1297        let expr = &ast[0];
1298        let tables = get_tables(expr);
1299        let names: Vec<&str> = tables
1300            .iter()
1301            .filter_map(|e| match e {
1302                Expression::Table(t) => Some(t.name.name.as_str()),
1303                _ => None,
1304            })
1305            .collect();
1306
1307        assert!(names.contains(&"src"));
1308        assert!(names.contains(&"dim"));
1309    }
1310
1311    // -- find_parent() tests --
1312
1313    #[test]
1314    fn test_find_parent_binary() {
1315        let left = make_column("a");
1316        let right = make_literal(1);
1317        let expr = Expression::Eq(Box::new(BinaryOp {
1318            left,
1319            right,
1320            left_comments: vec![],
1321            operator_comments: vec![],
1322            trailing_comments: vec![],
1323            inferred_type: None,
1324        }));
1325
1326        // Find the column child and get its parent
1327        let col = expr.find(is_column).unwrap();
1328        let parent = super::find_parent(&expr, col);
1329        assert!(parent.is_some());
1330        assert!(matches!(parent.unwrap(), Expression::Eq(_)));
1331    }
1332
1333    #[test]
1334    fn test_find_parent_root_has_none() {
1335        let lit = make_literal(1);
1336        let parent = super::find_parent(&lit, &lit);
1337        assert!(parent.is_none());
1338    }
1339
1340    // -- find_ancestor() tests --
1341
1342    #[test]
1343    fn test_find_ancestor_select() {
1344        let ast = crate::parser::Parser::parse_sql("SELECT a FROM t WHERE a > 1").unwrap();
1345        let expr = &ast[0];
1346
1347        // Find a column inside the WHERE clause
1348        let where_col = expr.dfs().find(|e| {
1349            if let Expression::Column(c) = e {
1350                c.name.name == "a"
1351            } else {
1352                false
1353            }
1354        });
1355        assert!(where_col.is_some());
1356
1357        // Find Select ancestor of that column
1358        let ancestor = super::find_ancestor(expr, where_col.unwrap(), is_select);
1359        assert!(ancestor.is_some());
1360        assert!(matches!(ancestor.unwrap(), Expression::Select(_)));
1361    }
1362
1363    #[test]
1364    fn test_find_ancestor_no_match() {
1365        let left = make_column("a");
1366        let right = make_literal(1);
1367        let expr = Expression::Eq(Box::new(BinaryOp {
1368            left,
1369            right,
1370            left_comments: vec![],
1371            operator_comments: vec![],
1372            trailing_comments: vec![],
1373            inferred_type: None,
1374        }));
1375
1376        let col = expr.find(is_column).unwrap();
1377        let ancestor = super::find_ancestor(&expr, col, is_select);
1378        assert!(ancestor.is_none());
1379    }
1380
1381    #[test]
1382    fn test_ancestors() {
1383        let col = make_column("a");
1384        let lit = make_literal(1);
1385        let inner = Expression::Add(Box::new(BinaryOp {
1386            left: col,
1387            right: lit,
1388            left_comments: vec![],
1389            operator_comments: vec![],
1390            trailing_comments: vec![],
1391            inferred_type: None,
1392        }));
1393        let outer = Expression::Eq(Box::new(BinaryOp {
1394            left: make_column("b"),
1395            right: inner,
1396            left_comments: vec![],
1397            operator_comments: vec![],
1398            trailing_comments: vec![],
1399            inferred_type: None,
1400        }));
1401
1402        let ctx = TreeContext::build(&outer);
1403
1404        // The inner Add's left child (column "a") should have ancestors
1405        // Node 0: Eq
1406        // Node 1: Column "b" (left of Eq)
1407        // Node 2: Add (right of Eq)
1408        // Node 3: Column "a" (left of Add)
1409        // Node 4: Literal (right of Add)
1410
1411        let ancestors = ctx.ancestors_of(3);
1412        assert_eq!(ancestors, vec![2, 0]); // Add, then Eq
1413    }
1414
1415    #[test]
1416    fn test_get_merge_target_and_source() {
1417        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1418
1419        // MERGE with aliased target and source tables
1420        let sql = "MERGE INTO orders o USING customers c ON o.customer_id = c.id WHEN MATCHED THEN UPDATE SET amount = amount + 100";
1421        let exprs = dialect.parse(sql).unwrap();
1422        let expr = &exprs[0];
1423
1424        assert!(is_merge(expr));
1425        assert!(is_query(expr));
1426
1427        let target = get_merge_target(expr).expect("should find target table");
1428        assert!(matches!(target, Expression::Table(_)));
1429        if let Expression::Table(t) = target {
1430            assert_eq!(t.name.name, "orders");
1431        }
1432
1433        let source = get_merge_source(expr).expect("should find source table");
1434        assert!(matches!(source, Expression::Table(_)));
1435        if let Expression::Table(t) = source {
1436            assert_eq!(t.name.name, "customers");
1437        }
1438    }
1439
1440    #[test]
1441    fn test_get_merge_source_subquery_returns_none() {
1442        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1443
1444        // MERGE with subquery source — get_merge_source should return None
1445        let sql = "MERGE INTO orders o USING (SELECT * FROM customers) c ON o.customer_id = c.id WHEN MATCHED THEN DELETE";
1446        let exprs = dialect.parse(sql).unwrap();
1447        let expr = &exprs[0];
1448
1449        assert!(get_merge_target(expr).is_some());
1450        assert!(get_merge_source(expr).is_none());
1451    }
1452
1453    #[test]
1454    fn test_get_merge_on_non_merge_returns_none() {
1455        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1456        let exprs = dialect.parse("SELECT 1").unwrap();
1457        assert!(get_merge_target(&exprs[0]).is_none());
1458        assert!(get_merge_source(&exprs[0]).is_none());
1459    }
1460
1461    #[test]
1462    fn test_get_tables_finds_tables_inside_in_subquery() {
1463        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1464        let sql = "SELECT id, name FROM customers WHERE id IN (SELECT customer_id FROM orders WHERE amount > 1000)";
1465        let exprs = dialect.parse(sql).unwrap();
1466        let tables = get_tables(&exprs[0]);
1467        let names: Vec<&str> = tables
1468            .iter()
1469            .filter_map(|e| {
1470                if let Expression::Table(t) = e {
1471                    Some(t.name.name.as_str())
1472                } else {
1473                    None
1474                }
1475            })
1476            .collect();
1477        assert!(names.contains(&"customers"), "should find outer table");
1478        assert!(names.contains(&"orders"), "should find subquery table");
1479    }
1480
1481    #[test]
1482    fn test_get_tables_finds_tables_inside_exists_subquery() {
1483        let dialect = crate::Dialect::get(crate::dialects::DialectType::Generic);
1484        let sql = "SELECT * FROM customers c WHERE EXISTS (SELECT 1 FROM orders o WHERE o.customer_id = c.id)";
1485        let exprs = dialect.parse(sql).unwrap();
1486        let tables = get_tables(&exprs[0]);
1487        let names: Vec<&str> = tables
1488            .iter()
1489            .filter_map(|e| {
1490                if let Expression::Table(t) = e {
1491                    Some(t.name.name.as_str())
1492                } else {
1493                    None
1494                }
1495            })
1496            .collect();
1497        assert!(names.contains(&"customers"), "should find outer table");
1498        assert!(
1499            names.contains(&"orders"),
1500            "should find EXISTS subquery table"
1501        );
1502    }
1503
1504    #[test]
1505    fn test_get_tables_finds_tables_in_correlated_subquery() {
1506        let dialect = crate::Dialect::get(crate::dialects::DialectType::TSQL);
1507        let sql = "SELECT id, name FROM customers WHERE id IN (SELECT customer_id FROM orders WHERE amount > 1000)";
1508        let exprs = dialect.parse(sql).unwrap();
1509        let tables = get_tables(&exprs[0]);
1510        let names: Vec<&str> = tables
1511            .iter()
1512            .filter_map(|e| {
1513                if let Expression::Table(t) = e {
1514                    Some(t.name.name.as_str())
1515                } else {
1516                    None
1517                }
1518            })
1519            .collect();
1520        assert!(
1521            names.contains(&"customers"),
1522            "TSQL: should find outer table"
1523        );
1524        assert!(
1525            names.contains(&"orders"),
1526            "TSQL: should find subquery table"
1527        );
1528    }
1529}
polyglot_sql/traversal.rs

polyglot_sql/
traversal.rs
