From c672b47553672af3fae1a0346dec8ccd7c333a84 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Mon, 1 Jan 2024 14:54:26 +1100 Subject: [PATCH 1/2] Support for pg ADD GENERATED in ALTER COLUMN statements Signed-off-by: Toby Hede --- src/ast/ddl.rs | 33 ++++++++++++++++++++++++++++++ src/keywords.rs | 1 + src/parser/mod.rs | 40 ++++++++++++++++++++++++++++++++----- tests/sqlparser_common.rs | 2 +- tests/sqlparser_postgres.rs | 14 +++++++++++++ 5 files changed, 84 insertions(+), 6 deletions(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index cd0cf1797..03fc5aabd 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -315,6 +315,13 @@ pub enum AlterColumnOperation { /// PostgreSQL specific using: Option, }, + /// `ADD GENERATED { ALWAYS | BY DEFAULT } AS IDENTITY [ ( sequence_options ) ]` + /// + /// Note: this is a PostgreSQL-specific operation. + AddGenerated { + generated_as: Option, + sequence_options: Option>, + }, } impl fmt::Display for AlterColumnOperation { @@ -335,6 +342,32 @@ impl fmt::Display for AlterColumnOperation { write!(f, "SET DATA TYPE {data_type}") } } + AlterColumnOperation::AddGenerated { + generated_as, + sequence_options, + } => { + let generated_as = match generated_as { + Some(GeneratedAs::Always) => " ALWAYS", + Some(GeneratedAs::ByDefault) => " BY DEFAULT", + _ => "", + }; + + write!(f, "ADD GENERATED{generated_as} AS IDENTITY",)?; + if let Some(options) = sequence_options { + if !options.is_empty() { + write!(f, " (")?; + } + + for sequence_option in options { + write!(f, "{sequence_option}")?; + } + + if !options.is_empty() { + write!(f, " )")?; + } + } + Ok(()) + } } } } diff --git a/src/keywords.rs b/src/keywords.rs index cba724019..e136f9905 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -552,6 +552,7 @@ define_keywords!( REPLICATION, RESET, RESPECT, + RESTART, RESTRICT, RESULT, RETAIN, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 92f4cf028..b56068ae7 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -4841,7 +4841,11 @@ impl<'a> Parser<'a> { let column_name = self.parse_identifier()?; let is_postgresql = dialect_of!(self is PostgreSqlDialect); - let op = if self.parse_keywords(&[Keyword::SET, Keyword::NOT, Keyword::NULL]) { + let op: AlterColumnOperation = if self.parse_keywords(&[ + Keyword::SET, + Keyword::NOT, + Keyword::NULL, + ]) { AlterColumnOperation::SetNotNull {} } else if self.parse_keywords(&[Keyword::DROP, Keyword::NOT, Keyword::NULL]) { AlterColumnOperation::DropNotNull {} @@ -4861,11 +4865,37 @@ impl<'a> Parser<'a> { None }; AlterColumnOperation::SetDataType { data_type, using } + } else if self.parse_keywords(&[Keyword::ADD, Keyword::GENERATED]) { + let generated_as = if self.parse_keyword(Keyword::ALWAYS) { + Some(GeneratedAs::Always) + } else if self.parse_keywords(&[Keyword::BY, Keyword::DEFAULT]) { + Some(GeneratedAs::ByDefault) + } else { + None + }; + + let _ = self.expect_keywords(&[Keyword::AS, Keyword::IDENTITY]); + + let mut sequence_options: Option> = None; + + if self.peek_token().token == Token::LParen { + self.expect_token(&Token::LParen)?; + sequence_options = self.parse_create_sequence_options().ok(); + self.expect_token(&Token::RParen)?; + } + + AlterColumnOperation::AddGenerated { + generated_as, + sequence_options, + } } else { - return self.expected( - "SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE after ALTER COLUMN", - self.peek_token(), - ); + let message = if is_postgresql { + "SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE, or ADD GENERATED after ALTER COLUMN" + } else { + "SET/DROP NOT NULL, SET DEFAULT, or SET DATA TYPE after ALTER COLUMN" + }; + + return self.expected(message, self.peek_token()); }; AlterTableOperation::AlterColumn { column_name, op } } else if self.parse_keyword(Keyword::SWAP) { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 6d393d615..eadc17696 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -3457,7 +3457,7 @@ fn parse_alter_table_alter_column_type() { let res = dialect.parse_sql_statements(&format!("{alter_stmt} ALTER COLUMN is_active TYPE TEXT")); assert_eq!( - ParserError::ParserError("Expected SET/DROP NOT NULL, SET DEFAULT, SET DATA TYPE after ALTER COLUMN, found: TYPE".to_string()), + ParserError::ParserError("Expected SET/DROP NOT NULL, SET DEFAULT, or SET DATA TYPE after ALTER COLUMN, found: TYPE".to_string()), res.unwrap_err() ); diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 0c1463f92..da3f06083 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -611,6 +611,20 @@ fn parse_alter_table_alter_column() { } } +#[test] +fn parse_alter_table_alter_column_add_generated() { + pg_and_generic() + .verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED ALWAYS AS IDENTITY"); + pg_and_generic() + .verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY"); + pg_and_generic().verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY"); + pg_and_generic().verified_stmt( + "ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( INCREMENT 1 MINVALUE 1 )", + ); + + pg_and_generic().verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( )"); +} + #[test] fn parse_alter_table_add_columns() { match pg().verified_stmt("ALTER TABLE IF EXISTS ONLY tab ADD COLUMN a TEXT, ADD COLUMN b INT") { From 28ef1926e4ae4c260e407f2dee13af150c90bec9 Mon Sep 17 00:00:00 2001 From: Toby Hede Date: Tue, 2 Jan 2024 09:53:15 +1100 Subject: [PATCH 2/2] Handle errors correctly Signed-off-by: Toby Hede --- src/parser/mod.rs | 4 ++-- tests/sqlparser_postgres.rs | 24 +++++++++++++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index b56068ae7..b43d7e7f2 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -4874,13 +4874,13 @@ impl<'a> Parser<'a> { None }; - let _ = self.expect_keywords(&[Keyword::AS, Keyword::IDENTITY]); + self.expect_keywords(&[Keyword::AS, Keyword::IDENTITY])?; let mut sequence_options: Option> = None; if self.peek_token().token == Token::LParen { self.expect_token(&Token::LParen)?; - sequence_options = self.parse_create_sequence_options().ok(); + sequence_options = Some(self.parse_create_sequence_options()?); self.expect_token(&Token::RParen)?; } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index da3f06083..f9e85e9c4 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -621,8 +621,30 @@ fn parse_alter_table_alter_column_add_generated() { pg_and_generic().verified_stmt( "ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( INCREMENT 1 MINVALUE 1 )", ); - pg_and_generic().verified_stmt("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( )"); + + let res = pg().parse_sql_statements( + "ALTER TABLE t ALTER COLUMN id ADD GENERATED ( INCREMENT 1 MINVALUE 1 )", + ); + assert_eq!( + ParserError::ParserError("Expected AS, found: (".to_string()), + res.unwrap_err() + ); + + let res = pg().parse_sql_statements( + "ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ( INCREMENT )", + ); + assert_eq!( + ParserError::ParserError("Expected a value, found: )".to_string()), + res.unwrap_err() + ); + + let res = + pg().parse_sql_statements("ALTER TABLE t ALTER COLUMN id ADD GENERATED AS IDENTITY ("); + assert_eq!( + ParserError::ParserError("Expected ), found: EOF".to_string()), + res.unwrap_err() + ); } #[test]