From e0b438b43d25665df8eaf9046245e46854e3ca33 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Tue, 27 Aug 2024 10:01:30 +1000 Subject: [PATCH 1/5] Support pg truncate: ONLY, IDENTITY, CASCADE --- src/ast/mod.rs | 48 ++++++++++++++++++++++++++++++++++++- src/keywords.rs | 1 + src/parser/mod.rs | 27 +++++++++++++++++++++ tests/sqlparser_postgres.rs | 32 +++++++++++++++++++++---- 4 files changed, 103 insertions(+), 5 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 8a56f3158..261252877 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -2016,6 +2016,15 @@ pub enum Statement { partitions: Option>, /// TABLE - optional keyword; table: bool, + /// Postgres-specific option + /// TRUNCATE [ TABLE ] [ ONLY ] name + only: bool, + /// Postgres-specific option + /// [ RESTART IDENTITY | CONTINUE IDENTITY ] + identity: Option, + /// Postgres-specific option + /// [ CASCADE | RESTRICT ] + cascade: Option, }, /// ```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)] diff --git a/src/keywords.rs b/src/keywords.rs index d2dcc57d1..222a45fb8 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -175,6 +175,7 @@ define_keywords!( CONNECTION, CONSTRAINT, CONTAINS, + CONTINUE, CONVERT, COPY, COPY_OPTIONS, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 302e5e660..e2bc3deea 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -681,17 +681,44 @@ impl<'a> Parser<'a> { pub fn parse_truncate(&mut self) -> Result { 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, }) } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 2f9fe86c9..d792b41df 100644 --- a/tests/sqlparser_postgres.rs +++ b/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(), ), From 1c3476f96bba06163a081164458bb7cc20357c53 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Tue, 27 Aug 2024 10:12:14 +1000 Subject: [PATCH 2/5] Support pg truncate: list of table names --- src/ast/mod.rs | 16 ++++++++++--- src/parser/mod.rs | 7 +++++- tests/sqlparser_postgres.rs | 46 +++++++++++++++++++++++++++++++++++-- 3 files changed, 63 insertions(+), 6 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 261252877..16c3c4aeb 100644 --- a/src/ast/mod.rs +++ b/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, partitions: Option>, /// 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::>() + .join(", "); + + write!(f, "TRUNCATE {table}{only}{table_names}")?; + if let Some(identity) = identity { match identity { TruncateIdentityOption::Restart => write!(f, " RESTART IDENTITY")?, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index e2bc3deea..59b991e97 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -682,7 +682,11 @@ impl<'a> Parser<'a> { pub fn parse_truncate(&mut self) -> Result { 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, diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index d792b41df..9885d045b 100644 --- a/tests/sqlparser_postgres.rs +++ b/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( From c87956814357d25b8f5721ad32fb5606a7a73d60 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Fri, 30 Aug 2024 10:41:07 +1000 Subject: [PATCH 3/5] Remove comment --- src/ast/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 16c3c4aeb..f89b9e1f0 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -3154,7 +3154,7 @@ impl fmt::Display for Statement { let table_names = table_names .iter() - .map(|table_name| table_name.to_string()) // replace `to_string()` with the appropriate method if necessary + .map(|table_name| table_name.to_string()) .collect::>() .join(", "); From c71d6a95e66f5b9d707a04c9b4ecc5dac1bd4cbb Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Mon, 2 Sep 2024 10:19:01 +1000 Subject: [PATCH 4/5] TruncateTableTarget to handle list of tables --- src/ast/mod.rs | 34 +++++++++++++++++++++++----------- src/parser/mod.rs | 10 +++++----- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index f89b9e1f0..ae59f7fcb 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -2011,9 +2011,7 @@ pub enum Statement { /// ``` /// Truncate (Hive) Truncate { - #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] - table_name: ObjectName, - table_names: Vec, + table_names: Vec, partitions: Option>, /// TABLE - optional keyword; table: bool, @@ -3141,7 +3139,6 @@ impl fmt::Display for Statement { Ok(()) } Statement::Truncate { - table_name: _, table_names, partitions, table, @@ -3152,13 +3149,11 @@ impl fmt::Display for Statement { let table = if *table { "TABLE " } else { "" }; let only = if *only { "ONLY " } else { "" }; - let table_names = table_names - .iter() - .map(|table_name| table_name.to_string()) - .collect::>() - .join(", "); - - write!(f, "TRUNCATE {table}{only}{table_names}")?; + write!( + f, + "TRUNCATE {table}{only}{table_names}", + table_names = display_comma_separated(table_names) + )?; if let Some(identity) = identity { match identity { @@ -4623,6 +4618,23 @@ impl fmt::Display for SequenceOptions { } } +/// Target of a `TRUNCATE TABLE` command +/// +/// Note this is its own struct because `visit_relation` requires an `ObjectName` (not a `Vec`) +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct TruncateTableTarget { + /// name of the table being truncated + #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] + pub name: ObjectName, +} + +impl fmt::Display for TruncateTableTarget { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.name) + } +} + /// PostgreSQL identity option for TRUNCATE table /// [ RESTART IDENTITY | CONTINUE IDENTITY ] #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 59b991e97..b79ba70c3 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -683,10 +683,11 @@ impl<'a> Parser<'a> { 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))?; - - // 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 table_names = self + .parse_comma_separated(|p| p.parse_object_name(false))? + .into_iter() + .map(|n| TruncateTableTarget { name: n }) + .collect(); let mut partitions = None; if self.parse_keyword(Keyword::PARTITION) { @@ -717,7 +718,6 @@ impl<'a> Parser<'a> { }; Ok(Statement::Truncate { - table_name, table_names, partitions, table, From 34da79d2c3d8fa3ba86c198dd7b4ca9b37485394 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Tue, 3 Sep 2024 09:57:00 +1000 Subject: [PATCH 5/5] derive visit traits for TruncateTableTarget --- src/ast/mod.rs | 1 + tests/sqlparser_postgres.rs | 22 ++++++++++++++-------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index ae59f7fcb..8c6fe5d82 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -4622,6 +4622,7 @@ impl fmt::Display for SequenceOptions { /// /// Note this is its own struct because `visit_relation` requires an `ObjectName` (not a `Vec`) #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct TruncateTableTarget { /// name of the table being truncated diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 9885d045b..1a5e4ebfb 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -3965,10 +3965,11 @@ fn parse_select_group_by_cube() { 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()]; + let table_names = vec![TruncateTableTarget { + name: table_name.clone(), + }]; assert_eq!( Statement::Truncate { - table_name, table_names, partitions: None, table: false, @@ -3986,11 +3987,12 @@ fn parse_truncate_with_options() { .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()]; + let table_names = vec![TruncateTableTarget { + name: table_name.clone(), + }]; assert_eq!( Statement::Truncate { - table_name, table_names, partitions: None, table: true, @@ -4008,16 +4010,20 @@ fn parse_truncate_with_table_list() { "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_name_a = ObjectName(vec![Ident::new("db"), Ident::new("table_name")]); + let table_name_b = ObjectName(vec![Ident::new("db"), Ident::new("other_table_name")]); let table_names = vec![ - table_name.clone(), - ObjectName(vec![Ident::new("db"), Ident::new("other_table_name")]), + TruncateTableTarget { + name: table_name_a.clone(), + }, + TruncateTableTarget { + name: table_name_b.clone(), + }, ]; assert_eq!( Statement::Truncate { - table_name, table_names, partitions: None, table: true,