Postgres: support for OWNER TO clause (#1314)

Co-authored-by: Andrew Lamb <[email protected]>

Author: gainings

src/ast/ddl.rs

@@ -157,6 +157,32 @@ pub enum AlterTableOperation {
     SwapWith { table_name: ObjectName },
     /// 'SET TBLPROPERTIES ( { property_key [ = ] property_val } [, ...] )'
     SetTblProperties { table_properties: Vec<SqlOption> },
+
+    /// `OWNER TO { <new_owner> | CURRENT_ROLE | CURRENT_USER | SESSION_USER }`
+    ///
+    /// Note: this is PostgreSQL-specific <https://www.postgresql.org/docs/current/sql-altertable.html>
+    OwnerTo { new_owner: Owner },
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum Owner {
+    Ident(Ident),
+    CurrentRole,
+    CurrentUser,
+    SessionUser,
+}
+
+impl fmt::Display for Owner {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Owner::Ident(ident) => write!(f, "{}", ident),
+            Owner::CurrentRole => write!(f, "CURRENT_ROLE"),
+            Owner::CurrentUser => write!(f, "CURRENT_USER"),
+            Owner::SessionUser => write!(f, "SESSION_USER"),
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
@@ -322,6 +348,9 @@ impl fmt::Display for AlterTableOperation {
             AlterTableOperation::SwapWith { table_name } => {
                 write!(f, "SWAP WITH {table_name}")
             }
+            AlterTableOperation::OwnerTo { new_owner } => {
+                write!(f, "OWNER TO {new_owner}")
+            }
             AlterTableOperation::SetTblProperties { table_properties } => {
                 write!(
                     f,

src/ast/mod.rs

@@ -34,8 +34,8 @@ pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue}
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
     ColumnOptionDef, ConstraintCharacteristics, DeferrableInitial, GeneratedAs,
-    GeneratedExpressionMode, IndexOption, IndexType, KeyOrIndexDisplay, Partition, ProcedureParam,
-    ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
+    GeneratedExpressionMode, IndexOption, IndexType, KeyOrIndexDisplay, Owner, Partition,
+    ProcedureParam, ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
     UserDefinedTypeRepresentation, ViewColumnDef,
 };
 pub use self::dml::{CreateIndex, CreateTable, Delete, Insert};

src/keywords.rs

@@ -527,6 +527,7 @@ define_keywords!(
     OVERLAY,
     OVERWRITE,
     OWNED,
+    OWNER,
     PARALLEL,
     PARAMETER,
     PARQUET,

src/parser/mod.rs

@@ -6447,6 +6447,25 @@ impl<'a> Parser<'a> {
             self.expect_keyword(Keyword::WITH)?;
             let table_name = self.parse_object_name(false)?;
             AlterTableOperation::SwapWith { table_name }
+        } else if dialect_of!(self is PostgreSqlDialect | GenericDialect)
+            && self.parse_keywords(&[Keyword::OWNER, Keyword::TO])
+        {
+            let new_owner = match self.parse_one_of_keywords( &[Keyword::CURRENT_USER, Keyword::CURRENT_ROLE, Keyword::SESSION_USER]) {
+                Some(Keyword::CURRENT_USER) => Owner::CurrentUser,
+                Some(Keyword::CURRENT_ROLE) => Owner::CurrentRole,
+                Some(Keyword::SESSION_USER) => Owner::SessionUser,
+                Some(_) => unreachable!(),
+                None => {
+                    match self.parse_identifier(false) {
+                        Ok(ident) => Owner::Ident(ident),
+                        Err(e) => {
+                            return Err(ParserError::ParserError(format!("Expected: CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO. {e}")))
+                        }
+                    }
+                },
+            };
+
+            AlterTableOperation::OwnerTo { new_owner }
         } else {
             let options: Vec<SqlOption> =
                 self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;

tests/sqlparser_postgres.rs

@@ -713,6 +713,78 @@ fn parse_alter_table_add_columns() {
     }
 }
 
+#[test]
+fn parse_alter_table_owner_to() {
+    struct TestCase {
+        sql: &'static str,
+        expected_owner: Owner,
+    }
+
+    let test_cases = vec![
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO new_owner",
+            expected_owner: Owner::Ident(Ident::new("new_owner".to_string())),
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO postgres",
+            expected_owner: Owner::Ident(Ident::new("postgres".to_string())),
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO CREATE", // treats CREATE as an identifier
+            expected_owner: Owner::Ident(Ident::new("CREATE".to_string())),
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO \"new_owner\"",
+            expected_owner: Owner::Ident(Ident::with_quote('\"', "new_owner".to_string())),
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO CURRENT_USER",
+            expected_owner: Owner::CurrentUser,
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO CURRENT_ROLE",
+            expected_owner: Owner::CurrentRole,
+        },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO SESSION_USER",
+            expected_owner: Owner::SessionUser,
+        },
+    ];
+
+    for case in test_cases {
+        match pg_and_generic().verified_stmt(case.sql) {
+            Statement::AlterTable {
+                name,
+                if_exists: _,
+                only: _,
+                operations,
+                location: _,
+            } => {
+                assert_eq!(name.to_string(), "tab");
+                assert_eq!(
+                    operations,
+                    vec![AlterTableOperation::OwnerTo {
+                        new_owner: case.expected_owner.clone()
+                    }]
+                );
+            }
+            _ => unreachable!("Expected an AlterTable statement"),
+        }
+    }
+
+    let res = pg().parse_sql_statements("ALTER TABLE tab OWNER TO CREATE FOO");
+    assert_eq!(
+        ParserError::ParserError("Expected: end of statement, found: FOO".to_string()),
+        res.unwrap_err()
+    );
+
+    let res = pg().parse_sql_statements("ALTER TABLE tab OWNER TO 4");
+    assert_eq!(
+        ParserError::ParserError("Expected: CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO. sql parser error: Expected: identifier, found: 4".to_string()),
+        res.unwrap_err()
+    );
+}
+
 #[test]
 fn parse_create_table_if_not_exists() {
     let sql = "CREATE TABLE IF NOT EXISTS uk_cities ()";