Support ADD PROJECTION syntax for ClickHouse (#1390)

This commit is contained in:
hulk 2024-08-23 00:33:44 +08:00 committed by GitHub
parent 11a6e6ff7b
commit 19e694aa91
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 231 additions and 65 deletions

View file

@ -26,7 +26,7 @@ use sqlparser_derive::{Visit, VisitMut};
use crate::ast::value::escape_single_quote_string; use crate::ast::value::escape_single_quote_string;
use crate::ast::{ use crate::ast::{
display_comma_separated, display_separated, DataType, Expr, Ident, MySQLColumnPosition, display_comma_separated, display_separated, DataType, Expr, Ident, MySQLColumnPosition,
ObjectName, SequenceOptions, SqlOption, ObjectName, ProjectionSelect, SequenceOptions, SqlOption,
}; };
use crate::tokenizer::Token; use crate::tokenizer::Token;
@ -48,6 +48,15 @@ pub enum AlterTableOperation {
/// MySQL `ALTER TABLE` only [FIRST | AFTER column_name] /// MySQL `ALTER TABLE` only [FIRST | AFTER column_name]
column_position: Option<MySQLColumnPosition>, 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` /// `DISABLE ROW LEVEL SECURITY`
/// ///
/// Note: this is a PostgreSQL-specific operation. /// Note: this is a PostgreSQL-specific operation.
@ -255,6 +264,17 @@ impl fmt::Display for AlterTableOperation {
Ok(()) 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 } => { AlterTableOperation::AlterColumn { column_name, op } => {
write!(f, "ALTER COLUMN {column_name} {op}") write!(f, "ALTER COLUMN {column_name} {op}")
} }

View file

@ -48,7 +48,7 @@ pub use self::query::{
InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonTableColumn, InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonTableColumn,
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern, JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
OffsetRows, OrderBy, OrderByExpr, PivotValueSource, Query, RenameSelectItem, OffsetRows, OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select, RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity, TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity,

View file

@ -68,16 +68,7 @@ impl fmt::Display for Query {
} }
write!(f, "{}", self.body)?; write!(f, "{}", self.body)?;
if let Some(ref order_by) = self.order_by { if let Some(ref order_by) = self.order_by {
write!(f, " 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")?,
}
}
} }
if let Some(ref limit) = self.limit { if let Some(ref limit) = self.limit {
write!(f, " LIMIT {limit}")?; write!(f, " LIMIT {limit}")?;
@ -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: /// A node in a tree, representing a "query body" expression, roughly:
/// `SELECT ... [ {UNION|EXCEPT|INTERSECT} SELECT ...]` /// `SELECT ... [ {UNION|EXCEPT|INTERSECT} SELECT ...]`
#[allow(clippy::large_enum_variant)] #[allow(clippy::large_enum_variant)]
@ -1717,6 +1735,22 @@ pub struct OrderBy {
pub interpolate: Option<Interpolate>, 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 /// An `ORDER BY` expression
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

View file

@ -576,6 +576,7 @@ define_keywords!(
PRIVILEGES, PRIVILEGES,
PROCEDURE, PROCEDURE,
PROGRAM, PROGRAM,
PROJECTION,
PURGE, PURGE,
QUALIFY, QUALIFY,
QUARTER, QUARTER,

View file

@ -6424,10 +6424,38 @@ impl<'a> Parser<'a> {
Ok(Partition::Partitions(partitions)) 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> { pub fn parse_alter_table_operation(&mut self) -> Result<AlterTableOperation, ParserError> {
let operation = if self.parse_keyword(Keyword::ADD) { let operation = if self.parse_keyword(Keyword::ADD) {
if let Some(constraint) = self.parse_optional_table_constraint()? { if let Some(constraint) = self.parse_optional_table_constraint()? {
AlterTableOperation::AddConstraint(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 { } else {
let if_not_exists = let if_not_exists =
self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
@ -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. /// Parse a possibly qualified, possibly quoted identifier, e.g.
/// `foo` or `myschema."table" /// `foo` or `myschema."table"
/// ///
@ -8264,21 +8352,7 @@ impl<'a> Parser<'a> {
} else { } else {
let body = self.parse_boxed_query_body(self.dialect.prec_unknown())?; 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 = self.parse_optional_order_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 mut limit = None; let mut limit = None;
let mut offset = None; let mut offset = None;
@ -8746,44 +8820,9 @@ impl<'a> Parser<'a> {
None None
}; };
let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) { let group_by = self
let expressions = if self.parse_keyword(Keyword::ALL) { .parse_optional_group_by()?
None .unwrap_or_else(|| GroupByExpr::Expressions(vec![], vec![]));
} 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 cluster_by = if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) { let cluster_by = if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
self.parse_comma_separated(Parser::parse_expr)? self.parse_comma_separated(Parser::parse_expr)?

View file

@ -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] #[test]
fn parse_optimize_table() { fn parse_optimize_table() {
clickhouse_and_generic().verified_stmt("OPTIMIZE TABLE t0"); clickhouse_and_generic().verified_stmt("OPTIMIZE TABLE t0");