Snowflake: Improve accuracy of lookahead in implicit LIMIT alias (#1941)

Author: yoavcloud

src/dialect/snowflake.rs

@@ -23,8 +23,8 @@ use crate::ast::helpers::stmt_data_loading::{
     FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
 };
 use crate::ast::{
-    ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, Ident,
-    IdentityParameters, IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind,
+    ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, DollarQuotedString,
+    Ident, IdentityParameters, IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind,
     IdentityPropertyOrder, ObjectName, ObjectNamePart, RowAccessPolicy, ShowObjects, SqlOption,
     Statement, TagsColumnOption, WrappedCollection,
 };
@@ -307,22 +307,22 @@ impl Dialect for SnowflakeDialect {
             // they are not followed by other tokens that may change their meaning
             // e.g. `SELECT * EXCEPT (col1) FROM tbl`
             Keyword::EXCEPT
-            // e.g. `SELECT 1 LIMIT 5`
-            | Keyword::LIMIT
-            // e.g. `SELECT 1 OFFSET 5 ROWS`
-            | Keyword::OFFSET
             // e.g. `INSERT INTO t SELECT 1 RETURNING *`
             | Keyword::RETURNING if !matches!(parser.peek_token_ref().token, Token::Comma | Token::EOF) =>
             {
                 false
             }
 
+            // e.g. `SELECT 1 LIMIT 5` - not an alias
+            // e.g. `SELECT 1 OFFSET 5 ROWS` - not an alias
+            Keyword::LIMIT | Keyword::OFFSET if peek_for_limit_options(parser) => false,
+
             // `FETCH` can be considered an alias as long as it's not followed by `FIRST`` or `NEXT`
             // which would give it a different meanings, for example: 
             // `SELECT 1 FETCH FIRST 10 ROWS` - not an alias
             // `SELECT 1 FETCH 10` - not an alias
             Keyword::FETCH if parser.peek_one_of_keywords(&[Keyword::FIRST, Keyword::NEXT]).is_some()
-                    || matches!(parser.peek_token().token, Token::Number(_, _)) =>
+                    || peek_for_limit_options(parser) =>
             {
                 false
             }
@@ -351,20 +351,23 @@ impl Dialect for SnowflakeDialect {
         match kw {
             // The following keywords can be considered an alias as long as
             // they are not followed by other tokens that may change their meaning
-            Keyword::LIMIT
-            | Keyword::RETURNING
+            Keyword::RETURNING
             | Keyword::INNER
             | Keyword::USING
             | Keyword::PIVOT
             | Keyword::UNPIVOT
             | Keyword::EXCEPT
             | Keyword::MATCH_RECOGNIZE
-            | Keyword::OFFSET
                 if !matches!(parser.peek_token_ref().token, Token::SemiColon | Token::EOF) =>
             {
                 false
             }
 
+            // `LIMIT` can be considered an alias as long as it's not followed by a value. For example:
+            // `SELECT * FROM tbl LIMIT WHERE 1=1` - alias
+            // `SELECT * FROM tbl LIMIT 3` - not an alias
+            Keyword::LIMIT | Keyword::OFFSET if peek_for_limit_options(parser) => false,
+
             // `FETCH` can be considered an alias as long as it's not followed by `FIRST`` or `NEXT`
             // which would give it a different meanings, for example:
             // `SELECT * FROM tbl FETCH FIRST 10 ROWS` - not an alias
@@ -373,7 +376,7 @@ impl Dialect for SnowflakeDialect {
                 if parser
                     .peek_one_of_keywords(&[Keyword::FIRST, Keyword::NEXT])
                     .is_some()
-                    || matches!(parser.peek_token().token, Token::Number(_, _)) =>
+                    || peek_for_limit_options(parser) =>
             {
                 false
             }
@@ -387,6 +390,7 @@ impl Dialect for SnowflakeDialect {
             {
                 false
             }
+
             Keyword::GLOBAL if parser.peek_keyword(Keyword::FULL) => false,
 
             // Reserved keywords by the Snowflake dialect, which seem to be less strictive
@@ -472,6 +476,18 @@ impl Dialect for SnowflakeDialect {
     }
 }
 
+// Peeks ahead to identify tokens that are expected after
+// a LIMIT/FETCH keyword.
+fn peek_for_limit_options(parser: &Parser) -> bool {
+    match &parser.peek_token_ref().token {
+        Token::Number(_, _) | Token::Placeholder(_) => true,
+        Token::SingleQuotedString(val) if val.is_empty() => true,
+        Token::DollarQuotedString(DollarQuotedString { value, .. }) if value.is_empty() => true,
+        Token::Word(w) if w.keyword == Keyword::NULL => true,
+        _ => false,
+    }
+}
+
 fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {
     let stage = parse_snowflake_stage_name(parser)?;
     let pattern = if parser.parse_keyword(Keyword::PATTERN) {

tests/sqlparser_snowflake.rs

@@ -3535,6 +3535,15 @@ fn test_sql_keywords_as_select_item_aliases() {
             .parse_sql_statements(&format!("SELECT 1 {kw}"))
             .is_err());
     }
+
+    // LIMIT is alias
+    snowflake().one_statement_parses_to("SELECT 1 LIMIT", "SELECT 1 AS LIMIT");
+    // LIMIT is not an alias
+    snowflake().verified_stmt("SELECT 1 LIMIT 1");
+    snowflake().verified_stmt("SELECT 1 LIMIT $1");
+    snowflake().verified_stmt("SELECT 1 LIMIT ''");
+    snowflake().verified_stmt("SELECT 1 LIMIT NULL");
+    snowflake().verified_stmt("SELECT 1 LIMIT $$$$");
 }
 
 #[test]
@@ -3586,6 +3595,15 @@ fn test_sql_keywords_as_table_aliases() {
             .parse_sql_statements(&format!("SELECT * FROM tbl {kw}"))
             .is_err());
     }
+
+    // LIMIT is alias
+    snowflake().one_statement_parses_to("SELECT * FROM tbl LIMIT", "SELECT * FROM tbl AS LIMIT");
+    // LIMIT is not an alias
+    snowflake().verified_stmt("SELECT * FROM tbl LIMIT 1");
+    snowflake().verified_stmt("SELECT * FROM tbl LIMIT $1");
+    snowflake().verified_stmt("SELECT * FROM tbl LIMIT ''");
+    snowflake().verified_stmt("SELECT * FROM tbl LIMIT NULL");
+    snowflake().verified_stmt("SELECT * FROM tbl LIMIT $$$$");
 }
 
 #[test]