From 903b24840d3d5ed2a383629b98cf3b7a83d3b1d6 Mon Sep 17 00:00:00 2001 From: Yoav Cohen Date: Tue, 3 Dec 2024 18:39:23 +0100 Subject: [PATCH 1/2] Parse Snowflake USE ROLE and USE SECONDARY ROLES --- src/ast/dcl.rs | 41 ++++++++++++++++++++++++++++++------ src/ast/mod.rs | 4 +++- src/ast/spans.rs | 17 ++++++++++----- src/keywords.rs | 2 ++ src/parser/mod.rs | 39 +++++++++++++++++++++++++++------- tests/sqlparser_snowflake.rs | 18 ++++++++++++++++ 6 files changed, 100 insertions(+), 21 deletions(-) diff --git a/src/ast/dcl.rs b/src/ast/dcl.rs index d47476ffa..735ab0cce 100644 --- a/src/ast/dcl.rs +++ b/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: +#[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), +} + +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)), + } + } +} diff --git a/src/ast/mod.rs b/src/ast/mod.rs index f782b3634..bc4dda349 100644 --- a/src/ast/mod.rs +++ b/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, diff --git a/src/ast/spans.rs b/src/ast/spans.rs index a54394174..cd3bda1c2 100644 --- a/src/ast/spans.rs +++ b/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(), } diff --git a/src/keywords.rs b/src/keywords.rs index be3910f8f..bd97c3c98 100644 --- a/src/keywords.rs +++ b/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, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 04a103c61..b5365b51d 100644 --- a/src/parser/mod.rs +++ b/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 { + 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 { let relation = self.parse_table_factor()?; // Note that for keywords to be properly handled here, they need to be diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index e31811c2b..fa81adc80 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -2709,6 +2709,20 @@ fn parse_use() { Ident::with_quote(quote, "my_schema") ]))) ); + std::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(), + )]))) + ); + std::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 @@ -2719,6 +2733,10 @@ fn parse_use() { 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] From 61199fe57089eae09b17f9cc216efa26fa5191c1 Mon Sep 17 00:00:00 2001 From: Yoav Cohen Date: Thu, 5 Dec 2024 17:08:21 +0100 Subject: [PATCH 2/2] Remove fully qualified import for asset_eq --- tests/sqlparser_common.rs | 14 +++++++------- tests/sqlparser_snowflake.rs | 20 ++++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 61c742dac..42616d51e 100644 --- a/tests/sqlparser_common.rs +++ b/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())), diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index fa81adc80..fb8a60cfa 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -2649,7 +2649,7 @@ 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() @@ -2657,7 +2657,7 @@ fn parse_use() { ); for "e in "e_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 "e in "e_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,35 +2688,35 @@ fn parse_use() { for "e in "e_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") ]))) ); - std::assert_eq!( + 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(), )]))) ); - std::assert_eq!( + assert_eq!( snowflake().verified_stmt(&format!("USE WAREHOUSE {0}my_wh{0}", quote)), Statement::Use(Use::Warehouse(ObjectName(vec![Ident::with_quote( quote, @@ -2728,7 +2728,7 @@ fn parse_use() { // 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()), );