Skip to content

Commit 7b50ac3

Browse files
yoavcloudyoavcloud
authored
Parse Snowflake USE ROLE and USE SECONDARY ROLES (#1578)
1 parent dd7ba72 commit 7b50ac3

File tree

7 files changed

+115
-36
lines changed

7 files changed

+115
-36
lines changed

‎src/ast/dcl.rs

+34-7
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ use serde::{Deserialize, Serialize};
2828
#[cfg(feature = "visitor")]
2929
use sqlparser_derive::{Visit, VisitMut};
3030

31-
use super::{Expr, Ident, Password};
31+
use super::{display_comma_separated, Expr, Ident, Password};
3232
use crate::ast::{display_separated, ObjectName};
3333

3434
/// An option in `ROLE` statement.
@@ -204,12 +204,14 @@ impl fmt::Display for AlterRoleOperation {
204204
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
205205
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
206206
pub enum Use {
207-
Catalog(ObjectName), // e.g. `USE CATALOG foo.bar`
208-
Schema(ObjectName), // e.g. `USE SCHEMA foo.bar`
209-
Database(ObjectName), // e.g. `USE DATABASE foo.bar`
210-
Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
211-
Object(ObjectName), // e.g. `USE foo.bar`
212-
Default, // e.g. `USE DEFAULT`
207+
Catalog(ObjectName), // e.g. `USE CATALOG foo.bar`
208+
Schema(ObjectName), // e.g. `USE SCHEMA foo.bar`
209+
Database(ObjectName), // e.g. `USE DATABASE foo.bar`
210+
Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
211+
Role(ObjectName), // e.g. `USE ROLE PUBLIC`
212+
SecondaryRoles(SecondaryRoles), // e.g. `USE SECONDARY ROLES ALL`
213+
Object(ObjectName), // e.g. `USE foo.bar`
214+
Default, // e.g. `USE DEFAULT`
213215
}
214216

215217
impl fmt::Display for Use {
@@ -220,8 +222,33 @@ impl fmt::Display for Use {
220222
Use::Schema(name) => write!(f, "SCHEMA {}", name),
221223
Use::Database(name) => write!(f, "DATABASE {}", name),
222224
Use::Warehouse(name) => write!(f, "WAREHOUSE {}", name),
225+
Use::Role(name) => write!(f, "ROLE {}", name),
226+
Use::SecondaryRoles(secondary_roles) => {
227+
write!(f, "SECONDARY ROLES {}", secondary_roles)
228+
}
223229
Use::Object(name) => write!(f, "{}", name),
224230
Use::Default => write!(f, "DEFAULT"),
225231
}
226232
}
227233
}
234+
235+
/// Snowflake `SECONDARY ROLES` USE variant
236+
/// See: <https://docs.snowflake.com/en/sql-reference/sql/use-secondary-roles>
237+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
238+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
239+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
240+
pub enum SecondaryRoles {
241+
All,
242+
None,
243+
List(Vec<Ident>),
244+
}
245+
246+
impl fmt::Display for SecondaryRoles {
247+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
248+
match self {
249+
SecondaryRoles::All => write!(f, "ALL"),
250+
SecondaryRoles::None => write!(f, "NONE"),
251+
SecondaryRoles::List(roles) => write!(f, "{}", display_comma_separated(roles)),
252+
}
253+
}
254+
}

‎src/ast/mod.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@ pub use self::data_type::{
4343
ArrayElemTypeDef, CharLengthUnits, CharacterLength, DataType, EnumMember, ExactNumberInfo,
4444
StructBracketKind, TimezoneInfo,
4545
};
46-
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
46+
pub use self::dcl::{
47+
AlterRoleOperation, ResetConfig, RoleOption, SecondaryRoles, SetConfigValue, Use,
48+
};
4749
pub use self::ddl::{
4850
AlterColumnOperation, AlterIndexOperation, AlterPolicyOperation, AlterTableOperation,
4951
ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef, ColumnPolicy, ColumnPolicyProperty,

‎src/ast/spans.rs

+12-5
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,11 @@ use core::iter;
33
use crate::tokenizer::Span;
44

55
use super::{
6-
AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array, Assignment,
7-
AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef,
8-
ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
9-
CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
10-
ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
6+
dcl::SecondaryRoles, AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Array,
7+
Assignment, AssignmentTarget, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption,
8+
ColumnOptionDef, ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex,
9+
CreateTable, CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem,
10+
Expr, ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
1111
FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
1212
IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator,
1313
JsonPath, JsonPathElem, LateralView, MatchRecognizePattern, Measure, NamedWindowDefinition,
@@ -484,6 +484,13 @@ impl Spanned for Use {
484484
Use::Schema(object_name) => object_name.span(),
485485
Use::Database(object_name) => object_name.span(),
486486
Use::Warehouse(object_name) => object_name.span(),
487+
Use::Role(object_name) => object_name.span(),
488+
Use::SecondaryRoles(secondary_roles) => {
489+
if let SecondaryRoles::List(roles) = secondary_roles {
490+
return union_spans(roles.iter().map(|i| i.span));
491+
}
492+
Span::empty()
493+
}
487494
Use::Object(object_name) => object_name.span(),
488495
Use::Default => Span::empty(),
489496
}

‎src/keywords.rs

+2
Original file line numberDiff line numberDiff line change
@@ -664,6 +664,7 @@ define_keywords!(
664664
RIGHT,
665665
RLIKE,
666666
ROLE,
667+
ROLES,
667668
ROLLBACK,
668669
ROLLUP,
669670
ROOT,
@@ -682,6 +683,7 @@ define_keywords!(
682683
SCROLL,
683684
SEARCH,
684685
SECOND,
686+
SECONDARY,
685687
SECRET,
686688
SECURITY,
687689
SELECT,

‎src/parser/mod.rs

+31-8
Original file line numberDiff line numberDiff line change
@@ -10093,23 +10093,46 @@ impl<'a> Parser<'a> {
1009310093
} else if dialect_of!(self is DatabricksDialect) {
1009410094
self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
1009510095
} else if dialect_of!(self is SnowflakeDialect) {
10096-
self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
10096+
self.parse_one_of_keywords(&[
10097+
Keyword::DATABASE,
10098+
Keyword::SCHEMA,
10099+
Keyword::WAREHOUSE,
10100+
Keyword::ROLE,
10101+
Keyword::SECONDARY,
10102+
])
1009710103
} else {
1009810104
None // No specific keywords for other dialects, including GenericDialect
1009910105
};
1010010106

10101-
let obj_name = self.parse_object_name(false)?;
10102-
let result = match parsed_keyword {
10103-
Some(Keyword::CATALOG) => Use::Catalog(obj_name),
10104-
Some(Keyword::DATABASE) => Use::Database(obj_name),
10105-
Some(Keyword::SCHEMA) => Use::Schema(obj_name),
10106-
Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
10107-
_ => Use::Object(obj_name),
10107+
let result = if matches!(parsed_keyword, Some(Keyword::SECONDARY)) {
10108+
self.parse_secondary_roles()?
10109+
} else {
10110+
let obj_name = self.parse_object_name(false)?;
10111+
match parsed_keyword {
10112+
Some(Keyword::CATALOG) => Use::Catalog(obj_name),
10113+
Some(Keyword::DATABASE) => Use::Database(obj_name),
10114+
Some(Keyword::SCHEMA) => Use::Schema(obj_name),
10115+
Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
10116+
Some(Keyword::ROLE) => Use::Role(obj_name),
10117+
_ => Use::Object(obj_name),
10118+
}
1010810119
};
1010910120

1011010121
Ok(Statement::Use(result))
1011110122
}
1011210123

10124+
fn parse_secondary_roles(&mut self) -> Result<Use, ParserError> {
10125+
self.expect_keyword(Keyword::ROLES)?;
10126+
if self.parse_keyword(Keyword::NONE) {
10127+
Ok(Use::SecondaryRoles(SecondaryRoles::None))
10128+
} else if self.parse_keyword(Keyword::ALL) {
10129+
Ok(Use::SecondaryRoles(SecondaryRoles::All))
10130+
} else {
10131+
let roles = self.parse_comma_separated(|parser| parser.parse_identifier(false))?;
10132+
Ok(Use::SecondaryRoles(SecondaryRoles::List(roles)))
10133+
}
10134+
}
10135+
1011310136
pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {
1011410137
let relation = self.parse_table_factor()?;
1011510138
// Note that for keywords to be properly handled here, they need to be

‎tests/sqlparser_common.rs

+7-7
Original file line numberDiff line numberDiff line change
@@ -4066,8 +4066,8 @@ fn test_alter_table_with_on_cluster() {
40664066
Statement::AlterTable {
40674067
name, on_cluster, ..
40684068
} => {
4069-
std::assert_eq!(name.to_string(), "t");
4070-
std::assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
4069+
assert_eq!(name.to_string(), "t");
4070+
assert_eq!(on_cluster, Some(Ident::with_quote('\'', "cluster")));
40714071
}
40724072
_ => unreachable!(),
40734073
}
@@ -4078,15 +4078,15 @@ fn test_alter_table_with_on_cluster() {
40784078
Statement::AlterTable {
40794079
name, on_cluster, ..
40804080
} => {
4081-
std::assert_eq!(name.to_string(), "t");
4082-
std::assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
4081+
assert_eq!(name.to_string(), "t");
4082+
assert_eq!(on_cluster, Some(Ident::new("cluster_name")));
40834083
}
40844084
_ => unreachable!(),
40854085
}
40864086

40874087
let res = all_dialects()
40884088
.parse_sql_statements("ALTER TABLE t ON CLUSTER 123 ADD CONSTRAINT bar PRIMARY KEY (baz)");
4089-
std::assert_eq!(
4089+
assert_eq!(
40904090
res.unwrap_err(),
40914091
ParserError::ParserError("Expected: identifier, found: 123".to_string())
40924092
)
@@ -11226,7 +11226,7 @@ fn test_group_by_nothing() {
1122611226
let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
1122711227
.verified_only_select("SELECT count(1) FROM t GROUP BY ()");
1122811228
{
11229-
std::assert_eq!(
11229+
assert_eq!(
1123011230
GroupByExpr::Expressions(vec![Expr::Tuple(vec![])], vec![]),
1123111231
group_by
1123211232
);
@@ -11235,7 +11235,7 @@ fn test_group_by_nothing() {
1123511235
let Select { group_by, .. } = all_dialects_where(|d| d.supports_group_by_expr())
1123611236
.verified_only_select("SELECT name, count(1) FROM t GROUP BY name, ()");
1123711237
{
11238-
std::assert_eq!(
11238+
assert_eq!(
1123911239
GroupByExpr::Expressions(
1124011240
vec![
1124111241
Identifier(Ident::new("name".to_string())),

‎tests/sqlparser_snowflake.rs

+26-8
Original file line numberDiff line numberDiff line change
@@ -2649,15 +2649,15 @@ fn parse_use() {
26492649
let quote_styles = ['\'', '"', '`'];
26502650
for object_name in &valid_object_names {
26512651
// Test single identifier without quotes
2652-
std::assert_eq!(
2652+
assert_eq!(
26532653
snowflake().verified_stmt(&format!("USE {}", object_name)),
26542654
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
26552655
object_name.to_string()
26562656
)])))
26572657
);
26582658
for &quote in &quote_styles {
26592659
// Test single identifier with different type of quotes
2660-
std::assert_eq!(
2660+
assert_eq!(
26612661
snowflake().verified_stmt(&format!("USE {}{}{}", quote, object_name, quote)),
26622662
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
26632663
quote,
@@ -2669,7 +2669,7 @@ fn parse_use() {
26692669

26702670
for &quote in &quote_styles {
26712671
// Test double identifier with different type of quotes
2672-
std::assert_eq!(
2672+
assert_eq!(
26732673
snowflake().verified_stmt(&format!("USE {0}CATALOG{0}.{0}my_schema{0}", quote)),
26742674
Statement::Use(Use::Object(ObjectName(vec![
26752675
Ident::with_quote(quote, "CATALOG"),
@@ -2678,7 +2678,7 @@ fn parse_use() {
26782678
);
26792679
}
26802680
// Test double identifier without quotes
2681-
std::assert_eq!(
2681+
assert_eq!(
26822682
snowflake().verified_stmt("USE mydb.my_schema"),
26832683
Statement::Use(Use::Object(ObjectName(vec![
26842684
Ident::new("mydb"),
@@ -2688,37 +2688,55 @@ fn parse_use() {
26882688

26892689
for &quote in &quote_styles {
26902690
// Test single and double identifier with keyword and different type of quotes
2691-
std::assert_eq!(
2691+
assert_eq!(
26922692
snowflake().verified_stmt(&format!("USE DATABASE {0}my_database{0}", quote)),
26932693
Statement::Use(Use::Database(ObjectName(vec![Ident::with_quote(
26942694
quote,
26952695
"my_database".to_string(),
26962696
)])))
26972697
);
2698-
std::assert_eq!(
2698+
assert_eq!(
26992699
snowflake().verified_stmt(&format!("USE SCHEMA {0}my_schema{0}", quote)),
27002700
Statement::Use(Use::Schema(ObjectName(vec![Ident::with_quote(
27012701
quote,
27022702
"my_schema".to_string(),
27032703
)])))
27042704
);
2705-
std::assert_eq!(
2705+
assert_eq!(
27062706
snowflake().verified_stmt(&format!("USE SCHEMA {0}CATALOG{0}.{0}my_schema{0}", quote)),
27072707
Statement::Use(Use::Schema(ObjectName(vec![
27082708
Ident::with_quote(quote, "CATALOG"),
27092709
Ident::with_quote(quote, "my_schema")
27102710
])))
27112711
);
2712+
assert_eq!(
2713+
snowflake().verified_stmt(&format!("USE ROLE {0}my_role{0}", quote)),
2714+
Statement::Use(Use::Role(ObjectName(vec![Ident::with_quote(
2715+
quote,
2716+
"my_role".to_string(),
2717+
)])))
2718+
);
2719+
assert_eq!(
2720+
snowflake().verified_stmt(&format!("USE WAREHOUSE {0}my_wh{0}", quote)),
2721+
Statement::Use(Use::Warehouse(ObjectName(vec![Ident::with_quote(
2722+
quote,
2723+
"my_wh".to_string(),
2724+
)])))
2725+
);
27122726
}
27132727

27142728
// Test invalid syntax - missing identifier
27152729
let invalid_cases = ["USE SCHEMA", "USE DATABASE", "USE WAREHOUSE"];
27162730
for sql in &invalid_cases {
2717-
std::assert_eq!(
2731+
assert_eq!(
27182732
snowflake().parse_sql_statements(sql).unwrap_err(),
27192733
ParserError::ParserError("Expected: identifier, found: EOF".to_string()),
27202734
);
27212735
}
2736+
2737+
snowflake().verified_stmt("USE SECONDARY ROLES ALL");
2738+
snowflake().verified_stmt("USE SECONDARY ROLES NONE");
2739+
snowflake().verified_stmt("USE SECONDARY ROLES r1, r2, r3");
27222740
}
27232741

27242742
#[test]

0 commit comments

Comments
 (0)