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

sqlparser/ast/
visitor.rs
