Implements ALTER POLICY syntax for PostgreSQL (apache#1446)

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

Author: git-hulk

src/ast/ddl.rs

@@ -234,6 +234,49 @@ pub enum AlterTableOperation {
     OwnerTo { new_owner: Owner },
 }
 
+/// An `ALTER Policy` (`Statement::AlterPolicy`) operation
+///
+/// [PostgreSQL Documentation](https://www.postgresql.org/docs/current/sql-altertable.html)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AlterPolicyOperation {
+    Rename {
+        new_name: Ident,
+    },
+    Apply {
+        to: Option<Vec<Owner>>,
+        using: Option<Expr>,
+        with_check: Option<Expr>,
+    },
+}
+
+impl fmt::Display for AlterPolicyOperation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            AlterPolicyOperation::Rename { new_name } => {
+                write!(f, " RENAME TO {new_name}")
+            }
+            AlterPolicyOperation::Apply {
+                to,
+                using,
+                with_check,
+            } => {
+                if let Some(to) = to {
+                    write!(f, " TO {}", display_comma_separated(to))?;
+                }
+                if let Some(using) = using {
+                    write!(f, " USING ({using})")?;
+                }
+                if let Some(with_check) = with_check {
+                    write!(f, " WITH CHECK ({with_check})")?;
+                }
+                Ok(())
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/ast/mod.rs

@@ -39,11 +39,12 @@ pub use self::data_type::{
 };
 pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
 pub use self::ddl::{
-    AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ClusteredBy, ColumnDef,
-    ColumnOption, ColumnOptionDef, ConstraintCharacteristics, Deduplicate, DeferrableInitial,
-    GeneratedAs, GeneratedExpressionMode, IdentityProperty, IndexOption, IndexType,
-    KeyOrIndexDisplay, Owner, Partition, ProcedureParam, ReferentialAction, TableConstraint,
-    UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation, ViewColumnDef,
+    AlterColumnOperation, AlterIndexOperation, AlterPolicyOperation, AlterTableOperation,
+    ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef, ConstraintCharacteristics, Deduplicate,
+    DeferrableInitial, GeneratedAs, GeneratedExpressionMode, IdentityProperty, IndexOption,
+    IndexType, KeyOrIndexDisplay, Owner, Partition, ProcedureParam, ReferentialAction,
+    TableConstraint, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
+    ViewColumnDef,
 };
 pub use self::dml::{CreateIndex, CreateTable, Delete, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};
@@ -2459,6 +2460,16 @@ pub enum Statement {
         operation: AlterRoleOperation,
     },
     /// ```sql
+    /// ALTER POLICY <NAME> ON <TABLE NAME> [<OPERATION>]
+    /// ```
+    /// (Postgresql-specific)
+    AlterPolicy {
+        name: Ident,
+        #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
+        table_name: ObjectName,
+        operation: AlterPolicyOperation,
+    },
+    /// ```sql
     /// ATTACH DATABASE 'path/to/file' AS alias
     /// ```
     /// (SQLite-specific)
@@ -4197,6 +4208,13 @@ impl fmt::Display for Statement {
             Statement::AlterRole { name, operation } => {
                 write!(f, "ALTER ROLE {name} {operation}")
             }
+            Statement::AlterPolicy {
+                name,
+                table_name,
+                operation,
+            } => {
+                write!(f, "ALTER POLICY {name} ON {table_name}{operation}")
+            }
             Statement::Drop {
                 object_type,
                 if_exists,

src/parser/alter.rs

@@ -17,7 +17,10 @@ use alloc::vec;
 
 use super::{Parser, ParserError};
 use crate::{
-    ast::{AlterRoleOperation, Expr, Password, ResetConfig, RoleOption, SetConfigValue, Statement},
+    ast::{
+        AlterPolicyOperation, AlterRoleOperation, Expr, Password, ResetConfig, RoleOption,
+        SetConfigValue, Statement,
+    },
     dialect::{MsSqlDialect, PostgreSqlDialect},
     keywords::Keyword,
     tokenizer::Token,
@@ -36,6 +39,66 @@ impl<'a> Parser<'a> {
         ))
     }
 
+    /// Parse ALTER POLICY statement
+    /// ```sql
+    /// ALTER POLICY policy_name ON table_name [ RENAME TO new_name ]
+    /// or
+    /// ALTER POLICY policy_name ON table_name
+    /// [ TO { role_name | PUBLIC | CURRENT_ROLE | CURRENT_USER | SESSION_USER } [, ...] ]
+    /// [ USING ( using_expression ) ]
+    /// [ WITH CHECK ( check_expression ) ]
+    /// ```
+    ///
+    /// [PostgreSQL](https://www.postgresql.org/docs/current/sql-alterpolicy.html)
+    pub fn parse_alter_policy(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_identifier(false)?;
+        self.expect_keyword(Keyword::ON)?;
+        let table_name = self.parse_object_name(false)?;
+
+        if self.parse_keyword(Keyword::RENAME) {
+            self.expect_keyword(Keyword::TO)?;
+            let new_name = self.parse_identifier(false)?;
+            Ok(Statement::AlterPolicy {
+                name,
+                table_name,
+                operation: AlterPolicyOperation::Rename { new_name },
+            })
+        } else {
+            let to = if self.parse_keyword(Keyword::TO) {
+                Some(self.parse_comma_separated(|p| p.parse_owner())?)
+            } else {
+                None
+            };
+
+            let using = if self.parse_keyword(Keyword::USING) {
+                self.expect_token(&Token::LParen)?;
+                let expr = self.parse_expr()?;
+                self.expect_token(&Token::RParen)?;
+                Some(expr)
+            } else {
+                None
+            };
+
+            let with_check = if self.parse_keywords(&[Keyword::WITH, Keyword::CHECK]) {
+                self.expect_token(&Token::LParen)?;
+                let expr = self.parse_expr()?;
+                self.expect_token(&Token::RParen)?;
+                Some(expr)
+            } else {
+                None
+            };
+            Ok(Statement::AlterPolicy {
+                name,
+                table_name,
+                operation: AlterPolicyOperation::Apply {
+                    to,
+                    using,
+                    with_check,
+                },
+            })
+        }
+    }
+
     fn parse_mssql_alter_role(&mut self) -> Result<Statement, ParserError> {
         let role_name = self.parse_identifier(false)?;
 

src/parser/mod.rs

@@ -7140,6 +7140,7 @@ impl<'a> Parser<'a> {
             Keyword::TABLE,
             Keyword::INDEX,
             Keyword::ROLE,
+            Keyword::POLICY,
         ])?;
         match object_type {
             Keyword::VIEW => self.parse_alter_view(),
@@ -7191,6 +7192,7 @@ impl<'a> Parser<'a> {
                 })
             }
             Keyword::ROLE => self.parse_alter_role(),
+            Keyword::POLICY => self.parse_alter_policy(),
             // unreachable because expect_one_of_keywords used above
             _ => unreachable!(),
         }

tests/sqlparser_common.rs

@@ -11167,3 +11167,84 @@ fn test_drop_policy() {
         "sql parser error: Expected: end of statement, found: WRONG"
     );
 }
+
+#[test]
+fn test_alter_policy() {
+    match verified_stmt("ALTER POLICY old_policy ON my_table RENAME TO new_policy") {
+        Statement::AlterPolicy {
+            name,
+            table_name,
+            operation,
+            ..
+        } => {
+            assert_eq!(name.to_string(), "old_policy");
+            assert_eq!(table_name.to_string(), "my_table");
+            assert_eq!(
+                operation,
+                AlterPolicyOperation::Rename {
+                    new_name: Ident::new("new_policy")
+                }
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    match verified_stmt(concat!(
+        "ALTER POLICY my_policy ON my_table TO CURRENT_USER ",
+        "USING ((SELECT c0)) WITH CHECK (c0 > 0)"
+    )) {
+        Statement::AlterPolicy {
+            name, table_name, ..
+        } => {
+            assert_eq!(name.to_string(), "my_policy");
+            assert_eq!(table_name.to_string(), "my_table");
+        }
+        _ => unreachable!(),
+    }
+
+    // omit TO / USING / WITH CHECK clauses is allowed
+    verified_stmt("ALTER POLICY my_policy ON my_table");
+
+    // mixing RENAME and APPLY expressions
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY old_policy ON my_table TO public RENAME TO new_policy")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: end of statement, found: RENAME"
+    );
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY old_policy ON my_table RENAME TO new_policy TO public")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: end of statement, found: TO"
+    );
+    // missing TO in RENAME TO
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY old_policy ON my_table RENAME")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: TO, found: EOF"
+    );
+    // missing new name in RENAME TO
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY old_policy ON my_table RENAME TO")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: identifier, found: EOF"
+    );
+
+    // missing the expression in USING
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY my_policy ON my_table USING")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: (, found: EOF"
+    );
+    // missing the expression in WITH CHECK
+    assert_eq!(
+        parse_sql_statements("ALTER POLICY my_policy ON my_table WITH CHECK")
+            .unwrap_err()
+            .to_string(),
+        "sql parser error: Expected: (, found: EOF"
+    );
+}