Add support for ClickHouse CSE.

src/ast/mod.rs

@@ -74,7 +74,7 @@ pub use self::ddl::{
 pub use self::dml::{Delete, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
-    AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
+    AfterMatchSkip, ConnectBy, Cse, Cte, CteAsMaterialized, CteOrCse, Distinct, EmptyMatchesMode,
     ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, ExprWithAliasAndOrderBy, Fetch, ForClause,
     ForJson, ForXml, FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias,
     IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join, JoinConstraint,

src/ast/query.rs

@@ -603,7 +603,7 @@ pub struct With {
     /// Token for the "WITH" keyword
     pub with_token: AttachedToken,
     pub recursive: bool,
-    pub cte_tables: Vec<Cte>,
+    pub cte_tables: Vec<CteOrCse>,
 }
 
 impl fmt::Display for With {
@@ -641,6 +641,56 @@ impl fmt::Display for CteAsMaterialized {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum CteOrCse {
+    Cte(Cte),
+    Cse(Cse),
+}
+
+impl CteOrCse {
+    pub fn cte(&self) -> Option<&Cte> {
+        match self {
+            CteOrCse::Cte(cte) => Some(cte),
+            CteOrCse::Cse(_) => None,
+        }
+    }
+
+    pub fn cse(&self) -> Option<&Cse> {
+        match self {
+            CteOrCse::Cte(_) => None,
+            CteOrCse::Cse(cse) => Some(cse),
+        }
+    }
+}
+
+impl fmt::Display for CteOrCse {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            CteOrCse::Cte(cte) => cte.fmt(f),
+            CteOrCse::Cse(cse) => cse.fmt(f),
+        }
+    }
+}
+
+/// A single CSE (used after `WITH`): `<expr> AS <ident>`.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct Cse {
+    pub expr: Expr,
+    pub ident: Ident,
+}
+
+impl fmt::Display for Cse {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.expr.fmt(f)?;
+        f.write_str(" AS ")?;
+        self.ident.fmt(f)?;
+        Ok(())
+    }
+}
 /// A single CTE (used after `WITH`): `<alias> [(col1, col2, ...)] AS <materialized> ( <query> )`
 /// The names in the column list before `AS`, when specified, replace the names
 /// of the columns returned by the query. The parser does not validate that the

src/ast/spans.rs

@@ -176,6 +176,23 @@ impl Spanned for With {
     }
 }
 
+impl Spanned for super::query::CteOrCse {
+    fn span(&self) -> Span {
+        match self {
+            super::query::CteOrCse::Cte(cte) => cte.span(),
+            super::query::CteOrCse::Cse(cse) => cse.span(),
+        }
+    }
+}
+
+impl Spanned for super::query::Cse {
+    fn span(&self) -> Span {
+        let super::query::Cse { expr, ident } = self;
+
+        union_spans(core::iter::once(expr.span()).chain(core::iter::once(ident.span)))
+    }
+}
+
 impl Spanned for Cte {
     fn span(&self) -> Span {
         let Cte {
@@ -2560,7 +2577,11 @@ pub mod tests {
 
         let query = test.0.parse_query().unwrap();
         let cte_span = query.clone().with.unwrap().cte_tables[0].span();
-        let cte_query_span = query.clone().with.unwrap().cte_tables[0].query.span();
+        let cte_query_span = query.clone().with.unwrap().cte_tables[0]
+            .cte()
+            .unwrap()
+            .query
+            .span();
         let body_span = query.body.span();
 
         // the WITH keyboard is part of the query

src/parser/mod.rs

@@ -11801,7 +11801,7 @@ impl<'a> Parser<'a> {
             Some(With {
                 with_token: with_token.clone().into(),
                 recursive: self.parse_keyword(Keyword::RECURSIVE),
-                cte_tables: self.parse_comma_separated(Parser::parse_cte)?,
+                cte_tables: self.parse_comma_separated(Parser::parse_cte_or_cse)?,
             })
         } else {
             None
@@ -12260,6 +12260,27 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a CTE or CSE.
+    pub fn parse_cte_or_cse(&mut self) -> Result<CteOrCse, ParserError> {
+        Ok(if dialect_of!(self is ClickHouseDialect) {
+            if let Some(cse) = self.maybe_parse(Parser::parse_cse)? {
+                CteOrCse::Cse(cse)
+            } else {
+                CteOrCse::Cte(self.parse_cte()?)
+            }
+        } else {
+            CteOrCse::Cte(self.parse_cte()?)
+        })
+    }
+
+    /// Parse a CSE (`<expr> AS <ident>`).
+    pub fn parse_cse(&mut self) -> Result<Cse, ParserError> {
+        let expr = self.parse_expr()?;
+        let _after_as = self.parse_keyword(Keyword::AS);
+        let ident = self.parse_identifier()?;
+        Ok(Cse { expr, ident })
+    }
+
     /// Parse a CTE (`alias [( col1, col2, ... )] AS (subquery)`)
     pub fn parse_cte(&mut self) -> Result<Cte, ParserError> {
         let name = self.parse_identifier()?;

tests/sqlparser_clickhouse.rs

@@ -1729,6 +1729,24 @@ fn test_parse_not_null_in_column_options() {
     );
 }
 
+#[test]
+fn parse_cses() {
+    clickhouse().verified_stmt("WITH x AS (SELECT 1) UPDATE t SET bar = (SELECT * FROM x)");
+
+    let with = concat!(
+        "WITH",
+        " toIntervalSecond(300) AS bucket_size,",
+        " toDateTime64(1735751460, 9) AS start_time,",
+        " toDateTime64(1735755060, 9) AS end_time ",
+        "SELECT",
+        " toStartOfInterval(EventTime, bucket_size) AS bucket,",
+        " count() AS count ",
+        "FROM logs",
+    );
+
+    clickhouse().verified_query(with);
+}
+
 fn clickhouse() -> TestedDialects {
     TestedDialects::new(vec![Box::new(ClickHouseDialect {})])
 }

tests/sqlparser_common.rs

@@ -7537,7 +7537,7 @@ fn parse_ctes() {
 
     fn assert_ctes_in_select(expected: &[&str], sel: &Query) {
         for (i, exp) in expected.iter().enumerate() {
-            let Cte { alias, query, .. } = &sel.with.as_ref().unwrap().cte_tables[i];
+            let Cte { alias, query, .. } = &sel.with.as_ref().unwrap().cte_tables[i].cte().unwrap();
             assert_eq!(*exp, query.to_string());
             assert_eq!(
                 if i == 0 {
@@ -7580,7 +7580,10 @@ fn parse_ctes() {
     // CTE in a CTE...
     let sql = &format!("WITH outer_cte AS ({with}) SELECT * FROM outer_cte");
     let select = verified_query(sql);
-    assert_ctes_in_select(&cte_sqls, &only(&select.with.unwrap().cte_tables).query);
+    assert_ctes_in_select(
+        &cte_sqls,
+        &only(&select.with.unwrap().cte_tables).cte().unwrap().query,
+    );
 }
 
 #[test]
@@ -7598,6 +7601,8 @@ fn parse_cte_renamed_columns() {
             .cte_tables
             .first()
             .unwrap()
+            .cte()
+            .unwrap()
             .alias
             .columns
     );
@@ -7628,7 +7633,7 @@ fn parse_recursive_cte() {
         materialized: None,
         closing_paren_token: AttachedToken::empty(),
     };
-    assert_eq!(with.cte_tables.first().unwrap(), &expected);
+    assert_eq!(with.cte_tables.first().unwrap().cte().unwrap(), &expected);
 }
 
 #[test]
@@ -17105,7 +17110,7 @@ fn test_parse_semantic_view_table_factor() {
     }
 
     let ast_sql = r#"SELECT * FROM SEMANTIC_VIEW(
-        my_model 
+        my_model
         DIMENSIONS DATE_PART('year', date_col), region_name
         METRICS orders.revenue, orders.count
         WHERE active = true