Repository: SeaQL/sea-query (language: rust)

Issue:
Unify SelectStatement and WithQuery
<!--

Welcome! Thanks for suggesting features!

Do you want to ask a question? Are you looking for support?
Please ask us on
- Discord: https://discord.com/invite/uCPdDXzbdv
- GitHub Discussions: https://github.com/SeaQL/sea-query/discussions/new

Make sure you have a clear feature specification before open an issue. Alternatively, please start an "Idea" thread on GitHub Discussions and let's formulate the solution together.

-->

## Motivation

WithQuery is currently a separate entity from SelectStatement which is unfortunate, because it restricts where you can use it, i.e. I'd like to use a WithQuery as a FROM with `from_subquery` in another query but I don't believe I can, because only SelectStatements are accepted.

https://github.com/SeaQL/sea-query/blob/53fd4e902910cf3db9f4332f1c1d1b6b6b3e49a6/src/query/select.rs#L1024

## Proposed Solutions

Move `with_clause` into `SelectStatement`/`UpdateStatement`/`DeleteStatement`/`InsertStatement` instead of having a separate code path for it.

## Additional Information

<!-- any other additional information that might be helpful -->



Discussion / hints:

Fix patch:
diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs
index 41cac862..089202ee 100644
--- a/src/backend/query_builder.rs
+++ b/src/backend/query_builder.rs
@@ -18,6 +18,10 @@ pub trait QueryBuilder:
 
     /// Translate [`InsertStatement`] into SQL statement.
     fn prepare_insert_statement(&self, insert: &InsertStatement, sql: &mut dyn SqlWriter) {
+        if let Some(with) = &insert.with {
+            self.prepare_with_clause(with, sql);
+        }
+
         self.prepare_insert(insert.replace, sql);
 
         if let Some(table) = &insert.table {
@@ -95,6 +99,10 @@ pub trait QueryBuilder:
 
     /// Translate [`SelectStatement`] into SQL statement.
     fn prepare_select_statement(&self, select: &SelectStatement, sql: &mut dyn SqlWriter) {
+        if let Some(with) = &select.with {
+            self.prepare_with_clause(with, sql);
+        }
+
         write!(sql, "SELECT ").unwrap();
 
         if let Some(distinct) = &select.distinct {
@@ -191,6 +199,10 @@ pub trait QueryBuilder:
 
     /// Translate [`UpdateStatement`] into SQL statement.
     fn prepare_update_statement(&self, update: &UpdateStatement, sql: &mut dyn SqlWriter) {
+        if let Some(with) = &update.with {
+            self.prepare_with_clause(with, sql);
+        }
+
         write!(sql, "UPDATE ").unwrap();
 
         if let Some(table) = &update.table {
@@ -245,6 +257,10 @@ pub trait QueryBuilder:
 
     /// Translate [`DeleteStatement`] into SQL statement.
     fn prepare_delete_statement(&self, delete: &DeleteStatement, sql: &mut dyn SqlWriter) {
+        if let Some(with) = &delete.with {
+            self.prepare_with_clause(with, sql);
+        }
+
         write!(sql, "DELETE ").unwrap();
 
         if let Some(table) = &delete.table {
diff --git a/src/query/delete.rs b/src/query/delete.rs
index 16b64b61..5fb727a8 100644
--- a/src/query/delete.rs
+++ b/src/query/delete.rs
@@ -44,6 +44,7 @@ pub struct DeleteStatement {
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
     pub(crate) returning: Option<ReturningClause>,
+    pub(crate) with: Option<WithClause>,
 }
 
 impl DeleteStatement {
@@ -226,6 +227,48 @@ impl DeleteStatement {
     pub fn with(self, clause: WithClause) -> WithQuery {
         clause.query(self)
     }
+
+    /// Create a Common Table Expression by specifying a [CommonTableExpression] or [WithClause] to execute this query with.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*};
+    ///
+    /// let select = SelectStatement::new()
+    ///         .columns([Glyph::Id])
+    ///         .from(Glyph::Table)
+    ///         .and_where(Expr::col(Glyph::Image).like("0%"))
+    ///         .to_owned();
+    ///     let cte = CommonTableExpression::new()
+    ///         .query(select)
+    ///         .column(Glyph::Id)
+    ///         .table_name(Alias::new("cte"))
+    ///         .to_owned();
+    ///     let with_clause = WithClause::new().cte(cte).to_owned();
+    ///     let query = DeleteStatement::new()
+    ///         .with_cte(with_clause)
+    ///         .from_table(Glyph::Table)
+    ///         .and_where(Expr::col(Glyph::Id).in_subquery(SelectStatement::new().column(Glyph::Id).from(Alias::new("cte")).to_owned()))
+    ///         .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"WITH `cte` (`id`) AS (SELECT `id` FROM `glyph` WHERE `image` LIKE '0%') DELETE FROM `glyph` WHERE `id` IN (SELECT `id` FROM `cte`)"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"WITH "cte" ("id") AS (SELECT "id" FROM "glyph" WHERE "image" LIKE '0%') DELETE FROM "glyph" WHERE "id" IN (SELECT "id" FROM "cte")"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"WITH "cte" ("id") AS (SELECT "id" FROM "glyph" WHERE "image" LIKE '0%') DELETE FROM "glyph" WHERE "id" IN (SELECT "id" FROM "cte")"#
+    /// );
+    /// ```
+    pub fn with_cte<C: Into<WithClause>>(&mut self, clause: C) -> &mut Self {
+        self.with = Some(clause.into());
+        self
+    }
 }
 
 #[inherent]
diff --git a/src/query/insert.rs b/src/query/insert.rs
index b82f5361..710516bb 100644
--- a/src/query/insert.rs
+++ b/src/query/insert.rs
@@ -51,6 +51,7 @@ pub struct InsertStatement {
     pub(crate) on_conflict: Option<OnConflict>,
     pub(crate) returning: Option<ReturningClause>,
     pub(crate) default_values: Option<u32>,
+    pub(crate) with: Option<WithClause>,
 }
 
 impl InsertStatement {
@@ -473,6 +474,55 @@ impl InsertStatement {
         clause.query(self)
     }
 
+    /// Create a Common Table Expression by specifying a [CommonTableExpression] or [WithClause] to execute this query with.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*};
+    ///
+    /// let select = SelectStatement::new()
+    ///         .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+    ///         .from(Glyph::Table)
+    ///         .to_owned();
+    ///     let cte = CommonTableExpression::new()
+    ///         .query(select)
+    ///         .column(Glyph::Id)
+    ///         .column(Glyph::Image)
+    ///         .column(Glyph::Aspect)
+    ///         .table_name(Alias::new("cte"))
+    ///         .to_owned();
+    ///     let with_clause = WithClause::new().cte(cte).to_owned();
+    ///     let select = SelectStatement::new()
+    ///         .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+    ///         .from(Alias::new("cte"))
+    ///         .to_owned();
+    ///     let mut query = Query::insert();
+    ///     query
+    ///         .with_cte(with_clause)
+    ///         .into_table(Glyph::Table)
+    ///         .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+    ///         .select_from(select)
+    ///         .unwrap();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"WITH `cte` (`id`, `image`, `aspect`) AS (SELECT `id`, `image`, `aspect` FROM `glyph`) INSERT INTO `glyph` (`id`, `image`, `aspect`) SELECT `id`, `image`, `aspect` FROM `cte`"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"WITH "cte" ("id", "image", "aspect") AS (SELECT "id", "image", "aspect" FROM "glyph") INSERT INTO "glyph" ("id", "image", "aspect") SELECT "id", "image", "aspect" FROM "cte""#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"WITH "cte" ("id", "image", "aspect") AS (SELECT "id", "image", "aspect" FROM "glyph") INSERT INTO "glyph" ("id", "image", "aspect") SELECT "id", "image", "aspect" FROM "cte""#
+    /// );
+    /// ```
+    pub fn with_cte<C: Into<WithClause>>(&mut self, clause: C) -> &mut Self {
+        self.with = Some(clause.into());
+        self
+    }
+
     /// Insert with default values if columns and values are not supplied.
     ///
     /// # Examples
diff --git a/src/query/select.rs b/src/query/select.rs
index debd0807..eee3c726 100644
--- a/src/query/select.rs
+++ b/src/query/select.rs
@@ -54,6 +54,7 @@ pub struct SelectStatement {
     pub(crate) offset: Option<Value>,
     pub(crate) lock: Option<LockClause>,
     pub(crate) window: Option<(DynIden, WindowStatement)>,
+    pub(crate) with: Option<WithClause>,
     #[cfg(feature = "backend-mysql")]
     pub(crate) index_hints: Vec<crate::extension::mysql::IndexHint>,
 }
@@ -162,6 +163,7 @@ impl SelectStatement {
             offset: self.offset.take(),
             lock: self.lock.take(),
             window: self.window.take(),
+            with: self.with.take(),
             #[cfg(feature = "backend-mysql")]
             index_hints: std::mem::take(&mut self.index_hints),
         }
@@ -2346,6 +2348,72 @@ impl SelectStatement {
         clause.query(self)
     }
 
+    /// Create a Common Table Expression by specifying a [CommonTableExpression] or [WithClause] to execute this query with.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*};
+    ///
+    /// let base_query = SelectStatement::new()
+    ///                     .column(Alias::new("id"))
+    ///                     .expr(1i32)
+    ///                     .column(Alias::new("next"))
+    ///                     .column(Alias::new("value"))
+    ///                     .from(Alias::new("table"))
+    ///                     .to_owned();
+    ///
+    /// let cte_referencing = SelectStatement::new()
+    ///                             .column(Alias::new("id"))
+    ///                             .expr(Expr::col(Alias::new("depth")).add(1i32))
+    ///                             .column(Alias::new("next"))
+    ///                             .column(Alias::new("value"))
+    ///                             .from(Alias::new("table"))
+    ///                             .join(
+    ///                                 JoinType::InnerJoin,
+    ///                                 Alias::new("cte_traversal"),
+    ///                                 Expr::col((Alias::new("cte_traversal"), Alias::new("next"))).equals((Alias::new("table"), Alias::new("id")))
+    ///                             )
+    ///                             .to_owned();
+    ///
+    /// let common_table_expression = CommonTableExpression::new()
+    ///             .query(
+    ///                 base_query.clone().union(UnionType::All, cte_referencing).to_owned()
+    ///             )
+    ///             .columns([Alias::new("id"), Alias::new("depth"), Alias::new("next"), Alias::new("value")])
+    ///             .table_name(Alias::new("cte_traversal"))
+    ///             .to_owned();
+    ///
+    /// let with_clause = WithClause::new()
+    ///         .recursive(true)
+    ///         .cte(common_table_expression)
+    ///         .cycle(Cycle::new_from_expr_set_using(SimpleExpr::Column(ColumnRef::Column(Alias::new("id").into_iden())), Alias::new("looped"), Alias::new("traversal_path")))
+    ///         .to_owned();
+    ///
+    /// let query = SelectStatement::new()
+    ///         .column(ColumnRef::Asterisk)
+    ///         .from(Alias::new("cte_traversal"))
+    ///         .with_cte(with_clause)
+    ///         .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"WITH RECURSIVE `cte_traversal` (`id`, `depth`, `next`, `value`) AS (SELECT `id`, 1, `next`, `value` FROM `table` UNION ALL (SELECT `id`, `depth` + 1, `next`, `value` FROM `table` INNER JOIN `cte_traversal` ON `cte_traversal`.`next` = `table`.`id`)) SELECT * FROM `cte_traversal`"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"WITH RECURSIVE "cte_traversal" ("id", "depth", "next", "value") AS (SELECT "id", 1, "next", "value" FROM "table" UNION ALL (SELECT "id", "depth" + 1, "next", "value" FROM "table" INNER JOIN "cte_traversal" ON "cte_traversal"."next" = "table"."id")) CYCLE "id" SET "looped" USING "traversal_path" SELECT * FROM "cte_traversal""#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"WITH RECURSIVE "cte_traversal" ("id", "depth", "next", "value") AS (SELECT "id", 1, "next", "value" FROM "table" UNION ALL SELECT "id", "depth" + 1, "next", "value" FROM "table" INNER JOIN "cte_traversal" ON "cte_traversal"."next" = "table"."id") SELECT * FROM "cte_traversal""#
+    /// );
+    /// ```
+    pub fn with_cte<C: Into<WithClause>>(&mut self, clause: C) -> &mut Self {
+        self.with = Some(clause.into());
+        self
+    }
+
     /// WINDOW
     ///
     /// # Examples:
diff --git a/src/query/update.rs b/src/query/update.rs
index a514127b..1b881569 100644
--- a/src/query/update.rs
+++ b/src/query/update.rs
@@ -44,6 +44,7 @@ pub struct UpdateStatement {
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
     pub(crate) returning: Option<ReturningClause>,
+    pub(crate) with: Option<WithClause>,
 }
 
 impl UpdateStatement {
@@ -307,6 +308,49 @@ impl UpdateStatement {
         clause.query(self)
     }
 
+    /// Create a Common Table Expression by specifying a [CommonTableExpression] or [WithClause] to execute this query with.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*};
+    ///
+    /// let select = SelectStatement::new()
+    ///         .columns([Glyph::Id])
+    ///         .from(Glyph::Table)
+    ///         .and_where(Expr::col(Glyph::Image).like("0%"))
+    ///         .to_owned();
+    ///     let cte = CommonTableExpression::new()
+    ///         .query(select)
+    ///         .column(Glyph::Id)
+    ///         .table_name(Alias::new("cte"))
+    ///         .to_owned();
+    ///     let with_clause = WithClause::new().cte(cte).to_owned();
+    ///     let query = UpdateStatement::new()
+    ///         .table(Glyph::Table)
+    ///         .and_where(Expr::col(Glyph::Id).in_subquery(SelectStatement::new().column(Glyph::Id).from(Alias::new("cte")).to_owned()))
+    ///         .value(Glyph::Aspect, Expr::cust("60 * 24 * 24"))
+    ///         .with_cte(with_clause)
+    ///         .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"WITH `cte` (`id`) AS (SELECT `id` FROM `glyph` WHERE `image` LIKE '0%') UPDATE `glyph` SET `aspect` = 60 * 24 * 24 WHERE `id` IN (SELECT `id` FROM `cte`)"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"WITH "cte" ("id") AS (SELECT "id" FROM "glyph" WHERE "image" LIKE '0%') UPDATE "glyph" SET "aspect" = 60 * 24 * 24 WHERE "id" IN (SELECT "id" FROM "cte")"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"WITH "cte" ("id") AS (SELECT "id" FROM "glyph" WHERE "image" LIKE '0%') UPDATE "glyph" SET "aspect" = 60 * 24 * 24 WHERE "id" IN (SELECT "id" FROM "cte")"#
+    /// );
+    /// ```
+    pub fn with_cte<C: Into<WithClause>>(&mut self, clause: C) -> &mut Self {
+        self.with = Some(clause.into());
+        self
+    }
+
     /// Get column values
     pub fn get_values(&self) -> &[(DynIden, Box<SimpleExpr>)] {
         &self.values
diff --git a/src/query/with.rs b/src/query/with.rs
index 94933f91..1157b5c8 100644
--- a/src/query/with.rs
+++ b/src/query/with.rs
@@ -485,6 +485,13 @@ impl WithClause {
         WithQuery::new().with_clause(self).query(query).to_owned()
     }
 }
+
+impl From<CommonTableExpression> for WithClause {
+    fn from(cte: CommonTableExpression) -> WithClause {
+        WithClause::new().cte(cte).to_owned()
+    }
+}
+
 /// A WITH query. A simple SQL query that has a WITH clause ([WithClause]).
 ///
 /// The [WithClause] can contain one or multiple common table expressions ([CommonTableExpression]).


Test patch:
diff --git a/tests/postgres/query.rs b/tests/postgres/query.rs
index 804818b1..1f90e9b9 100644
--- a/tests/postgres/query.rs
+++ b/tests/postgres/query.rs
@@ -1092,6 +1092,33 @@ fn select_62() {
     );
 }
 
+#[test]
+fn select_63() {
+    let select = SelectStatement::new()
+        .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+        .from(Glyph::Table)
+        .to_owned();
+    let cte = CommonTableExpression::new()
+        .query(select)
+        .table_name(Alias::new("cte"))
+        .to_owned();
+    let select = SelectStatement::new()
+        .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+        .from(Alias::new("cte"))
+        .with_cte(cte)
+        .to_owned();
+    assert_eq!(
+        select.to_string(PostgresQueryBuilder),
+        [
+            r#"WITH "cte" AS"#,
+            r#"(SELECT "id", "image", "aspect""#,
+            r#"FROM "glyph")"#,
+            r#"SELECT "id", "image", "aspect" FROM "cte""#,
+        ]
+        .join(" ")
+    );
+}
+
 #[test]
 #[allow(clippy::approx_constant)]
 fn insert_2() {
@@ -1275,6 +1302,40 @@ fn insert_10() {
     );
 }
 
+#[test]
+fn insert_11() -> error::Result<()> {
+    let select = SelectStatement::new()
+        .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+        .from(Glyph::Table)
+        .to_owned();
+    let cte = CommonTableExpression::new()
+        .query(select)
+        .column(Glyph::Id)
+        .column(Glyph::Image)
+        .column(Glyph::Aspect)
+        .table_name(Alias::new("cte"))
+        .to_owned();
+    let select = SelectStatement::new()
+        .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+        .from(Alias::new("cte"))
+        .to_owned();
+    let mut insert = Query::insert();
+    insert
+        .into_table(Glyph::Table)
+        .columns([Glyph::Id, Glyph::Image, Glyph::Aspect])
+        .with_cte(cte)
+        .select_from(select)?;
+    let sql = insert.to_string(PostgresQueryBuilder);
+    assert_eq!(
+        sql.as_str(),
+        [
+            r#"WITH "cte" ("id", "image", "aspect") AS (SELECT "id", "image", "aspect" FROM "glyph")"#,
+            r#"INSERT INTO "glyph" ("id", "image", "aspect") SELECT "id", "image", "aspect" FROM "cte""#,
+        ].join(" ")
+    );
+    Ok(())
+}
+
 #[test]
 #[allow(clippy::approx_constant)]
 fn insert_on_conflict_1() {