Support pg truncate: ONLY, IDENTITY, CASCADE

Author: tobyhede

src/ast/mod.rs

@@ -2016,6 +2016,15 @@ pub enum Statement {
         partitions: Option<Vec<Expr>>,
         /// TABLE - optional keyword;
         table: bool,
+        /// Postgres-specific option
+        /// TRUNCATE [ TABLE ] [ ONLY ] name
+        only: bool,
+        /// Postgres-specific option
+        /// [ RESTART IDENTITY | CONTINUE IDENTITY ]
+        identity: Option<TruncateIdentityOption>,
+        /// Postgres-specific option
+        /// [ CASCADE | RESTRICT ]
+        cascade: Option<TruncateCascadeOption>,
     },
     /// ```sql
     /// MSCK
@@ -3134,9 +3143,26 @@ impl fmt::Display for Statement {
                 table_name,
                 partitions,
                 table,
+                only,
+                identity,
+                cascade,
             } => {
                 let table = if *table { "TABLE " } else { "" };
-                write!(f, "TRUNCATE {table}{table_name}")?;
+                let only = if *only { "ONLY " } else { "" };
+                write!(f, "TRUNCATE {table}{only}{table_name}")?;
+                if let Some(identity) = identity {
+                    match identity {
+                        TruncateIdentityOption::Restart => write!(f, " RESTART IDENTITY")?,
+                        TruncateIdentityOption::Continue => write!(f, " CONTINUE IDENTITY")?,
+                    }
+                }
+                if let Some(cascade) = cascade {
+                    match cascade {
+                        TruncateCascadeOption::Cascade => write!(f, " CASCADE")?,
+                        TruncateCascadeOption::Restrict => write!(f, " RESTRICT")?,
+                    }
+                }
+
                 if let Some(ref parts) = partitions {
                     if !parts.is_empty() {
                         write!(f, " PARTITION ({})", display_comma_separated(parts))?;
@@ -4587,6 +4613,26 @@ impl fmt::Display for SequenceOptions {
     }
 }
 
+/// PostgreSQL identity option for TRUNCATE table
+/// [ RESTART IDENTITY | CONTINUE IDENTITY ]
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TruncateIdentityOption {
+    Restart,
+    Continue,
+}
+
+/// PostgreSQL cascade option for TRUNCATE table
+/// [ CASCADE | RESTRICT ]
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TruncateCascadeOption {
+    Cascade,
+    Restrict,
+}
+
 /// Can use to describe options in  create sequence or table column type identity
 /// [ MINVALUE minvalue | NO MINVALUE ] [ MAXVALUE maxvalue | NO MAXVALUE ]
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

src/keywords.rs

@@ -175,6 +175,7 @@ define_keywords!(
     CONNECTION,
     CONSTRAINT,
     CONTAINS,
+    CONTINUE,
     CONVERT,
     COPY,
     COPY_OPTIONS,

src/parser/mod.rs

@@ -681,17 +681,44 @@ impl<'a> Parser<'a> {
 
     pub fn parse_truncate(&mut self) -> Result<Statement, ParserError> {
         let table = self.parse_keyword(Keyword::TABLE);
+        let only = self.parse_keyword(Keyword::ONLY);
         let table_name = self.parse_object_name(false)?;
+
         let mut partitions = None;
         if self.parse_keyword(Keyword::PARTITION) {
             self.expect_token(&Token::LParen)?;
             partitions = Some(self.parse_comma_separated(Parser::parse_expr)?);
             self.expect_token(&Token::RParen)?;
         }
+
+        let mut identity = None;
+        let mut cascade = None;
+
+        if dialect_of!(self is PostgreSqlDialect | GenericDialect) {
+            identity = if self.parse_keywords(&[Keyword::RESTART, Keyword::IDENTITY]) {
+                Some(TruncateIdentityOption::Restart)
+            } else if self.parse_keywords(&[Keyword::CONTINUE, Keyword::IDENTITY]) {
+                Some(TruncateIdentityOption::Continue)
+            } else {
+                None
+            };
+
+            cascade = if self.parse_keyword(Keyword::CASCADE) {
+                Some(TruncateCascadeOption::Cascade)
+            } else if self.parse_keyword(Keyword::RESTRICT) {
+                Some(TruncateCascadeOption::Restrict)
+            } else {
+                None
+            };
+        };
+
         Ok(Statement::Truncate {
             table_name,
             partitions,
             table,
+            only,
+            identity,
+            cascade,
         })
     }
 

tests/sqlparser_postgres.rs

@@ -571,6 +571,10 @@ fn parse_alter_table_constraints_rename() {
 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");
+}
+
+#[test]
+fn parse_alter_table_disable_trigger() {
     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");
@@ -3957,7 +3961,27 @@ fn parse_truncate() {
         Statement::Truncate {
             table_name: ObjectName(vec![Ident::new("db"), Ident::new("table_name")]),
             partitions: None,
-            table: false
+            table: false,
+            only: false,
+            identity: None,
+            cascade: None,
+        },
+        truncate
+    );
+}
+
+#[test]
+fn parse_truncate_with_options() {
+    let truncate = pg_and_generic()
+        .verified_stmt("TRUNCATE TABLE ONLY db.table_name RESTART IDENTITY CASCADE");
+    assert_eq!(
+        Statement::Truncate {
+            table_name: ObjectName(vec![Ident::new("db"), Ident::new("table_name")]),
+            partitions: None,
+            table: true,
+            only: true,
+            identity: Some(TruncateIdentityOption::Restart),
+            cascade: Some(TruncateCascadeOption::Cascade)
         },
         truncate
     );
@@ -4731,12 +4755,12 @@ fn parse_trigger_related_functions() {
             IF NEW.salary IS NULL THEN
                 RAISE EXCEPTION '% cannot have null salary', NEW.empname;
             END IF;
-    
+
             -- Who works for us when they must pay for it?
             IF NEW.salary < 0 THEN
                 RAISE EXCEPTION '% cannot have a negative salary', NEW.empname;
             END IF;
-    
+
             -- Remember who changed the payroll when
             NEW.last_date := current_timestamp;
             NEW.last_user := current_user;
@@ -4868,7 +4892,7 @@ fn parse_trigger_related_functions() {
                     Expr::Value(
                         Value::DollarQuotedString(
                             DollarQuotedString {
-                                value: "\n        BEGIN\n            -- Check that empname and salary are given\n            IF NEW.empname IS NULL THEN\n                RAISE EXCEPTION 'empname cannot be null';\n            END IF;\n            IF NEW.salary IS NULL THEN\n                RAISE EXCEPTION '% cannot have null salary', NEW.empname;\n            END IF;\n    \n            -- Who works for us when they must pay for it?\n            IF NEW.salary < 0 THEN\n                RAISE EXCEPTION '% cannot have a negative salary', NEW.empname;\n            END IF;\n    \n            -- Remember who changed the payroll when\n            NEW.last_date := current_timestamp;\n            NEW.last_user := current_user;\n            RETURN NEW;\n        END;\n    ".to_owned(),
+                                value: "\n        BEGIN\n            -- Check that empname and salary are given\n            IF NEW.empname IS NULL THEN\n                RAISE EXCEPTION 'empname cannot be null';\n            END IF;\n            IF NEW.salary IS NULL THEN\n                RAISE EXCEPTION '% cannot have null salary', NEW.empname;\n            END IF;\n\n            -- Who works for us when they must pay for it?\n            IF NEW.salary < 0 THEN\n                RAISE EXCEPTION '% cannot have a negative salary', NEW.empname;\n            END IF;\n\n            -- Remember who changed the payroll when\n            NEW.last_date := current_timestamp;\n            NEW.last_user := current_user;\n            RETURN NEW;\n        END;\n    ".to_owned(),
                                 tag: Some(
                                     "emp_stamp".to_owned(),
                                 ),