From 3f4e5ab5a842acea01276c03cfb64815002f3f51 Mon Sep 17 00:00:00 2001 From: 05storm26 <05storm26@users.noreply.github.com> Date: Sun, 23 Jan 2022 15:45:55 +0100 Subject: [PATCH] With clause and with queries (#239) * With clause and with queries * Testing INSERT with CTE * Allow arbitrary subqueries --- src/backend/mysql/query.rs | 26 ++ src/backend/postgres/query.rs | 9 + src/backend/query_builder.rs | 180 +++++++++- src/backend/sqlite/query.rs | 18 + src/expr.rs | 2 +- src/prepare.rs | 6 +- src/query/delete.rs | 64 +++- src/query/insert.rs | 71 +++- src/query/mod.rs | 7 + src/query/select.rs | 92 +++++- src/query/traits.rs | 57 ++-- src/query/update.rs | 65 +++- src/query/with.rs | 601 ++++++++++++++++++++++++++++++++++ src/shim.rs | 9 +- tests/postgres/query.rs | 34 ++ 15 files changed, 1169 insertions(+), 72 deletions(-) create mode 100644 src/query/with.rs diff --git a/src/backend/mysql/query.rs b/src/backend/mysql/query.rs index bdbc0d4f69..6841e3f988 100644 --- a/src/backend/mysql/query.rs +++ b/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. + } } diff --git a/src/backend/postgres/query.rs b/src/backend/postgres/query.rs index 34d6598b48..a1fb8db14a 100644 --- a/src/backend/postgres/query.rs +++ b/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); + } } diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs index 7389293b20..b9baa6a4f5 100644 --- a/src/backend/query_builder.rs +++ b/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 { diff --git a/src/backend/sqlite/query.rs b/src/backend/sqlite/query.rs index a9088c68d8..e22b276be2 100644 --- a/src/backend/sqlite/query.rs +++ b/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. + } } diff --git a/src/expr.rs b/src/expr.rs index c6792d7f43..2991e49229 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -28,7 +28,7 @@ pub enum SimpleExpr { Unary(UnOper, Box), FunctionCall(Function, Vec), Binary(Box, BinOper, Box), - SubQuery(Box), + SubQuery(Box), Value(Value), Values(Vec), Custom(String), diff --git a/src/prepare.rs b/src/prepare.rs index de8b651a72..64a6ad1ae0 100644 --- a/src/prepare.rs +++ b/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), } } diff --git a/src/query/delete.rs b/src/query/delete.rs index 06bb3a16a1..acb7ac31ba 100644 --- a/src/query/delete.rs +++ b/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 { + 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 { diff --git a/src/query/insert.rs b/src/query/insert.rs index 1de694aa64..4e8bf8994f 100644 --- a/src/query/insert.rs +++ b/src/query/insert.rs @@ -1,6 +1,7 @@ use crate::{ backend::QueryBuilder, error::*, prepare::*, types::*, value::*, Expr, Query, - QueryStatementBuilder, SelectExpr, SelectStatement, SimpleExpr, + QueryStatementBuilder, QueryStatementBuilderGenerics, SelectExpr, SelectStatement, SimpleExpr, + WithClause, WithQuery, }; /// Represents a value source that can be used in an insert query. @@ -352,9 +353,67 @@ impl InsertStatement { { 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, 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 insert = Query::insert(); + /// insert + /// .into_table(Glyph::Table) + /// .columns([Glyph::Id, Glyph::Image, Glyph::Aspect]) + /// .select_from(select) + /// .unwrap(); + /// let query = insert.with(with_clause); + /// + /// 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(self, clause: WithClause) -> WithQuery { + clause.query(self) + } } impl QueryStatementBuilder for InsertStatement { + fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)) { + query_builder.prepare_insert_statement(self, sql, collector); + } + + fn box_clone(&self) -> Box { + Box::new(self.clone()) + } +} + +impl QueryStatementBuilderGenerics for InsertStatement { /// Build corresponding SQL statement for certain database backend and collect query parameters /// /// # Examples @@ -397,14 +456,4 @@ impl QueryStatementBuilder for InsertStatement { query_builder.prepare_insert_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_insert_statement(self, &mut sql, collector); - sql.result() - } } diff --git a/src/query/mod.rs b/src/query/mod.rs index c30ccc77aa..19d9b94a6f 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -15,6 +15,7 @@ mod select; mod shim; mod traits; mod update; +mod with; pub use condition::*; pub use delete::*; @@ -23,6 +24,7 @@ pub use ordered::*; pub use select::*; pub use traits::*; pub use update::*; +pub use with::*; /// Shorthand for constructing any table query #[derive(Debug, Clone)] @@ -57,4 +59,9 @@ impl Query { pub fn delete() -> DeleteStatement { DeleteStatement::new() } + + /// Construct [`WithClause`] + pub fn with() -> WithClause { + WithClause::new() + } } diff --git a/src/query/select.rs b/src/query/select.rs index 52d52bc55b..06764c6939 100644 --- a/src/query/select.rs +++ b/src/query/select.rs @@ -5,7 +5,7 @@ use crate::{ query::{condition::*, OrderedStatement}, types::*, value::*, - QueryStatementBuilder, + QueryStatementBuilder, QueryStatementBuilderGenerics, WithClause, WithQuery, }; /// Select rows from an existing table @@ -1672,9 +1672,89 @@ impl SelectStatement { self.unions.extend(unions); self } + + /// Create a [WithQuery] by specifying a [WithClause] to execute this query with. + /// + /// # Examples + /// + /// ``` + /// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*}; + /// + /// let base_query = SelectStatement::new() + /// .column(Alias::new("id")) + /// .expr(Expr::val(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::tbl(Alias::new("cte_traversal"), Alias::new("next")).equals(Alias::new("table"), Alias::new("id")).into_condition() + /// ) + /// .to_owned(); + /// + /// let common_table_expression = CommonTableExpression::new() + /// .query( + /// base_query.clone().union(UnionType::All, cte_referencing).to_owned() + /// ) + /// .column(Alias::new("id")) + /// .column(Alias::new("depth")) + /// .column(Alias::new("next")) + /// .column(Alias::new("value")) + /// .table_name(Alias::new("cte_traversal")) + /// .to_owned(); + /// + /// let select = SelectStatement::new() + /// .column(ColumnRef::Asterisk) + /// .from(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 = select.with(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(self, clause: WithClause) -> WithQuery { + clause.query(self) + } + } impl QueryStatementBuilder for SelectStatement { + fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)) { + query_builder.prepare_select_statement(self, sql, collector); + } + + fn box_clone(&self) -> Box { + Box::new(self.clone()) + } +} + +impl QueryStatementBuilderGenerics for SelectStatement { /// Build corresponding SQL statement for certain database backend and collect query parameters /// /// # Examples @@ -1716,16 +1796,6 @@ impl QueryStatementBuilder for SelectStatement { query_builder.prepare_select_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_select_statement(self, &mut sql, collector); - sql.result() - } } impl OrderedStatement for SelectStatement { diff --git a/src/query/traits.rs b/src/query/traits.rs index 34fce8c1f7..9ae525663f 100644 --- a/src/query/traits.rs +++ b/src/query/traits.rs @@ -1,10 +1,40 @@ -use crate::{ - backend::QueryBuilder, - prepare::inject_parameters, - value::{Value, Values}, -}; +use crate::{backend::QueryBuilder, prepare::inject_parameters, SqlWriter, value::{Value, Values}}; +use std::any::Any; +use std::fmt::Debug; -pub trait QueryStatementBuilder { +pub trait QueryStatementBuilder: Any + Debug { + /// Build corresponding SQL statement for certain database backend and collect query parameters into a vector + fn build_any(&self, query_builder: &dyn QueryBuilder) -> (String, Values) { + let mut values = Vec::new(); + let mut collector = |v| values.push(v); + let sql = self.build_collect_any(query_builder, &mut collector); + (sql, Values(values)) + } + + /// Build corresponding SQL statement for certain database backend and collect query parameters + fn build_collect_any( + &self, + query_builder: &dyn QueryBuilder, + collector: &mut dyn FnMut(Value), + ) -> String { + let mut sql = SqlWriter::new(); + self.build_collect_any_into(query_builder, &mut sql, collector); + sql.result() + } + + /// Build corresponding SQL statement into the SqlWriter for certain database backend and collect query parameters + fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)); + + fn box_clone(&self) -> Box; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.box_clone() + } +} + +pub trait QueryStatementBuilderGenerics: QueryStatementBuilder { /// Build corresponding SQL statement for certain database backend and return SQL string /// /// # Examples @@ -61,14 +91,6 @@ pub trait QueryStatementBuilder { (sql, Values(values)) } - /// Build corresponding SQL statement for certain database backend and collect query parameters into a vector - fn build_any(&self, query_builder: &dyn QueryBuilder) -> (String, Values) { - let mut values = Vec::new(); - let mut collector = |v| values.push(v); - let sql = self.build_collect_any(query_builder, &mut collector); - (sql, Values(values)) - } - /// Build corresponding SQL statement for certain database backend and collect query parameters /// /// # Examples @@ -106,11 +128,4 @@ pub trait QueryStatementBuilder { query_builder: T, collector: &mut dyn FnMut(Value), ) -> String; - - /// Build corresponding SQL statement for certain database backend and collect query parameters - fn build_collect_any( - &self, - query_builder: &dyn QueryBuilder, - collector: &mut dyn FnMut(Value), - ) -> String; } diff --git a/src/query/update.rs b/src/query/update.rs index f089d45493..34666a4ba9 100644 --- a/src/query/update.rs +++ b/src/query/update.rs @@ -5,7 +5,8 @@ use crate::{ query::{condition::*, OrderedStatement}, types::*, value::*, - Query, QueryStatementBuilder, SelectExpr, SelectStatement, + Query, QueryStatementBuilder, QueryStatementBuilderGenerics, SelectExpr, SelectStatement, + WithClause, WithQuery, }; /// Update existing rows in the table @@ -302,9 +303,61 @@ impl UpdateStatement { { 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 = UpdateStatement::new() + /// .table(Glyph::Table) + /// .and_where(Expr::col(Glyph::Id).in_subquery(SelectStatement::new().column(Glyph::Id).from(Alias::new("cte")).to_owned())) + /// .col_expr(Glyph::Aspect, Expr::cust("60 * 24 * 24")) + /// .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%') 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(self, clause: WithClause) -> WithQuery { + clause.query(self) + } } impl QueryStatementBuilder for UpdateStatement { + fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)) { + query_builder.prepare_update_statement(self, sql, collector); + } + + fn box_clone(&self) -> Box { + Box::new(self.clone()) + } +} + +impl QueryStatementBuilderGenerics for UpdateStatement { /// Build corresponding SQL statement for certain database backend and collect query parameters /// /// # Examples @@ -351,16 +404,6 @@ impl QueryStatementBuilder for UpdateStatement { query_builder.prepare_update_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_update_statement(self, &mut sql, collector); - sql.result() - } } impl OrderedStatement for UpdateStatement { diff --git a/src/query/with.rs b/src/query/with.rs new file mode 100644 index 0000000000..953011b6fd --- /dev/null +++ b/src/query/with.rs @@ -0,0 +1,601 @@ +use crate::{Alias, QueryBuilder}; +use crate::ColumnRef; +use crate::DynIden; +use crate::IntoIden; +use crate::QueryStatementBuilder; +use crate::QueryStatementBuilderGenerics; +use crate::SelectExpr; +use crate::SelectStatement; +use crate::SimpleExpr; +use crate::SqlWriter; +use crate::TableRef; +use crate::Value; +use std::ops::Deref; + +/// A table definition inside a WITH clause ([WithClause]). +/// +/// A WITH clause can contain one or multiple common table expressions ([CommonTableExpression]). +/// +/// These named queries can act as a "query local table" that are materialized during execution and +/// then can be used by the query prefixed with the WITH clause. +/// +/// A WITH clause can contain multiple of these [CommonTableExpression]. (Except in the case of +/// recursive WITH query which can only contain one [CommonTableExpression]). +/// +/// A [CommonTableExpression] is a name, column names and a query returning data for those columns. +/// +/// Some databases (like sqlite) restrict the acceptable kinds of queries inside of the WITH clause +/// common table expressions. These databases only allow [SelectStatement]s to form a common table +/// expression. +/// +/// Other databases like postgres allow modification queries (UPDATE, DELETE) inside of the WITH +/// clause but they have to return a table. (They must have a RETURNING clause). +/// +/// sea-query doesn't check this or restrict the kind of [CommonTableExpression] that you can create +/// in rust. This means that you can put an UPDATE or DELETE queries into WITH clause and sea-query +/// will succeed in generating that kind of sql query but the execution inside the database will +/// fail because they are invalid. +/// +/// It is your responsibility to ensure that the kind of WITH clause that you put together makes +/// sense and valid for that database that you are using. +/// +/// NOTE that for recursive WITH queries (in sql: "WITH RECURSIVE") you can only have a +/// single [CommonTableExpression] inside of the WITH clause. That query must match certain +/// requirements: +/// * It is a query of UNION or UNION ALL of two queries. +/// * The first part of the query (the left side of the UNION) must be executable first in itself. +/// It must be non-recursive. (Cannot contain self reference) +/// * The self reference must appear in the right hand side of the UNION. +/// * The query can only have a single self-reference. +/// * Recursive data-modifying statements are not supported, but you can use the results of a +/// recursive SELECT query in a data-modifying statement. (like so: WITH RECURSIVE +/// cte_name(a,b,c,d) AS (SELECT ... UNION SELECT ... FROM ... JOIN cte_name ON ... WHERE ...) +/// DELETE FROM table WHERE table.a = cte_name.a) +/// +/// It is mandatory to set the [Self::table_name] and the [Self::query]. +#[derive(Debug, Clone, Default)] +pub struct CommonTableExpression { + pub(crate) table_name: Option, + pub(crate) cols: Vec, + pub(crate) query: Option>, + pub(crate) materialized: Option, +} + +impl CommonTableExpression { + /// Construct a new [`CommonTableExpression`] + pub fn new() -> CommonTableExpression { + Self::default() + } + + /// Sets the CTE table name of the query. + pub fn table_name(&mut self, table_name: T) -> &mut Self + where + T: IntoIden, + { + self.table_name = Some(table_name.into_iden()); + self + } + + /// Adds a named column to the CTE table definition. + pub fn column(&mut self, col: C) -> &mut Self + where + C: IntoIden, + { + self.cols.push(col.into_iden()); + self + } + + /// Some databases allow you to put "MATERIALIZED" or "NOT MATERIALIZED" in the CTE definition. + /// This will affect how during the execution of [WithQuery] the CTE in the [WithClause] will be + /// executed. If the database doesn't support this syntax this option specified here will be + /// ignored and not appear in the generated sql. + pub fn materialized(&mut self, materialized: bool) -> &mut Self { + self.materialized = Some(materialized); + self + } + + /// Set the query generating the CTE content. The query's result must match the defined + /// columns. + pub fn query(&mut self, query: Q) -> &mut Self + where + Q: QueryStatementBuilder, + { + self.query = Some(Box::new(query)); + self + } + + /// Create a CTE from a [SelectStatement] if the selections are named columns then this will + /// return a [CommonTableExpression] that has the column names set. The [Self::table_name] is + /// not set. + pub fn from_select(select: SelectStatement) -> Self { + let mut cte = Self::default(); + cte.try_set_cols_from_selects(&select.selects); + if let Some(from) = &select.from { + match from.deref() { + TableRef::Table(iden) => cte.set_table_name_from_select(iden), + TableRef::SchemaTable(_, iden) => cte.set_table_name_from_select(iden), + TableRef::DatabaseSchemaTable(_, _, iden) => cte.set_table_name_from_select(iden), + TableRef::TableAlias(_, iden) => cte.set_table_name_from_select(iden), + TableRef::SchemaTableAlias(_, _, iden) => cte.set_table_name_from_select(iden), + TableRef::DatabaseSchemaTableAlias(_, _, _, iden) => { + cte.set_table_name_from_select(iden) + } + _ => {} + } + } + cte.query = Some(Box::new(select)); + cte + } + + fn set_table_name_from_select(&mut self, iden: &DynIden) { + self.table_name = Some(Alias::new(&format!("cte_{}", iden.to_string())).into_iden()) + } + + /// Set up the columns of the CTE to match the given [SelectStatement] selected columns. + /// This will fail if the select contains non named columns like expressions of wildcards. + /// + /// Returns true if the column setup from the select query was successful. If the returned + /// value is false the columns are untouched. + pub fn try_set_cols_from_select(&mut self, select: &SelectStatement) -> bool { + self.try_set_cols_from_selects(&select.selects) + } + + fn try_set_cols_from_selects(&mut self, selects: &[SelectExpr]) -> bool { + let vec: Option> = selects + .iter() + .map(|select| { + if let Some(ident) = &select.alias { + Some(ident.clone()) + } else { + match &select.expr { + SimpleExpr::Column(column) => match column { + ColumnRef::Column(iden) => Some(iden.clone()), + ColumnRef::TableColumn(table, column) => Some( + Alias::new(&format!( + "{}_{}", + table.to_string(), + column.to_string() + )) + .into_iden(), + ), + ColumnRef::SchemaTableColumn(schema, table, column) => Some( + Alias::new(&format!( + "{}_{}_{}", + schema.to_string(), + table.to_string(), + column.to_string() + )) + .into_iden(), + ), + _ => None, + }, + _ => None, + } + } + }) + .collect(); + + if let Some(c) = vec { + self.cols = c; + return true; + } + + false + } +} + +/// For recursive [WithQuery] [WithClause]s the traversing order can be specified in some databases +/// that support this functionality. +#[derive(Debug, Clone)] +pub enum SearchOrder { + /// Breadth first traversal during the execution of the recursive query. + BREADTH, + /// Depth first traversal during the execution of the recursive query. + DEPTH, +} + +/// For recursive [WithQuery] [WithClause]s the traversing order can be specified in some databases +/// that support this functionality. +/// +/// The clause contains the type of traversal: [SearchOrder] and the expression that is used to +/// construct the current path. +/// +/// A query can have both SEARCH and CYCLE clauses. +/// +/// Setting [Self::order] and [Self::expr] is mandatory. The [SelectExpr] used must specify an alias +/// which will be the name that you can use to order the result of the [CommonTableExpression]. +#[derive(Debug, Clone, Default)] +pub struct Search { + pub(crate) order: Option, + pub(crate) expr: Option, +} + +impl Search { + /// Create a complete [Search] specification from the [SearchOrder] and a [SelectExpr]. The + /// given [SelectExpr] must have an alias specified. + pub fn new_from_order_and_expr(order: SearchOrder, expr: EXPR) -> Self + where + EXPR: Into, + { + let expr = expr.into(); + expr.alias.as_ref().unwrap(); + Self { + order: Some(order), + expr: Some(expr), + } + } + + /// Constructs a new empty [Search]. + pub fn new() -> Self { + Self::default() + } + + /// The traversal order to be used. + pub fn order(&mut self, order: SearchOrder) -> &mut Self { + self.order = Some(order); + self + } + + /// The given [SelectExpr] must have an alias specified. + /// + /// The actual expression will be the one used to track the path in the graph. + /// + /// The alias of the given [SelectExpr] will be the name of the order column generated by this + /// clause. + pub fn expr(&mut self, expr: EXPR) -> &mut Self + where + EXPR: Into, + { + let expr = expr.into(); + expr.alias.as_ref().unwrap(); + self.expr = Some(expr); + self + } +} + +/// For recursive [WithQuery] [WithClause]s the CYCLE sql clause can be specified to avoid creating +/// an infinite traversals that loops on graph cycles indefinitely. You specify an expression that +/// identifies a node in the graph and that will be used to determine during the iteration of +/// the execution of the query when appending of new values whether the new values are distinct new +/// nodes or are already visited and therefore they should be added again into the result. +/// +/// A query can have both SEARCH and CYCLE clauses. +/// +/// Setting [Self::set], [Self::expr] and [Self::using] is mandatory. +#[derive(Debug, Clone, Default)] +pub struct Cycle { + pub(crate) expr: Option, + pub(crate) set_as: Option, + pub(crate) using: Option, +} + +impl Cycle { + /// Create a complete [Search] specification from the [SearchOrder] and a [SelectExpr]. The + /// given [SelectExpr] must have an alias specified. + pub fn new_from_expr_set_using(expr: EXPR, set: ID1, using: ID2) -> Self + where + EXPR: Into, + ID1: IntoIden, + ID2: IntoIden, + { + Self { + expr: Some(expr.into()), + set_as: Some(set.into_iden()), + using: Some(using.into_iden()), + } + } + + /// Constructs a new empty [Cycle]. + pub fn new() -> Self { + Self::default() + } + + /// The expression identifying nodes. + pub fn expr(&mut self, expr: EXPR) -> &mut Self + where + EXPR: Into, + { + self.expr = Some(expr.into()); + self + } + + /// The name of the boolean column containing whether we have completed a cycle or not yet + /// generated by this clause. + pub fn set(&mut self, set: ID) -> &mut Self + where + ID: IntoIden, + { + self.set_as = Some(set.into_iden()); + self + } + + /// The name of the array typed column that contains the node ids (generated using the + /// [Self::expr]) that specify the current nodes path that will be generated by this clause. + pub fn using(&mut self, using: ID) -> &mut Self + where + ID: IntoIden, + { + self.using = Some(using.into_iden()); + self + } +} + +/// A WITH clause can contain one or multiple common table expressions ([CommonTableExpression]). +/// +/// You can use this to generate [WithQuery] by calling [WithClause::query]. +/// +/// These named queries can act as a "query local table" that are materialized during execution and +/// then can be used by the query prefixed with the WITH clause. +/// +/// A WITH clause can contain multiple of these [CommonTableExpression]. (Except in the case of +/// recursive WITH query which can only contain one [CommonTableExpression]). +/// +/// A [CommonTableExpression] is a name, column names and a query returning data for those columns. +/// +/// Some databases (like sqlite) restrict the acceptable kinds of queries inside of the WITH clause +/// common table expressions. These databases only allow [SelectStatement]s to form a common table +/// expression. +/// +/// Other databases like postgres allow modification queries (UPDATE, DELETE) inside of the WITH +/// clause but they have to return a table. (They must have a RETURNING clause). +/// +/// sea-query doesn't check this or restrict the kind of [CommonTableExpression] that you can create +/// in rust. This means that you can put an UPDATE or DELETE queries into WITH clause and sea-query +/// will succeed in generating that kind of sql query but the execution inside the database will +/// fail because they are invalid. +/// +/// It is your responsibility to ensure that the kind of WITH clause that you put together makes +/// sense and valid for that database that you are using. +/// +/// NOTE that for recursive WITH queries (in sql: "WITH RECURSIVE") you can only have a +/// single [CommonTableExpression] inside of the WITH clause. That query must match certain +/// requirements: +/// * It is a query of UNION or UNION ALL of two queries. +/// * The first part of the query (the left side of the UNION) must be executable first in itself. +/// It must be non-recursive. (Cannot contain self reference) +/// * The self reference must appear in the right hand side of the UNION. +/// * The query can only have a single self-reference. +/// * Recursive data-modifying statements are not supported, but you can use the results of a +/// recursive SELECT query in a data-modifying statement. (like so: WITH RECURSIVE +/// cte_name(a,b,c,d) AS (SELECT ... UNION SELECT ... FROM ... JOIN cte_name ON ... WHERE ...) +/// DELETE FROM table WHERE table.a = cte_name.a) +/// +/// It is mandatory to set the [Self::cte]. With queries must have at least one CTE. +/// Recursive with query generation will panic if you specify more than one CTE. +/// +/// # Examples +/// +/// ``` +/// use sea_query::{*, IntoCondition, IntoIden, tests_cfg::*}; +/// +/// let base_query = SelectStatement::new() +/// .column(Alias::new("id")) +/// .expr(Expr::val(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::tbl(Alias::new("cte_traversal"), Alias::new("next")).equals(Alias::new("table"), Alias::new("id")).into_condition() +/// ) +/// .to_owned(); +/// +/// let common_table_expression = CommonTableExpression::new() +/// .query( +/// base_query.clone().union(UnionType::All, cte_referencing).to_owned() +/// ) +/// .column(Alias::new("id")) +/// .column(Alias::new("depth")) +/// .column(Alias::new("next")) +/// .column(Alias::new("value")) +/// .table_name(Alias::new("cte_traversal")) +/// .to_owned(); +/// +/// let select = SelectStatement::new() +/// .column(ColumnRef::Asterisk) +/// .from(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 = select.with(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""# +/// ); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct WithClause { + pub(crate) recursive: bool, + pub(crate) search: Option, + pub(crate) cycle: Option, + pub(crate) cte_expressions: Vec, +} + +impl WithClause { + /// Constructs a new [WithClause]. + pub fn new() -> Self { + Self::default() + } + + /// Sets whether this clause is a recursive with clause of not. + /// If set to true it will generate a 'WITH RECURSIVE' query. + /// + /// You can only specify a single [CommonTableExpression] containing a union query + /// if this is set to true. + pub fn recursive(&mut self, recursive: bool) -> &mut Self { + self.recursive = recursive; + self + } + + /// For recursive WITH queries you can specify the [Search] clause. + /// + /// This setting is not meaningful if the query is not recursive. + /// + /// Some databases don't support this clause. In that case this option will be silently ignored. + pub fn search(&mut self, search: Search) -> &mut Self { + self.search = Some(search); + self + } + + /// For recursive WITH queries you can specify the [Cycle] clause. + /// + /// This setting is not meaningful if the query is not recursive. + /// + /// Some databases don't support this clause. In that case this option will be silently ignored. + pub fn cycle(&mut self, cycle: Cycle) -> &mut Self { + self.cycle = Some(cycle); + self + } + + /// Add a [CommonTableExpression] to this with clause. + pub fn cte(&mut self, cte: CommonTableExpression) -> &mut Self { + self.cte_expressions.push(cte); + self + } + + /// You can turn this into a [WithQuery] using this function. The resulting WITH query will + /// execute the argument query with this WITH clause. + pub fn query(self, query: T) -> WithQuery + where + T: QueryStatementBuilder, + { + WithQuery::new().with_clause(self).query(query).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]). +/// +/// These named queries can act as a "query local table" that are materialized during execution and +/// then can be used by the query prefixed with the WITH clause. +/// +/// A WITH clause can contain multiple of these [CommonTableExpression]. (Except in the case of +/// recursive WITH query which can only contain one [CommonTableExpression]). +/// +/// A [CommonTableExpression] is a name, column names and a query returning data for those columns. +/// +/// Some databases (like sqlite) restrict the acceptable kinds of queries inside of the WITH clause +/// common table expressions. These databases only allow [SelectStatement]s to form a common table +/// expression. +/// +/// Other databases like postgres allow modification queries (UPDATE, DELETE) inside of the WITH +/// clause but they have to return a table. (They must have a RETURNING clause). +/// +/// sea-query doesn't check this or restrict the kind of [CommonTableExpression] that you can create +/// in rust. This means that you can put an UPDATE or DELETE queries into WITH clause and sea-query +/// will succeed in generating that kind of sql query but the execution inside the database will +/// fail because they are invalid. +/// +/// It is your responsibility to ensure that the kind of WITH clause that you put together makes +/// sense and valid for that database that you are using. +/// +/// NOTE that for recursive WITH queries (in sql: "WITH RECURSIVE") you can only have a +/// single [CommonTableExpression] inside of the WITH clause. That query must match certain +/// requirements: +/// * It is a query of UNION or UNION ALL of two queries. +/// * The first part of the query (the left side of the UNION) must be executable first in itself. +/// It must be non-recursive. (Cannot contain self reference) +/// * The self reference must appear in the right hand side of the UNION. +/// * The query can only have a single self-reference. +/// * Recursive data-modifying statements are not supported, but you can use the results of a +/// recursive SELECT query in a data-modifying statement. (like so: WITH RECURSIVE +/// cte_name(a,b,c,d) AS (SELECT ... UNION SELECT ... FROM ... JOIN cte_name ON ... WHERE ...) +/// DELETE FROM table WHERE table.a = cte_name.a) +/// +/// It is mandatory to set the [Self::cte] and the [Self::query]. +#[derive(Debug, Clone, Default)] +pub struct WithQuery { + pub(crate) with_clause: WithClause, + pub(crate) query: Option>, +} + +impl WithQuery { + /// Constructs a new empty [WithQuery]. + pub fn new() -> Self { + Self::default() + } + + /// Set the whole [WithClause]. + pub fn with_clause(&mut self, with_clause: WithClause) -> &mut Self { + self.with_clause = with_clause; + self + } + + /// Set the [WithClause::recursive]. See that method for more information. + pub fn recursive(&mut self, recursive: bool) -> &mut Self { + self.with_clause.recursive = recursive; + self + } + + /// Add the [WithClause::search]. See that method for more information. + pub fn search(&mut self, search: Search) -> &mut Self { + self.with_clause.search = Some(search); + self + } + + /// Set the [WithClause::cycle]. See that method for more information. + pub fn cycle(&mut self, cycle: Cycle) -> &mut Self { + self.with_clause.cycle = Some(cycle); + self + } + + /// Add a [CommonTableExpression] to the with clause. See [WithClause::cte]. + pub fn cte(&mut self, cte: CommonTableExpression) -> &mut Self { + self.with_clause.cte_expressions.push(cte); + self + } + + /// Set the query that you execute with the [WithClause]. + pub fn query(&mut self, query: T) -> &mut Self + where + T: QueryStatementBuilder, + { + self.query = Some(Box::new(query)); + self + } +} + +impl QueryStatementBuilder for WithQuery { + fn build_collect_any_into(&self, query_builder: &dyn QueryBuilder, sql: &mut SqlWriter, collector: &mut dyn FnMut(Value)) { + query_builder.prepare_with_query(self, sql, collector); + } + + fn box_clone(&self) -> Box { + Box::new(self.clone()) + } +} + +impl QueryStatementBuilderGenerics for WithQuery { + fn build_collect( + &self, + query_builder: T, + collector: &mut dyn FnMut(Value), + ) -> String { + let mut sql = SqlWriter::new(); + query_builder.prepare_with_query(self, &mut sql, collector); + sql.result() + } +} diff --git a/src/shim.rs b/src/shim.rs index 14a5d09c24..0530a53077 100644 --- a/src/shim.rs +++ b/src/shim.rs @@ -27,15 +27,18 @@ macro_rules! impl_query_statement_builder { ( $mod_name: ident, $struct_name: ident ) => { mod $mod_name { - use crate::{$struct_name, QueryBuilder, QueryStatementBuilder, Values}; + use crate::{ + $struct_name, QueryBuilder, QueryStatementBuilder, QueryStatementBuilderGenerics, + Values, + }; impl $struct_name { pub fn to_string(&self, query_builder: T) -> String { - ::to_string(self, query_builder) + ::to_string(self, query_builder) } pub fn build(&self, query_builder: T) -> (String, Values) { - ::build(self, query_builder) + ::build(self, query_builder) } pub fn build_any(&self, query_builder: &dyn QueryBuilder) -> (String, Values) { diff --git a/tests/postgres/query.rs b/tests/postgres/query.rs index 9818ef7f54..dfa0d34185 100644 --- a/tests/postgres/query.rs +++ b/tests/postgres/query.rs @@ -965,6 +965,40 @@ fn insert_from_select() { ); } +#[test] +fn insert_6() -> sea_query::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 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 insert = Query::insert(); + insert + .into_table(Glyph::Table) + .columns([Glyph::Id, Glyph::Image, Glyph::Aspect]) + .select_from(select)?; + let sql = insert.with(with_clause).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] fn update_1() { assert_eq!(