Skip to content

Commit 1ea8858

Browse files
gruuyagruuya
authored
Table time travel clause support, add visit_table_factor to Visitor (#951)
1 parent 9500649 commit 1ea8858

15 files changed

+225
-12
lines changed

‎src/ast/mod.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ pub use self::query::{
4040
JoinConstraint, JoinOperator, LateralView, LockClause, LockType, NamedWindowDefinition,
4141
NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem, ReplaceSelectElement,
4242
ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table,
43-
TableAlias, TableFactor, TableWithJoins, Top, Values, WildcardAdditionalOptions, With,
43+
TableAlias, TableFactor, TableVersion, TableWithJoins, Top, Values, WildcardAdditionalOptions,
44+
With,
4445
};
4546
pub use self::value::{
4647
escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,

‎src/ast/query.rs

+24
Original file line numberDiff line numberDiff line change
@@ -646,6 +646,7 @@ impl fmt::Display for TableWithJoins {
646646
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
647647
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
648648
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
649+
#[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
649650
pub enum TableFactor {
650651
Table {
651652
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
@@ -661,6 +662,9 @@ pub enum TableFactor {
661662
args: Option<Vec<FunctionArg>>,
662663
/// MSSQL-specific `WITH (...)` hints such as NOLOCK.
663664
with_hints: Vec<Expr>,
665+
/// Optional version qualifier to facilitate table time-travel, as
666+
/// supported by BigQuery and MSSQL.
667+
version: Option<TableVersion>,
664668
},
665669
Derived {
666670
lateral: bool,
@@ -720,6 +724,7 @@ impl fmt::Display for TableFactor {
720724
alias,
721725
args,
722726
with_hints,
727+
version,
723728
} => {
724729
write!(f, "{name}")?;
725730
if let Some(args) = args {
@@ -731,6 +736,9 @@ impl fmt::Display for TableFactor {
731736
if !with_hints.is_empty() {
732737
write!(f, " WITH ({})", display_comma_separated(with_hints))?;
733738
}
739+
if let Some(version) = version {
740+
write!(f, "{version}")?;
741+
}
734742
Ok(())
735743
}
736744
TableFactor::Derived {
@@ -835,6 +843,22 @@ impl fmt::Display for TableAlias {
835843
}
836844
}
837845

846+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
847+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
848+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
849+
pub enum TableVersion {
850+
ForSystemTimeAsOf(Expr),
851+
}
852+
853+
impl Display for TableVersion {
854+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
855+
match self {
856+
TableVersion::ForSystemTimeAsOf(e) => write!(f, " FOR SYSTEM_TIME AS OF {e}")?,
857+
}
858+
Ok(())
859+
}
860+
}
861+
838862
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
839863
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
840864
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

‎src/ast/visitor.rs

+73-9
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
//! Recursive visitors for ast Nodes. See [`Visitor`] for more details.
1414
15-
use crate::ast::{Expr, ObjectName, Statement};
15+
use crate::ast::{Expr, ObjectName, Statement, TableFactor};
1616
use core::ops::ControlFlow;
1717

1818
/// A type that can be visited by a [`Visitor`]. See [`Visitor`] for
@@ -115,8 +115,8 @@ visit_noop!(bigdecimal::BigDecimal);
115115

116116
/// A visitor that can be used to walk an AST tree.
117117
///
118-
/// `previst_` methods are invoked before visiting all children of the
119-
/// node and `postvisit_` methods are invoked after visiting all
118+
/// `pre_visit_` methods are invoked before visiting all children of the
119+
/// node and `post_visit_` methods are invoked after visiting all
120120
/// children of the node.
121121
///
122122
/// # See also
@@ -139,7 +139,7 @@ visit_noop!(bigdecimal::BigDecimal);
139139
/// }
140140
///
141141
/// // Visit relations and exprs before children are visited (depth first walk)
142-
/// // Note you can also visit statements and visit exprs after children have been visitoed
142+
/// // Note you can also visit statements and visit exprs after children have been visited
143143
/// impl Visitor for V {
144144
/// type Break = ();
145145
///
@@ -189,6 +189,16 @@ pub trait Visitor {
189189
ControlFlow::Continue(())
190190
}
191191

192+
/// Invoked for any table factors that appear in the AST before visiting children
193+
fn pre_visit_table_factor(&mut self, _table_factor: &TableFactor) -> ControlFlow<Self::Break> {
194+
ControlFlow::Continue(())
195+
}
196+
197+
/// Invoked for any table factors that appear in the AST after visiting children
198+
fn post_visit_table_factor(&mut self, _table_factor: &TableFactor) -> ControlFlow<Self::Break> {
199+
ControlFlow::Continue(())
200+
}
201+
192202
/// Invoked for any expressions that appear in the AST before visiting children
193203
fn pre_visit_expr(&mut self, _expr: &Expr) -> ControlFlow<Self::Break> {
194204
ControlFlow::Continue(())
@@ -212,8 +222,8 @@ pub trait Visitor {
212222

213223
/// A visitor that can be used to mutate an AST tree.
214224
///
215-
/// `previst_` methods are invoked before visiting all children of the
216-
/// node and `postvisit_` methods are invoked after visiting all
225+
/// `pre_visit_` methods are invoked before visiting all children of the
226+
/// node and `post_visit_` methods are invoked after visiting all
217227
/// children of the node.
218228
///
219229
/// # See also
@@ -267,6 +277,22 @@ pub trait VisitorMut {
267277
ControlFlow::Continue(())
268278
}
269279

280+
/// Invoked for any table factors that appear in the AST before visiting children
281+
fn pre_visit_table_factor(
282+
&mut self,
283+
_table_factor: &mut TableFactor,
284+
) -> ControlFlow<Self::Break> {
285+
ControlFlow::Continue(())
286+
}
287+
288+
/// Invoked for any table factors that appear in the AST after visiting children
289+
fn post_visit_table_factor(
290+
&mut self,
291+
_table_factor: &mut TableFactor,
292+
) -> ControlFlow<Self::Break> {
293+
ControlFlow::Continue(())
294+
}
295+
270296
/// Invoked for any expressions that appear in the AST before visiting children
271297
fn pre_visit_expr(&mut self, _expr: &mut Expr) -> ControlFlow<Self::Break> {
272298
ControlFlow::Continue(())
@@ -609,6 +635,24 @@ mod tests {
609635
ControlFlow::Continue(())
610636
}
611637

638+
fn pre_visit_table_factor(
639+
&mut self,
640+
table_factor: &TableFactor,
641+
) -> ControlFlow<Self::Break> {
642+
self.visited
643+
.push(format!("PRE: TABLE FACTOR: {table_factor}"));
644+
ControlFlow::Continue(())
645+
}
646+
647+
fn post_visit_table_factor(
648+
&mut self,
649+
table_factor: &TableFactor,
650+
) -> ControlFlow<Self::Break> {
651+
self.visited
652+
.push(format!("POST: TABLE FACTOR: {table_factor}"));
653+
ControlFlow::Continue(())
654+
}
655+
612656
fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
613657
self.visited.push(format!("PRE: EXPR: {expr}"));
614658
ControlFlow::Continue(())
@@ -647,22 +691,28 @@ mod tests {
647691
fn test_sql() {
648692
let tests = vec![
649693
(
650-
"SELECT * from table_name",
694+
"SELECT * from table_name as my_table",
651695
vec![
652-
"PRE: STATEMENT: SELECT * FROM table_name",
696+
"PRE: STATEMENT: SELECT * FROM table_name AS my_table",
697+
"PRE: TABLE FACTOR: table_name AS my_table",
653698
"PRE: RELATION: table_name",
654699
"POST: RELATION: table_name",
655-
"POST: STATEMENT: SELECT * FROM table_name",
700+
"POST: TABLE FACTOR: table_name AS my_table",
701+
"POST: STATEMENT: SELECT * FROM table_name AS my_table",
656702
],
657703
),
658704
(
659705
"SELECT * from t1 join t2 on t1.id = t2.t1_id",
660706
vec![
661707
"PRE: STATEMENT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
708+
"PRE: TABLE FACTOR: t1",
662709
"PRE: RELATION: t1",
663710
"POST: RELATION: t1",
711+
"POST: TABLE FACTOR: t1",
712+
"PRE: TABLE FACTOR: t2",
664713
"PRE: RELATION: t2",
665714
"POST: RELATION: t2",
715+
"POST: TABLE FACTOR: t2",
666716
"PRE: EXPR: t1.id = t2.t1_id",
667717
"PRE: EXPR: t1.id",
668718
"POST: EXPR: t1.id",
@@ -676,13 +726,17 @@ mod tests {
676726
"SELECT * from t1 where EXISTS(SELECT column from t2)",
677727
vec![
678728
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
729+
"PRE: TABLE FACTOR: t1",
679730
"PRE: RELATION: t1",
680731
"POST: RELATION: t1",
732+
"POST: TABLE FACTOR: t1",
681733
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
682734
"PRE: EXPR: column",
683735
"POST: EXPR: column",
736+
"PRE: TABLE FACTOR: t2",
684737
"PRE: RELATION: t2",
685738
"POST: RELATION: t2",
739+
"POST: TABLE FACTOR: t2",
686740
"POST: EXPR: EXISTS (SELECT column FROM t2)",
687741
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
688742
],
@@ -691,13 +745,17 @@ mod tests {
691745
"SELECT * from t1 where EXISTS(SELECT column from t2)",
692746
vec![
693747
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
748+
"PRE: TABLE FACTOR: t1",
694749
"PRE: RELATION: t1",
695750
"POST: RELATION: t1",
751+
"POST: TABLE FACTOR: t1",
696752
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
697753
"PRE: EXPR: column",
698754
"POST: EXPR: column",
755+
"PRE: TABLE FACTOR: t2",
699756
"PRE: RELATION: t2",
700757
"POST: RELATION: t2",
758+
"POST: TABLE FACTOR: t2",
701759
"POST: EXPR: EXISTS (SELECT column FROM t2)",
702760
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
703761
],
@@ -706,16 +764,22 @@ mod tests {
706764
"SELECT * from t1 where EXISTS(SELECT column from t2) UNION SELECT * from t3",
707765
vec![
708766
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
767+
"PRE: TABLE FACTOR: t1",
709768
"PRE: RELATION: t1",
710769
"POST: RELATION: t1",
770+
"POST: TABLE FACTOR: t1",
711771
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
712772
"PRE: EXPR: column",
713773
"POST: EXPR: column",
774+
"PRE: TABLE FACTOR: t2",
714775
"PRE: RELATION: t2",
715776
"POST: RELATION: t2",
777+
"POST: TABLE FACTOR: t2",
716778
"POST: EXPR: EXISTS (SELECT column FROM t2)",
779+
"PRE: TABLE FACTOR: t3",
717780
"PRE: RELATION: t3",
718781
"POST: RELATION: t3",
782+
"POST: TABLE FACTOR: t3",
719783
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
720784
],
721785
),

‎src/parser/mod.rs

+18
Original file line numberDiff line numberDiff line change
@@ -6210,6 +6210,9 @@ impl<'a> Parser<'a> {
62106210
} else {
62116211
let name = self.parse_object_name()?;
62126212

6213+
// Parse potential version qualifier
6214+
let version = self.parse_table_version()?;
6215+
62136216
// Postgres, MSSQL: table-valued functions:
62146217
let args = if self.consume_token(&Token::LParen) {
62156218
Some(self.parse_optional_args()?)
@@ -6240,10 +6243,25 @@ impl<'a> Parser<'a> {
62406243
alias,
62416244
args,
62426245
with_hints,
6246+
version,
62436247
})
62446248
}
62456249
}
62466250

6251+
/// Parse a given table version specifier.
6252+
///
6253+
/// For now it only supports timestamp versioning for BigQuery and MSSQL dialects.
6254+
pub fn parse_table_version(&mut self) -> Result<Option<TableVersion>, ParserError> {
6255+
if dialect_of!(self is BigQueryDialect | MsSqlDialect)
6256+
&& self.parse_keywords(&[Keyword::FOR, Keyword::SYSTEM_TIME, Keyword::AS, Keyword::OF])
6257+
{
6258+
let expr = self.parse_expr()?;
6259+
Ok(Some(TableVersion::ForSystemTimeAsOf(expr)))
6260+
} else {
6261+
Ok(None)
6262+
}
6263+
}
6264+
62476265
pub fn parse_derived_table_factor(
62486266
&mut self,
62496267
lateral: IsLateral,

‎src/test_utils.rs

+1
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
221221
alias: None,
222222
args: None,
223223
with_hints: vec![],
224+
version: None,
224225
}
225226
}
226227

‎tests/sqlparser_bigquery.rs

+26
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ fn parse_table_identifiers() {
9595
alias: None,
9696
args: None,
9797
with_hints: vec![],
98+
version: None,
9899
},
99100
joins: vec![]
100101
},]
@@ -143,6 +144,31 @@ fn parse_table_identifiers() {
143144
test_table_ident("abc5.GROUP", vec![Ident::new("abc5"), Ident::new("GROUP")]);
144145
}
145146

147+
#[test]
148+
fn parse_table_time_travel() {
149+
let version = "2023-08-18 23:08:18".to_string();
150+
let sql = format!("SELECT 1 FROM t1 FOR SYSTEM_TIME AS OF '{version}'");
151+
let select = bigquery().verified_only_select(&sql);
152+
assert_eq!(
153+
select.from,
154+
vec![TableWithJoins {
155+
relation: TableFactor::Table {
156+
name: ObjectName(vec![Ident::new("t1")]),
157+
alias: None,
158+
args: None,
159+
with_hints: vec![],
160+
version: Some(TableVersion::ForSystemTimeAsOf(Expr::Value(
161+
Value::SingleQuotedString(version)
162+
))),
163+
},
164+
joins: vec![]
165+
},]
166+
);
167+
168+
let sql = "SELECT 1 FROM t1 FOR SYSTEM TIME AS OF 'some_timestamp'".to_string();
169+
assert!(bigquery().parse_sql_statements(&sql).is_err());
170+
}
171+
146172
#[test]
147173
fn parse_join_constraint_unnest_alias() {
148174
assert_eq!(

‎tests/sqlparser_clickhouse.rs

+3
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ fn parse_map_access_expr() {
6262
alias: None,
6363
args: None,
6464
with_hints: vec![],
65+
version: None,
6566
},
6667
joins: vec![]
6768
}],
@@ -169,11 +170,13 @@ fn parse_delimited_identifiers() {
169170
alias,
170171
args,
171172
with_hints,
173+
version,
172174
} => {
173175
assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
174176
assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);
175177
assert!(args.is_none());
176178
assert!(with_hints.is_empty());
179+
assert!(version.is_none());
177180
}
178181
_ => panic!("Expecting TableFactor::Table"),
179182
}

0 commit comments

Comments
 (0)