Support pg truncate: list of table names

Author: tobyhede

src/ast/mod.rs

@@ -2013,11 +2013,12 @@ pub enum Statement {
     Truncate {
         #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
         table_name: ObjectName,
+        table_names: Vec<ObjectName>,
         partitions: Option<Vec<Expr>>,
         /// TABLE - optional keyword;
         table: bool,
         /// Postgres-specific option
-        /// TRUNCATE [ TABLE ] [ ONLY ] name
+        /// [ TRUNCATE TABLE ONLY ]
         only: bool,
         /// Postgres-specific option
         /// [ RESTART IDENTITY | CONTINUE IDENTITY ]
@@ -3140,7 +3141,8 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Statement::Truncate {
-                table_name,
+                table_name: _,
+                table_names,
                 partitions,
                 table,
                 only,
@@ -3149,7 +3151,15 @@ impl fmt::Display for Statement {
             } => {
                 let table = if *table { "TABLE " } else { "" };
                 let only = if *only { "ONLY " } else { "" };
-                write!(f, "TRUNCATE {table}{only}{table_name}")?;
+
+                let table_names = table_names
+                    .iter()
+                    .map(|table_name| table_name.to_string()) // replace `to_string()` with the appropriate method if necessary
+                    .collect::<Vec<String>>()
+                    .join(", ");
+
+                write!(f, "TRUNCATE {table}{only}{table_names}")?;
+
                 if let Some(identity) = identity {
                     match identity {
                         TruncateIdentityOption::Restart => write!(f, " RESTART IDENTITY")?,

src/parser/mod.rs

@@ -682,7 +682,11 @@ 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 table_names = self.parse_comma_separated(|p| p.parse_object_name(false))?;
+
+        // Unwrap is safe - the preceding parse  fails if there is not at least one table name
+        let table_name = table_names.first().unwrap().clone();
 
         let mut partitions = None;
         if self.parse_keyword(Keyword::PARTITION) {
@@ -714,6 +718,7 @@ impl<'a> Parser<'a> {
 
         Ok(Statement::Truncate {
             table_name,
+            table_names,
             partitions,
             table,
             only,

tests/sqlparser_postgres.rs

@@ -593,6 +593,13 @@ fn parse_alter_table_enable() {
     pg_and_generic().verified_stmt("ALTER TABLE tab ENABLE TRIGGER trigger_name");
 }
 
+#[test]
+fn parse_truncate_table() {
+    pg_and_generic()
+        .verified_stmt("TRUNCATE TABLE \"users\", \"orders\" RESTART IDENTITY RESTRICT");
+    pg_and_generic().verified_stmt("TRUNCATE users, orders RESTART IDENTITY");
+}
+
 #[test]
 fn parse_create_extension() {
     pg_and_generic().verified_stmt("CREATE EXTENSION extension_name");
@@ -3957,9 +3964,12 @@ fn parse_select_group_by_cube() {
 #[test]
 fn parse_truncate() {
     let truncate = pg_and_generic().verified_stmt("TRUNCATE db.table_name");
+    let table_name = ObjectName(vec![Ident::new("db"), Ident::new("table_name")]);
+    let table_names = vec![table_name.clone()];
     assert_eq!(
         Statement::Truncate {
-            table_name: ObjectName(vec![Ident::new("db"), Ident::new("table_name")]),
+            table_name,
+            table_names,
             partitions: None,
             table: false,
             only: false,
@@ -3974,9 +3984,14 @@ fn parse_truncate() {
 fn parse_truncate_with_options() {
     let truncate = pg_and_generic()
         .verified_stmt("TRUNCATE TABLE ONLY db.table_name RESTART IDENTITY CASCADE");
+
+    let table_name = ObjectName(vec![Ident::new("db"), Ident::new("table_name")]);
+    let table_names = vec![table_name.clone()];
+
     assert_eq!(
         Statement::Truncate {
-            table_name: ObjectName(vec![Ident::new("db"), Ident::new("table_name")]),
+            table_name,
+            table_names,
             partitions: None,
             table: true,
             only: true,
@@ -3987,6 +4002,33 @@ fn parse_truncate_with_options() {
     );
 }
 
+#[test]
+fn parse_truncate_with_table_list() {
+    let truncate = pg().verified_stmt(
+        "TRUNCATE TABLE db.table_name, db.other_table_name RESTART IDENTITY CASCADE",
+    );
+
+    let table_name = ObjectName(vec![Ident::new("db"), Ident::new("table_name")]);
+
+    let table_names = vec![
+        table_name.clone(),
+        ObjectName(vec![Ident::new("db"), Ident::new("other_table_name")]),
+    ];
+
+    assert_eq!(
+        Statement::Truncate {
+            table_name,
+            table_names,
+            partitions: None,
+            table: true,
+            only: false,
+            identity: Some(TruncateIdentityOption::Restart),
+            cascade: Some(TruncateCascadeOption::Cascade)
+        },
+        truncate
+    );
+}
+
 #[test]
 fn parse_select_regexp_as_column_name() {
     pg_and_generic().verified_only_select(