Support 'SET ROLE' statement (apache#455)

src/ast/mod.rs

@@ -1013,6 +1013,16 @@ pub enum Statement {
  /// Note: this is a PostgreSQL-specific statement,
  /// but may also compatible with other SQL.
  Discard { object_type: DiscardObject },
+ /// SET [ SESSION | LOCAL ] ROLE role_name
+ ///
+ /// Note: this is a PostgreSQL-specific statement,
+ /// but may also compatible with other SQL.
+ SetRole {
+ local: bool,
+ // SESSION is the default if neither SESSION nor LOCAL appears.
+ session: bool,
+ role_name: Option<Ident>,
+ },
  /// SET <variable>
  ///
  /// Note: this is not a standard SQL statement, but it is supported by at
@@ -1755,6 +1765,24 @@ impl fmt::Display for Statement {
  write!(f, "DISCARD {object_type}", object_type = object_type)?;
  Ok(())
  }
+ Statement::SetRole {
+ local,
+ session,
+ role_name,
+ } => {
+ write!(
+ f,
+ "SET {local}{session}ROLE",
+ local = if *local { "LOCAL " } else { "" },
+ session = if *session { "SESSION " } else { "" },
+ )?;
+ if let Some(role_name) = role_name {
+ write!(f, " {}", role_name)?;
+ } else {
+ f.write_str(" NONE")?;
+ }
+ Ok(())
+ }
  Statement::SetVariable { key_values } => {
  f.write_str("SET ")?;
 

src/parser.rs

@@ -3681,6 +3681,17 @@ impl<'a> Parser<'a> {
  .to_vec(),
  });
  }
+ } else if self.parse_keyword(Keyword::ROLE) {
+ let role_name = if self.parse_keyword(Keyword::NONE) {
+ None
+ } else {
+ Some(self.parse_identifier()?)
+ };
+ return Ok(Statement::SetRole {
+ local: modifier == Some(Keyword::LOCAL),
+ session: modifier == Some(Keyword::SESSION),
+ role_name,
+ });
  }
 
  let mut key_values: Vec<SetVariableKeyValue> = vec![];

tests/sqlparser_postgres.rs

@@ -906,6 +906,45 @@ fn parse_set() {
  );
 }
 
+#[test]
+fn parse_set_role() {
+ let stmt = pg_and_generic().verified_stmt("SET SESSION ROLE NONE");
+ assert_eq!(
+ stmt,
+ Statement::SetRole {
+ local: false,
+ session: true,
+ role_name: None,
+ }
+ );
+
+ let stmt = pg_and_generic().verified_stmt("SET LOCAL ROLE \"rolename\"");
+ assert_eq!(
+ stmt,
+ Statement::SetRole {
+ local: true,
+ session: false,
+ role_name: Some(Ident {
+ value: "rolename".to_string(),
+ quote_style: Some('\"'),
+ }),
+ }
+ );
+
+ let stmt = pg_and_generic().verified_stmt("SET ROLE 'rolename'");
+ assert_eq!(
+ stmt,
+ Statement::SetRole {
+ local: false,
+ session: false,
+ role_name: Some(Ident {
+ value: "rolename".to_string(),
+ quote_style: Some('\''),
+ }),
+ }
+ );
+}
+
 #[test]
 fn parse_show() {
  let stmt = pg_and_generic().verified_stmt("SHOW a a");