Support 'SET ROLE' statement (#455)

apache · Apr 13, 2022 · 8f207db · 8f207db
1 parent 7480e89
commit 8f207db
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 0 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -841,6 +841,16 @@ pub enum Statement {
         /// deleted along with the dropped table
         purge: bool,
     },
+    /// 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
@@ -1453,6 +1463,24 @@ impl fmt::Display for Statement {
                 if *cascade { " CASCADE" } else { "" },
                 if *purge { " PURGE" } else { "" }
             ),
+            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 {
                 local,
                 variable,

diff --git a/src/parser.rs b/src/parser.rs
@@ -3116,7 +3116,19 @@ impl<'a> Parser<'a> {
             self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]);
         if let Some(Keyword::HIVEVAR) = modifier {
             self.expect_token(&Token::Colon)?;
+        } 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 variable = self.parse_identifier()?;
         if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
             let mut values = vec![];

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -855,6 +855,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");