diff --git a/src/ast/table_constraints.rs b/src/ast/table_constraints.rs index 9ba196a81..2415e81b5 100644 --- a/src/ast/table_constraints.rs +++ b/src/ast/table_constraints.rs @@ -443,6 +443,8 @@ pub struct PrimaryKeyConstraint { pub index_type: Option, /// Identifiers of the columns that form the primary key. pub columns: Vec, + /// INCLUDE clause: + pub include: Vec, /// Optional index options such as `USING`. pub index_options: Vec, /// Optional characteristics like `DEFERRABLE`. @@ -461,6 +463,10 @@ impl fmt::Display for PrimaryKeyConstraint { display_comma_separated(&self.columns), )?; + if !self.include.is_empty() { + write!(f, " INCLUDE ({})", display_comma_separated(&self.include))?; + } + if !self.index_options.is_empty() { write!(f, " {}", display_separated(&self.index_options, " "))?; } @@ -482,6 +488,7 @@ impl crate::ast::Spanned for PrimaryKeyConstraint { .map(|i| i.span) .chain(self.index_name.iter().map(|i| i.span)) .chain(self.columns.iter().map(|i| i.span())) + .chain(self.include.iter().map(|i| i.span)) .chain(self.characteristics.iter().map(|i| i.span())), ) } @@ -506,6 +513,8 @@ pub struct UniqueConstraint { pub index_type: Option, /// Identifiers of the columns that are unique. pub columns: Vec, + /// INCLUDE clause: + pub include: Vec, /// Optional index options such as `USING`. pub index_options: Vec, /// Optional characteristics like `DEFERRABLE`. @@ -528,6 +537,10 @@ impl fmt::Display for UniqueConstraint { display_comma_separated(&self.columns), )?; + if !self.include.is_empty() { + write!(f, " INCLUDE ({})", display_comma_separated(&self.include))?; + } + if !self.index_options.is_empty() { write!(f, " {}", display_separated(&self.index_options, " "))?; } @@ -549,6 +562,7 @@ impl crate::ast::Spanned for UniqueConstraint { .map(|i| i.span) .chain(self.index_name.iter().map(|i| i.span)) .chain(self.columns.iter().map(|i| i.span())) + .chain(self.include.iter().map(|i| i.span)) .chain(self.characteristics.iter().map(|i| i.span())), ) } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index d7dd28034..19001180f 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9378,6 +9378,7 @@ impl<'a> Parser<'a> { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics, } @@ -9398,6 +9399,7 @@ impl<'a> Parser<'a> { index_type_display, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics, nulls_distinct: NullsDistinctOption::None, @@ -9414,6 +9416,7 @@ impl<'a> Parser<'a> { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics, } @@ -9841,6 +9844,7 @@ impl<'a> Parser<'a> { let index_type = self.parse_optional_using_then_index_type()?; let columns = self.parse_parenthesized_index_column_list()?; + let include = self.parse_optional_include_columns()?; let index_options = self.parse_index_options()?; let characteristics = self.parse_constraint_characteristics()?; Ok(Some( @@ -9850,6 +9854,7 @@ impl<'a> Parser<'a> { index_type_display, index_type, columns, + include, index_options, characteristics, nulls_distinct, @@ -9874,6 +9879,7 @@ impl<'a> Parser<'a> { let index_type = self.parse_optional_using_then_index_type()?; let columns = self.parse_parenthesized_index_column_list()?; + let include = self.parse_optional_include_columns()?; let index_options = self.parse_index_options()?; let characteristics = self.parse_constraint_characteristics()?; Ok(Some( @@ -9882,6 +9888,7 @@ impl<'a> Parser<'a> { index_name, index_type, columns, + include, index_options, characteristics, } @@ -10157,6 +10164,18 @@ impl<'a> Parser<'a> { } } + /// Parse an optional `INCLUDE (col, ...)` clause on a table constraint. + pub fn parse_optional_include_columns(&mut self) -> Result, ParserError> { + if self.parse_keyword(Keyword::INCLUDE) { + self.expect_token(&Token::LParen)?; + let columns = self.parse_comma_separated(|p| p.parse_identifier())?; + self.expect_token(&Token::RParen)?; + Ok(columns) + } else { + Ok(vec![]) + } + } + /// Parse a single `SqlOption` used by various dialect-specific DDL statements. pub fn parse_sql_option(&mut self) -> Result { let is_mssql = dialect_of!(self is MsSqlDialect|GenericDialect); diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 627016427..d47addf39 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -3944,6 +3944,7 @@ fn parse_create_table() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }), @@ -3960,6 +3961,7 @@ fn parse_create_table() { index_type_display: KeyOrIndexDisplay::None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, nulls_distinct: NullsDistinctOption::None, @@ -4298,6 +4300,7 @@ fn parse_create_table_column_constraint_characteristics() { index_type_display: KeyOrIndexDisplay::None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: expected_value, nulls_distinct: NullsDistinctOption::None, diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs index f513d3670..ffd5852ad 100644 --- a/tests/sqlparser_mysql.rs +++ b/tests/sqlparser_mysql.rs @@ -655,6 +655,7 @@ fn parse_create_table_auto_increment() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }), @@ -706,6 +707,7 @@ fn table_constraint_unique_primary_ctor( index_type_display, index_type, columns, + include: vec![], index_options, characteristics, nulls_distinct: NullsDistinctOption::None, @@ -716,6 +718,7 @@ fn table_constraint_unique_primary_ctor( index_name, index_type, columns, + include: vec![], index_options, characteristics, } @@ -764,6 +767,7 @@ fn parse_create_table_primary_and_unique_key() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }), @@ -1435,6 +1439,7 @@ fn parse_quote_identifiers() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }), diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 274988be0..92b62f02f 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -705,6 +705,50 @@ fn parse_alter_table_constraint_using_index() { ); } +#[test] +fn parse_constraint_include_columns() { + // INCLUDE covering columns on PRIMARY KEY / UNIQUE table constraints. + // https://www.postgresql.org/docs/current/sql-createtable.html + pg_and_generic().verified_stmt( + "CREATE TABLE t (id INT, payload TEXT, CONSTRAINT t_pk PRIMARY KEY (id) INCLUDE (payload))", + ); + pg_and_generic().verified_stmt( + "CREATE TABLE t (id INT, email TEXT, payload TEXT, CONSTRAINT t_uk UNIQUE (email) INCLUDE (payload))", + ); + pg_and_generic().verified_stmt( + "CREATE TABLE t (a INT, b INT, c INT, d INT, CONSTRAINT t_pk PRIMARY KEY (a, b) INCLUDE (c, d))", + ); + pg_and_generic() + .verified_stmt("ALTER TABLE t ADD CONSTRAINT t_pk PRIMARY KEY (id) INCLUDE (payload)"); + pg_and_generic() + .verified_stmt("ALTER TABLE t ADD CONSTRAINT t_uk UNIQUE (email) INCLUDE (payload)"); + pg_and_generic().verified_stmt( + "ALTER TABLE t ADD CONSTRAINT t_pk PRIMARY KEY (id) INCLUDE (payload) DEFERRABLE INITIALLY DEFERRED", + ); + + match pg_and_generic().verified_stmt( + "ALTER TABLE t ADD CONSTRAINT t_pk PRIMARY KEY (id) INCLUDE (payload, extra)", + ) { + Statement::AlterTable(alter_table) => match &alter_table.operations[0] { + AlterTableOperation::AddConstraint { + constraint: TableConstraint::PrimaryKey(pk), + .. + } => { + assert_eq!(pk.name.as_ref().unwrap().to_string(), "t_pk"); + assert_eq!( + pk.include + .iter() + .map(|i| i.value.clone()) + .collect::>(), + vec!["payload".to_string(), "extra".to_string()] + ); + } + _ => unreachable!(), + }, + _ => unreachable!(), + } +} + #[test] fn parse_alter_table_disable() { pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE ROW LEVEL SECURITY"); diff --git a/tests/sqlparser_sqlite.rs b/tests/sqlparser_sqlite.rs index f9536bc28..5b936dde0 100644 --- a/tests/sqlparser_sqlite.rs +++ b/tests/sqlparser_sqlite.rs @@ -223,6 +223,7 @@ fn parse_create_table_auto_increment() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }), @@ -255,6 +256,7 @@ fn parse_create_table_primary_key_asc_desc() { index_name: None, index_type: None, columns: vec![], + include: vec![], index_options: vec![], characteristics: None, }),