Grammar: Tidy up Snowflake/MySQL/HIVE USE statements (#2187)

* Grammar: Tidy up USE statements

* Update src/sqlfluff/dialects/dialect_snowflake.py

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

* Update src/sqlfluff/dialects/dialect_ansi.py

* black format

* Update dialect_snowflake.py

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

src/sqlfluff/dialects/dialect_ansi.py

@@ -3095,24 +3095,12 @@ class DescribeStatementSegment(BaseSegment):
 
 @ansi_dialect.segment()
 class UseStatementSegment(BaseSegment):
-    """A `USE` statement.
-
-    USE [ ROLE ] <name>
-
-    USE [ WAREHOUSE ] <name>
-
-    USE [ DATABASE ] <name>
-
-    USE [ SCHEMA ] [<db_name>.]<name>
-    """
+    """A `USE` statement."""
 
     type = "use_statement"
-    match_grammar = StartsWith("USE")
-
-    parse_grammar = Sequence(
+    match_grammar = Sequence(
         "USE",
-        OneOf("ROLE", "WAREHOUSE", "DATABASE", "SCHEMA", optional=True),
-        Ref("ObjectReferenceSegment"),
+        Ref("DatabaseReferenceSegment"),
     )
 
 

src/sqlfluff/dialects/dialect_hive.py

@@ -496,17 +496,6 @@ class TruncateStatementSegment(BaseSegment):
     )
 
 
-@hive_dialect.segment(replace=True)
-class UseStatementSegment(BaseSegment):
-    """An `USE` statement."""
-
-    type = "use_statement"
-    match_grammar = Sequence(
-        "USE",
-        Ref("DatabaseReferenceSegment"),
-    )
-
-
 @hive_dialect.segment(replace=True)
 class StatementSegment(ansi_dialect.get_segment("StatementSegment")):  # type: ignore
     """Overriding StatementSegment to allow for additional segment parsing."""

src/sqlfluff/dialects/dialect_snowflake.py

@@ -539,7 +539,6 @@ class StatementSegment(ansi_dialect.get_segment("StatementSegment")):  # type: i
 
     parse_grammar = ansi_dialect.get_segment("StatementSegment").parse_grammar.copy(
         insert=[
-            Ref("UseStatementSegment"),
             Ref("CreateStatementSegment"),
             Ref("CreateTaskSegment"),
             Ref("CreateCloneStatementSegment"),
@@ -3519,3 +3518,36 @@ class CommentStatementSegment(BaseSegment):
         "IS",
         Ref("QuotedLiteralSegment"),
     )
+
+
+@snowflake_dialect.segment(replace=True)
+class UseStatementSegment(BaseSegment):
+    """A `USE` statement.
+
+    https://docs.snowflake.com/en/sql-reference/sql/use.html
+    """
+
+    type = "use_statement"
+    match_grammar = Sequence(
+        "USE",
+        OneOf(
+            Sequence("ROLE", Ref("ObjectReferenceSegment")),
+            Sequence("WAREHOUSE", Ref("ObjectReferenceSegment")),
+            Sequence(
+                Ref.keyword("DATABASE", optional=True),
+                Ref("DatabaseReferenceSegment"),
+            ),
+            Sequence(
+                Ref.keyword("SCHEMA", optional=True),
+                Ref("SchemaReferenceSegment"),
+            ),
+            Sequence(
+                "SECONDARY",
+                "ROLES",
+                OneOf(
+                    "ALL",
+                    "NONE",
+                ),
+            ),
+        ),
+    )

src/sqlfluff/dialects/dialect_snowflake_keywords.py

@@ -328,6 +328,7 @@
 SCHEDULE
 SCHEMAS
 SECOND
+SECONDARY
 SECURE
 SECURITY
 SEED

test/fixtures/dialects/mysql/use_database.yml

@@ -3,11 +3,11 @@
 # 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: 33fd349f0a269928e8467c2d5e344540bd7451fd757a92fc2497cf14c1dea31a
+_hash: 0a7521fef2b73d430f586f5eb3db1363de28f5d740a86a1a4644f0fb7f289566
 file:
   statement:
     use_statement:
       keyword: use
-      object_reference:
+      database_reference:
         identifier: my_db
   statement_terminator: ;

test/fixtures/dialects/mysql/use_statement.sql

@@ -0,0 +1 @@
+USE db;

test/fixtures/dialects/mysql/use_statement.yml

@@ -0,0 +1,13 @@
+# YML test files are auto-generated from SQL files and should not be edited by
+# hand. To help enforce this, the "hash" field in the file must match a hash
+# 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: 2224c350cd615fa21509176099d7565e9e369a63294787cea56a3010591bcedc
+file:
+  statement:
+    use_statement:
+      keyword: USE
+      database_reference:
+        identifier: db
+  statement_terminator: ;

test/fixtures/dialects/snowflake/snowflake_use.sql

@@ -19,3 +19,7 @@ USE SCHEMA "MY_DATABASE"."MY_SCHEMA";
 USE SCHEMA "MY_SCHEMA";
 
 USE "MY_DATABASE"."MY_SCHEMA";
+
+USE SECONDARY ROLES ALL;
+
+USE SECONDARY ROLES NONE;

test/fixtures/dialects/snowflake/snowflake_use.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: ac5358b10e8c24fcb3d4d780a0a7926c3d2499f6cf85759aa9610a217e5a2dba
+_hash: 304ece91d8540fd10a09ebec75402ff76d951690615eac803dee97975aa9c73f
 file:
 - statement:
     use_statement:
@@ -23,14 +23,14 @@ file:
     use_statement:
     - keyword: use
     - keyword: database
-    - object_reference:
+    - database_reference:
         identifier: my_database
 - statement_terminator: ;
 - statement:
     use_statement:
     - keyword: use
     - keyword: schema
-    - object_reference:
+    - schema_reference:
         identifier: my_schema
 - statement_terminator: ;
 - statement:
@@ -51,20 +51,20 @@ file:
     use_statement:
     - keyword: USE
     - keyword: DATABASE
-    - object_reference:
+    - database_reference:
         identifier: '"MY_DATABASE"'
 - statement_terminator: ;
 - statement:
     use_statement:
       keyword: USE
-      object_reference:
+      database_reference:
         identifier: '"MY_DATABASE"'
 - statement_terminator: ;
 - statement:
     use_statement:
     - keyword: USE
     - keyword: SCHEMA
-    - object_reference:
+    - schema_reference:
       - identifier: '"MY_DATABASE"'
       - dot: .
       - identifier: '"MY_SCHEMA"'
@@ -73,14 +73,28 @@ file:
     use_statement:
     - keyword: USE
     - keyword: SCHEMA
-    - object_reference:
+    - schema_reference:
         identifier: '"MY_SCHEMA"'
 - statement_terminator: ;
 - statement:
     use_statement:
       keyword: USE
-      object_reference:
+      database_reference:
       - identifier: '"MY_DATABASE"'
       - dot: .
       - identifier: '"MY_SCHEMA"'
 - statement_terminator: ;
+- statement:
+    use_statement:
+    - keyword: USE
+    - keyword: SECONDARY
+    - keyword: ROLES
+    - keyword: ALL
+- statement_terminator: ;
+- statement:
+    use_statement:
+    - keyword: USE
+    - keyword: SECONDARY
+    - keyword: ROLES
+    - keyword: NONE
+- statement_terminator: ;