snowflake: fix rendering of SELECT TOP (#1070)

Author: jmhain

src/ast/mod.rs

@@ -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,

src/ast/query.rs

@@ -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}")
         }

src/parser/mod.rs

@@ -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);

tests/sqlparser_mssql.rs

@@ -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);
 }
 
@@ -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);
 }
 
@@ -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);
 }
 
@@ -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]

tests/sqlparser_snowflake.rs

@@ -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",
+    );
+}