Extended returning support

.github/workflows/rust.yml

@@ -22,7 +22,7 @@ jobs:
           toolchain: stable
           components: clippy
           override: true
-          
+
       - uses: actions-rs/clippy-check@v1
         with:
           token: ${{ secrets.GITHUB_TOKEN }}

Cargo.toml

@@ -68,6 +68,7 @@ with-rust_decimal = ["rust_decimal"]
 with-bigdecimal = ["bigdecimal"]
 with-uuid = ["uuid"]
 
+
 [[test]]
 name = "test-derive"
 path = "tests/derive/mod.rs"

src/backend/postgres/query.rs

@@ -6,24 +6,6 @@ impl QueryBuilder for PostgresQueryBuilder {
         ("$", true)
     }
 
-    fn prepare_returning(
-        &self,
-        returning: &[SelectExpr],
-        sql: &mut SqlWriter,
-        collector: &mut dyn FnMut(Value),
-    ) {
-        if !returning.is_empty() {
-            write!(sql, " RETURNING ").unwrap();
-            returning.iter().fold(true, |first, expr| {
-                if !first {
-                    write!(sql, ", ").unwrap()
-                }
-                self.prepare_select_expr(expr, sql, collector);
-                false
-            });
-        }
-    }
-
     fn if_null_function(&self) -> &str {
         "COALESCE"
     }

src/backend/query_builder.rs

@@ -756,10 +756,31 @@ pub trait QueryBuilder: QuotedBuilder {
     /// Hook to insert "RETURNING" statements.
     fn prepare_returning(
         &self,
-        _returning: &[SelectExpr],
-        _sql: &mut SqlWriter,
+        returning: &Returning,
+        sql: &mut SqlWriter,
         _collector: &mut dyn FnMut(Value),
     ) {
+        match returning {
+            Returning::All => write!(sql, " RETURNING *").unwrap(),
+            Returning::Columns(cols) => {
+                write!(sql, " RETURNING ").unwrap();
+                cols.into_iter().fold(true, |first, column_ref| {
+                    if !first {
+                        write!(sql, ", ").unwrap()
+                    }
+                    match column_ref {
+                        ColumnRef::Column(column) => column.prepare(sql, self.quote()),
+                        ColumnRef::TableColumn(table, column) => {
+                            table.prepare(sql, self.quote());
+                            write!(sql, ".").unwrap();
+                            column.prepare(sql, self.quote());
+                        }
+                    };
+                    false
+                });
+            }
+            Returning::Nothing => return,
+        }
     }
 
     #[doc(hidden)]

src/query/delete.rs

@@ -4,7 +4,7 @@ use crate::{
     query::{condition::*, OrderedStatement},
     types::*,
     value::*,
-    Query, QueryStatementBuilder, SelectExpr, SelectStatement,
+    ColumnRef, QueryStatementBuilder, Returning,
 };
 
 /// Delete existing rows from the table
@@ -39,7 +39,7 @@ pub struct DeleteStatement {
     pub(crate) wherei: ConditionHolder,
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
-    pub(crate) returning: Vec<SelectExpr>,
+    pub(crate) returning: Returning,
 }
 
 impl Default for DeleteStatement {
@@ -56,7 +56,7 @@ impl DeleteStatement {
             wherei: ConditionHolder::new(),
             orders: Vec::new(),
             limit: None,
-            returning: Vec::new(),
+            returning: Returning::default(),
         }
     }
 
@@ -100,36 +100,36 @@ impl DeleteStatement {
         self
     }
 
-    /// RETURNING expressions. Postgres only.
+    /// RETURNING expressions. Supported fully by postgres, version deppendant on other databases
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
     ///
     /// let query = Query::delete()
     ///     .from_table(Glyph::Table)
     ///     .and_where(Expr::col(Glyph::Id).eq(1))
-    ///     .returning(Query::select().column(Glyph::Id).take())
+    ///     .returning(Returning::Columns(vec![Glyph::Id.into_column_ref()]))
     ///     .to_owned();
     ///
     /// assert_eq!(
     ///     query.to_string(MysqlQueryBuilder),
-    ///     r#"DELETE FROM `glyph` WHERE `id` = 1"#
+    ///     r#"DELETE FROM `glyph` WHERE `id` = 1 RETURNING `id`"#
     /// );
     /// assert_eq!(
     ///     query.to_string(PostgresQueryBuilder),
     ///     r#"DELETE FROM "glyph" WHERE "id" = 1 RETURNING "id""#
     /// );
     /// assert_eq!(
     ///     query.to_string(SqliteQueryBuilder),
-    ///     r#"DELETE FROM `glyph` WHERE `id` = 1"#
+    ///     r#"DELETE FROM `glyph` WHERE `id` = 1 RETURNING `id`"#
     /// );
     /// ```
-    pub fn returning(&mut self, select: SelectStatement) -> &mut Self {
-        self.returning = select.selects;
+    pub fn returning(&mut self, returning_cols: Returning) -> &mut Self {
+        self.returning = returning_cols;
         self
     }
 
-    /// RETURNING a column after delete. Postgres only.
+    /// RETURNING a column after delete. Supported fully by postgres, version deppendant on other databases
     /// Wrapper over [`DeleteStatement::returning()`].
     ///
     /// ```
@@ -143,22 +143,22 @@ impl DeleteStatement {
     ///
     /// assert_eq!(
     ///     query.to_string(MysqlQueryBuilder),
-    ///     r#"DELETE FROM `glyph` WHERE `id` = 1"#
+    ///     r#"DELETE FROM `glyph` WHERE `id` = 1 RETURNING `id`"#
     /// );
     /// assert_eq!(
     ///     query.to_string(PostgresQueryBuilder),
     ///     r#"DELETE FROM "glyph" WHERE "id" = 1 RETURNING "id""#
     /// );
     /// assert_eq!(
     ///     query.to_string(SqliteQueryBuilder),
-    ///     r#"DELETE FROM `glyph` WHERE `id` = 1"#
+    ///     r#"DELETE FROM `glyph` WHERE `id` = 1 RETURNING `id`"#
     /// );
     /// ```
     pub fn returning_col<C>(&mut self, col: C) -> &mut Self
     where
         C: IntoIden,
     {
-        self.returning(Query::select().column(col.into_iden()).take())
+        self.returning(Returning::Columns(vec![ColumnRef::Column(col.into_iden())]))
     }
 }
 

src/query/insert.rs

@@ -1,6 +1,6 @@
 use crate::{
-    backend::QueryBuilder, error::*, prepare::*, types::*, value::*, Expr, Query,
-    QueryStatementBuilder, SelectExpr, SelectStatement, SimpleExpr,
+    backend::QueryBuilder, error::*, prepare::*, types::*, value::*, Expr, QueryStatementBuilder,
+    Returning, SimpleExpr,
 };
 
 /// Insert any new rows into an existing table
@@ -35,7 +35,7 @@ pub struct InsertStatement {
     pub(crate) table: Option<Box<TableRef>>,
     pub(crate) columns: Vec<DynIden>,
     pub(crate) values: Vec<Vec<SimpleExpr>>,
-    pub(crate) returning: Vec<SelectExpr>,
+    pub(crate) returning: Returning,
 }
 
 impl InsertStatement {
@@ -179,7 +179,7 @@ impl InsertStatement {
         self.exprs(values).unwrap()
     }
 
-    /// RETURNING expressions. Postgres only.
+    /// RETURNING expressions. Supported fully by postgres, version deppendant on other databases
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
@@ -188,28 +188,28 @@ impl InsertStatement {
     ///     .into_table(Glyph::Table)
     ///     .columns(vec![Glyph::Image])
     ///     .values_panic(vec!["12A".into()])
-    ///     .returning(Query::select().column(Glyph::Id).take())
+    ///     .returning(Returning::Columns(vec![Glyph::Id.into_column_ref()]))
     ///     .to_owned();
     ///
     /// assert_eq!(
     ///     query.to_string(MysqlQueryBuilder),
-    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A') RETURNING `id`"
     /// );
     /// assert_eq!(
     ///     query.to_string(PostgresQueryBuilder),
     ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING "id""#
     /// );
     /// assert_eq!(
     ///     query.to_string(SqliteQueryBuilder),
-    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A') RETURNING `id`"
     /// );
     /// ```
-    pub fn returning(&mut self, select: SelectStatement) -> &mut Self {
-        self.returning = select.selects;
+    pub fn returning(&mut self, returning: Returning) -> &mut Self {
+        self.returning = returning;
         self
     }
 
-    /// RETURNING a column after insertion. Postgres only. This is equivalent to MySQL's LAST_INSERT_ID.
+    /// RETURNING a column after insertion. Supported fully by postgres, version deppendant on other databases
     /// Wrapper over [`InsertStatement::returning()`].
     ///
     /// ```
@@ -224,22 +224,22 @@ impl InsertStatement {
     ///
     /// assert_eq!(
     ///     query.to_string(MysqlQueryBuilder),
-    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A') RETURNING `id`"
     /// );
     /// assert_eq!(
     ///     query.to_string(PostgresQueryBuilder),
     ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING "id""#
     /// );
     /// assert_eq!(
     ///     query.to_string(SqliteQueryBuilder),
-    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A') RETURNING `id`"
     /// );
     /// ```
     pub fn returning_col<C>(&mut self, col: C) -> &mut Self
     where
         C: IntoIden,
     {
-        self.returning(Query::select().column(col.into_iden()).take())
+        self.returning(Returning::Columns(vec![ColumnRef::Column(col.into_iden())]))
     }
 }
 

src/query/mod.rs

@@ -11,6 +11,7 @@ mod condition;
 mod delete;
 mod insert;
 mod ordered;
+mod returning;
 mod select;
 mod shim;
 mod traits;
@@ -20,6 +21,7 @@ pub use condition::*;
 pub use delete::*;
 pub use insert::*;
 pub use ordered::*;
+pub use returning::*;
 pub use select::*;
 pub use traits::*;
 pub use update::*;

src/query/returning.rs

@@ -0,0 +1,27 @@
+use std::default::Default;
+
+use crate::{ColumnRef, IntoColumnRef};
+
+#[derive(Clone, Debug)]
+pub enum Returning {
+    All,
+    Columns(Vec<ColumnRef>),
+    Nothing,
+}
+
+impl Returning {
+    pub fn cols<T, I>(cols: I) -> Self
+    where
+        T: IntoColumnRef,
+        I: IntoIterator<Item = T>,
+    {
+        let cols: Vec<_> = cols.into_iter().map(|c| c.into_column_ref()).collect();
+        Self::Columns(cols)
+    }
+}
+
+impl Default for Returning {
+    fn default() -> Self {
+        Self::Nothing
+    }
+}