From 19e694aa91d2211557cfc0fa76a408ec1c664896 Mon Sep 17 00:00:00 2001 From: hulk Date: Fri, 23 Aug 2024 00:33:44 +0800 Subject: [PATCH] Support `ADD PROJECTION` syntax for ClickHouse (#1390) --- src/ast/ddl.rs | 22 +++++- src/ast/mod.rs | 2 +- src/ast/query.rs | 54 ++++++++++--- src/keywords.rs | 1 + src/parser/mod.rs | 145 +++++++++++++++++++++------------- tests/sqlparser_clickhouse.rs | 72 +++++++++++++++++ 6 files changed, 231 insertions(+), 65 deletions(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 265ae7727..31e216a9f 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -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; @@ -48,6 +48,15 @@ pub enum AlterTableOperation { /// MySQL `ALTER TABLE` only [FIRST | AFTER column_name] column_position: Option, }, + /// `ADD PROJECTION [IF NOT EXISTS] name ( SELECT [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. @@ -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}") } diff --git a/src/ast/mod.rs b/src/ast/mod.rs index c4533ef57..4f9aac885 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -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, diff --git a/src/ast/query.rs b/src/ast/query.rs index cda7430be..c52d01105 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -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}")?; @@ -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 [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, + pub order_by: Option, + pub group_by: Option, +} + +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)] @@ -1717,6 +1735,22 @@ pub struct OrderBy { pub interpolate: Option, } +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))] diff --git a/src/keywords.rs b/src/keywords.rs index 538c2d380..acb913d57 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -576,6 +576,7 @@ define_keywords!( PRIVILEGES, PROCEDURE, PROGRAM, + PROJECTION, PURGE, QUALIFY, QUARTER, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 2632d807a..1eff8f7d5 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -6424,10 +6424,38 @@ impl<'a> Parser<'a> { Ok(Partition::Partitions(partitions)) } + pub fn parse_projection_select(&mut self) -> Result { + 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 { + 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 { 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]); @@ -7672,6 +7700,66 @@ impl<'a> Parser<'a> { } } + pub fn parse_optional_group_by(&mut self) -> Result, 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, 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" /// @@ -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; @@ -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)? diff --git a/tests/sqlparser_clickhouse.rs b/tests/sqlparser_clickhouse.rs index 931899ff9..fe255cda5 100644 --- a/tests/sqlparser_clickhouse.rs +++ b/tests/sqlparser_clickhouse.rs @@ -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");