Skip to content

Commit

Permalink
Add support for various Snowflake grantees (#1640)
Browse files Browse the repository at this point in the history
  • Loading branch information
yoavcloud authored Jan 5, 2025
1 parent 02d60cc commit e23877c
Show file tree
Hide file tree
Showing 4 changed files with 128 additions and 2 deletions.
62 changes: 61 additions & 1 deletion src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3159,7 +3159,7 @@ pub enum Statement {
Grant {
privileges: Privileges,
objects: GrantObjects,
grantees: Vec<Ident>,
grantees: Vec<Grantee>,
with_grant_option: bool,
granted_by: Option<Ident>,
},
Expand Down Expand Up @@ -5366,6 +5366,66 @@ impl fmt::Display for Action {
}
}

/// The principal that receives the privileges
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Grantee {
pub grantee_type: GranteesType,
pub name: Option<ObjectName>,
}

impl fmt::Display for Grantee {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self.grantee_type {
GranteesType::Role => {
write!(f, "ROLE ")?;
}
GranteesType::Share => {
write!(f, "SHARE ")?;
}
GranteesType::User => {
write!(f, "USER ")?;
}
GranteesType::Group => {
write!(f, "GROUP ")?;
}
GranteesType::Public => {
write!(f, "PUBLIC ")?;
}
GranteesType::DatabaseRole => {
write!(f, "DATABASE ROLE ")?;
}
GranteesType::Application => {
write!(f, "APPLICATION ")?;
}
GranteesType::ApplicationRole => {
write!(f, "APPLICATION ROLE ")?;
}
GranteesType::None => (),
}
if let Some(ref name) = self.name {
write!(f, "{}", name)?;
}
Ok(())
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum GranteesType {
Role,
Share,
User,
Group,
Public,
DatabaseRole,
Application,
ApplicationRole,
None,
}

/// Objects on which privileges are granted in a GRANT statement.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -615,6 +615,7 @@ define_keywords!(
PROCEDURE,
PROGRAM,
PROJECTION,
PUBLIC,
PURGE,
QUALIFY,
QUARTER,
Expand Down
58 changes: 57 additions & 1 deletion src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11611,7 +11611,7 @@ impl<'a> Parser<'a> {
let (privileges, objects) = self.parse_grant_revoke_privileges_objects()?;

self.expect_keyword_is(Keyword::TO)?;
let grantees = self.parse_comma_separated(|p| p.parse_identifier())?;
let grantees = self.parse_grantees()?;

let with_grant_option =
self.parse_keywords(&[Keyword::WITH, Keyword::GRANT, Keyword::OPTION]);
Expand All @@ -11629,6 +11629,62 @@ impl<'a> Parser<'a> {
})
}

fn parse_grantees(&mut self) -> Result<Vec<Grantee>, ParserError> {
let mut values = vec![];
let mut grantee_type = GranteesType::None;
loop {
grantee_type = if self.parse_keyword(Keyword::ROLE) {
GranteesType::Role
} else if self.parse_keyword(Keyword::USER) {
GranteesType::User
} else if self.parse_keyword(Keyword::SHARE) {
GranteesType::Share
} else if self.parse_keyword(Keyword::GROUP) {
GranteesType::Group
} else if self.parse_keyword(Keyword::PUBLIC) {
GranteesType::Public
} else if self.parse_keywords(&[Keyword::DATABASE, Keyword::ROLE]) {
GranteesType::DatabaseRole
} else if self.parse_keywords(&[Keyword::APPLICATION, Keyword::ROLE]) {
GranteesType::ApplicationRole
} else if self.parse_keyword(Keyword::APPLICATION) {
GranteesType::Application
} else {
grantee_type // keep from previous iteraton, if not specified
};

let grantee = if grantee_type == GranteesType::Public {
Grantee {
grantee_type: grantee_type.clone(),
name: None,
}
} else {
let mut name = self.parse_object_name(false)?;
if self.consume_token(&Token::Colon) {
// Redshift supports namespace prefix for extenrnal users and groups:
// <Namespace>:<GroupName> or <Namespace>:<UserName>
// https://docs.aws.amazon.com/redshift/latest/mgmt/redshift-iam-access-control-native-idp.html
let ident = self.parse_identifier()?;
if let Some(n) = name.0.first() {
name = ObjectName(vec![Ident::new(format!("{}:{}", n.value, ident.value))]);
};
}
Grantee {
grantee_type: grantee_type.clone(),
name: Some(name),
}
};

values.push(grantee);

if !self.consume_token(&Token::Comma) {
break;
}
}

Ok(values)
}

pub fn parse_grant_revoke_privileges_objects(
&mut self,
) -> Result<(Privileges, GrantObjects), ParserError> {
Expand Down
9 changes: 9 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8497,6 +8497,15 @@ fn parse_grant() {
},
_ => unreachable!(),
}

verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO ROLE role1");
verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO ROLE role1 WITH GRANT OPTION");
verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO DATABASE ROLE role1");
verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO APPLICATION role1");
verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO APPLICATION ROLE role1");
verified_stmt("GRANT SELECT ON ALL TABLES IN SCHEMA db1.sc1 TO SHARE share1");
verified_stmt("GRANT USAGE ON SCHEMA sc1 TO a:b");
verified_stmt("GRANT USAGE ON SCHEMA sc1 TO GROUP group1");
}

#[test]
Expand Down

0 comments on commit e23877c

Please sign in to comment.