From 51cbd5a3e69f9cc903f91d7df335658094b5e811 Mon Sep 17 00:00:00 2001 From: hulk Date: Sun, 29 Sep 2024 18:17:31 +0800 Subject: [PATCH] Implements ALTER POLICY syntax for PostgreSQL (#1446) Co-authored-by: Andrew Lamb --- src/ast/ddl.rs | 43 +++++++++++++++++++++ src/ast/mod.rs | 28 +++++++++++--- src/parser/alter.rs | 65 ++++++++++++++++++++++++++++++- src/parser/mod.rs | 2 + tests/sqlparser_common.rs | 81 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 213 insertions(+), 6 deletions(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index e0441d07a..4578ae8f1 100644 --- a/src/ast/ddl.rs +++ b/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>, + using: Option, + with_check: Option, + }, +} + +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))] diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 19354ab96..a11fa63f4 100644 --- a/src/ast/mod.rs +++ b/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 ON [] + /// ``` + /// (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, diff --git a/src/parser/alter.rs b/src/parser/alter.rs index 7bf99af66..28fdaf764 100644 --- a/src/parser/alter.rs +++ b/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 { + 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 { let role_name = self.parse_identifier(false)?; diff --git a/src/parser/mod.rs b/src/parser/mod.rs index e5fd53554..a60dea3b8 100644 --- a/src/parser/mod.rs +++ b/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!(), } diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index e066451df..29352981a 100644 --- a/tests/sqlparser_common.rs +++ b/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" + ); +}