Skip to content

Commit

Permalink
ClickHouse: support alter table update ...
Browse files Browse the repository at this point in the history
  • Loading branch information
fangxu.hu committed Oct 16, 2024
1 parent 7c20d4a commit 17080f1
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 2 deletions.
23 changes: 23 additions & 0 deletions src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ use crate::ast::{
use crate::keywords::Keyword;
use crate::tokenizer::Token;

use super::Assignment;

/// An `ALTER TABLE` (`Statement::AlterTable`) operation
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
Expand Down Expand Up @@ -143,6 +145,17 @@ pub enum AlterTableOperation {
partition: Partition,
with_name: Option<Ident>,
},
/// `UPDATE <column> = <expression> [, ...] [IN PARTITION partition_id] WHERE <filter_expr>`
/// Note: this is a ClickHouse-specific operation, please refer to
/// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/alter/update)
Update {
/// Column assignments
assignments: Vec<Assignment>,
/// PARTITION
partition_id: Option<Ident>,
/// WHERE
selection: Option<Expr>,
},
/// `DROP PRIMARY KEY`
///
/// Note: this is a MySQL-specific operation.
Expand Down Expand Up @@ -546,6 +559,16 @@ impl fmt::Display for AlterTableOperation {
}
Ok(())
}
AlterTableOperation::Update { assignments, partition_id, selection } => {
write!(f, "UPDATE {}", display_comma_separated(assignments))?;
if let Some(partition_id) = partition_id {
write!(f, " IN PARTITION {}", partition_id)?;
}
if let Some(selection) = selection {
write!(f, " WHERE {}", selection)?;
}
Ok(())
}
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions src/dialect/clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,10 @@ impl Dialect for ClickHouseDialect {
self.is_identifier_start(ch) || ch.is_ascii_digit()
}

fn supports_alter_table_update(&self) -> bool {
true
}

fn supports_string_literal_backslash_escape(&self) -> bool {
true
}
Expand Down
4 changes: 4 additions & 0 deletions src/dialect/generic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,10 @@ impl Dialect for GenericDialect {
|| ch == '_'
}

fn supports_alter_table_update(&self) -> bool {
true
}

fn supports_unicode_string_literal(&self) -> bool {
true
}
Expand Down
7 changes: 6 additions & 1 deletion src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ use alloc::boxed::Box;

/// Convenience check if a [`Parser`] uses a certain dialect.
///
/// Note: when possible please the new style, adding a method to the [`Dialect`]
/// Note: when possible please use the new style, adding a method to the [`Dialect`]
/// trait rather than using this macro.
///
/// The benefits of adding a method on `Dialect` over this macro are:
Expand Down Expand Up @@ -149,6 +149,11 @@ pub trait Dialect: Debug + Any {
false
}

/// Determine if the dialect supports `ALTER TABLE ... UPDATE ...` statements.
fn supports_alter_table_update(&self) -> bool {
false
}

/// Determine if the dialect supports escaping characters via '\' in string literals.
///
/// Some dialects like BigQuery and Snowflake support this while others like
Expand Down
29 changes: 28 additions & 1 deletion src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1081,7 +1081,7 @@ impl<'a> Parser<'a> {
self.parse_bigquery_struct_literal()
}
Keyword::PRIOR if matches!(self.state, ParserState::ConnectBy) => {
let expr = self.parse_subexpr(self.dialect.prec_value(Precedence::PlusMinus))?;
let expr = self.parse_subexpr(self.dialect.prec_value(Precedence::Pipe))?;
Ok(Expr::Prior(Box::new(expr)))
}
Keyword::MAP if self.peek_token() == Token::LBrace && self.dialect.support_map_literal_syntax() => {
Expand Down Expand Up @@ -7156,6 +7156,33 @@ impl<'a> Parser<'a> {
partition,
with_name,
}
} else if self.dialect.supports_alter_table_update() && self.parse_keyword(Keyword::UPDATE)
{
let mut assignments = vec![];
loop {
let target = self.parse_assignment_target()?;
self.expect_token(&Token::Eq)?;
let value = self.parse_subexpr(self.dialect.prec_value(Precedence::Between))?;
assignments.push(Assignment { target, value });
if self.is_parse_comma_separated_end() {
break;
}
}
let partition_id = if self.parse_keywords(&[Keyword::IN, Keyword::PARTITION]) {
Some(self.parse_identifier(false)?)
} else {
None
};
let selection = if self.parse_keyword(Keyword::WHERE) {
Some(self.parse_expr()?)
} else {
None
};
AlterTableOperation::Update {
assignments,
partition_id,
selection,
}
} else {
let options: Vec<SqlOption> =
self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;
Expand Down
21 changes: 21 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11432,3 +11432,24 @@ fn test_any_some_all_comparison() {
verified_stmt("SELECT c1 FROM tbl WHERE c1 <> SOME(SELECT c2 FROM tbl)");
verified_stmt("SELECT 1 = ANY(WITH x AS (SELECT 1) SELECT * FROM x)");
}

#[test]
fn test_parse_alter_table_update() {
let dialects = all_dialects_where(|d| d.supports_alter_table_update());
let cases = [
(
"ALTER TABLE t UPDATE col1 = 1, col2 = col3 + col4 WHERE cod4 = 1",
true,
),
("ALTER TABLE t UPDATE c = 0 IN PARTITION abc", true),
("ALTER TABLE t UPDATE", false),
("ALTER TABLE t UPDATE c WHERE 1 = 1", false),
];
for (sql, is_valid) in cases {
if is_valid {
dialects.verified_stmt(sql);
} else {
assert!(dialects.parse_sql_statements(sql).is_err());
}
}
}

0 comments on commit 17080f1

Please sign in to comment.