ClickHouse: support alter table update ...

Author: fangxu.hu

src/ast/ddl.rs

@@ -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))]
@@ -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)
+    UpdateData {
+        /// Column assignments
+        assignments: Vec<Assignment>,
+        /// PARTITION
+        partition_id: Option<Ident>,
+        /// WHERE
+        selection: Option<Expr>,
+    },
     /// `DROP PRIMARY KEY`
     ///
     /// Note: this is a MySQL-specific operation.
@@ -546,6 +559,16 @@ impl fmt::Display for AlterTableOperation {
                 }
                 Ok(())
             }
+            AlterTableOperation::UpdateData { 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(())
+            }
         }
     }
 }

src/parser/mod.rs

@@ -7156,6 +7156,25 @@ impl<'a> Parser<'a> {
                 partition,
                 with_name,
             }
+        } else if dialect_of!(self is ClickHouseDialect|GenericDialect)
+            && self.parse_keyword(Keyword::UPDATE)
+        {
+            let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
+            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::UpdateData {
+                assignments,
+                partition_id,
+                selection,
+            }
         } else {
             let options: Vec<SqlOption> =
                 self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;

tests/sqlparser_clickhouse.rs

@@ -1422,6 +1422,38 @@ fn parse_create_table_on_commit_and_as_query() {
     }
 }
 
+#[test]
+fn parse_update_data() {
+    let sql = "ALTER TABLE t UPDATE col1 = 1, col2 = col3 + col4 WHERE cod4 = 1";
+    match clickhouse_and_generic().verified_stmt(sql) {
+        Statement::AlterTable { operations, .. } => {
+            assert_eq!(operations.len(), 1);
+            if let AlterTableOperation::UpdateData {
+                assignments,
+                partition_id,
+                selection,
+            } = &operations[0]
+            {
+                assert_eq!(assignments.len(), 2);
+                assert!(partition_id.is_none());
+                if let Some(Expr::BinaryOp { op, .. }) = selection {
+                    assert_eq!(*op, BinaryOperator::Eq);
+                } else {
+                    unreachable!()
+                }
+            } else {
+                unreachable!();
+            }
+        }
+        _ => unreachable!(),
+    }
+
+    let fails = ["ALTER TABLE t UPDATE", "ALTER TABLE t UPDATE c WHERE 1 = 1"];
+    for f in fails {
+        assert!(clickhouse_and_generic().parse_sql_statements(f).is_err());
+    }
+}
+
 #[test]
 fn parse_freeze_and_unfreeze_partition() {
     // test cases without `WITH NAME`