Skip to content

Commit b2ff6ae

Browse files
rs-sacrs-sac
authored
Add visitors for ORDER BY and GROUP BY (#2406)
1 parent 57a62ae commit b2ff6ae

2 files changed

Lines changed: 191 additions & 1 deletion

File tree

‎src/ast/query.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2888,6 +2888,7 @@ pub enum OrderByKind {
28882888
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
28892889
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
28902890
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
2891+
#[cfg_attr(feature = "visitor", visit(with = "visit_order_by"))]
28912892
/// Represents an `ORDER BY` clause with its kind and optional `INTERPOLATE`.
28922893
pub struct OrderBy {
28932894
/// The kind of ordering (expressions or `ALL`).
@@ -2924,6 +2925,7 @@ impl fmt::Display for OrderBy {
29242925
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
29252926
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
29262927
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
2928+
#[cfg_attr(feature = "visitor", visit(with = "visit_order_by_expr"))]
29272929
pub struct OrderByExpr {
29282930
/// The expression to order by.
29292931
pub expr: Expr,
@@ -3776,6 +3778,7 @@ impl fmt::Display for GroupByWithModifier {
37763778
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
37773779
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
37783780
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
3781+
#[cfg_attr(feature = "visitor", visit(with = "visit_group_by"))]
37793782
/// Represents the two syntactic forms that `GROUP BY` can take, including
37803783
/// `GROUP BY ALL` with optional modifiers and ordinary `GROUP BY <exprs>`.
37813784
pub enum GroupByExpr {

‎src/ast/visitor.rs‎

Lines changed: 188 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@
2121
use alloc::{boxed::Box, string::String, vec::Vec};
2222
use core::ops::ControlFlow;
2323

24-
use crate::ast::{Expr, Ident, ObjectName, Query, Select, Statement, TableFactor, ValueWithSpan};
24+
use crate::ast::{
25+
Expr, GroupByExpr, Ident, ObjectName, OrderBy, OrderByExpr, Query, Select, Statement,
26+
TableFactor, ValueWithSpan,
27+
};
2528

2629
/// A type that can be visited by a [`Visitor`]. See [`Visitor`] for
2730
/// recursively visiting parsed SQL statements.
@@ -279,6 +282,42 @@ pub trait Visitor {
279282
fn post_visit_ident(&mut self, _ident: &Ident) -> ControlFlow<Self::Break> {
280283
ControlFlow::Continue(())
281284
}
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+
}
282321
}
283322

284323
/// A visitor that can be used to mutate an AST tree.
@@ -417,6 +456,42 @@ pub trait VisitorMut {
417456
fn post_visit_ident(&mut self, _ident: &mut Ident) -> ControlFlow<Self::Break> {
418457
ControlFlow::Continue(())
419458
}
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+
}
420495
}
421496

422497
struct RelationVisitor<F>(F);
@@ -809,6 +884,44 @@ mod tests {
809884
self.visited.push(format!("POST: STATEMENT: {statement}"));
810885
ControlFlow::Continue(())
811886
}
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+
}
812925
}
813926

814927
fn do_visit<V: Visitor<Break = ()>>(sql: &str, visitor: &mut V) -> Statement {
@@ -837,6 +950,8 @@ mod tests {
837950
"PRE: RELATION: table_name",
838951
"POST: RELATION: table_name",
839952
"POST: TABLE FACTOR: table_name AS my_table",
953+
"PRE: GROUP BY: GROUP BY ",
954+
"POST: GROUP BY: GROUP BY ",
840955
"POST: SELECT: SELECT * FROM table_name AS my_table",
841956
"POST: QUERY: SELECT * FROM table_name AS my_table",
842957
"POST: STATEMENT: SELECT * FROM table_name AS my_table",
@@ -862,6 +977,8 @@ mod tests {
862977
"PRE: EXPR: t2.t1_id",
863978
"POST: EXPR: t2.t1_id",
864979
"POST: EXPR: t1.id = t2.t1_id",
980+
"PRE: GROUP BY: GROUP BY ",
981+
"POST: GROUP BY: GROUP BY ",
865982
"POST: SELECT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
866983
"POST: QUERY: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
867984
"POST: STATEMENT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
@@ -886,9 +1003,13 @@ mod tests {
8861003
"PRE: RELATION: t2",
8871004
"POST: RELATION: t2",
8881005
"POST: TABLE FACTOR: t2",
1006+
"PRE: GROUP BY: GROUP BY ",
1007+
"POST: GROUP BY: GROUP BY ",
8891008
"POST: SELECT: SELECT column FROM t2",
8901009
"POST: QUERY: SELECT column FROM t2",
8911010
"POST: EXPR: EXISTS (SELECT column FROM t2)",
1011+
"PRE: GROUP BY: GROUP BY ",
1012+
"POST: GROUP BY: GROUP BY ",
8921013
"POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
8931014
"POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
8941015
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
@@ -913,9 +1034,13 @@ mod tests {
9131034
"PRE: RELATION: t2",
9141035
"POST: RELATION: t2",
9151036
"POST: TABLE FACTOR: t2",
1037+
"PRE: GROUP BY: GROUP BY ",
1038+
"POST: GROUP BY: GROUP BY ",
9161039
"POST: SELECT: SELECT column FROM t2",
9171040
"POST: QUERY: SELECT column FROM t2",
9181041
"POST: EXPR: EXISTS (SELECT column FROM t2)",
1042+
"PRE: GROUP BY: GROUP BY ",
1043+
"POST: GROUP BY: GROUP BY ",
9191044
"POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
9201045
"POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
9211046
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
@@ -940,15 +1065,21 @@ mod tests {
9401065
"PRE: RELATION: t2",
9411066
"POST: RELATION: t2",
9421067
"POST: TABLE FACTOR: t2",
1068+
"PRE: GROUP BY: GROUP BY ",
1069+
"POST: GROUP BY: GROUP BY ",
9431070
"POST: SELECT: SELECT column FROM t2",
9441071
"POST: QUERY: SELECT column FROM t2",
9451072
"POST: EXPR: EXISTS (SELECT column FROM t2)",
1073+
"PRE: GROUP BY: GROUP BY ",
1074+
"POST: GROUP BY: GROUP BY ",
9461075
"POST: SELECT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
9471076
"PRE: SELECT: SELECT * FROM t3",
9481077
"PRE: TABLE FACTOR: t3",
9491078
"PRE: RELATION: t3",
9501079
"POST: RELATION: t3",
9511080
"POST: TABLE FACTOR: t3",
1081+
"PRE: GROUP BY: GROUP BY ",
1082+
"POST: GROUP BY: GROUP BY ",
9521083
"POST: SELECT: SELECT * FROM t3",
9531084
"POST: QUERY: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
9541085
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
@@ -984,9 +1115,15 @@ mod tests {
9841115
"PRE: EXPR: 'APR'",
9851116
"POST: EXPR: 'APR'",
9861117
"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 ",
9871120
"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",
9881123
"PRE: EXPR: EMPID",
9891124
"POST: EXPR: EMPID",
1125+
"POST: ORDER BY EXPR: EMPID",
1126+
"POST: ORDER BY: ORDER BY EMPID",
9901127
"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",
9911128
"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",
9921129
]
@@ -1000,6 +1137,56 @@ mod tests {
10001137
"POST: STATEMENT: SHOW COLUMNS FROM t1",
10011138
],
10021139
),
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+
),
10031190
];
10041191
for (sql, expected) in tests {
10051192
let mut visitor = TestVisitor::default();

0 commit comments

Comments
 (0)