From 9a39afbe07b6eecd8819fe1f7d9af2b4aad3a6d7 Mon Sep 17 00:00:00 2001 From: Forbes Lindesay Date: Thu, 17 Aug 2023 16:47:11 +0100 Subject: [PATCH] feat: support more Postgres index syntax (#943) --- src/ast/mod.rs | 35 ++++++- src/keywords.rs | 2 + src/parser/mod.rs | 38 ++++++- tests/sqlparser_common.rs | 12 ++- tests/sqlparser_postgres.rs | 202 ++++++++++++++++++++++++++++++++++++ 5 files changed, 279 insertions(+), 10 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index edab3c63..47419a89 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1348,13 +1348,17 @@ pub enum Statement { /// CREATE INDEX CreateIndex { /// index name - name: ObjectName, + name: Option, #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] table_name: ObjectName, using: Option, columns: Vec, unique: bool, + concurrently: bool, if_not_exists: bool, + include: Vec, + nulls_distinct: Option, + predicate: Option, }, /// CREATE ROLE /// See [postgres](https://www.postgresql.org/docs/current/sql-createrole.html) @@ -2464,20 +2468,41 @@ impl fmt::Display for Statement { using, columns, unique, + concurrently, if_not_exists, + include, + nulls_distinct, + predicate, } => { write!( f, - "CREATE {unique}INDEX {if_not_exists}{name} ON {table_name}", + "CREATE {unique}INDEX {concurrently}{if_not_exists}", unique = if *unique { "UNIQUE " } else { "" }, + concurrently = if *concurrently { "CONCURRENTLY " } else { "" }, if_not_exists = if *if_not_exists { "IF NOT EXISTS " } else { "" }, - name = name, - table_name = table_name )?; + if let Some(value) = name { + write!(f, "{value} ")?; + } + write!(f, "ON {table_name}")?; if let Some(value) = using { write!(f, " USING {value} ")?; } - write!(f, "({})", display_separated(columns, ",")) + write!(f, "({})", display_separated(columns, ","))?; + if !include.is_empty() { + write!(f, " INCLUDE ({})", display_separated(include, ","))?; + } + if let Some(value) = nulls_distinct { + if *value { + write!(f, " NULLS DISTINCT")?; + } else { + write!(f, " NULLS NOT DISTINCT")?; + } + } + if let Some(predicate) = predicate { + write!(f, " WHERE {predicate}")?; + } + Ok(()) } Statement::CreateRole { names, diff --git a/src/keywords.rs b/src/keywords.rs index 5f3e4402..c73535fc 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -153,6 +153,7 @@ define_keywords!( COMMITTED, COMPRESSION, COMPUTE, + CONCURRENTLY, CONDITION, CONFLICT, CONNECT, @@ -310,6 +311,7 @@ define_keywords!( ILIKE, IMMUTABLE, IN, + INCLUDE, INCREMENT, INDEX, INDICATOR, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index d6f45359..94814627 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -3370,9 +3370,15 @@ impl<'a> Parser<'a> { } pub fn parse_create_index(&mut self, unique: bool) -> Result { + let concurrently = self.parse_keyword(Keyword::CONCURRENTLY); let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let index_name = self.parse_object_name()?; - self.expect_keyword(Keyword::ON)?; + let index_name = if if_not_exists || !self.parse_keyword(Keyword::ON) { + let index_name = self.parse_object_name()?; + self.expect_keyword(Keyword::ON)?; + Some(index_name) + } else { + None + }; let table_name = self.parse_object_name()?; let using = if self.parse_keyword(Keyword::USING) { Some(self.parse_identifier()?) @@ -3382,13 +3388,41 @@ impl<'a> Parser<'a> { self.expect_token(&Token::LParen)?; let columns = self.parse_comma_separated(Parser::parse_order_by_expr)?; self.expect_token(&Token::RParen)?; + + let include = if self.parse_keyword(Keyword::INCLUDE) { + self.expect_token(&Token::LParen)?; + let columns = self.parse_comma_separated(Parser::parse_identifier)?; + self.expect_token(&Token::RParen)?; + columns + } else { + vec![] + }; + + let nulls_distinct = if self.parse_keyword(Keyword::NULLS) { + let not = self.parse_keyword(Keyword::NOT); + self.expect_keyword(Keyword::DISTINCT)?; + Some(!not) + } else { + None + }; + + let predicate = if self.parse_keyword(Keyword::WHERE) { + Some(self.parse_expr()?) + } else { + None + }; + Ok(Statement::CreateIndex { name: index_name, table_name, using, columns, unique, + concurrently, if_not_exists, + include, + nulls_distinct, + predicate, }) } diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index ca30e951..8dfcc6e7 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -5981,7 +5981,7 @@ fn parse_create_index() { ]; match verified_stmt(sql) { Statement::CreateIndex { - name, + name: Some(name), table_name, columns, unique, @@ -6015,19 +6015,25 @@ fn test_create_index_with_using_function() { ]; match verified_stmt(sql) { Statement::CreateIndex { - name, + name: Some(name), table_name, using, columns, unique, + concurrently, if_not_exists, + include, + nulls_distinct: None, + predicate: None, } => { assert_eq!("idx_name", name.to_string()); assert_eq!("test", table_name.to_string()); assert_eq!("btree", using.unwrap().to_string()); assert_eq!(indexed_columns, columns); assert!(unique); - assert!(if_not_exists) + assert!(!concurrently); + assert!(if_not_exists); + assert!(include.is_empty()); } _ => unreachable!(), } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 09196db4..a62c41e4 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -1761,6 +1761,208 @@ fn parse_array_index_expr() { ); } +#[test] +fn parse_create_index() { + let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2)"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + nulls_distinct: None, + include, + predicate: None, + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + } + _ => unreachable!(), + } +} + +#[test] +fn parse_create_anonymous_index() { + let sql = "CREATE INDEX ON my_table(col1,col2)"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name, + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: None, + predicate: None, + } => { + assert_eq!(None, name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(!if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + } + _ => unreachable!(), + } +} + +#[test] +fn parse_create_index_concurrently() { + let sql = "CREATE INDEX CONCURRENTLY IF NOT EXISTS my_index ON my_table(col1,col2)"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: None, + predicate: None, + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + } + _ => unreachable!(), + } +} + +#[test] +fn parse_create_index_with_predicate() { + let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2) WHERE col3 IS NULL"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: None, + predicate: Some(_), + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + } + _ => unreachable!(), + } +} + +#[test] +fn parse_create_index_with_include() { + let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2) INCLUDE (col3)"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: None, + predicate: None, + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert_eq_vec(&["col3"], &include); + } + _ => unreachable!(), + } +} + +#[test] +fn parse_create_index_with_nulls_distinct() { + let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2) NULLS NOT DISTINCT"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: Some(nulls_distinct), + predicate: None, + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + assert!(!nulls_distinct); + } + _ => unreachable!(), + } + + let sql = "CREATE INDEX IF NOT EXISTS my_index ON my_table(col1,col2) NULLS DISTINCT"; + match pg().verified_stmt(sql) { + Statement::CreateIndex { + name: Some(ObjectName(name)), + table_name: ObjectName(table_name), + using, + columns, + unique, + concurrently, + if_not_exists, + include, + nulls_distinct: Some(nulls_distinct), + predicate: None, + } => { + assert_eq_vec(&["my_index"], &name); + assert_eq_vec(&["my_table"], &table_name); + assert_eq!(None, using); + assert!(!unique); + assert!(!concurrently); + assert!(if_not_exists); + assert_eq_vec(&["col1", "col2"], &columns); + assert!(include.is_empty()); + assert!(nulls_distinct); + } + _ => unreachable!(), + } +} + #[test] fn parse_array_subquery_expr() { let sql = "SELECT ARRAY(SELECT 1 UNION SELECT 2)";