use reifydb_core::{common::IndexType, sort::SortDirection};
use crate::{
Result,
ast::{
ast::{AstCreate, AstCreateIndex, AstIndexColumn},
identifier::MaybeQualifiedIndexIdentifier,
parse::{Parser, Precedence},
},
bump::BumpBox,
token::{
keyword::Keyword::{Asc, Desc, Filter, Index, Map, On, Unique},
operator::Operator,
separator::Separator::Comma,
token::{Token, TokenKind},
},
};
impl<'bump> Parser<'bump> {
pub(crate) fn peek_is_index_creation(&mut self) -> Result<bool> {
Ok(matches!(self.current()?.kind, TokenKind::Keyword(Index) | TokenKind::Keyword(Unique)))
}
pub(crate) fn parse_create_index(&mut self, create_token: Token<'bump>) -> Result<AstCreate<'bump>> {
let index_type = self.parse_index_type()?;
let name_token = self.consume(TokenKind::Identifier)?;
self.consume_keyword(On)?;
let mut segments = self.parse_double_colon_separated_identifiers()?;
let table_fragment = segments.pop().unwrap().into_fragment();
let namespace: Vec<_> = segments.into_iter().map(|s| s.into_fragment()).collect();
let index =
MaybeQualifiedIndexIdentifier::new(table_fragment, name_token.fragment).with_shape(namespace);
let columns = self.parse_index_columns()?;
let mut filters = Vec::new();
while self.consume_if(TokenKind::Keyword(Filter))?.is_some() {
filters.push(BumpBox::new_in(self.parse_node(Precedence::None)?, self.bump()));
}
let map = if self.consume_if(TokenKind::Keyword(Map))?.is_some() {
Some(BumpBox::new_in(self.parse_node(Precedence::None)?, self.bump()))
} else {
None
};
Ok(AstCreate::Index(AstCreateIndex {
token: create_token,
index_type,
index,
columns,
filters,
map,
}))
}
fn parse_index_type(&mut self) -> Result<IndexType> {
if self.consume_if(TokenKind::Keyword(Unique))?.is_some() {
self.consume_keyword(Index)?;
Ok(IndexType::Unique)
} else {
self.consume_keyword(Index)?;
Ok(IndexType::Index)
}
}
fn parse_index_columns(&mut self) -> Result<Vec<AstIndexColumn<'bump>>> {
let mut columns = Vec::new();
self.consume_operator(Operator::OpenCurly)?;
loop {
self.skip_new_line()?;
if self.current()?.is_operator(Operator::CloseCurly) {
break;
}
let column = self.parse_column_identifier()?;
let order = if self.consume_if(TokenKind::Operator(Operator::Colon))?.is_some() {
if self.consume_if(TokenKind::Keyword(Asc))?.is_some() {
Some(SortDirection::Asc)
} else if self.consume_if(TokenKind::Keyword(Desc))?.is_some() {
Some(SortDirection::Desc)
} else {
None
}
} else {
None
};
columns.push(AstIndexColumn {
column,
order,
});
if self.consume_if(TokenKind::Separator(Comma))?.is_none() {
break;
}
}
self.consume_operator(Operator::CloseCurly)?;
Ok(columns)
}
}
#[cfg(test)]
pub mod tests {
use reifydb_core::{common::IndexType, sort::SortDirection};
use crate::{
ast::{
ast::{AstCreate, AstCreateIndex},
parse::Parser,
tokenize,
},
bump::Bump,
};
#[test]
fn test_create_index() {
let bump = Bump::new();
let source = r#"create index idx_email on test::users {email}"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
index_type,
index,
columns,
filters,
..
}) => {
assert_eq!(*index_type, IndexType::Index);
assert_eq!(index.name.text(), "idx_email");
assert_eq!(index.namespace[0].text(), "test");
assert_eq!(index.table.text(), "users");
assert_eq!(columns.len(), 1);
assert_eq!(columns[0].column.name.text(), "email");
assert!(columns[0].order.is_none());
assert_eq!(filters.len(), 0);
}
_ => unreachable!(),
}
}
#[test]
fn test_create_unique_index() {
let bump = Bump::new();
let source = r#"create unique index idx_email on test::users {email}"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
index_type,
index,
columns,
filters,
..
}) => {
assert_eq!(*index_type, IndexType::Unique);
assert_eq!(index.name.text(), "idx_email");
assert_eq!(index.namespace[0].text(), "test");
assert_eq!(index.table.text(), "users");
assert_eq!(columns.len(), 1);
assert_eq!(columns[0].column.name.text(), "email");
assert_eq!(filters.len(), 0);
}
_ => unreachable!(),
}
}
#[test]
fn test_create_composite_index() {
let bump = Bump::new();
let source = r#"create index idx_name on test::users {last_name, first_name}"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
columns,
filters,
..
}) => {
assert_eq!(columns.len(), 2);
assert_eq!(columns[0].column.name.text(), "last_name");
assert_eq!(columns[1].column.name.text(), "first_name");
assert_eq!(filters.len(), 0);
}
_ => unreachable!(),
}
}
#[test]
fn test_create_index_with_ordering() {
let bump = Bump::new();
let source = r#"create index idx_status on test::users {created_at:desc, status:asc}"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
columns,
filters,
..
}) => {
assert_eq!(columns.len(), 2);
assert_eq!(columns[0].column.name.text(), "created_at");
assert_eq!(columns[0].order, Some(SortDirection::Desc));
assert_eq!(columns[1].column.name.text(), "status");
assert_eq!(columns[1].order, Some(SortDirection::Asc));
assert_eq!(filters.len(), 0);
}
_ => unreachable!(),
}
}
#[test]
fn test_create_index_with_single_filter() {
let bump = Bump::new();
let source = r#"create index idx_active_email on test::users {email} filter active == true"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
columns,
filters,
..
}) => {
assert_eq!(columns.len(), 1);
assert_eq!(columns[0].column.name.text(), "email");
assert_eq!(filters.len(), 1);
assert!(filters[0].is_infix());
}
_ => unreachable!(),
}
}
#[test]
fn test_create_index_with_multiple_filters() {
let bump = Bump::new();
let source = r#"create index idx_filtered on test::users {email} filter active == true filter age > 18 filter country == "US""#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
columns,
filters,
..
}) => {
assert_eq!(columns.len(), 1);
assert_eq!(columns[0].column.name.text(), "email");
assert_eq!(filters.len(), 3);
assert!(filters[0].is_infix());
assert!(filters[1].is_infix());
assert!(filters[2].is_infix());
}
_ => unreachable!(),
}
}
#[test]
fn test_create_index_with_filters_and_map() {
let bump = Bump::new();
let source = r#"create index idx_comptokenize on test::users {email} filter active == true filter age > 18 map email"#;
let tokens = tokenize(&bump, source).unwrap().into_iter().collect();
let mut parser = Parser::new(&bump, source, tokens);
let mut result = parser.parse().unwrap();
assert_eq!(result.len(), 1);
let result = result.pop().unwrap();
let create = result.first_unchecked().as_create();
match create {
AstCreate::Index(AstCreateIndex {
columns,
filters,
map,
..
}) => {
assert_eq!(columns.len(), 1);
assert_eq!(columns[0].column.name.text(), "email");
assert_eq!(filters.len(), 2);
assert!(filters[0].is_infix());
assert!(filters[1].is_infix());
assert!(map.is_some());
}
_ => unreachable!(),
}
}
}