feat: Add support for SET SESSION AUTHORIZATION apache#2086 (apache#2087)

Author: watford-ep

src/ast/mod.rs

@@ -2953,6 +2953,15 @@ pub enum Set {
     /// MySQL-style
     /// SET a = 1, b = 2, ..;
     MultipleAssignments { assignments: Vec<SetAssignment> },
+    /// Session authorization for Postgres/Redshift
+    ///
+    /// ```sql
+    /// SET SESSION AUTHORIZATION { user_name | DEFAULT }
+    /// ```
+    ///
+    /// See <https://www.postgresql.org/docs/current/sql-set-session-authorization.html>
+    /// See <https://docs.aws.amazon.com/redshift/latest/dg/r_SET_SESSION_AUTHORIZATION.html>
+    SetSessionAuthorization(SetSessionAuthorizationParam),
     /// MS-SQL session
     ///
     /// See <https://learn.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql>
@@ -3027,6 +3036,7 @@ impl Display for Set {
                     modifier = context_modifier.map(|m| format!("{m}")).unwrap_or_default()
                 )
             }
+            Self::SetSessionAuthorization(kind) => write!(f, "SET SESSION AUTHORIZATION {kind}"),
             Self::SetSessionParam(kind) => write!(f, "SET {kind}"),
             Self::SetTransaction {
                 modes,
@@ -9856,6 +9866,42 @@ impl fmt::Display for TableObject {
     }
 }
 
+/// Represents a SET SESSION AUTHORIZATION statement
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SetSessionAuthorizationParam {
+    pub scope: ContextModifier,
+    pub kind: SetSessionAuthorizationParamKind,
+}
+
+impl fmt::Display for SetSessionAuthorizationParam {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.kind)
+    }
+}
+
+/// Represents the parameter kind for SET SESSION AUTHORIZATION
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum SetSessionAuthorizationParamKind {
+    /// Default authorization
+    Default,
+
+    /// User name
+    User(Ident),
+}
+
+impl fmt::Display for SetSessionAuthorizationParamKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SetSessionAuthorizationParamKind::Default => write!(f, "DEFAULT"),
+            SetSessionAuthorizationParamKind::User(name) => write!(f, "{}", name),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/parser/mod.rs

@@ -13143,6 +13143,18 @@ impl<'a> Parser<'a> {
                 session: false,
             }
             .into());
+        } else if self.parse_keyword(Keyword::AUTHORIZATION) {
+            let auth_value = if self.parse_keyword(Keyword::DEFAULT) {
+                SetSessionAuthorizationParamKind::Default
+            } else {
+                let value = self.parse_identifier()?;
+                SetSessionAuthorizationParamKind::User(value)
+            };
+            return Ok(Set::SetSessionAuthorization(SetSessionAuthorizationParam {
+                scope: scope.expect("SET ... AUTHORIZATION must have a scope"),
+                kind: auth_value,
+            })
+            .into());
         }
 
         if self.dialect.supports_comma_separated_set_assignments() {

tests/sqlparser_common.rs

@@ -17698,3 +17698,28 @@ fn parse_reset_statement() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn test_parse_set_session_authorization() {
+    let stmt = verified_stmt("SET SESSION AUTHORIZATION DEFAULT");
+    assert_eq!(
+        stmt,
+        Statement::Set(Set::SetSessionAuthorization(SetSessionAuthorizationParam {
+            scope: ContextModifier::Session,
+            kind: SetSessionAuthorizationParamKind::Default,
+        }))
+    );
+
+    let stmt = verified_stmt("SET SESSION AUTHORIZATION 'username'");
+    assert_eq!(
+        stmt,
+        Statement::Set(Set::SetSessionAuthorization(SetSessionAuthorizationParam {
+            scope: ContextModifier::Session,
+            kind: SetSessionAuthorizationParamKind::User(Ident {
+                value: "username".to_string(),
+                quote_style: Some('\''),
+                span: Span::empty(),
+            }),
+        }))
+    );
+}