Add support for ENABLE and DISABLE on ALTER TABLE for pg (#1077)

Signed-off-by: Toby Hede <toby@cipherstash.com>
This commit is contained in:
Toby Hede 2024-01-01 01:14:52 +11:00 committed by GitHub
parent 593c090b21
commit a75778c8c7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 143 additions and 7 deletions

View file

@ -1 +1 @@
rust 1.73.0
rust 1.75.0

View file

@ -45,6 +45,18 @@ pub enum AlterTableOperation {
/// <column_def>.
column_def: ColumnDef,
},
/// `DISABLE ROW LEVEL SECURITY`
///
/// Note: this is a PostgreSQL-specific operation.
DisableRowLevelSecurity,
/// `DISABLE RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
DisableRule { name: Ident },
/// `DISABLE TRIGGER [ trigger_name | ALL | USER ]`
///
/// Note: this is a PostgreSQL-specific operation.
DisableTrigger { name: Ident },
/// `DROP CONSTRAINT [ IF EXISTS ] <name>`
DropConstraint {
if_exists: bool,
@ -61,6 +73,34 @@ pub enum AlterTableOperation {
///
/// Note: this is a MySQL-specific operation.
DropPrimaryKey,
/// `ENABLE ALWAYS RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableAlwaysRule { name: Ident },
/// `ENABLE ALWAYS TRIGGER trigger_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableAlwaysTrigger { name: Ident },
/// `ENABLE REPLICA RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableReplicaRule { name: Ident },
/// `ENABLE REPLICA TRIGGER trigger_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableReplicaTrigger { name: Ident },
/// `ENABLE ROW LEVEL SECURITY`
///
/// Note: this is a PostgreSQL-specific operation.
EnableRowLevelSecurity,
/// `ENABLE RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableRule { name: Ident },
/// `ENABLE TRIGGER [ trigger_name | ALL | USER ]`
///
/// Note: this is a PostgreSQL-specific operation.
EnableTrigger { name: Ident },
/// `RENAME TO PARTITION (partition=val)`
RenamePartitions {
old_partitions: Vec<Expr>,
@ -143,6 +183,15 @@ impl fmt::Display for AlterTableOperation {
AlterTableOperation::AlterColumn { column_name, op } => {
write!(f, "ALTER COLUMN {column_name} {op}")
}
AlterTableOperation::DisableRowLevelSecurity => {
write!(f, "DISABLE ROW LEVEL SECURITY")
}
AlterTableOperation::DisableRule { name } => {
write!(f, "DISABLE RULE {name}")
}
AlterTableOperation::DisableTrigger { name } => {
write!(f, "DISABLE TRIGGER {name}")
}
AlterTableOperation::DropPartitions {
partitions,
if_exists,
@ -177,6 +226,27 @@ impl fmt::Display for AlterTableOperation {
column_name,
if *cascade { " CASCADE" } else { "" }
),
AlterTableOperation::EnableAlwaysRule { name } => {
write!(f, "ENABLE ALWAYS RULE {name}")
}
AlterTableOperation::EnableAlwaysTrigger { name } => {
write!(f, "ENABLE ALWAYS TRIGGER {name}")
}
AlterTableOperation::EnableReplicaRule { name } => {
write!(f, "ENABLE REPLICA RULE {name}")
}
AlterTableOperation::EnableReplicaTrigger { name } => {
write!(f, "ENABLE REPLICA TRIGGER {name}")
}
AlterTableOperation::EnableRowLevelSecurity => {
write!(f, "ENABLE ROW LEVEL SECURITY")
}
AlterTableOperation::EnableRule { name } => {
write!(f, "ENABLE RULE {name}")
}
AlterTableOperation::EnableTrigger { name } => {
write!(f, "ENABLE TRIGGER {name}")
}
AlterTableOperation::RenamePartitions {
old_partitions,
new_partitions,

View file

@ -349,6 +349,7 @@ mod tests {
}
}
#[allow(clippy::needless_raw_string_hashes)]
let statement = r#"SELECT 'Wayne\'s World'"#;
let res1 = Parser::parse_sql(&MySqlDialect {}, statement);
let res2 = Parser::parse_sql(&WrappedDialect(MySqlDialect {}), statement);

View file

@ -223,6 +223,7 @@ define_keywords!(
DETAIL,
DETERMINISTIC,
DIRECTORY,
DISABLE,
DISCARD,
DISCONNECT,
DISTINCT,
@ -241,6 +242,7 @@ define_keywords!(
ELEMENTS,
ELSE,
EMPTY,
ENABLE,
ENCODING,
ENCRYPTION,
END,
@ -546,6 +548,7 @@ define_keywords!(
REPAIR,
REPEATABLE,
REPLACE,
REPLICA,
REPLICATION,
RESET,
RESPECT,
@ -566,6 +569,7 @@ define_keywords!(
ROWID,
ROWS,
ROW_NUMBER,
RULE,
RUN,
SAFE_CAST,
SAVEPOINT,
@ -574,6 +578,7 @@ define_keywords!(
SCROLL,
SEARCH,
SECOND,
SECURITY,
SELECT,
SEMI,
SENSITIVE,

View file

@ -4724,6 +4724,48 @@ impl<'a> Parser<'a> {
new_column_name,
}
}
} else if self.parse_keyword(Keyword::DISABLE) {
if self.parse_keywords(&[Keyword::ROW, Keyword::LEVEL, Keyword::SECURITY]) {
AlterTableOperation::DisableRowLevelSecurity {}
} else if self.parse_keyword(Keyword::RULE) {
let name = self.parse_identifier()?;
AlterTableOperation::DisableRule { name }
} else if self.parse_keyword(Keyword::TRIGGER) {
let name = self.parse_identifier()?;
AlterTableOperation::DisableTrigger { name }
} else {
return self.expected(
"ROW LEVEL SECURITY, RULE, or TRIGGER after DISABLE",
self.peek_token(),
);
}
} else if self.parse_keyword(Keyword::ENABLE) {
if self.parse_keywords(&[Keyword::ALWAYS, Keyword::RULE]) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableAlwaysRule { name }
} else if self.parse_keywords(&[Keyword::ALWAYS, Keyword::TRIGGER]) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableAlwaysTrigger { name }
} else if self.parse_keywords(&[Keyword::ROW, Keyword::LEVEL, Keyword::SECURITY]) {
AlterTableOperation::EnableRowLevelSecurity {}
} else if self.parse_keywords(&[Keyword::REPLICA, Keyword::RULE]) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableReplicaRule { name }
} else if self.parse_keywords(&[Keyword::REPLICA, Keyword::TRIGGER]) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableReplicaTrigger { name }
} else if self.parse_keyword(Keyword::RULE) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableRule { name }
} else if self.parse_keyword(Keyword::TRIGGER) {
let name = self.parse_identifier()?;
AlterTableOperation::EnableTrigger { name }
} else {
return self.expected(
"ALWAYS, REPLICA, ROW LEVEL SECURITY, RULE, or TRIGGER after ENABLE",
self.peek_token(),
);
}
} else if self.parse_keyword(Keyword::DROP) {
if self.parse_keywords(&[Keyword::IF, Keyword::EXISTS, Keyword::PARTITION]) {
self.expect_token(&Token::LParen)?;

View file

@ -727,10 +727,7 @@ impl<'a> Tokenizer<'a> {
// match binary literal that starts with 0x
if s == "0" && chars.peek() == Some(&'x') {
chars.next();
let s2 = peeking_take_while(
chars,
|ch| matches!(ch, '0'..='9' | 'A'..='F' | 'a'..='f'),
);
let s2 = peeking_take_while(chars, |ch| ch.is_ascii_hexdigit());
return Ok(Some(Token::HexStringLiteral(s2)));
}
@ -1077,7 +1074,7 @@ impl<'a> Tokenizer<'a> {
match chars.peek() {
Some('$') => {
chars.next();
for (_, c) in value.chars().enumerate() {
for c in value.chars() {
let next_char = chars.next();
if Some(c) != next_char {
return self.tokenizer_error(

View file

@ -563,6 +563,27 @@ fn parse_alter_table_constraints_rename() {
}
}
#[test]
fn parse_alter_table_disable() {
pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE ROW LEVEL SECURITY");
pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE RULE rule_name");
pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE TRIGGER ALL");
pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE TRIGGER USER");
pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE TRIGGER trigger_name");
}
#[test]
fn parse_alter_table_enable() {
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE ALWAYS RULE rule_name");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE ALWAYS TRIGGER trigger_name");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE REPLICA TRIGGER trigger_name");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE REPLICA RULE rule_name");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE ROW LEVEL SECURITY");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE RULE rule_name");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER ALL");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER USER");
pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER trigger_name");
}
#[test]
fn parse_alter_table_alter_column() {
pg().one_statement_parses_to(
@ -3256,7 +3277,7 @@ fn parse_dollar_quoted_string() {
let stmt = pg().parse_sql_statements(sql).unwrap();
let projection = match stmt.get(0).unwrap() {
let projection = match stmt.first().unwrap() {
Statement::Query(query) => match &*query.body {
SetExpr::Select(select) => &select.projection,
_ => unreachable!(),