Skip to content

Commit

Permalink
Implements ALTER POLICY syntax for PostgreSQL (#1446)
Browse files Browse the repository at this point in the history
Co-authored-by: Andrew Lamb <[email protected]>
  • Loading branch information
git-hulk and alamb authored Sep 29, 2024
1 parent 73dc8a3 commit 51cbd5a
Show file tree
Hide file tree
Showing 5 changed files with 213 additions and 6 deletions.
43 changes: 43 additions & 0 deletions src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))]
Expand Down
28 changes: 23 additions & 5 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
65 changes: 64 additions & 1 deletion src/parser/alter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)?;

Expand Down
2 changes: 2 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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!(),
}
Expand Down
81 changes: 81 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"
);
}

0 comments on commit 51cbd5a

Please sign in to comment.