Skip to content

Commit

Permalink
Support ADD PROJECTION syntax for ClickHouse (#1390)
Browse files Browse the repository at this point in the history
  • Loading branch information
git-hulk authored Aug 22, 2024
1 parent 11a6e6f commit 19e694a
Show file tree
Hide file tree
Showing 6 changed files with 231 additions and 65 deletions.
22 changes: 21 additions & 1 deletion src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use sqlparser_derive::{Visit, VisitMut};
use crate::ast::value::escape_single_quote_string;
use crate::ast::{
display_comma_separated, display_separated, DataType, Expr, Ident, MySQLColumnPosition,
ObjectName, SequenceOptions, SqlOption,
ObjectName, ProjectionSelect, SequenceOptions, SqlOption,
};
use crate::tokenizer::Token;

Expand All @@ -48,6 +48,15 @@ pub enum AlterTableOperation {
/// MySQL `ALTER TABLE` only [FIRST | AFTER column_name]
column_position: Option<MySQLColumnPosition>,
},
/// `ADD PROJECTION [IF NOT EXISTS] name ( SELECT <COLUMN LIST EXPR> [GROUP BY] [ORDER BY])`
///
/// Note: this is a ClickHouse-specific operation.
/// Please refer to [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/alter/projection#add-projection)
AddProjection {
if_not_exists: bool,
name: Ident,
select: ProjectionSelect,
},
/// `DISABLE ROW LEVEL SECURITY`
///
/// Note: this is a PostgreSQL-specific operation.
Expand Down Expand Up @@ -255,6 +264,17 @@ impl fmt::Display for AlterTableOperation {

Ok(())
}
AlterTableOperation::AddProjection {
if_not_exists,
name,
select: query,
} => {
write!(f, "ADD PROJECTION")?;
if *if_not_exists {
write!(f, " IF NOT EXISTS")?;
}
write!(f, " {} ({})", name, query)
}
AlterTableOperation::AlterColumn { column_name, op } => {
write!(f, "ALTER COLUMN {column_name} {op}")
}
Expand Down
2 changes: 1 addition & 1 deletion src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ pub use self::query::{
InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonTableColumn,
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
OffsetRows, OrderBy, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
OffsetRows, OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity,
Expand Down
54 changes: 44 additions & 10 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,16 +68,7 @@ impl fmt::Display for Query {
}
write!(f, "{}", self.body)?;
if let Some(ref order_by) = self.order_by {
write!(f, " ORDER BY")?;
if !order_by.exprs.is_empty() {
write!(f, " {}", display_comma_separated(&order_by.exprs))?;
}
if let Some(ref interpolate) = order_by.interpolate {
match &interpolate.exprs {
Some(exprs) => write!(f, " INTERPOLATE ({})", display_comma_separated(exprs))?,
None => write!(f, " INTERPOLATE")?,
}
}
write!(f, " {order_by}")?;
}
if let Some(ref limit) = self.limit {
write!(f, " LIMIT {limit}")?;
Expand Down Expand Up @@ -107,6 +98,33 @@ impl fmt::Display for Query {
}
}

/// Query syntax for ClickHouse ADD PROJECTION statement.
/// Its syntax is similar to SELECT statement, but it is used to add a new projection to a table.
/// Syntax is `SELECT <COLUMN LIST EXPR> [GROUP BY] [ORDER BY]`
///
/// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/alter/projection#add-projection)
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct ProjectionSelect {
pub projection: Vec<SelectItem>,
pub order_by: Option<OrderBy>,
pub group_by: Option<GroupByExpr>,
}

impl fmt::Display for ProjectionSelect {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "SELECT {}", display_comma_separated(&self.projection))?;
if let Some(ref group_by) = self.group_by {
write!(f, " {group_by}")?;
}
if let Some(ref order_by) = self.order_by {
write!(f, " {order_by}")?;
}
Ok(())
}
}

/// A node in a tree, representing a "query body" expression, roughly:
/// `SELECT ... [ {UNION|EXCEPT|INTERSECT} SELECT ...]`
#[allow(clippy::large_enum_variant)]
Expand Down Expand Up @@ -1717,6 +1735,22 @@ pub struct OrderBy {
pub interpolate: Option<Interpolate>,
}

impl fmt::Display for OrderBy {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "ORDER BY")?;
if !self.exprs.is_empty() {
write!(f, " {}", display_comma_separated(&self.exprs))?;
}
if let Some(ref interpolate) = self.interpolate {
match &interpolate.exprs {
Some(exprs) => write!(f, " INTERPOLATE ({})", display_comma_separated(exprs))?,
None => write!(f, " INTERPOLATE")?,
}
}
Ok(())
}
}

/// An `ORDER BY` expression
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -576,6 +576,7 @@ define_keywords!(
PRIVILEGES,
PROCEDURE,
PROGRAM,
PROJECTION,
PURGE,
QUALIFY,
QUARTER,
Expand Down
145 changes: 92 additions & 53 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6424,10 +6424,38 @@ impl<'a> Parser<'a> {
Ok(Partition::Partitions(partitions))
}

pub fn parse_projection_select(&mut self) -> Result<ProjectionSelect, ParserError> {
self.expect_token(&Token::LParen)?;
self.expect_keyword(Keyword::SELECT)?;
let projection = self.parse_projection()?;
let group_by = self.parse_optional_group_by()?;
let order_by = self.parse_optional_order_by()?;
self.expect_token(&Token::RParen)?;
Ok(ProjectionSelect {
projection,
group_by,
order_by,
})
}
pub fn parse_alter_table_add_projection(&mut self) -> Result<AlterTableOperation, ParserError> {
let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
let name = self.parse_identifier(false)?;
let query = self.parse_projection_select()?;
Ok(AlterTableOperation::AddProjection {
if_not_exists,
name,
select: query,
})
}

pub fn parse_alter_table_operation(&mut self) -> Result<AlterTableOperation, ParserError> {
let operation = if self.parse_keyword(Keyword::ADD) {
if let Some(constraint) = self.parse_optional_table_constraint()? {
AlterTableOperation::AddConstraint(constraint)
} else if dialect_of!(self is ClickHouseDialect|GenericDialect)
&& self.parse_keyword(Keyword::PROJECTION)
{
return self.parse_alter_table_add_projection();
} else {
let if_not_exists =
self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
Expand Down Expand Up @@ -7672,6 +7700,66 @@ impl<'a> Parser<'a> {
}
}

pub fn parse_optional_group_by(&mut self) -> Result<Option<GroupByExpr>, ParserError> {
if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
let expressions = if self.parse_keyword(Keyword::ALL) {
None
} else {
Some(self.parse_comma_separated(Parser::parse_group_by_expr)?)
};

let mut modifiers = vec![];
if dialect_of!(self is ClickHouseDialect | GenericDialect) {
loop {
if !self.parse_keyword(Keyword::WITH) {
break;
}
let keyword = self.expect_one_of_keywords(&[
Keyword::ROLLUP,
Keyword::CUBE,
Keyword::TOTALS,
])?;
modifiers.push(match keyword {
Keyword::ROLLUP => GroupByWithModifier::Rollup,
Keyword::CUBE => GroupByWithModifier::Cube,
Keyword::TOTALS => GroupByWithModifier::Totals,
_ => {
return parser_err!(
"BUG: expected to match GroupBy modifier keyword",
self.peek_token().location
)
}
});
}
}
let group_by = match expressions {
None => GroupByExpr::All(modifiers),
Some(exprs) => GroupByExpr::Expressions(exprs, modifiers),
};
Ok(Some(group_by))
} else {
Ok(None)
}
}

pub fn parse_optional_order_by(&mut self) -> Result<Option<OrderBy>, ParserError> {
if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
let order_by_exprs = self.parse_comma_separated(Parser::parse_order_by_expr)?;
let interpolate = if dialect_of!(self is ClickHouseDialect | GenericDialect) {
self.parse_interpolations()?
} else {
None
};

Ok(Some(OrderBy {
exprs: order_by_exprs,
interpolate,
}))
} else {
Ok(None)
}
}

/// Parse a possibly qualified, possibly quoted identifier, e.g.
/// `foo` or `myschema."table"
///
Expand Down Expand Up @@ -8264,21 +8352,7 @@ impl<'a> Parser<'a> {
} else {
let body = self.parse_boxed_query_body(self.dialect.prec_unknown())?;

let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
let order_by_exprs = self.parse_comma_separated(Parser::parse_order_by_expr)?;
let interpolate = if dialect_of!(self is ClickHouseDialect | GenericDialect) {
self.parse_interpolations()?
} else {
None
};

Some(OrderBy {
exprs: order_by_exprs,
interpolate,
})
} else {
None
};
let order_by = self.parse_optional_order_by()?;

let mut limit = None;
let mut offset = None;
Expand Down Expand Up @@ -8746,44 +8820,9 @@ impl<'a> Parser<'a> {
None
};

let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
let expressions = if self.parse_keyword(Keyword::ALL) {
None
} else {
Some(self.parse_comma_separated(Parser::parse_group_by_expr)?)
};

let mut modifiers = vec![];
if dialect_of!(self is ClickHouseDialect | GenericDialect) {
loop {
if !self.parse_keyword(Keyword::WITH) {
break;
}
let keyword = self.expect_one_of_keywords(&[
Keyword::ROLLUP,
Keyword::CUBE,
Keyword::TOTALS,
])?;
modifiers.push(match keyword {
Keyword::ROLLUP => GroupByWithModifier::Rollup,
Keyword::CUBE => GroupByWithModifier::Cube,
Keyword::TOTALS => GroupByWithModifier::Totals,
_ => {
return parser_err!(
"BUG: expected to match GroupBy modifier keyword",
self.peek_token().location
)
}
});
}
}
match expressions {
None => GroupByExpr::All(modifiers),
Some(exprs) => GroupByExpr::Expressions(exprs, modifiers),
}
} else {
GroupByExpr::Expressions(vec![], vec![])
};
let group_by = self
.parse_optional_group_by()?
.unwrap_or_else(|| GroupByExpr::Expressions(vec![], vec![]));

let cluster_by = if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
self.parse_comma_separated(Parser::parse_expr)?
Expand Down
72 changes: 72 additions & 0 deletions tests/sqlparser_clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -287,6 +287,78 @@ fn parse_alter_table_attach_and_detach_partition() {
}
}

#[test]
fn parse_alter_table_add_projection() {
match clickhouse_and_generic().verified_stmt(concat!(
"ALTER TABLE t0 ADD PROJECTION IF NOT EXISTS my_name",
" (SELECT a, b GROUP BY a ORDER BY b)",
)) {
Statement::AlterTable {
name, operations, ..
} => {
assert_eq!(name, ObjectName(vec!["t0".into()]));
assert_eq!(1, operations.len());
assert_eq!(
operations[0],
AlterTableOperation::AddProjection {
if_not_exists: true,
name: "my_name".into(),
select: ProjectionSelect {
projection: vec![
UnnamedExpr(Identifier(Ident::new("a"))),
UnnamedExpr(Identifier(Ident::new("b"))),
],
group_by: Some(GroupByExpr::Expressions(
vec![Identifier(Ident::new("a"))],
vec![]
)),
order_by: Some(OrderBy {
exprs: vec![OrderByExpr {
expr: Identifier(Ident::new("b")),
asc: None,
nulls_first: None,
with_fill: None,
}],
interpolate: None,
}),
}
}
)
}
_ => unreachable!(),
}

// leave out IF NOT EXISTS is allowed
clickhouse_and_generic()
.verified_stmt("ALTER TABLE t0 ADD PROJECTION my_name (SELECT a, b GROUP BY a ORDER BY b)");
// leave out GROUP BY is allowed
clickhouse_and_generic()
.verified_stmt("ALTER TABLE t0 ADD PROJECTION my_name (SELECT a, b ORDER BY b)");
// leave out ORDER BY is allowed
clickhouse_and_generic()
.verified_stmt("ALTER TABLE t0 ADD PROJECTION my_name (SELECT a, b GROUP BY a)");

// missing select query is not allowed
assert_eq!(
clickhouse_and_generic()
.parse_sql_statements("ALTER TABLE t0 ADD PROJECTION my_name")
.unwrap_err(),
ParserError("Expected: (, found: EOF".to_string())
);
assert_eq!(
clickhouse_and_generic()
.parse_sql_statements("ALTER TABLE t0 ADD PROJECTION my_name ()")
.unwrap_err(),
ParserError("Expected: SELECT, found: )".to_string())
);
assert_eq!(
clickhouse_and_generic()
.parse_sql_statements("ALTER TABLE t0 ADD PROJECTION my_name (SELECT)")
.unwrap_err(),
ParserError("Expected: an expression:, found: )".to_string())
);
}

#[test]
fn parse_optimize_table() {
clickhouse_and_generic().verified_stmt("OPTIMIZE TABLE t0");
Expand Down

0 comments on commit 19e694a

Please sign in to comment.