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

Signed-off-by: Toby Hede <toby@cipherstash.com>

.tool-versions

@@ -1 +1 @@
-rust 1.73.0
+rust 1.75.0

src/ast/ddl.rs

@@ -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,

src/dialect/mod.rs

@@ -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);

src/keywords.rs

@@ -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,

src/parser/mod.rs

@@ -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)?;

src/tokenizer.rs

@@ -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(
+                        let s2 = peeking_take_while(chars, |ch| ch.is_ascii_hexdigit());
-                            chars,
-                            |ch| matches!(ch, '0'..='9' | 'A'..='F' | 'a'..='f'),
-                        );
                         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(

tests/sqlparser_postgres.rs

@@ -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!(),