From c3e9e0762dd35eaae5774989bee758ac2c408d6f Mon Sep 17 00:00:00 2001
From: slhmy <31381093+slhmy@users.noreply.github.com>
Date: Thu, 14 Apr 2022 02:45:53 +0800
Subject: [PATCH] Support 'SET ROLE' statement (#455)

---
 src/ast/mod.rs              | 28 ++++++++++++++++++++++++++
 src/parser.rs               | 11 +++++++++++
 tests/sqlparser_postgres.rs | 39 +++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index e8f15a5a9..feec8bda1 100644
--- a/src/ast/mod.rs
+++ b/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 ")?;
 
diff --git a/src/parser.rs b/src/parser.rs
index 0a179a23b..8a2074f5f 100644
--- a/src/parser.rs
+++ b/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![];
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index ac8d283f9..68033e367 100644
--- a/tests/sqlparser_postgres.rs
+++ b/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");