sqlparser/ast/
visitor.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Recursive visitors for ast Nodes. See [`Visitor`] for more details.
19
20#[cfg(not(feature = "std"))]
21use alloc::{boxed::Box, string::String, vec::Vec};
22use core::ops::ControlFlow;
23
24use crate::ast::{
25    Expr, GroupByExpr, Ident, ObjectName, OrderBy, OrderByExpr, Query, Select, Statement,
26    TableFactor, ValueWithSpan,
27};
28
29/// A type that can be visited by a [`Visitor`]. See [`Visitor`] for
30/// recursively visiting parsed SQL statements.
31///
32/// # Note
33///
34/// This trait should be automatically derived for sqlparser AST nodes
35/// using the [Visit](sqlparser_derive::Visit) proc macro.
36///
37/// ```text
38/// #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
39/// ```
40pub trait Visit {
41    /// Visit this node with the provided [`Visitor`].
42    ///
43    /// Implementations should call the appropriate visitor hooks to traverse
44    /// child nodes and return a `ControlFlow` value to allow early exit.
45    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break>;
46}
47
48/// A type that can be visited by a [`VisitorMut`]. See [`VisitorMut`] for
49/// recursively visiting parsed SQL statements.
50///
51/// # Note
52///
53/// This trait should be automatically derived for sqlparser AST nodes
54/// using the [VisitMut](sqlparser_derive::VisitMut) proc macro.
55///
56/// ```text
57/// #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
58/// ```
59pub trait VisitMut {
60    /// Mutably visit this node with the provided [`VisitorMut`].
61    ///
62    /// Implementations should call the appropriate mutable visitor hooks to
63    /// traverse and allow in-place mutation of child nodes. Returning a
64    /// `ControlFlow` value permits early termination of the traversal.
65    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break>;
66}
67
68impl<T: Visit> Visit for Option<T> {
69    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
70        if let Some(s) = self {
71            s.visit(visitor)?;
72        }
73        ControlFlow::Continue(())
74    }
75}
76
77impl<T: Visit> Visit for Vec<T> {
78    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
79        for v in self {
80            v.visit(visitor)?;
81        }
82        ControlFlow::Continue(())
83    }
84}
85
86impl<T: Visit> Visit for Box<T> {
87    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
88        T::visit(self, visitor)
89    }
90}
91
92impl<T: VisitMut> VisitMut for Option<T> {
93    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
94        if let Some(s) = self {
95            s.visit(visitor)?;
96        }
97        ControlFlow::Continue(())
98    }
99}
100
101impl<T: VisitMut> VisitMut for Vec<T> {
102    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
103        for v in self {
104            v.visit(visitor)?;
105        }
106        ControlFlow::Continue(())
107    }
108}
109
110impl<T: VisitMut> VisitMut for Box<T> {
111    fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
112        T::visit(self, visitor)
113    }
114}
115
116macro_rules! visit_noop {
117    ($($t:ty),+) => {
118        $(impl Visit for $t {
119            fn visit<V: Visitor>(&self, _visitor: &mut V) -> ControlFlow<V::Break> {
120               ControlFlow::Continue(())
121            }
122        })+
123        $(impl VisitMut for $t {
124            fn visit<V: VisitorMut>(&mut self, _visitor: &mut V) -> ControlFlow<V::Break> {
125               ControlFlow::Continue(())
126            }
127        })+
128    };
129}
130
131visit_noop!(u8, u16, u32, u64, i8, i16, i32, i64, char, bool, String);
132
133#[cfg(feature = "bigdecimal")]
134visit_noop!(bigdecimal::BigDecimal);
135
136/// A visitor that can be used to walk an AST tree.
137///
138/// `pre_visit_` methods are invoked before visiting all children of the
139/// node and `post_visit_` methods are invoked after visiting all
140/// children of the node.
141///
142/// # See also
143///
144/// These methods provide a more concise way of visiting nodes of a certain type:
145/// * [visit_relations]
146/// * [visit_expressions]
147/// * [visit_statements]
148///
149/// # Example
150/// ```
151/// # use sqlparser::parser::Parser;
152/// # use sqlparser::dialect::GenericDialect;
153/// # use sqlparser::ast::{Visit, Visitor, ObjectName, Expr};
154/// # use core::ops::ControlFlow;
155/// // A structure that records statements and relations
156/// #[derive(Default)]
157/// struct V {
158///    visited: Vec<String>,
159/// }
160///
161/// // Visit relations and exprs before children are visited (depth first walk)
162/// // Note you can also visit statements and visit exprs after children have been visited
163/// impl Visitor for V {
164///   type Break = ();
165///
166///   fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
167///     self.visited.push(format!("PRE: RELATION: {}", relation));
168///     ControlFlow::Continue(())
169///   }
170///
171///   fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
172///     self.visited.push(format!("PRE: EXPR: {}", expr));
173///     ControlFlow::Continue(())
174///   }
175/// }
176///
177/// let sql = "SELECT a FROM foo where x IN (SELECT y FROM bar)";
178/// let statements = Parser::parse_sql(&GenericDialect{}, sql)
179///    .unwrap();
180///
181/// // Drive the visitor through the AST
182/// let mut visitor = V::default();
183/// statements.visit(&mut visitor);
184///
185/// // The visitor has visited statements and expressions in pre-traversal order
186/// let expected : Vec<_> = [
187///   "PRE: EXPR: a",
188///   "PRE: RELATION: foo",
189///   "PRE: EXPR: x IN (SELECT y FROM bar)",
190///   "PRE: EXPR: x",
191///   "PRE: EXPR: y",
192///   "PRE: RELATION: bar",
193/// ]
194///   .into_iter().map(|s| s.to_string()).collect();
195///
196/// assert_eq!(visitor.visited, expected);
197/// ```
198pub trait Visitor {
199    /// Type returned when the recursion returns early.
200    ///
201    /// Important note: The `Break` type should be kept as small as possible to prevent
202    /// stack overflow during recursion. If you need to return an error, consider
203    /// boxing it with `Box` to minimize stack usage.
204    type Break;
205
206    /// Invoked for any queries that appear in the AST before visiting children
207    fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow<Self::Break> {
208        ControlFlow::Continue(())
209    }
210
211    /// Invoked for any queries that appear in the AST after visiting children
212    fn post_visit_query(&mut self, _query: &Query) -> ControlFlow<Self::Break> {
213        ControlFlow::Continue(())
214    }
215
216    /// Invoked for any [Select] that appear in the AST before visiting children
217    fn pre_visit_select(&mut self, _select: &Select) -> ControlFlow<Self::Break> {
218        ControlFlow::Continue(())
219    }
220
221    /// Invoked for any [Select] that appear in the AST after visiting children
222    fn post_visit_select(&mut self, _select: &Select) -> ControlFlow<Self::Break> {
223        ControlFlow::Continue(())
224    }
225
226    /// Invoked for any relations (e.g. tables) that appear in the AST before visiting children
227    fn pre_visit_relation(&mut self, _relation: &ObjectName) -> ControlFlow<Self::Break> {
228        ControlFlow::Continue(())
229    }
230
231    /// Invoked for any relations (e.g. tables) that appear in the AST after visiting children
232    fn post_visit_relation(&mut self, _relation: &ObjectName) -> ControlFlow<Self::Break> {
233        ControlFlow::Continue(())
234    }
235
236    /// Invoked for any table factors that appear in the AST before visiting children
237    fn pre_visit_table_factor(&mut self, _table_factor: &TableFactor) -> ControlFlow<Self::Break> {
238        ControlFlow::Continue(())
239    }
240
241    /// Invoked for any table factors that appear in the AST after visiting children
242    fn post_visit_table_factor(&mut self, _table_factor: &TableFactor) -> ControlFlow<Self::Break> {
243        ControlFlow::Continue(())
244    }
245
246    /// Invoked for any expressions that appear in the AST before visiting children
247    fn pre_visit_expr(&mut self, _expr: &Expr) -> ControlFlow<Self::Break> {
248        ControlFlow::Continue(())
249    }
250
251    /// Invoked for any expressions that appear in the AST
252    fn post_visit_expr(&mut self, _expr: &Expr) -> ControlFlow<Self::Break> {
253        ControlFlow::Continue(())
254    }
255
256    /// Invoked for any statements that appear in the AST before visiting children
257    fn pre_visit_statement(&mut self, _statement: &Statement) -> ControlFlow<Self::Break> {
258        ControlFlow::Continue(())
259    }
260
261    /// Invoked for any statements that appear in the AST after visiting children
262    fn post_visit_statement(&mut self, _statement: &Statement) -> ControlFlow<Self::Break> {
263        ControlFlow::Continue(())
264    }
265
266    /// Invoked for any Value that appear in the AST before visiting children
267    fn pre_visit_value(&mut self, _value: &ValueWithSpan) -> ControlFlow<Self::Break> {
268        ControlFlow::Continue(())
269    }
270
271    /// Invoked for any Value that appear in the AST after visiting children
272    fn post_visit_value(&mut self, _value: &ValueWithSpan) -> ControlFlow<Self::Break> {
273        ControlFlow::Continue(())
274    }
275
276    /// Invoked for any identifiers that appear in the AST before visiting children
277    fn pre_visit_ident(&mut self, _ident: &Ident) -> ControlFlow<Self::Break> {
278        ControlFlow::Continue(())
279    }
280
281    /// Invoked for any identifiers that appear in the AST after visiting children
282    fn post_visit_ident(&mut self, _ident: &Ident) -> ControlFlow<Self::Break> {
283        ControlFlow::Continue(())
284    }
285
286    /// Invoked for any `ORDER BY` clauses that appear in the AST before visiting children
287    fn pre_visit_order_by(&mut self, _order_by: &OrderBy) -> ControlFlow<Self::Break> {
288        ControlFlow::Continue(())
289    }
290
291    /// Invoked for any `ORDER BY` clauses that appear in the AST after visiting children
292    fn post_visit_order_by(&mut self, _order_by: &OrderBy) -> ControlFlow<Self::Break> {
293        ControlFlow::Continue(())
294    }
295
296    /// Invoked for any `ORDER BY` expressions that appear in the AST before visiting children
297    fn pre_visit_order_by_expr(
298        &mut self,
299        _order_by_expr: &OrderByExpr,
300    ) -> ControlFlow<Self::Break> {
301        ControlFlow::Continue(())
302    }
303
304    /// Invoked for any `ORDER BY` expressions that appear in the AST after visiting children
305    fn post_visit_order_by_expr(
306        &mut self,
307        _order_by_expr: &OrderByExpr,
308    ) -> ControlFlow<Self::Break> {
309        ControlFlow::Continue(())
310    }
311
312    /// Invoked for any `GROUP BY` clauses that appear in the AST before visiting children
313    fn pre_visit_group_by(&mut self, _group_by: &GroupByExpr) -> ControlFlow<Self::Break> {
314        ControlFlow::Continue(())
315    }
316
317    /// Invoked for any `GROUP BY` clauses that appear in the AST after visiting children
318    fn post_visit_group_by(&mut self, _group_by: &GroupByExpr) -> ControlFlow<Self::Break> {
319        ControlFlow::Continue(())
320    }
321}
322
323/// A visitor that can be used to mutate an AST tree.
324///
325/// `pre_visit_` methods are invoked before visiting all children of the
326/// node and `post_visit_` methods are invoked after visiting all
327/// children of the node.
328///
329/// # See also
330///
331/// These methods provide a more concise way of visiting nodes of a certain type:
332/// * [visit_relations_mut]
333/// * [visit_expressions_mut]
334/// * [visit_statements_mut]
335///
336/// # Example
337/// ```
338/// # use sqlparser::parser::Parser;
339/// # use sqlparser::dialect::GenericDialect;
340/// # use sqlparser::ast::{VisitMut, VisitorMut, ObjectName, Expr, Ident};
341/// # use core::ops::ControlFlow;
342///
343/// // A visitor that replaces "to_replace" with "replaced" in all expressions
344/// struct Replacer;
345///
346/// // Visit each expression after its children have been visited
347/// impl VisitorMut for Replacer {
348///   type Break = ();
349///
350///   fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
351///     if let Expr::Identifier(Ident{ value, ..}) = expr {
352///         *value = value.replace("to_replace", "replaced")
353///     }
354///     ControlFlow::Continue(())
355///   }
356/// }
357///
358/// let sql = "SELECT to_replace FROM foo where to_replace IN (SELECT to_replace FROM bar)";
359/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql).unwrap();
360///
361/// // Drive the visitor through the AST
362/// statements.visit(&mut Replacer);
363///
364/// assert_eq!(statements[0].to_string(), "SELECT replaced FROM foo WHERE replaced IN (SELECT replaced FROM bar)");
365/// ```
366pub trait VisitorMut {
367    /// Type returned when the recursion returns early.
368    ///
369    /// Important note: The `Break` type should be kept as small as possible to prevent
370    /// stack overflow during recursion. If you need to return an error, consider
371    /// boxing it with `Box` to minimize stack usage.
372    type Break;
373
374    /// Invoked for any queries that appear in the AST before visiting children
375    fn pre_visit_query(&mut self, _query: &mut Query) -> ControlFlow<Self::Break> {
376        ControlFlow::Continue(())
377    }
378
379    /// Invoked for any queries that appear in the AST after visiting children
380    fn post_visit_query(&mut self, _query: &mut Query) -> ControlFlow<Self::Break> {
381        ControlFlow::Continue(())
382    }
383
384    /// Invoked for any [Select] that appear in the AST before visiting children
385    fn pre_visit_select(&mut self, _select: &mut Select) -> ControlFlow<Self::Break> {
386        ControlFlow::Continue(())
387    }
388
389    /// Invoked for any [Select] that appear in the AST after visiting children
390    fn post_visit_select(&mut self, _select: &mut Select) -> ControlFlow<Self::Break> {
391        ControlFlow::Continue(())
392    }
393
394    /// Invoked for any relations (e.g. tables) that appear in the AST before visiting children
395    fn pre_visit_relation(&mut self, _relation: &mut ObjectName) -> ControlFlow<Self::Break> {
396        ControlFlow::Continue(())
397    }
398
399    /// Invoked for any relations (e.g. tables) that appear in the AST after visiting children
400    fn post_visit_relation(&mut self, _relation: &mut ObjectName) -> ControlFlow<Self::Break> {
401        ControlFlow::Continue(())
402    }
403
404    /// Invoked for any table factors that appear in the AST before visiting children
405    fn pre_visit_table_factor(
406        &mut self,
407        _table_factor: &mut TableFactor,
408    ) -> ControlFlow<Self::Break> {
409        ControlFlow::Continue(())
410    }
411
412    /// Invoked for any table factors that appear in the AST after visiting children
413    fn post_visit_table_factor(
414        &mut self,
415        _table_factor: &mut TableFactor,
416    ) -> ControlFlow<Self::Break> {
417        ControlFlow::Continue(())
418    }
419
420    /// Invoked for any expressions that appear in the AST before visiting children
421    fn pre_visit_expr(&mut self, _expr: &mut Expr) -> ControlFlow<Self::Break> {
422        ControlFlow::Continue(())
423    }
424
425    /// Invoked for any expressions that appear in the AST
426    fn post_visit_expr(&mut self, _expr: &mut Expr) -> ControlFlow<Self::Break> {
427        ControlFlow::Continue(())
428    }
429
430    /// Invoked for any statements that appear in the AST before visiting children
431    fn pre_visit_statement(&mut self, _statement: &mut Statement) -> ControlFlow<Self::Break> {
432        ControlFlow::Continue(())
433    }
434
435    /// Invoked for any statements that appear in the AST after visiting children
436    fn post_visit_statement(&mut self, _statement: &mut Statement) -> ControlFlow<Self::Break> {
437        ControlFlow::Continue(())
438    }
439
440    /// Invoked for any value that appear in the AST before visiting children
441    fn pre_visit_value(&mut self, _value: &mut ValueWithSpan) -> ControlFlow<Self::Break> {
442        ControlFlow::Continue(())
443    }
444
445    /// Invoked for any statements that appear in the AST after visiting children
446    fn post_visit_value(&mut self, _value: &mut ValueWithSpan) -> ControlFlow<Self::Break> {
447        ControlFlow::Continue(())
448    }
449
450    /// Invoked for any identifiers that appear in the AST before visiting children
451    fn pre_visit_ident(&mut self, _ident: &mut Ident) -> ControlFlow<Self::Break> {
452        ControlFlow::Continue(())
453    }
454
455    /// Invoked for any identifiers that appear in the AST after visiting children
456    fn post_visit_ident(&mut self, _ident: &mut Ident) -> ControlFlow<Self::Break> {
457        ControlFlow::Continue(())
458    }
459
460    /// Invoked for any `ORDER BY` clauses that appear in the AST before visiting children
461    fn pre_visit_order_by(&mut self, _order_by: &mut OrderBy) -> ControlFlow<Self::Break> {
462        ControlFlow::Continue(())
463    }
464
465    /// Invoked for any `ORDER BY` clauses that appear in the AST after visiting children
466    fn post_visit_order_by(&mut self, _order_by: &mut OrderBy) -> ControlFlow<Self::Break> {
467        ControlFlow::Continue(())
468    }
469
470    /// Invoked for any `ORDER BY` expressions that appear in the AST before visiting children
471    fn pre_visit_order_by_expr(
472        &mut self,
473        _order_by_expr: &mut OrderByExpr,
474    ) -> ControlFlow<Self::Break> {
475        ControlFlow::Continue(())
476    }
477
478    /// Invoked for any `ORDER BY` expressions that appear in the AST after visiting children
479    fn post_visit_order_by_expr(
480        &mut self,
481        _order_by_expr: &mut OrderByExpr,
482    ) -> ControlFlow<Self::Break> {
483        ControlFlow::Continue(())
484    }
485
486    /// Invoked for any `GROUP BY` clauses that appear in the AST before visiting children
487    fn pre_visit_group_by(&mut self, _group_by: &mut GroupByExpr) -> ControlFlow<Self::Break> {
488        ControlFlow::Continue(())
489    }
490
491    /// Invoked for any `GROUP BY` clauses that appear in the AST after visiting children
492    fn post_visit_group_by(&mut self, _group_by: &mut GroupByExpr) -> ControlFlow<Self::Break> {
493        ControlFlow::Continue(())
494    }
495}
496
497struct RelationVisitor<F>(F);
498
499impl<E, F: FnMut(&ObjectName) -> ControlFlow<E>> Visitor for RelationVisitor<F> {
500    type Break = E;
501
502    fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
503        self.0(relation)
504    }
505}
506
507impl<E, F: FnMut(&mut ObjectName) -> ControlFlow<E>> VisitorMut for RelationVisitor<F> {
508    type Break = E;
509
510    fn post_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow<Self::Break> {
511        self.0(relation)
512    }
513}
514
515/// Invokes the provided closure on all relations (e.g. table names) present in `v`
516///
517/// # Example
518/// ```
519/// # use sqlparser::parser::Parser;
520/// # use sqlparser::dialect::GenericDialect;
521/// # use sqlparser::ast::{visit_relations};
522/// # use core::ops::ControlFlow;
523/// let sql = "SELECT a FROM foo where x IN (SELECT y FROM bar)";
524/// let statements = Parser::parse_sql(&GenericDialect{}, sql)
525///    .unwrap();
526///
527/// // visit statements, capturing relations (table names)
528/// let mut visited = vec![];
529/// visit_relations(&statements, |relation| {
530///   visited.push(format!("RELATION: {}", relation));
531///   ControlFlow::<()>::Continue(())
532/// });
533///
534/// let expected : Vec<_> = [
535///   "RELATION: foo",
536///   "RELATION: bar",
537/// ]
538///   .into_iter().map(|s| s.to_string()).collect();
539///
540/// assert_eq!(visited, expected);
541/// ```
542pub fn visit_relations<V, E, F>(v: &V, f: F) -> ControlFlow<E>
543where
544    V: Visit,
545    F: FnMut(&ObjectName) -> ControlFlow<E>,
546{
547    let mut visitor = RelationVisitor(f);
548    v.visit(&mut visitor)?;
549    ControlFlow::Continue(())
550}
551
552/// Invokes the provided closure with a mutable reference to all relations (e.g. table names)
553/// present in `v`.
554///
555/// When the closure mutates its argument, the new mutated relation will not be visited again.
556///
557/// # Example
558/// ```
559/// # use sqlparser::parser::Parser;
560/// # use sqlparser::dialect::GenericDialect;
561/// # use sqlparser::ast::{ObjectName, ObjectNamePart, Ident, visit_relations_mut};
562/// # use core::ops::ControlFlow;
563/// let sql = "SELECT a FROM foo";
564/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql)
565///    .unwrap();
566///
567/// // visit statements, renaming table foo to bar
568/// visit_relations_mut(&mut statements, |table| {
569///   table.0[0] = ObjectNamePart::Identifier(Ident::new("bar"));
570///   ControlFlow::<()>::Continue(())
571/// });
572///
573/// assert_eq!(statements[0].to_string(), "SELECT a FROM bar");
574/// ```
575pub fn visit_relations_mut<V, E, F>(v: &mut V, f: F) -> ControlFlow<E>
576where
577    V: VisitMut,
578    F: FnMut(&mut ObjectName) -> ControlFlow<E>,
579{
580    let mut visitor = RelationVisitor(f);
581    v.visit(&mut visitor)?;
582    ControlFlow::Continue(())
583}
584
585struct ExprVisitor<F>(F);
586
587impl<E, F: FnMut(&Expr) -> ControlFlow<E>> Visitor for ExprVisitor<F> {
588    type Break = E;
589
590    fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
591        self.0(expr)
592    }
593}
594
595impl<E, F: FnMut(&mut Expr) -> ControlFlow<E>> VisitorMut for ExprVisitor<F> {
596    type Break = E;
597
598    fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
599        self.0(expr)
600    }
601}
602
603/// Invokes the provided closure on all expressions (e.g. `1 + 2`) present in `v`
604///
605/// # Example
606/// ```
607/// # use sqlparser::parser::Parser;
608/// # use sqlparser::dialect::GenericDialect;
609/// # use sqlparser::ast::{visit_expressions};
610/// # use core::ops::ControlFlow;
611/// let sql = "SELECT a FROM foo where x IN (SELECT y FROM bar)";
612/// let statements = Parser::parse_sql(&GenericDialect{}, sql)
613///    .unwrap();
614///
615/// // visit all expressions
616/// let mut visited = vec![];
617/// visit_expressions(&statements, |expr| {
618///   visited.push(format!("EXPR: {}", expr));
619///   ControlFlow::<()>::Continue(())
620/// });
621///
622/// let expected : Vec<_> = [
623///   "EXPR: a",
624///   "EXPR: x IN (SELECT y FROM bar)",
625///   "EXPR: x",
626///   "EXPR: y",
627/// ]
628///   .into_iter().map(|s| s.to_string()).collect();
629///
630/// assert_eq!(visited, expected);
631/// ```
632pub fn visit_expressions<V, E, F>(v: &V, f: F) -> ControlFlow<E>
633where
634    V: Visit,
635    F: FnMut(&Expr) -> ControlFlow<E>,
636{
637    let mut visitor = ExprVisitor(f);
638    v.visit(&mut visitor)?;
639    ControlFlow::Continue(())
640}
641
642/// Invokes the provided closure iteratively with a mutable reference to all expressions
643/// present in `v`.
644///
645/// This performs a depth-first search, so if the closure mutates the expression
646///
647/// # Example
648///
649/// ## Remove all select limits in sub-queries
650/// ```
651/// # use sqlparser::parser::Parser;
652/// # use sqlparser::dialect::GenericDialect;
653/// # use sqlparser::ast::{Expr, visit_expressions_mut, visit_statements_mut};
654/// # use core::ops::ControlFlow;
655/// let sql = "SELECT (SELECT y FROM z LIMIT 9) FROM t LIMIT 3";
656/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql).unwrap();
657///
658/// // Remove all select limits in sub-queries
659/// visit_expressions_mut(&mut statements, |expr| {
660///   if let Expr::Subquery(q) = expr {
661///      q.limit_clause = None;
662///   }
663///   ControlFlow::<()>::Continue(())
664/// });
665///
666/// assert_eq!(statements[0].to_string(), "SELECT (SELECT y FROM z) FROM t LIMIT 3");
667/// ```
668///
669/// ## Wrap column name in function call
670///
671/// This demonstrates how to effectively replace an expression with another more complicated one
672/// that references the original. This example avoids unnecessary allocations by using the
673/// [`std::mem`] family of functions.
674///
675/// ```
676/// # use sqlparser::parser::Parser;
677/// # use sqlparser::dialect::GenericDialect;
678/// # use sqlparser::ast::*;
679/// # use core::ops::ControlFlow;
680/// let sql = "SELECT x, y FROM t";
681/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql).unwrap();
682///
683/// visit_expressions_mut(&mut statements, |expr| {
684///   if matches!(expr, Expr::Identifier(col_name) if col_name.value == "x") {
685///     let old_expr = std::mem::replace(expr, Expr::value(Value::Null));
686///     *expr = Expr::Function(Function {
687///           name: ObjectName::from(vec![Ident::new("f")]),
688///           uses_odbc_syntax: false,
689///           args: FunctionArguments::List(FunctionArgumentList {
690///               duplicate_treatment: None,
691///               args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(old_expr))],
692///               clauses: vec![],
693///           }),
694///           null_treatment: None,
695///           filter: None,
696///           over: None,
697///           parameters: FunctionArguments::None,
698///           within_group: vec![],
699///      });
700///   }
701///   ControlFlow::<()>::Continue(())
702/// });
703///
704/// assert_eq!(statements[0].to_string(), "SELECT f(x), y FROM t");
705/// ```
706pub fn visit_expressions_mut<V, E, F>(v: &mut V, f: F) -> ControlFlow<E>
707where
708    V: VisitMut,
709    F: FnMut(&mut Expr) -> ControlFlow<E>,
710{
711    v.visit(&mut ExprVisitor(f))?;
712    ControlFlow::Continue(())
713}
714
715struct StatementVisitor<F>(F);
716
717impl<E, F: FnMut(&Statement) -> ControlFlow<E>> Visitor for StatementVisitor<F> {
718    type Break = E;
719
720    fn pre_visit_statement(&mut self, statement: &Statement) -> ControlFlow<Self::Break> {
721        self.0(statement)
722    }
723}
724
725impl<E, F: FnMut(&mut Statement) -> ControlFlow<E>> VisitorMut for StatementVisitor<F> {
726    type Break = E;
727
728    fn post_visit_statement(&mut self, statement: &mut Statement) -> ControlFlow<Self::Break> {
729        self.0(statement)
730    }
731}
732
733/// Invokes the provided closure iteratively with a mutable reference to all statements
734/// present in `v` (e.g. `SELECT`, `CREATE TABLE`, etc).
735///
736/// # Example
737/// ```
738/// # use sqlparser::parser::Parser;
739/// # use sqlparser::dialect::GenericDialect;
740/// # use sqlparser::ast::{visit_statements};
741/// # use core::ops::ControlFlow;
742/// let sql = "SELECT a FROM foo where x IN (SELECT y FROM bar); CREATE TABLE baz(q int)";
743/// let statements = Parser::parse_sql(&GenericDialect{}, sql)
744///    .unwrap();
745///
746/// // visit all statements
747/// let mut visited = vec![];
748/// visit_statements(&statements, |stmt| {
749///   visited.push(format!("STATEMENT: {}", stmt));
750///   ControlFlow::<()>::Continue(())
751/// });
752///
753/// let expected : Vec<_> = [
754///   "STATEMENT: SELECT a FROM foo WHERE x IN (SELECT y FROM bar)",
755///   "STATEMENT: CREATE TABLE baz (q INT)"
756/// ]
757///   .into_iter().map(|s| s.to_string()).collect();
758///
759/// assert_eq!(visited, expected);
760/// ```
761pub fn visit_statements<V, E, F>(v: &V, f: F) -> ControlFlow<E>
762where
763    V: Visit,
764    F: FnMut(&Statement) -> ControlFlow<E>,
765{
766    let mut visitor = StatementVisitor(f);
767    v.visit(&mut visitor)?;
768    ControlFlow::Continue(())
769}
770
771/// Invokes the provided closure on all statements (e.g. `SELECT`, `CREATE TABLE`, etc) present in `v`
772///
773/// # Example
774/// ```
775/// # use sqlparser::parser::Parser;
776/// # use sqlparser::dialect::GenericDialect;
777/// # use sqlparser::ast::{Statement, visit_statements_mut};
778/// # use core::ops::ControlFlow;
779/// let sql = "SELECT x FROM foo LIMIT 9+$limit; SELECT * FROM t LIMIT f()";
780/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql).unwrap();
781///
782/// // Remove all select limits in outer statements (not in sub-queries)
783/// visit_statements_mut(&mut statements, |stmt| {
784///   if let Statement::Query(q) = stmt {
785///      q.limit_clause = None;
786///   }
787///   ControlFlow::<()>::Continue(())
788/// });
789///
790/// assert_eq!(statements[0].to_string(), "SELECT x FROM foo");
791/// assert_eq!(statements[1].to_string(), "SELECT * FROM t");
792/// ```
793pub fn visit_statements_mut<V, E, F>(v: &mut V, f: F) -> ControlFlow<E>
794where
795    V: VisitMut,
796    F: FnMut(&mut Statement) -> ControlFlow<E>,
797{
798    v.visit(&mut StatementVisitor(f))?;
799    ControlFlow::Continue(())
800}
801
802#[cfg(test)]
803mod tests {
804    use super::*;
805    use crate::ast::Statement;
806    use crate::dialect::GenericDialect;
807    use crate::parser::Parser;
808    use crate::tokenizer::Tokenizer;
809
810    #[derive(Default)]
811    struct TestVisitor {
812        visited: Vec<String>,
813    }
814
815    impl Visitor for TestVisitor {
816        type Break = ();
817
818        /// Invoked for any queries that appear in the AST before visiting children
819        fn pre_visit_query(&mut self, query: &Query) -> ControlFlow<Self::Break> {
820            self.visited.push(format!("PRE: QUERY: {query}"));
821            ControlFlow::Continue(())
822        }
823
824        /// Invoked for any queries that appear in the AST after visiting children
825        fn post_visit_query(&mut self, query: &Query) -> ControlFlow<Self::Break> {
826            self.visited.push(format!("POST: QUERY: {query}"));
827            ControlFlow::Continue(())
828        }
829
830        fn pre_visit_select(&mut self, select: &Select) -> ControlFlow<Self::Break> {
831            self.visited.push(format!("PRE: SELECT: {select}"));
832            ControlFlow::Continue(())
833        }
834
835        fn post_visit_select(&mut self, select: &Select) -> ControlFlow<Self::Break> {
836            self.visited.push(format!("POST: SELECT: {select}"));
837            ControlFlow::Continue(())
838        }
839
840        fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
841            self.visited.push(format!("PRE: RELATION: {relation}"));
842            ControlFlow::Continue(())
843        }
844
845        fn post_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
846            self.visited.push(format!("POST: RELATION: {relation}"));
847            ControlFlow::Continue(())
848        }
849
850        fn pre_visit_table_factor(
851            &mut self,
852            table_factor: &TableFactor,
853        ) -> ControlFlow<Self::Break> {
854            self.visited
855                .push(format!("PRE: TABLE FACTOR: {table_factor}"));
856            ControlFlow::Continue(())
857        }
858
859        fn post_visit_table_factor(
860            &mut self,
861            table_factor: &TableFactor,
862        ) -> ControlFlow<Self::Break> {
863            self.visited
864                .push(format!("POST: TABLE FACTOR: {table_factor}"));
865            ControlFlow::Continue(())
866        }
867
868        fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
869            self.visited.push(format!("PRE: EXPR: {expr}"));
870            ControlFlow::Continue(())
871        }
872
873        fn post_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
874            self.visited.push(format!("POST: EXPR: {expr}"));
875            ControlFlow::Continue(())
876        }
877
878        fn pre_visit_statement(&mut self, statement: &Statement) -> ControlFlow<Self::Break> {
879            self.visited.push(format!("PRE: STATEMENT: {statement}"));
880            ControlFlow::Continue(())
881        }
882
883        fn post_visit_statement(&mut self, statement: &Statement) -> ControlFlow<Self::Break> {
884            self.visited.push(format!("POST: STATEMENT: {statement}"));
885            ControlFlow::Continue(())
886        }
887
888        fn pre_visit_order_by(&mut self, order_by: &OrderBy) -> ControlFlow<Self::Break> {
889            self.visited.push(format!("PRE: ORDER BY: {order_by}"));
890            ControlFlow::Continue(())
891        }
892
893        fn post_visit_order_by(&mut self, order_by: &OrderBy) -> ControlFlow<Self::Break> {
894            self.visited.push(format!("POST: ORDER BY: {order_by}"));
895            ControlFlow::Continue(())
896        }
897
898        fn pre_visit_order_by_expr(
899            &mut self,
900            order_by_expr: &OrderByExpr,
901        ) -> ControlFlow<Self::Break> {
902            self.visited
903                .push(format!("PRE: ORDER BY EXPR: {order_by_expr}"));
904            ControlFlow::Continue(())
905        }
906
907        fn post_visit_order_by_expr(
908            &mut self,
909            order_by_expr: &OrderByExpr,
910        ) -> ControlFlow<Self::Break> {
911            self.visited
912                .push(format!("POST: ORDER BY EXPR: {order_by_expr}"));
913            ControlFlow::Continue(())
914        }
915
916        fn pre_visit_group_by(&mut self, group_by: &GroupByExpr) -> ControlFlow<Self::Break> {
917            self.visited.push(format!("PRE: GROUP BY: {group_by}"));
918            ControlFlow::Continue(())
919        }
920
921        fn post_visit_group_by(&mut self, group_by: &GroupByExpr) -> ControlFlow<Self::Break> {
922            self.visited.push(format!("POST: GROUP BY: {group_by}"));
923            ControlFlow::Continue(())
924        }
925    }
926
927    fn do_visit<V: Visitor<Break = ()>>(sql: &str, visitor: &mut V) -> Statement {
928        let dialect = GenericDialect {};
929        let tokens = Tokenizer::new(&dialect, sql).tokenize().unwrap();
930        let s = Parser::new(&dialect)
931            .with_tokens(tokens)
932            .parse_statement()
933            .unwrap();
934
935        let flow = s.visit(visitor);
936        assert_eq!(flow, ControlFlow::Continue(()));
937        s
938    }
939
940    #[test]
941    fn test_sql() {
942        let tests = vec![
943            (
944                "SELECT * from table_name as my_table",
945                vec![
946                    "PRE: STATEMENT: SELECT * FROM table_name AS my_table",
947                    "PRE: QUERY: SELECT * FROM table_name AS my_table",
948                    "PRE: SELECT: SELECT * FROM table_name AS my_table",
949                    "PRE: TABLE FACTOR: table_name AS my_table",
950                    "PRE: RELATION: table_name",
951                    "POST: RELATION: table_name",
952                    "POST: TABLE FACTOR: table_name AS my_table",
953                    "PRE: GROUP BY: GROUP BY ",
954                    "POST: GROUP BY: GROUP BY ",
955                    "POST: SELECT: SELECT * FROM table_name AS my_table",
956                    "POST: QUERY: SELECT * FROM table_name AS my_table",
957                    "POST: STATEMENT: SELECT * FROM table_name AS my_table",
958                ],
959            ),
960            (
961                "SELECT * from t1 join t2 on t1.id = t2.t1_id",
962                vec![
963                    "PRE: STATEMENT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
964                    "PRE: QUERY: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
965                    "PRE: SELECT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
966                    "PRE: TABLE FACTOR: t1",
967                    "PRE: RELATION: t1",
968                    "POST: RELATION: t1",
969                    "POST: TABLE FACTOR: t1",
970                    "PRE: TABLE FACTOR: t2",
971                    "PRE: RELATION: t2",
972                    "POST: RELATION: t2",
973                    "POST: TABLE FACTOR: t2",
974                    "PRE: EXPR: t1.id = t2.t1_id",
975                    "PRE: EXPR: t1.id",
976                    "POST: EXPR: t1.id",
977                    "PRE: EXPR: t2.t1_id",
978                    "POST: EXPR: t2.t1_id",
979                    "POST: EXPR: t1.id = t2.t1_id",
980                    "PRE: GROUP BY: GROUP BY ",
981                    "POST: GROUP BY: GROUP BY ",
982                    "POST: SELECT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
983                    "POST: QUERY: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
984                    "POST: STATEMENT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
985                ],
986            ),
987            (
988                "SELECT * from t1 where EXISTS(SELECT column from t2)",
989                vec![
990                    "PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
991                    "PRE: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
992                    "PRE: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
993                    "PRE: TABLE FACTOR: t1",
994                    "PRE: RELATION: t1",
995                    "POST: RELATION: t1",
996                    "POST: TABLE FACTOR: t1",
997                    "PRE: EXPR: EXISTS (SELECT column FROM t2)",
998                    "PRE: QUERY: SELECT column FROM t2",
999                    "PRE: SELECT: SELECT column FROM t2",
1000                    "PRE: EXPR: column",
1001                    "POST: EXPR: column",
1002                    "PRE: TABLE FACTOR: t2",
1003                    "PRE: RELATION: t2",
1004                    "POST: RELATION: t2",
1005                    "POST: TABLE FACTOR: t2",
1006                    "PRE: GROUP BY: GROUP BY ",
1007                    "POST: GROUP BY: GROUP BY ",
1008                    "POST: SELECT: SELECT column FROM t2",
1009                    "POST: QUERY: SELECT column FROM t2",
1010                    "POST: EXPR: EXISTS (SELECT column FROM t2)",
1011                    "PRE: GROUP BY: GROUP BY ",
1012                    "POST: GROUP BY: GROUP BY ",
1013                    "POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1014                    "POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1015                    "POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1016                ],
1017            ),
1018            (
1019                "SELECT * from t1 where EXISTS(SELECT column from t2)",
1020                vec![
1021                    "PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1022                    "PRE: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1023                    "PRE: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1024                    "PRE: TABLE FACTOR: t1",
1025                    "PRE: RELATION: t1",
1026                    "POST: RELATION: t1",
1027                    "POST: TABLE FACTOR: t1",
1028                    "PRE: EXPR: EXISTS (SELECT column FROM t2)",
1029                    "PRE: QUERY: SELECT column FROM t2",
1030                    "PRE: SELECT: SELECT column FROM t2",
1031                    "PRE: EXPR: column",
1032                    "POST: EXPR: column",
1033                    "PRE: TABLE FACTOR: t2",
1034                    "PRE: RELATION: t2",
1035                    "POST: RELATION: t2",
1036                    "POST: TABLE FACTOR: t2",
1037                    "PRE: GROUP BY: GROUP BY ",
1038                    "POST: GROUP BY: GROUP BY ",
1039                    "POST: SELECT: SELECT column FROM t2",
1040                    "POST: QUERY: SELECT column FROM t2",
1041                    "POST: EXPR: EXISTS (SELECT column FROM t2)",
1042                    "PRE: GROUP BY: GROUP BY ",
1043                    "POST: GROUP BY: GROUP BY ",
1044                    "POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1045                    "POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1046                    "POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1047                ],
1048            ),
1049            (
1050                "SELECT * from t1 where EXISTS(SELECT column from t2) UNION SELECT * from t3",
1051                vec![
1052                    "PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
1053                    "PRE: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
1054                    "PRE: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1055                    "PRE: TABLE FACTOR: t1",
1056                    "PRE: RELATION: t1",
1057                    "POST: RELATION: t1",
1058                    "POST: TABLE FACTOR: t1",
1059                    "PRE: EXPR: EXISTS (SELECT column FROM t2)",
1060                    "PRE: QUERY: SELECT column FROM t2",
1061                    "PRE: SELECT: SELECT column FROM t2",
1062                    "PRE: EXPR: column",
1063                    "POST: EXPR: column",
1064                    "PRE: TABLE FACTOR: t2",
1065                    "PRE: RELATION: t2",
1066                    "POST: RELATION: t2",
1067                    "POST: TABLE FACTOR: t2",
1068                    "PRE: GROUP BY: GROUP BY ",
1069                    "POST: GROUP BY: GROUP BY ",
1070                    "POST: SELECT: SELECT column FROM t2",
1071                    "POST: QUERY: SELECT column FROM t2",
1072                    "POST: EXPR: EXISTS (SELECT column FROM t2)",
1073                    "PRE: GROUP BY: GROUP BY ",
1074                    "POST: GROUP BY: GROUP BY ",
1075                    "POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
1076                    "PRE: SELECT: SELECT * FROM t3",
1077                    "PRE: TABLE FACTOR: t3",
1078                    "PRE: RELATION: t3",
1079                    "POST: RELATION: t3",
1080                    "POST: TABLE FACTOR: t3",
1081                    "PRE: GROUP BY: GROUP BY ",
1082                    "POST: GROUP BY: GROUP BY ",
1083                    "POST: SELECT: SELECT * FROM t3",
1084                    "POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
1085                    "POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
1086                ],
1087            ),
1088            (
1089                concat!(
1090                    "SELECT * FROM monthly_sales ",
1091                    "PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
1092                    "ORDER BY EMPID"
1093                ),
1094                vec![
1095                    "PRE: STATEMENT: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ORDER BY EMPID",
1096                    "PRE: QUERY: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ORDER BY EMPID",
1097                    "PRE: SELECT: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
1098                    "PRE: TABLE FACTOR: monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
1099                    "PRE: TABLE FACTOR: monthly_sales",
1100                    "PRE: RELATION: monthly_sales",
1101                    "POST: RELATION: monthly_sales",
1102                    "POST: TABLE FACTOR: monthly_sales",
1103                    "PRE: EXPR: SUM(a.amount)",
1104                    "PRE: EXPR: a.amount",
1105                    "POST: EXPR: a.amount",
1106                    "POST: EXPR: SUM(a.amount)",
1107                    "PRE: EXPR: a.MONTH",
1108                    "POST: EXPR: a.MONTH",
1109                    "PRE: EXPR: 'JAN'",
1110                    "POST: EXPR: 'JAN'",
1111                    "PRE: EXPR: 'FEB'",
1112                    "POST: EXPR: 'FEB'",
1113                    "PRE: EXPR: 'MAR'",
1114                    "POST: EXPR: 'MAR'",
1115                    "PRE: EXPR: 'APR'",
1116                    "POST: EXPR: 'APR'",
1117                    "POST: TABLE FACTOR: monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
1118                    "PRE: GROUP BY: GROUP BY ",
1119                    "POST: GROUP BY: GROUP BY ",
1120                    "POST: SELECT: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
1121                    "PRE: ORDER BY: ORDER BY EMPID",
1122                    "PRE: ORDER BY EXPR: EMPID",
1123                    "PRE: EXPR: EMPID",
1124                    "POST: EXPR: EMPID",
1125                    "POST: ORDER BY EXPR: EMPID",
1126                    "POST: ORDER BY: ORDER BY EMPID",
1127                    "POST: QUERY: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ORDER BY EMPID",
1128                    "POST: STATEMENT: SELECT * FROM monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ORDER BY EMPID",
1129                ]
1130            ),
1131            (
1132                "SHOW COLUMNS FROM t1",
1133                vec![
1134                    "PRE: STATEMENT: SHOW COLUMNS FROM t1",
1135                    "PRE: RELATION: t1",
1136                    "POST: RELATION: t1",
1137                    "POST: STATEMENT: SHOW COLUMNS FROM t1",
1138                ],
1139            ),
1140            (
1141                "SELECT * FROM t1 ORDER BY a DESC, b",
1142                vec![
1143                    "PRE: STATEMENT: SELECT * FROM t1 ORDER BY a DESC, b",
1144                    "PRE: QUERY: SELECT * FROM t1 ORDER BY a DESC, b",
1145                    "PRE: SELECT: SELECT * FROM t1",
1146                    "PRE: TABLE FACTOR: t1",
1147                    "PRE: RELATION: t1",
1148                    "POST: RELATION: t1",
1149                    "POST: TABLE FACTOR: t1",
1150                    "PRE: GROUP BY: GROUP BY ",
1151                    "POST: GROUP BY: GROUP BY ",
1152                    "POST: SELECT: SELECT * FROM t1",
1153                    "PRE: ORDER BY: ORDER BY a DESC, b",
1154                    "PRE: ORDER BY EXPR: a DESC",
1155                    "PRE: EXPR: a",
1156                    "POST: EXPR: a",
1157                    "POST: ORDER BY EXPR: a DESC",
1158                    "PRE: ORDER BY EXPR: b",
1159                    "PRE: EXPR: b",
1160                    "POST: EXPR: b",
1161                    "POST: ORDER BY EXPR: b",
1162                    "POST: ORDER BY: ORDER BY a DESC, b",
1163                    "POST: QUERY: SELECT * FROM t1 ORDER BY a DESC, b",
1164                    "POST: STATEMENT: SELECT * FROM t1 ORDER BY a DESC, b",
1165                ],
1166            ),
1167            (
1168                "SELECT a FROM t GROUP BY a, b",
1169                vec![
1170                    "PRE: STATEMENT: SELECT a FROM t GROUP BY a, b",
1171                    "PRE: QUERY: SELECT a FROM t GROUP BY a, b",
1172                    "PRE: SELECT: SELECT a FROM t GROUP BY a, b",
1173                    "PRE: EXPR: a",
1174                    "POST: EXPR: a",
1175                    "PRE: TABLE FACTOR: t",
1176                    "PRE: RELATION: t",
1177                    "POST: RELATION: t",
1178                    "POST: TABLE FACTOR: t",
1179                    "PRE: GROUP BY: GROUP BY a, b",
1180                    "PRE: EXPR: a",
1181                    "POST: EXPR: a",
1182                    "PRE: EXPR: b",
1183                    "POST: EXPR: b",
1184                    "POST: GROUP BY: GROUP BY a, b",
1185                    "POST: SELECT: SELECT a FROM t GROUP BY a, b",
1186                    "POST: QUERY: SELECT a FROM t GROUP BY a, b",
1187                    "POST: STATEMENT: SELECT a FROM t GROUP BY a, b",
1188                ],
1189            ),
1190        ];
1191        for (sql, expected) in tests {
1192            let mut visitor = TestVisitor::default();
1193            let _ = do_visit(sql, &mut visitor);
1194            let actual: Vec<_> = visitor.visited.iter().map(|x| x.as_str()).collect();
1195            assert_eq!(actual, expected)
1196        }
1197    }
1198
1199    struct QuickVisitor; // [`TestVisitor`] is too slow to iterate over thousands of nodes
1200
1201    impl Visitor for QuickVisitor {
1202        type Break = ();
1203    }
1204
1205    #[test]
1206    fn overflow() {
1207        let cond = (0..1000)
1208            .map(|n| format!("X = {n}"))
1209            .collect::<Vec<_>>()
1210            .join(" OR ");
1211        let sql = format!("SELECT x where {cond}");
1212
1213        let dialect = GenericDialect {};
1214        let tokens = Tokenizer::new(&dialect, sql.as_str()).tokenize().unwrap();
1215        let s = Parser::new(&dialect)
1216            .with_tokens(tokens)
1217            .parse_statement()
1218            .unwrap();
1219
1220        let mut visitor = QuickVisitor {};
1221        let flow = s.visit(&mut visitor);
1222        assert_eq!(flow, ControlFlow::Continue(()));
1223    }
1224
1225    #[derive(Default)]
1226    struct IdentVisitor {
1227        idents: Vec<String>,
1228    }
1229
1230    impl Visitor for IdentVisitor {
1231        type Break = ();
1232
1233        fn pre_visit_ident(&mut self, ident: &Ident) -> ControlFlow<Self::Break> {
1234            self.idents.push(ident.value.clone());
1235            ControlFlow::Continue(())
1236        }
1237    }
1238
1239    #[test]
1240    fn test_pre_visit_ident() {
1241        let mut visitor = IdentVisitor::default();
1242        do_visit("SELECT a, b FROM t", &mut visitor);
1243        assert_eq!(visitor.idents, vec!["a", "b", "t"]);
1244    }
1245
1246    #[test]
1247    fn visits_function_clauses_in_source_order() {
1248        #[derive(Default)]
1249        struct ExprVisitor {
1250            idents: Vec<String>,
1251        }
1252
1253        impl Visitor for ExprVisitor {
1254            type Break = ();
1255
1256            fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
1257                if let Expr::Identifier(ident) = expr {
1258                    self.idents.push(ident.value.clone());
1259                }
1260                ControlFlow::Continue(())
1261            }
1262        }
1263
1264        let mut visitor = ExprVisitor::default();
1265        do_visit(
1266            "SELECT LISTAGG(value) WITHIN GROUP (ORDER BY order_key) FILTER (WHERE filter_key)",
1267            &mut visitor,
1268        );
1269        assert_eq!(visitor.idents, ["value", "order_key", "filter_key"]);
1270    }
1271
1272    #[derive(Default)]
1273    struct RelationVisitor {
1274        relations: Vec<String>,
1275    }
1276
1277    impl Visitor for RelationVisitor {
1278        type Break = ();
1279
1280        fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
1281            self.relations.push(relation.to_string());
1282            ControlFlow::Continue(())
1283        }
1284    }
1285
1286    #[test]
1287    fn test_visit_create_view_name_as_relation() {
1288        let mut visitor = RelationVisitor::default();
1289        do_visit("CREATE VIEW db1.v AS SELECT * FROM t", &mut visitor);
1290        assert_eq!(visitor.relations, vec!["db1.v", "t"]);
1291    }
1292}
1293
1294#[cfg(test)]
1295mod visit_mut_tests {
1296    use crate::ast::{Ident, Statement, Value, ValueWithSpan, VisitMut, VisitorMut};
1297    use crate::dialect::GenericDialect;
1298    use crate::parser::Parser;
1299    use crate::tokenizer::Tokenizer;
1300    use core::ops::ControlFlow;
1301
1302    #[derive(Default)]
1303    struct MutatorVisitor {
1304        index: u64,
1305    }
1306
1307    impl VisitorMut for MutatorVisitor {
1308        type Break = ();
1309
1310        fn pre_visit_value(&mut self, value: &mut ValueWithSpan) -> ControlFlow<Self::Break> {
1311            self.index += 1;
1312            value.value = Value::SingleQuotedString(format!("REDACTED_{}", self.index));
1313            ControlFlow::Continue(())
1314        }
1315
1316        fn post_visit_value(&mut self, _value: &mut ValueWithSpan) -> ControlFlow<Self::Break> {
1317            ControlFlow::Continue(())
1318        }
1319    }
1320
1321    fn do_visit_mut<V: VisitorMut<Break = ()>>(sql: &str, visitor: &mut V) -> Statement {
1322        let dialect = GenericDialect {};
1323        let tokens = Tokenizer::new(&dialect, sql).tokenize().unwrap();
1324        let mut s = Parser::new(&dialect)
1325            .with_tokens(tokens)
1326            .parse_statement()
1327            .unwrap();
1328
1329        let flow = s.visit(visitor);
1330        assert_eq!(flow, ControlFlow::Continue(()));
1331        s
1332    }
1333
1334    #[test]
1335    fn test_value_redact() {
1336        let tests = vec![
1337            (
1338                concat!(
1339                    "SELECT * FROM monthly_sales ",
1340                    "PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
1341                    "ORDER BY EMPID"
1342                ),
1343                concat!(
1344                    "SELECT * FROM monthly_sales ",
1345                    "PIVOT(SUM(a.amount) FOR a.MONTH IN ('REDACTED_1', 'REDACTED_2', 'REDACTED_3', 'REDACTED_4')) AS p (c, d) ",
1346                    "ORDER BY EMPID"
1347                ),
1348            ),
1349        ];
1350
1351        for (sql, expected) in tests {
1352            let mut visitor = MutatorVisitor::default();
1353            let mutated = do_visit_mut(sql, &mut visitor);
1354            assert_eq!(mutated.to_string(), expected)
1355        }
1356    }
1357
1358    #[derive(Default)]
1359    struct IdentMutator;
1360
1361    impl VisitorMut for IdentMutator {
1362        type Break = ();
1363
1364        fn pre_visit_ident(&mut self, ident: &mut Ident) -> ControlFlow<Self::Break> {
1365            ident.value = ident.value.to_uppercase();
1366            ControlFlow::Continue(())
1367        }
1368    }
1369
1370    #[test]
1371    fn test_pre_visit_ident_mut() {
1372        let mut visitor = IdentMutator;
1373        let mutated = do_visit_mut("SELECT a, b FROM t", &mut visitor);
1374        assert_eq!(mutated.to_string(), "SELECT A, B FROM T");
1375    }
1376}
sqlparser/ast/visitor.rs

sqlparser/ast/
visitor.rs
