perf(quaint): support insert common table expressions

Signed-off-by: Alexey Orlenko's AI Agent <robot@aqrln.net>

Author: tensordreams

quaint/src/ast.rs

@@ -34,6 +34,7 @@ pub use column::{Column, DefaultValue, TypeDataLength, TypeFamily};
 pub use compare::{Comparable, Compare, JsonCompare, JsonType};
 pub use conditions::ConditionTree;
 pub use conjunctive::Conjunctive;
+pub(crate) use cte::CommonTableExpressionBody;
 pub use cte::{CommonTableExpression, IntoCommonTableExpression};
 pub use delete::Delete;
 pub use enums::{EnumName, EnumVariant};

quaint/src/ast/cte.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 
-use super::SelectQuery;
+use super::{Insert, Select, SelectQuery, Union};
 
 /// A builder for a common table expression (CTE) statement, to be used in the
 /// `WITH` block of a `SELECT` statement.
@@ -12,7 +12,13 @@ use super::SelectQuery;
 pub struct CommonTableExpression<'a> {
     pub(crate) identifier: Cow<'a, str>,
     pub(crate) columns: Vec<Cow<'a, str>>,
-    pub(crate) selection: SelectQuery<'a>,
+    pub(crate) body: CommonTableExpressionBody<'a>,
+}
+
+#[derive(Debug, PartialEq, Clone)]
+pub(crate) enum CommonTableExpressionBody<'a> {
+    Selection(SelectQuery<'a>),
+    Insert(Insert<'a>),
 }
 
 impl<'a> CommonTableExpression<'a> {
@@ -30,14 +36,40 @@ impl<'a> CommonTableExpression<'a> {
 ///
 /// [`Select#with`]: struct.Select.html#method.with
 pub trait IntoCommonTableExpression<'a> {
-    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a>
-    where
-        Self: Into<SelectQuery<'a>>,
-    {
-        CommonTableExpression {
-            identifier: identifier.into(),
-            columns: Vec::new(),
-            selection: self.into(),
-        }
+    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a>;
+}
+
+fn common_table_expression<'a>(
+    identifier: impl Into<Cow<'a, str>>,
+    body: CommonTableExpressionBody<'a>,
+) -> CommonTableExpression<'a> {
+    CommonTableExpression {
+        identifier: identifier.into(),
+        columns: Vec::new(),
+        body,
+    }
+}
+
+impl<'a> IntoCommonTableExpression<'a> for Select<'a> {
+    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a> {
+        common_table_expression(identifier, CommonTableExpressionBody::Selection(self.into()))
+    }
+}
+
+impl<'a> IntoCommonTableExpression<'a> for Union<'a> {
+    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a> {
+        common_table_expression(identifier, CommonTableExpressionBody::Selection(self.into()))
+    }
+}
+
+impl<'a> IntoCommonTableExpression<'a> for SelectQuery<'a> {
+    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a> {
+        common_table_expression(identifier, CommonTableExpressionBody::Selection(self))
+    }
+}
+
+impl<'a> IntoCommonTableExpression<'a> for Insert<'a> {
+    fn into_cte(self, identifier: impl Into<Cow<'a, str>>) -> CommonTableExpression<'a> {
+        common_table_expression(identifier, CommonTableExpressionBody::Insert(self))
     }
 }

quaint/src/ast/query.rs

@@ -1,8 +1,6 @@
 use crate::ast::{Delete, Insert, Merge, Select, Union, Update};
 use std::borrow::Cow;
 
-use super::IntoCommonTableExpression;
-
 /// A database query
 #[derive(Debug, PartialEq)]
 pub enum Query<'a> {
@@ -107,5 +105,3 @@ impl<'a> From<SelectQuery<'a>> for Query<'a> {
         }
     }
 }
-
-impl<'a> IntoCommonTableExpression<'a> for SelectQuery<'a> {}

quaint/src/ast/select.rs

@@ -722,8 +722,6 @@ impl<'a> Select<'a> {
     }
 }
 
-impl<'a> IntoCommonTableExpression<'a> for Select<'a> {}
-
 impl<'a> Extend<Expression<'a>> for Select<'a> {
     fn extend<T: IntoIterator<Item = Expression<'a>>>(&mut self, iter: T) {
         self.columns.extend(iter);

quaint/src/ast/union.rs

@@ -2,7 +2,6 @@ use crate::ast::{Expression, Query, Select};
 use std::{collections::BTreeSet, fmt};
 
 use super::CommonTableExpression;
-use super::IntoCommonTableExpression;
 
 #[derive(Debug, PartialEq, Clone, Copy)]
 pub(crate) enum UnionType {
@@ -145,5 +144,3 @@ impl<'a> Union<'a> {
         }
     }
 }
-
-impl<'a> IntoCommonTableExpression<'a> for Union<'a> {}

quaint/src/visitor.rs

@@ -1350,8 +1350,11 @@ pub trait Visitor<'a> {
 
         self.write(" AS ")?;
 
-        let selection = cte.selection;
-        self.surround_with("(", ")", |ref mut s| s.visit_selection(selection))
+        let body = cte.body;
+        self.surround_with("(", ")", |ref mut s| match body {
+            CommonTableExpressionBody::Selection(selection) => s.visit_selection(selection),
+            CommonTableExpressionBody::Insert(insert) => s.visit_insert(insert),
+        })
     }
 
     fn visit_comment(&mut self, comment: Cow<'a, str>) -> Result {

quaint/src/visitor/postgres.rs

@@ -934,6 +934,28 @@ mod tests {
         assert_eq!(expected.1, params);
     }
 
+    #[test]
+    #[cfg(feature = "postgresql")]
+    fn test_insert_common_table_expression() {
+        let expected = expected_values(
+            "WITH \"inserted\" AS (INSERT INTO \"relations\" (\"parent_id\",\"child_id\") SELECT \"parent_id\", \"child_id\" FROM \"children\" WHERE \"child_id\" = $1 ON CONFLICT DO NOTHING RETURNING \"child_id\") SELECT \"child_id\" FROM \"inserted\"",
+            vec![10],
+        );
+        let selection = Select::from_table("children")
+            .columns(vec!["parent_id", "child_id"])
+            .so_that("child_id".equals(10));
+        let insert = Insert::expression_into("relations", vec!["parent_id", "child_id"], selection)
+            .on_conflict(OnConflict::DoNothing)
+            .returning(vec!["child_id"]);
+        let query = Select::from_table("inserted")
+            .column("child_id")
+            .with(insert.into_cte("inserted"));
+        let (sql, params) = Postgres::build(query).unwrap();
+
+        assert_eq!(expected.0, sql);
+        assert_eq!(expected.1, params);
+    }
+
     #[test]
     #[cfg(feature = "postgresql")]
     fn test_insert_on_conflict_update() {