Snowflake: Complete LIMIT grammar (#2784)

* Snowflake: Complete LIMIT grammar

* Update src/sqlfluff/dialects/dialect_snowflake.py

Co-authored-by: Barry Pollard <barry_pollard@hotmail.com>

* Update src/sqlfluff/dialects/dialect_snowflake.py

Co-authored-by: Barry Pollard <barry_pollard@hotmail.com>

* Implement review feedback

* Black

* use match_grammar.terminator insert

* Use Ref.keyword for terminator insert

Co-authored-by: Barry Pollard <barry_pollard@hotmail.com>

src/sqlfluff/dialects/dialect_snowflake.py

@@ -219,7 +219,15 @@
             # Can `GROUP BY coalesce(col, 1)`
             Ref("ExpressionSegment"),
         ),
-        terminator=OneOf("ORDER", "LIMIT", "HAVING", "QUALIFY", "WINDOW"),
+        terminator=OneOf(
+            "ORDER", "LIMIT", "FETCH", "OFFSET", "HAVING", "QUALIFY", "WINDOW"
+        ),
+    ),
+    LimitLiteralGrammar=OneOf(
+        Ref("NumericLiteralSegment"),
+        "NULL",
+        # '' and $$$$ are allowed as alternatives to NULL.
+        Ref("QuotedLiteralSegment"),
     ),
 )
 
@@ -320,6 +328,41 @@
         Sequence("ILIKE", Ref.keyword("ANY", optional=True)),
         "REGEXP",
     ),
+    SelectClauseElementTerminatorGrammar=OneOf(
+        "FROM",
+        "WHERE",
+        Sequence("ORDER", "BY"),
+        "LIMIT",
+        "FETCH",
+        "OFFSET",
+        Ref("CommaSegment"),
+        Ref("SetOperatorSegment"),
+    ),
+    FromClauseTerminatorGrammar=OneOf(
+        "WHERE",
+        "LIMIT",
+        "FETCH",
+        "OFFSET",
+        Sequence("GROUP", "BY"),
+        Sequence("ORDER", "BY"),
+        "HAVING",
+        "QUALIFY",
+        "WINDOW",
+        Ref("SetOperatorSegment"),
+        Ref("WithNoSchemaBindingClauseSegment"),
+        Ref("WithDataClauseSegment"),
+    ),
+    WhereClauseTerminatorGrammar=OneOf(
+        "LIMIT",
+        "FETCH",
+        "OFFSET",
+        Sequence("GROUP", "BY"),
+        Sequence("ORDER", "BY"),
+        "HAVING",
+        "QUALIFY",
+        "WINDOW",
+        "OVERLAPS",
+    ),
 )
 
 # Add all Snowflake keywords
@@ -473,7 +516,9 @@ class GroupByClauseSegment(BaseSegment):
     type = "groupby_clause"
     match_grammar = StartsWith(
         Sequence("GROUP", "BY"),
-        terminator=OneOf("ORDER", "LIMIT", "HAVING", "QUALIFY", "WINDOW"),
+        terminator=OneOf(
+            "ORDER", "LIMIT", "FETCH", "OFFSET", "HAVING", "QUALIFY", "WINDOW"
+        ),
         enforce_whitespace_preceding_terminator=True,
     )
     parse_grammar = Sequence(
@@ -860,12 +905,14 @@ class QualifyClauseSegment(BaseSegment):
     https://docs.snowflake.com/en/sql-reference/constructs/qualify.html
     """
 
-    type = "having_clause"
+    type = "qualify_clause"
     match_grammar = StartsWith(
         "QUALIFY",
         terminator=OneOf(
             Sequence("ORDER", "BY"),
             "LIMIT",
+            "FETCH",
+            "OFFSET",
         ),
     )
     parse_grammar = Sequence(
@@ -3840,3 +3887,117 @@ class CallStatementSegment(BaseSegment):
             ),
         ),
     )
+
+
+@snowflake_dialect.segment(replace=True)
+class LimitClauseSegment(BaseSegment):
+    """A `LIMIT` clause.
+
+    https://docs.snowflake.com/en/sql-reference/constructs/limit.html
+    """
+
+    type = "limit_clause"
+    match_grammar = OneOf(
+        Sequence(
+            "LIMIT",
+            Indent,
+            Ref("LimitLiteralGrammar"),
+            Dedent,
+            Sequence(
+                "OFFSET",
+                Indent,
+                Ref("LimitLiteralGrammar"),
+                Dedent,
+                optional=True,
+            ),
+        ),
+        Sequence(
+            Sequence(
+                "OFFSET",
+                Indent,
+                Ref("LimitLiteralGrammar"),
+                OneOf(
+                    "ROW",
+                    "ROWS",
+                    optional=True,
+                ),
+                Dedent,
+                optional=True,
+            ),
+            "FETCH",
+            Indent,
+            OneOf(
+                "FIRST",
+                "NEXT",
+                optional=True,
+            ),
+            Ref("LimitLiteralGrammar"),
+            OneOf(
+                "ROW",
+                "ROWS",
+                optional=True,
+            ),
+            Ref.keyword("ONLY", optional=True),
+            Dedent,
+        ),
+    )
+
+
+@snowflake_dialect.segment(replace=True)
+class SelectClauseSegment(BaseSegment):
+    """A group of elements in a select target statement."""
+
+    type = "select_clause"
+    match_grammar = ansi_dialect.get_segment("SelectClauseSegment").match_grammar.copy()
+    match_grammar.terminator = match_grammar.terminator.copy(
+        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET")],
+    )
+    parse_grammar = ansi_dialect.get_segment("SelectClauseSegment").parse_grammar.copy()
+
+
+@snowflake_dialect.segment(replace=True)
+class OrderByClauseSegment(BaseSegment):
+    """An `ORDER BY` clause.
+
+    https://docs.snowflake.com/en/sql-reference/constructs/order-by.html
+    """
+
+    type = "orderby_clause"
+    match_grammar = ansi_dialect.get_segment(
+        "OrderByClauseSegment"
+    ).match_grammar.copy()
+    match_grammar.terminator = match_grammar.terminator.copy(
+        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET")],
+    )
+    parse_grammar = Sequence(
+        "ORDER",
+        "BY",
+        Indent,
+        Delimited(
+            Sequence(
+                OneOf(
+                    Ref("ColumnReferenceSegment"),
+                    # Can `ORDER BY 1`
+                    Ref("NumericLiteralSegment"),
+                    # Can order by an expression
+                    Ref("ExpressionSegment"),
+                ),
+                OneOf("ASC", "DESC", optional=True),
+                Sequence("NULLS", OneOf("FIRST", "LAST"), optional=True),
+            ),
+            terminator=OneOf("LIMIT", "FETCH", "OFFSET", Ref("FrameClauseUnitGrammar")),
+        ),
+        Dedent,
+    )
+
+
+@snowflake_dialect.segment(replace=True)
+class HavingClauseSegment(BaseSegment):
+    """A `HAVING` clause."""
+
+    type = "having_clause"
+    match_grammar = ansi_dialect.get_segment("HavingClauseSegment").match_grammar.copy()
+    match_grammar.terminator = match_grammar.terminator.copy(
+        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET")],
+    )
+    parse_grammar = ansi_dialect.get_segment("HavingClauseSegment").parse_grammar.copy()

src/sqlfluff/dialects/dialect_snowflake_keywords.py

@@ -183,6 +183,7 @@
 EXPLAIN
 EXTENSION
 EXTERNAL
+FETCH
 FILE
 FILE_FORMAT
 FILES
@@ -250,6 +251,7 @@
 NAME
 NAN
 NETWORK
+NEXT
 NEXTVAL
 NO
 NOCACHE
@@ -263,6 +265,7 @@
 OBJECTS
 OFFSET
 ON_ERROR
+ONLY
 OPERATE
 OPTIMIZATION
 OPTION

test/fixtures/dialects/snowflake/snowflake_limit.sql

@@ -0,0 +1,29 @@
+select c1 from testtable order by c1 limit 3;
+
+select c1 from testtable order by c1 limit 3 offset 3;
+
+select * from demo1 order by i limit null offset null;
+
+select * from demo1 order by i limit '' offset '';
+
+select * from demo1 order by i limit $$$$ offset $$$$;
+
+select c1 from testtable order by c1 fetch 3;
+
+select c1 from testtable order by c1 fetch first 3;
+
+select c1 from testtable order by c1 fetch next 3;
+
+select c1 from testtable order by c1 fetch 1 row;
+
+select c1 from testtable order by c1 fetch 3 rows;
+
+select c1 from testtable order by c1 fetch 3 only;
+
+select c1 from testtable order by c1 offset 3 fetch 3;
+
+select c1 from testtable order by c1 offset 1 row fetch 1 row;
+
+select c1 from testtable order by c1 offset 3 rows fetch 3 rows;
+
+select c1 from testtable offset 3 fetch 3;