diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 330af3d33..8772389bf 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -36,8 +36,8 @@ pub use self::ddl::{ pub use self::operator::{BinaryOperator, UnaryOperator}; pub use self::query::{ Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, Offset, OffsetRows, - OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, TableAlias, - TableFactor, TableWithJoins, Top, Values, With, + OrderByExpr, Query, SamplingMethod, Select, SelectInto, SelectItem, SetExpr, SetOperator, + TableAlias, TableFactor, TableSample, TableWithJoins, Top, Values, With, }; pub use self::value::{DateTimeField, TrimWhereField, Value}; diff --git a/src/ast/query.rs b/src/ast/query.rs index 011d4658b..4aa47299d 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -344,6 +344,7 @@ pub enum TableFactor { args: Option>, /// MSSQL-specific `WITH (...)` hints such as NOLOCK. with_hints: Vec, + tablesample: Option, }, Derived { lateral: bool, @@ -389,6 +390,7 @@ impl fmt::Display for TableFactor { alias, args, with_hints, + tablesample, } => { write!(f, "{}", name)?; if let Some(args) = args { @@ -400,6 +402,9 @@ impl fmt::Display for TableFactor { if !with_hints.is_empty() { write!(f, " WITH ({})", display_comma_separated(with_hints))?; } + if let Some(tablesample) = tablesample { + write!(f, " {}", tablesample)?; + } Ok(()) } TableFactor::Derived { @@ -455,6 +460,44 @@ impl fmt::Display for TableFactor { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct TableSample { + pub sampling_method: SamplingMethod, + pub sampling_fraction: Box, + pub repeatable_seed: Option>, +} + +impl fmt::Display for TableSample { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "TABLESAMPLE {} ({})", + self.sampling_method, self.sampling_fraction + )?; + if let Some(ref repeatable_seed) = self.repeatable_seed { + write!(f, " REPEATABLE ({})", repeatable_seed)?; + } + Ok(()) + } +} +/// Stores the different Sampling Methods +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum SamplingMethod { + BERNOULLI, + SYSTEM, +} + +impl fmt::Display for SamplingMethod { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + SamplingMethod::BERNOULLI => write!(f, " BERNOULLI"), + SamplingMethod::SYSTEM => write!(f, "SYSTEM"), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct TableAlias { diff --git a/src/keywords.rs b/src/keywords.rs index 35d98eef7..fe56b0453 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -99,6 +99,7 @@ define_keywords!( BEGIN, BEGIN_FRAME, BEGIN_PARTITION, + BERNOULLI, BETWEEN, BIGINT, BINARY, diff --git a/src/parser.rs b/src/parser.rs index e1f0c87a1..e65c0b783 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -4466,11 +4466,14 @@ impl<'a> Parser<'a> { self.prev_token(); } }; + let tablesample = self.parse_tablesample_args()?; + Ok(TableFactor::Table { name, alias, args, with_hints, + tablesample, }) } } @@ -4492,6 +4495,33 @@ impl<'a> Parser<'a> { }) } + pub fn parse_tablesample_args(&mut self) -> Result, ParserError> { + if self.parse_keyword(Keyword::TABLESAMPLE) { + let sampling_method = + match self.expect_one_of_keywords(&[Keyword::BERNOULLI, Keyword::SYSTEM])? { + Keyword::BERNOULLI => SamplingMethod::BERNOULLI, + Keyword::SYSTEM => SamplingMethod::SYSTEM, + _ => unreachable!(), + }; + self.expect_token(&Token::LParen)?; + let sampling_fraction = Box::new(self.parse_expr()?); + self.expect_token(&Token::RParen)?; + let mut repeatable_seed = None; + if self.parse_keyword(Keyword::REPEATABLE) { + self.expect_token(&Token::LParen)?; + repeatable_seed = Some(Box::new(self.parse_expr()?)); + self.expect_token(&Token::RParen)?; + } + Ok(Some(TableSample { + sampling_method, + sampling_fraction, + repeatable_seed, + })) + } else { + Ok(None) + } + } + pub fn parse_join_constraint(&mut self, natural: bool) -> Result { if natural { Ok(JoinConstraint::Natural) diff --git a/src/test_utils.rs b/src/test_utils.rs index c3d60ee62..e2c8ae6f7 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -186,6 +186,7 @@ pub fn table(name: impl Into) -> TableFactor { alias: None, args: None, with_hints: vec![], + tablesample: None, } } diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs index 0a606c3ec..1ab35d031 100644 --- a/tests/sqlparser_bigquery.rs +++ b/tests/sqlparser_bigquery.rs @@ -46,6 +46,7 @@ fn parse_table_identifiers() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![] },] diff --git a/tests/sqlparser_clickhouse.rs b/tests/sqlparser_clickhouse.rs index a61df73cc..1602f5e01 100644 --- a/tests/sqlparser_clickhouse.rs +++ b/tests/sqlparser_clickhouse.rs @@ -61,6 +61,7 @@ fn parse_map_access_expr() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![] }], diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index e304f38c3..93171b2f7 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -206,6 +206,7 @@ fn parse_update_with_table_alias() { }), args: None, with_hints: vec![], + tablesample: None, }, joins: vec![] }, @@ -259,7 +260,8 @@ fn parse_delete_statement() { name: ObjectName(vec![Ident::with_quote('"', "table")]), alias: None, args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, }, table_name ); @@ -284,7 +286,8 @@ fn parse_where_delete_statement() { name: ObjectName(vec![Ident::new("foo")]), alias: None, args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, }, table_name, ); @@ -322,7 +325,8 @@ fn parse_where_delete_with_alias_statement() { columns: vec![] }), args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None }, table_name, ); @@ -335,7 +339,8 @@ fn parse_where_delete_with_alias_statement() { columns: vec![] }), args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, }), using ); @@ -3315,11 +3320,13 @@ fn parse_delimited_identifiers() { alias, args, with_hints, + tablesample, } => { assert_eq!(vec![Ident::with_quote('"', "a table")], name.0); assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name); assert!(args.is_none()); assert!(with_hints.is_empty()); + assert!(tablesample.is_none()); } _ => panic!("Expecting TableFactor::Table"), } @@ -3456,6 +3463,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], }, @@ -3465,6 +3473,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], } @@ -3482,6 +3491,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![Join { relation: TableFactor::Table { @@ -3489,6 +3499,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, join_operator: JoinOperator::Inner(JoinConstraint::Natural), }] @@ -3499,6 +3510,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![Join { relation: TableFactor::Table { @@ -3506,6 +3518,7 @@ fn parse_implicit_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, join_operator: JoinOperator::Inner(JoinConstraint::Natural), }] @@ -3526,6 +3539,7 @@ fn parse_cross_join() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, join_operator: JoinOperator::CrossJoin }, @@ -3546,6 +3560,7 @@ fn parse_joins_on() { alias, args: None, with_hints: vec![], + tablesample: None, }, join_operator: f(JoinConstraint::On(Expr::BinaryOp { left: Box::new(Expr::Identifier("c1".into())), @@ -3599,6 +3614,7 @@ fn parse_joins_using() { alias, args: None, with_hints: vec![], + tablesample: None, }, join_operator: f(JoinConstraint::Using(vec!["c1".into()])), } @@ -3644,6 +3660,7 @@ fn parse_natural_join() { alias, args: None, with_hints: vec![], + tablesample: None, }, join_operator: f(JoinConstraint::Natural), } @@ -3911,6 +3928,7 @@ fn parse_derived_tables() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, join_operator: JoinOperator::Inner(JoinConstraint::Natural), }], @@ -5111,7 +5129,8 @@ fn parse_merge() { columns: vec![] }), args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, } ); assert_eq!(table, table_no_into); @@ -5132,7 +5151,8 @@ fn parse_merge() { name: ObjectName(vec![Ident::new("s"), Ident::new("foo")]), alias: None, args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, }, joins: vec![] }], diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs index f46d5d23e..1864fd581 100644 --- a/tests/sqlparser_mysql.rs +++ b/tests/sqlparser_mysql.rs @@ -810,10 +810,11 @@ fn parse_update_with_joins() { name: ObjectName(vec![Ident::new("orders")]), alias: Some(TableAlias { name: Ident::new("o"), - columns: vec![] + columns: vec![], }), args: None, with_hints: vec![], + tablesample: None, }, joins: vec![Join { relation: TableFactor::Table { @@ -824,6 +825,7 @@ fn parse_update_with_joins() { }), args: None, with_hints: vec![], + tablesample: None, }, join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp { left: Box::new(Expr::CompoundIdentifier(vec![ @@ -930,7 +932,8 @@ fn parse_substring_in_select() { }]), alias: None, args: None, - with_hints: vec![] + with_hints: vec![], + tablesample: None, }, joins: vec![] }], diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 171624c0b..84464a5d7 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -414,6 +414,7 @@ fn parse_update_set_from() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], }, @@ -440,6 +441,7 @@ fn parse_update_set_from() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], }], diff --git a/tests/sqlparser_redshift.rs b/tests/sqlparser_redshift.rs index 6f77cf335..34c8637f9 100644 --- a/tests/sqlparser_redshift.rs +++ b/tests/sqlparser_redshift.rs @@ -45,6 +45,7 @@ fn test_square_brackets_over_db_schema_table_name() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], } @@ -89,6 +90,7 @@ fn test_double_quotes_over_db_schema_table_name() { alias: None, args: None, with_hints: vec![], + tablesample: None, }, joins: vec![], }