Snowflake: fix connect by prior selects (#5996)

Co-authored-by: Gergely Havlicsek <ghavlicsek@apple.com>
sqlfluff · Jul 3, 2024 · 77d40d2 · 77d40d2
1 parent 9d65692
commit 77d40d2
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 13 deletions.
diff --git a/src/sqlfluff/dialects/dialect_snowflake.py b/src/sqlfluff/dialects/dialect_snowflake.py
@@ -1102,19 +1102,40 @@ class ConnectByClauseSegment(BaseSegment):
     """
 
     type = "connectby_clause"
-    match_grammar = Sequence(
-        "START",
-        "WITH",
-        Ref("ExpressionSegment"),
-        "CONNECT",
-        "BY",
-        Delimited(
+    match_grammar = OneOf(
+        Sequence(
+            "START",
+            "WITH",
+            Ref("ExpressionSegment"),
+            "CONNECT",
+            "BY",
+            Delimited(
+                Sequence(
+                    Ref.keyword("PRIOR", optional=True),
+                    Ref("ColumnReferenceSegment"),
+                    Ref("EqualsSegment"),
+                    Ref.keyword("PRIOR", optional=True),
+                    Ref("ColumnReferenceSegment"),
+                ),
+            ),
+        ),
+        Sequence(
+            "CONNECT",
+            "BY",
+            Delimited(
+                Sequence(
+                    Ref.keyword("PRIOR", optional=True),
+                    Ref("ColumnReferenceSegment"),
+                    Ref("EqualsSegment"),
+                    Ref("ColumnReferenceSegment"),
+                ),
+                delimiter="AND",
+            ),
             Sequence(
-                Ref.keyword("PRIOR", optional=True),
-                Ref("ColumnReferenceSegment"),
-                Ref("EqualsSegment"),
-                Ref.keyword("PRIOR", optional=True),
-                Ref("ColumnReferenceSegment"),
+                "START",
+                "WITH",
+                Ref("ExpressionSegment"),
+                optional=True,
             ),
         ),
     )

diff --git a/test/fixtures/dialects/snowflake/connect_by.sql b/test/fixtures/dialects/snowflake/connect_by.sql
@@ -32,3 +32,13 @@ start with title = 'President'
 connect by
     manager_id = prior employee_id
 order by employee_id;
+
+select
+  description,
+  quantity,
+  component_id,
+  parent_component_id,
+  component_type
+from components c
+connect by prior c.parent_component_id = c.component_id AND PRIOR c.component_type = c.component_type
+order by quantity;
diff --git a/test/fixtures/dialects/snowflake/connect_by.yml b/test/fixtures/dialects/snowflake/connect_by.yml
@@ -3,7 +3,7 @@
 # computed by SQLFluff when running the tests. Please run
 # `python test/generate_parse_fixture_yml.py`  to generate them after adding or
 # altering SQL files.
-_hash: b2ea3dbc494b78abe541194309260d24ae0a2a77bf63f9da74da0b73cdb74e19
+_hash: 6cf317f99f2d83dd1f2fa16620d22f7640c0da6447e2991be8b07f194ce94671
 file:
 - statement:
     select_statement:
@@ -232,3 +232,67 @@ file:
       - column_reference:
           naked_identifier: employee_id
 - statement_terminator: ;
+- statement:
+    select_statement:
+      select_clause:
+      - keyword: select
+      - select_clause_element:
+          column_reference:
+            naked_identifier: description
+      - comma: ','
+      - select_clause_element:
+          column_reference:
+            naked_identifier: quantity
+      - comma: ','
+      - select_clause_element:
+          column_reference:
+            naked_identifier: component_id
+      - comma: ','
+      - select_clause_element:
+          column_reference:
+            naked_identifier: parent_component_id
+      - comma: ','
+      - select_clause_element:
+          column_reference:
+            naked_identifier: component_type
+      from_clause:
+        keyword: from
+        from_expression:
+          from_expression_element:
+            table_expression:
+              table_reference:
+                naked_identifier: components
+            alias_expression:
+              naked_identifier: c
+          connectby_clause:
+          - keyword: connect
+          - keyword: by
+          - keyword: prior
+          - column_reference:
+            - naked_identifier: c
+            - dot: .
+            - naked_identifier: parent_component_id
+          - comparison_operator:
+              raw_comparison_operator: '='
+          - column_reference:
+            - naked_identifier: c
+            - dot: .
+            - naked_identifier: component_id
+          - keyword: AND
+          - keyword: PRIOR
+          - column_reference:
+            - naked_identifier: c
+            - dot: .
+            - naked_identifier: component_type
+          - comparison_operator:
+              raw_comparison_operator: '='
+          - column_reference:
+            - naked_identifier: c
+            - dot: .
+            - naked_identifier: component_type
+      orderby_clause:
+      - keyword: order
+      - keyword: by
+      - column_reference:
+          naked_identifier: quantity
+- statement_terminator: ;