feat: support more Postgres index syntax (#943)

This commit is contained in:
Forbes Lindesay 2023-08-17 16:47:11 +01:00 committed by GitHub
parent a49ea1908d
commit 9a39afbe07
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 279 additions and 10 deletions

View file

@ -1348,13 +1348,17 @@ pub enum Statement {
/// CREATE INDEX
CreateIndex {
/// index name
name: ObjectName,
name: Option<ObjectName>,
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
table_name: ObjectName,
using: Option<Ident>,
columns: Vec<OrderByExpr>,
unique: bool,
concurrently: bool,
if_not_exists: bool,
include: Vec<Ident>,
nulls_distinct: Option<bool>,
predicate: Option<Expr>,
},
/// 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,

View file

@ -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,

View file

@ -3370,9 +3370,15 @@ impl<'a> Parser<'a> {
}
pub fn parse_create_index(&mut self, unique: bool) -> Result<Statement, ParserError> {
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,
})
}

View file

@ -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!(),
}

View file

@ -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)";