Parse Snowflake USE ROLE and USE SECONDARY ROLES (#1578)

src/ast/dcl.rs

@@ -28,7 +28,7 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
-use super::{Expr, Ident, Password};
+use super::{display_comma_separated, Expr, Ident, Password};
 use crate::ast::{display_separated, ObjectName};
 
 /// An option in `ROLE` statement.
@@ -204,12 +204,14 @@ impl fmt::Display for AlterRoleOperation {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum Use {
-    Catalog(ObjectName),   // e.g. `USE CATALOG foo.bar`
-    Schema(ObjectName),    // e.g. `USE SCHEMA foo.bar`
-    Database(ObjectName),  // e.g. `USE DATABASE foo.bar`
-    Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
-    Object(ObjectName),    // e.g. `USE foo.bar`
-    Default,               // e.g. `USE DEFAULT`
+    Catalog(ObjectName),            // e.g. `USE CATALOG foo.bar`
+    Schema(ObjectName),             // e.g. `USE SCHEMA foo.bar`
+    Database(ObjectName),           // e.g. `USE DATABASE foo.bar`
+    Warehouse(ObjectName),          // e.g. `USE WAREHOUSE foo.bar`
+    Role(ObjectName),               // e.g. `USE ROLE PUBLIC`
+    SecondaryRoles(SecondaryRoles), // e.g. `USE SECONDARY ROLES ALL`
+    Object(ObjectName),             // e.g. `USE foo.bar`
+    Default,                        // e.g. `USE DEFAULT`
 }
 
 impl fmt::Display for Use {
@@ -220,8 +222,33 @@ impl fmt::Display for Use {
             Use::Schema(name) => write!(f, "SCHEMA {}", name),
             Use::Database(name) => write!(f, "DATABASE {}", name),
             Use::Warehouse(name) => write!(f, "WAREHOUSE {}", name),
+            Use::Role(name) => write!(f, "ROLE {}", name),
+            Use::SecondaryRoles(secondary_roles) => {
+                write!(f, "SECONDARY ROLES {}", secondary_roles)
+            }
             Use::Object(name) => write!(f, "{}", name),
             Use::Default => write!(f, "DEFAULT"),
         }
     }
 }
+
+/// Snowflake `SECONDARY ROLES` USE variant
+/// See: <https://docs.snowflake.com/en/sql-reference/sql/use-secondary-roles>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum SecondaryRoles {
+    All,
+    None,
+    List(Vec<Ident>),
+}
+
+impl fmt::Display for SecondaryRoles {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SecondaryRoles::All => write!(f, "ALL"),
+            SecondaryRoles::None => write!(f, "NONE"),
+            SecondaryRoles::List(roles) => write!(f, "{}", display_comma_separated(roles)),
+        }
+    }
+}

src/ast/mod.rs

@@ -43,7 +43,9 @@ pub use self::data_type::{
     ArrayElemTypeDef, CharLengthUnits, CharacterLength, DataType, EnumMember, ExactNumberInfo,
     StructBracketKind, TimezoneInfo,
 };
-pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
+pub use self::dcl::{
+    AlterRoleOperation, ResetConfig, RoleOption, SecondaryRoles, SetConfigValue, Use,
+};
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterPolicyOperation, AlterTableOperation,
     ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef, ColumnPolicy, ColumnPolicyProperty,

src/ast/spans.rs

@@ -3,11 +3,11 @@ use core::iter;
 use crate::tokenizer::Span;
 
 use super::{
-    AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array, Assignment,
-    AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef,
-    ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
-    CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
-    ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
+    dcl::SecondaryRoles, AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array,
+    Assignment, AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption,
+    ColumnOptionDef, ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex,
+    CreateTable, CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem,
+    Expr, ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
     FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
     IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator,
     JsonPath, JsonPathElem, LateralView, MatchRecognizePattern, Measure, NamedWindowDefinition,
@@ -484,6 +484,13 @@ impl Spanned for Use {
             Use::Schema(object_name) => object_name.span(),
             Use::Database(object_name) => object_name.span(),
             Use::Warehouse(object_name) => object_name.span(),
+            Use::Role(object_name) => object_name.span(),
+            Use::SecondaryRoles(secondary_roles) => {
+                if let SecondaryRoles::List(roles) = secondary_roles {
+                    return union_spans(roles.iter().map(|i| i.span));
+                }
+                Span::empty()
+            }
             Use::Object(object_name) => object_name.span(),
             Use::Default => Span::empty(),
         }

src/keywords.rs

@@ -664,6 +664,7 @@ define_keywords!(
     RIGHT,
     RLIKE,
     ROLE,
+    ROLES,
     ROLLBACK,
     ROLLUP,
     ROOT,
@@ -682,6 +683,7 @@ define_keywords!(
     SCROLL,
     SEARCH,
     SECOND,
+    SECONDARY,
     SECRET,
     SECURITY,
     SELECT,

src/parser/mod.rs

@@ -10093,23 +10093,46 @@ impl<'a> Parser<'a> {
         } else if dialect_of!(self is DatabricksDialect) {
             self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
         } else if dialect_of!(self is SnowflakeDialect) {
-            self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
+            self.parse_one_of_keywords(&[
+                Keyword::DATABASE,
+                Keyword::SCHEMA,
+                Keyword::WAREHOUSE,
+                Keyword::ROLE,
+                Keyword::SECONDARY,
+            ])
         } else {
             None // No specific keywords for other dialects, including GenericDialect
         };
 
-        let obj_name = self.parse_object_name(false)?;
-        let result = match parsed_keyword {
-            Some(Keyword::CATALOG) => Use::Catalog(obj_name),
-            Some(Keyword::DATABASE) => Use::Database(obj_name),
-            Some(Keyword::SCHEMA) => Use::Schema(obj_name),
-            Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
-            _ => Use::Object(obj_name),
+        let result = if matches!(parsed_keyword, Some(Keyword::SECONDARY)) {
+            self.parse_secondary_roles()?
+        } else {
+            let obj_name = self.parse_object_name(false)?;
+            match parsed_keyword {
+                Some(Keyword::CATALOG) => Use::Catalog(obj_name),
+                Some(Keyword::DATABASE) => Use::Database(obj_name),
+                Some(Keyword::SCHEMA) => Use::Schema(obj_name),
+                Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
+                Some(Keyword::ROLE) => Use::Role(obj_name),
+                _ => Use::Object(obj_name),
+            }
         };
 
         Ok(Statement::Use(result))
     }
 
+    fn parse_secondary_roles(&mut self) -> Result<Use, ParserError> {
+        self.expect_keyword(Keyword::ROLES)?;
+        if self.parse_keyword(Keyword::NONE) {
+            Ok(Use::SecondaryRoles(SecondaryRoles::None))
+        } else if self.parse_keyword(Keyword::ALL) {
+            Ok(Use::SecondaryRoles(SecondaryRoles::All))
+        } else {
+            let roles = self.parse_comma_separated(|parser| parser.parse_identifier(false))?;
+            Ok(Use::SecondaryRoles(SecondaryRoles::List(roles)))
+        }
+    }
+
     pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {
         let relation = self.parse_table_factor()?;
         // Note that for keywords to be properly handled here, they need to be

tests/sqlparser_common.rs

@@ -4066,8 +4066,8 @@ fn test_alter_table_with_on_cluster() {
         Statement::AlterTable {
             name, on_cluster, ..
         } => {
-            std::assert_eq!(name.to_string(), "t");
-            std::assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
+            assert_eq!(name.to_string(), "t");
+            assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
         }
         _ => unreachable!(),
     }
@@ -4078,15 +4078,15 @@ fn test_alter_table_with_on_cluster() {
         Statement::AlterTable {
             name, on_cluster, ..
         } => {
-            std::assert_eq!(name.to_string(), "t");
-            std::assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
+            assert_eq!(name.to_string(), "t");
+            assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
         }
         _ => unreachable!(),
     }
 
     let res = all_dialects()
         .parse_sql_statements("ALTER TABLE t ON CLUSTER 123 ADD CONSTRAINT bar PRIMARY KEY (baz)");
-    std::assert_eq!(
+    assert_eq!(
         res.unwrap_err(),
         ParserError::ParserError("Expected: identifier, found: 123".to_string())
     )
@@ -11226,7 +11226,7 @@ fn test_group_by_nothing() {
     let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
         .verified_only_select("SELECT count(1) FROM t GROUP BY ()");
     {
-        std::assert_eq!(
+        assert_eq!(
             GroupByExpr::Expressions(vec![Expr::Tuple(vec![])], vec![]),
             group_by
         );
@@ -11235,7 +11235,7 @@ fn test_group_by_nothing() {
     let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
         .verified_only_select("SELECT name, count(1) FROM t GROUP BY name, ()");
     {
-        std::assert_eq!(
+        assert_eq!(
             GroupByExpr::Expressions(
                 vec![
                     Identifier(Ident::new("name".to_string())),

tests/sqlparser_snowflake.rs

@@ -2649,15 +2649,15 @@ fn parse_use() {
     let quote_styles = ['\'', '"', '`'];
     for object_name in &valid_object_names {
         // Test single identifier without quotes
-        std::assert_eq!(
+        assert_eq!(
             snowflake().verified_stmt(&format!("USE {}", object_name)),
             Statement::Use(Use::Object(ObjectName(vec![Ident::new(
                 object_name.to_string()
             )])))
         );
         for &quote in &quote_styles {
             // Test single identifier with different type of quotes
-            std::assert_eq!(
+            assert_eq!(
                 snowflake().verified_stmt(&format!("USE {}{}{}", quote, object_name, quote)),
                 Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
                     quote,
@@ -2669,7 +2669,7 @@ fn parse_use() {
 
     for &quote in &quote_styles {
         // Test double identifier with different type of quotes
-        std::assert_eq!(
+        assert_eq!(
             snowflake().verified_stmt(&format!("USE {0}CATALOG{0}.{0}my_schema{0}", quote)),
             Statement::Use(Use::Object(ObjectName(vec![
                 Ident::with_quote(quote, "CATALOG"),
@@ -2678,7 +2678,7 @@ fn parse_use() {
         );
     }
     // Test double identifier without quotes
-    std::assert_eq!(
+    assert_eq!(
         snowflake().verified_stmt("USE mydb.my_schema"),
         Statement::Use(Use::Object(ObjectName(vec![
             Ident::new("mydb"),
@@ -2688,37 +2688,55 @@ fn parse_use() {
 
     for &quote in &quote_styles {
         // Test single and double identifier with keyword and different type of quotes
-        std::assert_eq!(
+        assert_eq!(
             snowflake().verified_stmt(&format!("USE DATABASE {0}my_database{0}", quote)),
             Statement::Use(Use::Database(ObjectName(vec![Ident::with_quote(
                 quote,
                 "my_database".to_string(),
             )])))
         );
-        std::assert_eq!(
+        assert_eq!(
             snowflake().verified_stmt(&format!("USE SCHEMA {0}my_schema{0}", quote)),
             Statement::Use(Use::Schema(ObjectName(vec![Ident::with_quote(
                 quote,
                 "my_schema".to_string(),
             )])))
         );
-        std::assert_eq!(
+        assert_eq!(
             snowflake().verified_stmt(&format!("USE SCHEMA {0}CATALOG{0}.{0}my_schema{0}", quote)),
             Statement::Use(Use::Schema(ObjectName(vec![
                 Ident::with_quote(quote, "CATALOG"),
                 Ident::with_quote(quote, "my_schema")
             ])))
         );
+        assert_eq!(
+            snowflake().verified_stmt(&format!("USE ROLE {0}my_role{0}", quote)),
+            Statement::Use(Use::Role(ObjectName(vec![Ident::with_quote(
+                quote,
+                "my_role".to_string(),
+            )])))
+        );
+        assert_eq!(
+            snowflake().verified_stmt(&format!("USE WAREHOUSE {0}my_wh{0}", quote)),
+            Statement::Use(Use::Warehouse(ObjectName(vec![Ident::with_quote(
+                quote,
+                "my_wh".to_string(),
+            )])))
+        );
     }
 
     // Test invalid syntax - missing identifier
     let invalid_cases = ["USE SCHEMA", "USE DATABASE", "USE WAREHOUSE"];
     for sql in &invalid_cases {
-        std::assert_eq!(
+        assert_eq!(
             snowflake().parse_sql_statements(sql).unwrap_err(),
             ParserError::ParserError("Expected: identifier, found: EOF".to_string()),
         );
     }
+
+    snowflake().verified_stmt("USE SECONDARY ROLES ALL");
+    snowflake().verified_stmt("USE SECONDARY ROLES NONE");
+    snowflake().verified_stmt("USE SECONDARY ROLES r1, r2, r3");
 }
 
 #[test]