Postgres: Apply ONLY keyword per table in TRUNCATE stmt (apache#1872)

Author: MohamedAbdeen21

src/ast/mod.rs

@@ -3014,9 +3014,6 @@ pub enum Statement {
         /// TABLE - optional keyword;
         table: bool,
         /// Postgres-specific option
-        /// [ TRUNCATE TABLE ONLY ]
-        only: bool,
-        /// Postgres-specific option
         /// [ RESTART IDENTITY | CONTINUE IDENTITY ]
         identity: Option<TruncateIdentityOption>,
         /// Postgres-specific option
@@ -4425,17 +4422,15 @@ impl fmt::Display for Statement {
                 table_names,
                 partitions,
                 table,
-                only,
                 identity,
                 cascade,
                 on_cluster,
             } => {
                 let table = if *table { "TABLE " } else { "" };
-                let only = if *only { "ONLY " } else { "" };
 
                 write!(
                     f,
-                    "TRUNCATE {table}{only}{table_names}",
+                    "TRUNCATE {table}{table_names}",
                     table_names = display_comma_separated(table_names)
                 )?;
 
@@ -6106,10 +6101,17 @@ pub struct TruncateTableTarget {
     /// name of the table being truncated
     #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
     pub name: ObjectName,
+    /// Postgres-specific option
+    /// [ TRUNCATE TABLE ONLY ]
+    /// <https://www.postgresql.org/docs/current/sql-truncate.html>
+    pub only: bool,
 }
 
 impl fmt::Display for TruncateTableTarget {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.only {
+            write!(f, "ONLY ")?;
+        };
         write!(f, "{}", self.name)
     }
 }

src/ast/spans.rs

@@ -311,7 +311,6 @@ impl Spanned for Statement {
                 table_names,
                 partitions,
                 table: _,
-                only: _,
                 identity: _,
                 cascade: _,
                 on_cluster: _,

src/parser/mod.rs

@@ -960,12 +960,13 @@ 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_names = self
-            .parse_comma_separated(|p| p.parse_object_name(false))?
+            .parse_comma_separated(|p| {
+                Ok((p.parse_keyword(Keyword::ONLY), p.parse_object_name(false)?))
+            })?
             .into_iter()
-            .map(|n| TruncateTableTarget { name: n })
+            .map(|(only, name)| TruncateTableTarget { name, only })
             .collect();
 
         let mut partitions = None;
@@ -996,7 +997,6 @@ impl<'a> Parser<'a> {
             table_names,
             partitions,
             table,
-            only,
             identity,
             cascade,
             on_cluster,

tests/sqlparser_common.rs

@@ -15302,3 +15302,31 @@ fn test_open() {
         })
     );
 }
+
+#[test]
+fn parse_truncate_only() {
+    let truncate = all_dialects().verified_stmt("TRUNCATE TABLE employee, ONLY dept");
+
+    let table_names = vec![
+        TruncateTableTarget {
+            name: ObjectName::from(vec![Ident::new("employee")]),
+            only: false,
+        },
+        TruncateTableTarget {
+            name: ObjectName::from(vec![Ident::new("dept")]),
+            only: true,
+        },
+    ];
+
+    assert_eq!(
+        Statement::Truncate {
+            table_names,
+            partitions: None,
+            table: true,
+            identity: None,
+            cascade: None,
+            on_cluster: None,
+        },
+        truncate
+    );
+}

tests/sqlparser_postgres.rs

@@ -4788,13 +4788,13 @@ fn parse_truncate() {
     let table_name = ObjectName::from(vec![Ident::new("db"), Ident::new("table_name")]);
     let table_names = vec![TruncateTableTarget {
         name: table_name.clone(),
+        only: false,
     }];
     assert_eq!(
         Statement::Truncate {
             table_names,
             partitions: None,
             table: false,
-            only: false,
             identity: None,
             cascade: None,
             on_cluster: None,
@@ -4811,14 +4811,14 @@ fn parse_truncate_with_options() {
     let table_name = ObjectName::from(vec![Ident::new("db"), Ident::new("table_name")]);
     let table_names = vec![TruncateTableTarget {
         name: table_name.clone(),
+        only: true,
     }];
 
     assert_eq!(
         Statement::Truncate {
             table_names,
             partitions: None,
             table: true,
-            only: true,
             identity: Some(TruncateIdentityOption::Restart),
             cascade: Some(CascadeOption::Cascade),
             on_cluster: None,
@@ -4839,9 +4839,11 @@ fn parse_truncate_with_table_list() {
     let table_names = vec![
         TruncateTableTarget {
             name: table_name_a.clone(),
+            only: false,
         },
         TruncateTableTarget {
             name: table_name_b.clone(),
+            only: false,
         },
     ];
 
@@ -4850,7 +4852,6 @@ fn parse_truncate_with_table_list() {
             table_names,
             partitions: None,
             table: true,
-            only: false,
             identity: Some(TruncateIdentityOption::Restart),
             cascade: Some(CascadeOption::Cascade),
             on_cluster: None,