With clause and with queries (SeaQL#239)

* With clause and with queries

* Testing INSERT with CTE

* Allow arbitrary subqueries

src/backend/mysql/query.rs

@@ -30,4 +30,30 @@ impl QueryBuilder for MysqlQueryBuilder {
         write!(sql, " ").unwrap();
         self.prepare_order(&order_expr.order, sql, collector);
     }
+
+    fn prepare_query_statement(
+        &self,
+        query: &dyn QueryStatementBuilder,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        query.build_collect_any_into(self, sql, collector);
+    }
+
+    fn prepare_with_clause_recursive_options(
+        &self,
+        _: &WithClause,
+        _: &mut SqlWriter,
+        _: &mut dyn FnMut(Value),
+    ) {
+        // MySQL doesn't support sql recursive with query 'SEARCH' and 'CYCLE' options.
+    }
+
+    fn prepare_with_query_clause_materialization(
+        &self,
+        _: &CommonTableExpression,
+        _: &mut SqlWriter,
+    ) {
+        // MySQL doesn't support declaring materialization in SQL for with query.
+    }
 }

src/backend/postgres/query.rs

@@ -90,4 +90,13 @@ impl QueryBuilder for PostgresQueryBuilder {
             Some(NullOrdering::First) => write!(sql, " NULLS FIRST").unwrap(),
         }
     }
+
+    fn prepare_query_statement(
+        &self,
+        query: &dyn QueryStatementBuilder,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        query.build_collect_any_into(self, sql, collector);
+    }
 }

src/backend/query_builder.rs

@@ -309,7 +309,7 @@ pub trait QueryBuilder: QuotedBuilder {
             }
             SimpleExpr::SubQuery(sel) => {
                 write!(sql, "(").unwrap();
-                self.prepare_select_statement(sel, sql, collector);
+                self.prepare_query_statement(sel.deref(), sql, collector);
                 write!(sql, ")").unwrap();
             }
             SimpleExpr::Value(val) => {
@@ -635,6 +635,173 @@ pub trait QueryBuilder: QuotedBuilder {
         }
     }
 
+    /// Translate [`QueryStatement`] into SQL statement.
+    fn prepare_query_statement(
+        &self,
+        query: &dyn QueryStatementBuilder,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    );
+
+    fn prepare_with_query(
+        &self,
+        query: &WithQuery,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        self.prepare_with_clause(&query.with_clause, sql, collector);
+        self.prepare_query_statement(query.query.as_ref().unwrap().deref(), sql, collector);
+    }
+
+    fn prepare_with_clause(
+        &self,
+        with_clause: &WithClause,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        self.prepare_with_clause_start(with_clause, sql);
+        self.prepare_with_clause_common_tables(with_clause, sql, collector);
+        if with_clause.recursive {
+            self.prepare_with_clause_recursive_options(with_clause, sql, collector);
+        }
+    }
+
+    fn prepare_with_clause_recursive_options(
+        &self,
+        with_clause: &WithClause,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        if with_clause.recursive {
+            if let Some(search) = &with_clause.search {
+                write!(
+                    sql,
+                    "SEARCH {} FIRST BY ",
+                    match &search.order.as_ref().unwrap() {
+                        SearchOrder::BREADTH => "BREADTH",
+                        SearchOrder::DEPTH => "DEPTH",
+                    }
+                )
+                .unwrap();
+
+                self.prepare_simple_expr(&search.expr.as_ref().unwrap().expr, sql, collector);
+
+                write!(sql, " SET ").unwrap();
+
+                search
+                    .expr
+                    .as_ref()
+                    .unwrap()
+                    .alias
+                    .as_ref()
+                    .unwrap()
+                    .prepare(sql, self.quote());
+                write!(sql, " ").unwrap();
+            }
+            if let Some(cycle) = &with_clause.cycle {
+                write!(sql, "CYCLE ").unwrap();
+
+                self.prepare_simple_expr(cycle.expr.as_ref().unwrap(), sql, collector);
+
+                write!(sql, " SET ").unwrap();
+
+                cycle.set_as.as_ref().unwrap().prepare(sql, self.quote());
+                write!(sql, " USING ").unwrap();
+                cycle.using.as_ref().unwrap().prepare(sql, self.quote());
+                write!(sql, " ").unwrap();
+            }
+        }
+    }
+
+    fn prepare_with_clause_common_tables(
+        &self,
+        with_clause: &WithClause,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        let mut cte_first = true;
+        assert_ne!(
+            with_clause.cte_expressions.len(),
+            0,
+            "Cannot build a with query that has no common table expression!"
+        );
+
+        if with_clause.recursive {
+            assert_eq!(
+                with_clause.cte_expressions.len(),
+                1,
+                "Cannot build a recursive query with more than one common table! \
+                A recursive with query must have a single cte inside it that has a union query of \
+                two queries!"
+            );
+        }
+        for cte in &with_clause.cte_expressions {
+            if !cte_first {
+                write!(sql, ", ").unwrap();
+            }
+            cte_first = false;
+
+            self.prepare_with_query_clause_common_table(cte, sql, collector);
+        }
+    }
+
+    fn prepare_with_query_clause_common_table(
+        &self,
+        cte: &CommonTableExpression,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        cte.table_name.as_ref().unwrap().prepare(sql, self.quote());
+
+        if !cte.cols.is_empty() {
+            write!(sql, " (").unwrap();
+
+            let mut col_first = true;
+            for col in &cte.cols {
+                if !col_first {
+                    write!(sql, ", ").unwrap();
+                }
+                col_first = false;
+                col.prepare(sql, self.quote());
+            }
+
+            write!(sql, ") ").unwrap();
+        }
+
+        write!(sql, "AS ").unwrap();
+
+        self.prepare_with_query_clause_materialization(cte, sql);
+
+        write!(sql, "(").unwrap();
+
+        self.prepare_query_statement(cte.query.as_ref().unwrap().deref(), sql, collector);
+
+        write!(sql, ") ").unwrap();
+    }
+
+    fn prepare_with_query_clause_materialization(
+        &self,
+        cte: &CommonTableExpression,
+        sql: &mut SqlWriter,
+    ) {
+        if let Some(materialized) = cte.materialized {
+            write!(
+                sql,
+                "{} MATERIALIZED ",
+                if materialized { "" } else { "NOT" }
+            )
+            .unwrap()
+        }
+    }
+
+    fn prepare_with_clause_start(&self, with_clause: &WithClause, sql: &mut SqlWriter) {
+        write!(sql, "WITH ").unwrap();
+
+        if with_clause.recursive {
+            write!(sql, "RECURSIVE ").unwrap();
+        }
+    }
+
     fn prepare_function(
         &self,
         function: &Function,
@@ -995,7 +1162,16 @@ pub trait QueryBuilder: QuotedBuilder {
 
 pub(crate) struct CommonSqlQueryBuilder;
 
-impl QueryBuilder for CommonSqlQueryBuilder {}
+impl QueryBuilder for CommonSqlQueryBuilder {
+    fn prepare_query_statement(
+        &self,
+        query: &dyn QueryStatementBuilder,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        query.build_collect_any_into(self, sql, collector);
+    }
+}
 
 impl QuotedBuilder for CommonSqlQueryBuilder {
     fn quote(&self) -> char {

src/backend/sqlite/query.rs

@@ -29,4 +29,22 @@ impl QueryBuilder for SqliteQueryBuilder {
             Some(NullOrdering::First) => write!(sql, " NULLS FIRST").unwrap(),
         }
     }
+
+    fn prepare_query_statement(
+        &self,
+        query: &dyn QueryStatementBuilder,
+        sql: &mut SqlWriter,
+        collector: &mut dyn FnMut(Value),
+    ) {
+        query.build_collect_any_into(self, sql, collector);
+    }
+
+    fn prepare_with_clause_recursive_options(
+        &self,
+        _: &WithClause,
+        _: &mut SqlWriter,
+        _: &mut dyn FnMut(Value),
+    ) {
+        // Sqlite doesn't support sql recursive with query 'SEARCH' and 'CYCLE' options.
+    }
 }

src/expr.rs

@@ -28,7 +28,7 @@ pub enum SimpleExpr {
     Unary(UnOper, Box<SimpleExpr>),
     FunctionCall(Function, Vec<SimpleExpr>),
     Binary(Box<SimpleExpr>, BinOper, Box<SimpleExpr>),
-    SubQuery(Box<SelectStatement>),
+    SubQuery(Box<dyn QueryStatementBuilder>),
     Value(Value),
     Values(Vec<Value>),
     Custom(String),

src/prepare.rs

@@ -50,8 +50,12 @@ where
 
 impl SqlWriter {
     pub fn new() -> Self {
+        Self::new_params_from(0)
+    }
+
+    pub fn new_params_from(start_params_from: usize) -> Self {
         Self {
-            counter: 0,
+            counter: start_params_from,
             string: String::with_capacity(256),
         }
     }

src/query/delete.rs

@@ -4,7 +4,8 @@ use crate::{
     query::{condition::*, OrderedStatement},
     types::*,
     value::*,
-    Query, QueryStatementBuilder, SelectExpr, SelectStatement,
+    Query, QueryStatementBuilder, QueryStatementBuilderGenerics, SelectExpr, SelectStatement,
+    WithClause, WithQuery,
 };
 
 /// Delete existing rows from the table
@@ -174,9 +175,60 @@ impl DeleteStatement {
     {
         self.returning(Query::select().column(col.into_iden()).take())
     }
+
+    /// Create a [WithQuery] by specifying a [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 update = DeleteStatement::new()
+    ///         .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();
+    ///     let query = update.with(with_clause);
+    ///
+    /// 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(self, clause: WithClause) -> WithQuery {
+        clause.query(self)
+    }
 }
 
 impl QueryStatementBuilder for DeleteStatement {
+    fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)) {
+        query_builder.prepare_delete_statement(self, sql, collector);
+    }
+
+    fn box_clone(&self) -> Box<dyn QueryStatementBuilder> {
+        Box::new(self.clone())
+    }
+}
+
+impl QueryStatementBuilderGenerics for DeleteStatement {
     /// Build corresponding SQL statement for certain database backend and collect query parameters
     ///
     /// # Examples
@@ -212,16 +264,6 @@ impl QueryStatementBuilder for DeleteStatement {
         query_builder.prepare_delete_statement(self, &mut sql, collector);
         sql.result()
     }
-
-    fn build_collect_any(
-        &self,
-        query_builder: &dyn QueryBuilder,
-        collector: &mut dyn FnMut(Value),
-    ) -> String {
-        let mut sql = SqlWriter::new();
-        query_builder.prepare_delete_statement(self, &mut sql, collector);
-        sql.result()
-    }
 }
 
 impl OrderedStatement for DeleteStatement {