Skip to content

Commit

Permalink
support for snowflake SELECT TOP
Browse files Browse the repository at this point in the history
  • Loading branch information
jmhain committed Dec 21, 2023
1 parent 1baec96 commit 98bfd63
Show file tree
Hide file tree
Showing 5 changed files with 53 additions and 11 deletions.
4 changes: 2 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,8 @@ pub use self::query::{
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, NamedWindowDefinition,
NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem, ReplaceSelectElement,
ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table,
TableAlias, TableFactor, TableVersion, TableWithJoins, Top, Values, WildcardAdditionalOptions,
With,
TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, Values,
WildcardAdditionalOptions, With,
};
pub use self::value::{
escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,
Expand Down
21 changes: 19 additions & 2 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1278,17 +1278,34 @@ impl fmt::Display for Distinct {
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Top {
/// SQL semantic equivalent of LIMIT but with same structure as FETCH.
/// MSSQL only.
pub with_ties: bool,
/// MSSQL only.
pub percent: bool,
pub quantity: Option<Expr>,
pub quantity: Option<TopQuantity>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum TopQuantity {
// A parenthesized expression. MSSQL only.
Expr(Expr),
// An unparenthesized integer constant.
Constant(u64),
}

impl fmt::Display for Top {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let extension = if self.with_ties { " WITH TIES" } else { "" };
if let Some(ref quantity) = self.quantity {
let percent = if self.percent { " PERCENT" } else { "" };
write!(f, "TOP ({quantity}){percent}{extension}")
match quantity {
TopQuantity::Expr(quantity) => write!(f, "TOP ({quantity}){percent}{extension}"),
TopQuantity::Constant(quantity) => {
write!(f, "TOP {quantity}{percent}{extension}")
}
}
} else {
write!(f, "TOP{extension}")
}
Expand Down
9 changes: 7 additions & 2 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7843,9 +7843,14 @@ impl<'a> Parser<'a> {
let quantity = if self.consume_token(&Token::LParen) {
let quantity = self.parse_expr()?;
self.expect_token(&Token::RParen)?;
Some(quantity)
Some(TopQuantity::Expr(quantity))
} else {
Some(Expr::Value(self.parse_number_value()?))
let next_token = self.next_token();
let quantity = match next_token.token {
Token::Number(s, _) => s.parse::<u64>().expect("literal int"),
_ => self.expected("literal int", next_token)?,
};
Some(TopQuantity::Constant(quantity))
};

let percent = self.parse_keyword(Keyword::PERCENT);
Expand Down
22 changes: 17 additions & 5 deletions tests/sqlparser_mssql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,10 @@ fn parse_mssql_top_paren() {
let sql = "SELECT TOP (5) * FROM foo";
let select = ms_and_generic().verified_only_select(sql);
let top = select.top.unwrap();
assert_eq!(Some(Expr::Value(number("5"))), top.quantity);
assert_eq!(
Some(TopQuantity::Expr(Expr::Value(number("5")))),
top.quantity
);
assert!(!top.percent);
}

Expand All @@ -189,7 +192,10 @@ fn parse_mssql_top_percent() {
let sql = "SELECT TOP (5) PERCENT * FROM foo";
let select = ms_and_generic().verified_only_select(sql);
let top = select.top.unwrap();
assert_eq!(Some(Expr::Value(number("5"))), top.quantity);
assert_eq!(
Some(TopQuantity::Expr(Expr::Value(number("5")))),
top.quantity
);
assert!(top.percent);
}

Expand All @@ -198,7 +204,10 @@ fn parse_mssql_top_with_ties() {
let sql = "SELECT TOP (5) WITH TIES * FROM foo";
let select = ms_and_generic().verified_only_select(sql);
let top = select.top.unwrap();
assert_eq!(Some(Expr::Value(number("5"))), top.quantity);
assert_eq!(
Some(TopQuantity::Expr(Expr::Value(number("5")))),
top.quantity
);
assert!(top.with_ties);
}

Expand All @@ -207,14 +216,17 @@ fn parse_mssql_top_percent_with_ties() {
let sql = "SELECT TOP (10) PERCENT WITH TIES * FROM foo";
let select = ms_and_generic().verified_only_select(sql);
let top = select.top.unwrap();
assert_eq!(Some(Expr::Value(number("10"))), top.quantity);
assert_eq!(
Some(TopQuantity::Expr(Expr::Value(number("10")))),
top.quantity
);
assert!(top.percent);
}

#[test]
fn parse_mssql_top() {
let sql = "SELECT TOP 5 bar, baz FROM foo";
let _ = ms_and_generic().one_statement_parses_to(sql, "SELECT TOP (5) bar, baz FROM foo");
let _ = ms_and_generic().one_statement_parses_to(sql, "SELECT TOP 5 bar, baz FROM foo");
}

#[test]
Expand Down
8 changes: 8 additions & 0 deletions tests/sqlparser_snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1147,3 +1147,11 @@ fn parse_pivot_of_table_factor_derived() {
"SELECT * FROM (SELECT place_id, weekday, open FROM times AS p) PIVOT(max(open) FOR weekday IN (0, 1, 2, 3, 4, 5, 6)) AS p (place_id, open_sun, open_mon, open_tue, open_wed, open_thu, open_fri, open_sat)"
);
}

#[test]
fn parse_top() {
snowflake().one_statement_parses_to(
"SELECT TOP 4 c1 FROM testtable",
"SELECT TOP 4 c1 FROM testtable",
);
}

0 comments on commit 98bfd63

Please sign in to comment.