Skip to content

Commit e7c6d55

Browse files
gruuyaserprex
authored andcommitted
Table time travel clause support, add visit_table_factor to Visitor (apache#951)
1 parent 1ff2e70 commit e7c6d55

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
@@ -650,6 +650,7 @@ impl fmt::Display for TableWithJoins {
650650
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
651651
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
652652
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
653+
#[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
653654
pub enum TableFactor {
654655
Table {
655656
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
@@ -665,6 +666,9 @@ pub enum TableFactor {
665666
args: Option<Vec<FunctionArg>>,
666667
/// MSSQL-specific `WITH (...)` hints such as NOLOCK.
667668
with_hints: Vec<Expr>,
669+
/// Optional version qualifier to facilitate table time-travel, as
670+
/// supported by BigQuery and MSSQL.
671+
version: Option<TableVersion>,
668672
},
669673
Derived {
670674
lateral: bool,
@@ -724,6 +728,7 @@ impl fmt::Display for TableFactor {
724728
alias,
725729
args,
726730
with_hints,
731+
version,
727732
} => {
728733
write!(f, "{name}")?;
729734
if let Some(args) = args {
@@ -735,6 +740,9 @@ impl fmt::Display for TableFactor {
735740
if !with_hints.is_empty() {
736741
write!(f, " WITH ({})", display_comma_separated(with_hints))?;
737742
}
743+
if let Some(version) = version {
744+
write!(f, "{version}")?;
745+
}
738746
Ok(())
739747
}
740748
TableFactor::Derived {
@@ -839,6 +847,22 @@ impl fmt::Display for TableAlias {
839847
}
840848
}
841849

850+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
851+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
852+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
853+
pub enum TableVersion {
854+
ForSystemTimeAsOf(Expr),
855+
}
856+
857+
impl Display for TableVersion {
858+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
859+
match self {
860+
TableVersion::ForSystemTimeAsOf(e) => write!(f, " FOR SYSTEM_TIME AS OF {e}")?,
861+
}
862+
Ok(())
863+
}
864+
}
865+
842866
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
843867
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
844868
#[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, FunctionArgExpr, ObjectName, SetExpr, Statement};
15+
use crate::ast::{Expr, FunctionArgExpr, ObjectName, SetExpr, 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(())
@@ -231,8 +241,8 @@ pub trait Visitor {
231241

232242
/// A visitor that can be used to mutate an AST tree.
233243
///
234-
/// `previst_` methods are invoked before visiting all children of the
235-
/// node and `postvisit_` methods are invoked after visiting all
244+
/// `pre_visit_` methods are invoked before visiting all children of the
245+
/// node and `post_visit_` methods are invoked after visiting all
236246
/// children of the node.
237247
///
238248
/// # See also
@@ -286,6 +296,22 @@ pub trait VisitorMut {
286296
ControlFlow::Continue(())
287297
}
288298

299+
/// Invoked for any table factors that appear in the AST before visiting children
300+
fn pre_visit_table_factor(
301+
&mut self,
302+
_table_factor: &mut TableFactor,
303+
) -> ControlFlow<Self::Break> {
304+
ControlFlow::Continue(())
305+
}
306+
307+
/// Invoked for any table factors that appear in the AST after visiting children
308+
fn post_visit_table_factor(
309+
&mut self,
310+
_table_factor: &mut TableFactor,
311+
) -> ControlFlow<Self::Break> {
312+
ControlFlow::Continue(())
313+
}
314+
289315
/// Invoked for any expressions that appear in the AST before visiting children
290316
fn pre_visit_expr(&mut self, _expr: &mut Expr) -> ControlFlow<Self::Break> {
291317
ControlFlow::Continue(())
@@ -719,6 +745,24 @@ mod tests {
719745
ControlFlow::Continue(())
720746
}
721747

748+
fn pre_visit_table_factor(
749+
&mut self,
750+
table_factor: &TableFactor,
751+
) -> ControlFlow<Self::Break> {
752+
self.visited
753+
.push(format!("PRE: TABLE FACTOR: {table_factor}"));
754+
ControlFlow::Continue(())
755+
}
756+
757+
fn post_visit_table_factor(
758+
&mut self,
759+
table_factor: &TableFactor,
760+
) -> ControlFlow<Self::Break> {
761+
self.visited
762+
.push(format!("POST: TABLE FACTOR: {table_factor}"));
763+
ControlFlow::Continue(())
764+
}
765+
722766
fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
723767
self.visited.push(format!("PRE: EXPR: {expr}"));
724768
ControlFlow::Continue(())
@@ -757,22 +801,28 @@ mod tests {
757801
fn test_sql() {
758802
let tests = vec![
759803
(
760-
"SELECT * from table_name",
804+
"SELECT * from table_name as my_table",
761805
vec![
762-
"PRE: STATEMENT: SELECT * FROM table_name",
806+
"PRE: STATEMENT: SELECT * FROM table_name AS my_table",
807+
"PRE: TABLE FACTOR: table_name AS my_table",
763808
"PRE: RELATION: table_name",
764809
"POST: RELATION: table_name",
765-
"POST: STATEMENT: SELECT * FROM table_name",
810+
"POST: TABLE FACTOR: table_name AS my_table",
811+
"POST: STATEMENT: SELECT * FROM table_name AS my_table",
766812
],
767813
),
768814
(
769815
"SELECT * from t1 join t2 on t1.id = t2.t1_id",
770816
vec![
771817
"PRE: STATEMENT: SELECT * FROM t1 JOIN t2 ON t1.id = t2.t1_id",
818+
"PRE: TABLE FACTOR: t1",
772819
"PRE: RELATION: t1",
773820
"POST: RELATION: t1",
821+
"POST: TABLE FACTOR: t1",
822+
"PRE: TABLE FACTOR: t2",
774823
"PRE: RELATION: t2",
775824
"POST: RELATION: t2",
825+
"POST: TABLE FACTOR: t2",
776826
"PRE: EXPR: t1.id = t2.t1_id",
777827
"PRE: EXPR: t1.id",
778828
"POST: EXPR: t1.id",
@@ -786,13 +836,17 @@ mod tests {
786836
"SELECT * from t1 where EXISTS(SELECT column from t2)",
787837
vec![
788838
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
839+
"PRE: TABLE FACTOR: t1",
789840
"PRE: RELATION: t1",
790841
"POST: RELATION: t1",
842+
"POST: TABLE FACTOR: t1",
791843
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
792844
"PRE: EXPR: column",
793845
"POST: EXPR: column",
846+
"PRE: TABLE FACTOR: t2",
794847
"PRE: RELATION: t2",
795848
"POST: RELATION: t2",
849+
"POST: TABLE FACTOR: t2",
796850
"POST: EXPR: EXISTS (SELECT column FROM t2)",
797851
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
798852
],
@@ -801,13 +855,17 @@ mod tests {
801855
"SELECT * from t1 where EXISTS(SELECT column from t2)",
802856
vec![
803857
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
858+
"PRE: TABLE FACTOR: t1",
804859
"PRE: RELATION: t1",
805860
"POST: RELATION: t1",
861+
"POST: TABLE FACTOR: t1",
806862
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
807863
"PRE: EXPR: column",
808864
"POST: EXPR: column",
865+
"PRE: TABLE FACTOR: t2",
809866
"PRE: RELATION: t2",
810867
"POST: RELATION: t2",
868+
"POST: TABLE FACTOR: t2",
811869
"POST: EXPR: EXISTS (SELECT column FROM t2)",
812870
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2)",
813871
],
@@ -816,16 +874,22 @@ mod tests {
816874
"SELECT * from t1 where EXISTS(SELECT column from t2) UNION SELECT * from t3",
817875
vec![
818876
"PRE: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
877+
"PRE: TABLE FACTOR: t1",
819878
"PRE: RELATION: t1",
820879
"POST: RELATION: t1",
880+
"POST: TABLE FACTOR: t1",
821881
"PRE: EXPR: EXISTS (SELECT column FROM t2)",
822882
"PRE: EXPR: column",
823883
"POST: EXPR: column",
884+
"PRE: TABLE FACTOR: t2",
824885
"PRE: RELATION: t2",
825886
"POST: RELATION: t2",
887+
"POST: TABLE FACTOR: t2",
826888
"POST: EXPR: EXISTS (SELECT column FROM t2)",
889+
"PRE: TABLE FACTOR: t3",
827890
"PRE: RELATION: t3",
828891
"POST: RELATION: t3",
892+
"POST: TABLE FACTOR: t3",
829893
"POST: STATEMENT: SELECT * FROM t1 WHERE EXISTS (SELECT column FROM t2) UNION SELECT * FROM t3",
830894
],
831895
),

src/parser/mod.rs

+18
Original file line numberDiff line numberDiff line change
@@ -6257,6 +6257,9 @@ impl<'a> Parser<'a> {
62576257
} else {
62586258
let name = self.parse_object_name()?;
62596259

6260+
// Parse potential version qualifier
6261+
let version = self.parse_table_version()?;
6262+
62606263
// Postgres, MSSQL: table-valued functions:
62616264
let args = if self.consume_token(&Token::LParen) {
62626265
Some(self.parse_optional_args()?)
@@ -6287,10 +6290,25 @@ impl<'a> Parser<'a> {
62876290
alias,
62886291
args,
62896292
with_hints,
6293+
version,
62906294
})
62916295
}
62926296
}
62936297

6298+
/// Parse a given table version specifier.
6299+
///
6300+
/// For now it only supports timestamp versioning for BigQuery and MSSQL dialects.
6301+
pub fn parse_table_version(&mut self) -> Result<Option<TableVersion>, ParserError> {
6302+
if dialect_of!(self is BigQueryDialect | MsSqlDialect)
6303+
&& self.parse_keywords(&[Keyword::FOR, Keyword::SYSTEM_TIME, Keyword::AS, Keyword::OF])
6304+
{
6305+
let expr = self.parse_expr()?;
6306+
Ok(Some(TableVersion::ForSystemTimeAsOf(expr)))
6307+
} else {
6308+
Ok(None)
6309+
}
6310+
}
6311+
62946312
pub fn parse_derived_table_factor(
62956313
&mut self,
62966314
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)