Allow returning many on insert and delete

src/executor/delete.rs

@@ -1,7 +1,12 @@
-use crate::{error::*, ActiveModelTrait, ConnectionTrait, DeleteMany, DeleteOne, EntityTrait};
-use sea_query::DeleteStatement;
+use crate::{
+    error::*, ActiveModelTrait, ColumnTrait, ConnectionTrait, DeleteMany, DeleteOne, EntityTrait,
+    Iterable,
+};
+use sea_query::{DeleteStatement, Query};
 use std::future::Future;
 
+use super::{SelectModel, SelectorRaw};
+
 /// Handles DELETE operations in a ActiveModel using [DeleteStatement]
 #[derive(Clone, Debug)]
 pub struct Deleter {
@@ -27,6 +32,17 @@ where
         // so that self is dropped before entering await
         exec_delete_only(self.query, db)
     }
+
+    /// Execute an delete operation and return the deleted model (use `RETURNING` syntax if supported)
+    pub fn exec_with_returning<C>(
+        self,
+        db: &'a C,
+    ) -> impl Future<Output = Result<Option<<A::Entity as EntityTrait>::Model>, DbErr>> + '_
+    where
+        C: ConnectionTrait,
+    {
+        exec_delete_with_returning_one::<A::Entity, _>(self.query, db)
+    }
 }
 
 impl<'a, E> DeleteMany<E>
@@ -41,6 +57,18 @@ where
         // so that self is dropped before entering await
         exec_delete_only(self.query, db)
     }
+
+    /// Execute an delete operation and return the deleted model (use `RETURNING` syntax if supported)
+    pub fn exec_with_returning<C>(
+        self,
+        db: &C,
+    ) -> impl Future<Output = Result<Vec<E::Model>, DbErr>> + '_
+    where
+        E: EntityTrait,
+        C: ConnectionTrait,
+    {
+        exec_delete_with_returning_many::<E, _>(self.query, db)
+    }
 }
 
 impl Deleter {
@@ -56,6 +84,18 @@ impl Deleter {
     {
         exec_delete(self.query, db)
     }
+
+    /// Execute an delete operation and return the deleted model (use `RETURNING` syntax if supported)
+    pub fn exec_with_returning<E, C>(
+        self,
+        db: &C,
+    ) -> impl Future<Output = Result<Vec<E::Model>, DbErr>> + '_
+    where
+        E: EntityTrait,
+        C: ConnectionTrait,
+    {
+        exec_delete_with_returning_many::<E, _>(self.query, db)
+    }
 }
 
 async fn exec_delete_only<C>(query: DeleteStatement, db: &C) -> Result<DeleteResult, DbErr>
@@ -77,3 +117,49 @@ where
         rows_affected: result.rows_affected(),
     })
 }
+
+async fn exec_delete_with_returning_one<E, C>(
+    mut query: DeleteStatement,
+    db: &C,
+) -> Result<Option<E::Model>, DbErr>
+where
+    E: EntityTrait,
+    C: ConnectionTrait,
+{
+    let models = match db.support_returning() {
+        true => {
+            let db_backend = db.get_database_backend();
+            let delete_statement = db_backend.build(&query.returning_all().to_owned());
+            SelectorRaw::<SelectModel<<E>::Model>>::from_statement(delete_statement)
+                .one(db)
+                .await?
+        }
+        false => unimplemented!("Database backend doesn't support RETURNING"),
+    };
+    Ok(models)
+}
+
+async fn exec_delete_with_returning_many<E, C>(
+    mut query: DeleteStatement,
+    db: &C,
+) -> Result<Vec<E::Model>, DbErr>
+where
+    E: EntityTrait,
+    C: ConnectionTrait,
+{
+    let models = match db.support_returning() {
+        true => {
+            let db_backend = db.get_database_backend();
+            let returning = Query::returning().exprs(
+                E::Column::iter().map(|c| c.select_enum_as(c.into_returning_expr(db_backend))),
+            );
+            let query = query.returning(returning);
+            let delete_statement = db_backend.build(&query.to_owned());
+            SelectorRaw::<SelectModel<<E>::Model>>::from_statement(delete_statement)
+                .all(db)
+                .await?
+        }
+        false => unimplemented!("Database backend doesn't support RETURNING"),
+    };
+    Ok(models)
+}

src/executor/insert.rs

@@ -102,6 +102,28 @@ where
             Err(err) => Err(err),
         }
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if database supported)
+    pub async fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> Result<TryInsertResult<Vec<<A::Entity as EntityTrait>::Model>>, DbErr>
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        if self.insert_struct.columns.is_empty() {
+            return Ok(TryInsertResult::Empty);
+        }
+
+        let res = self.insert_struct.exec_with_returning_many(db).await;
+        match res {
+            Ok(res) => Ok(TryInsertResult::Inserted(res)),
+            Err(DbErr::RecordNotInserted) => Ok(TryInsertResult::Conflicted),
+            Err(err) => Err(err),
+        }
+    }
 }
 
 impl<A> Insert<A>
@@ -155,6 +177,25 @@ where
     {
         Inserter::<A>::new(self.primary_key, self.query).exec_with_returning(db)
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if supported)
+    pub async fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        if self.columns.is_empty() {
+            Ok(vec![])
+        } else {
+            Inserter::<A>::new(self.primary_key, self.query)
+                .exec_with_returning_many(db)
+                .await
+        }
+    }
 }
 
 impl<A> Inserter<A>
@@ -203,6 +244,19 @@ where
     {
         exec_insert_with_returning::<A, _>(self.primary_key, self.query, db)
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if supported)
+    pub fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> impl Future<Output = Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>> + '_
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        exec_insert_many_with_returning::<A, _>(self.query, db)
+    }
 }
 
 async fn exec_insert<A, C>(
@@ -313,3 +367,33 @@ where
         )),
     }
 }
+
+async fn exec_insert_many_with_returning<A, C>(
+    mut insert_statement: InsertStatement,
+    db: &C,
+) -> Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>
+where
+    <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+    C: ConnectionTrait,
+    A: ActiveModelTrait,
+{
+    match db.support_returning() {
+        true => {
+            let db_backend = db.get_database_backend();
+            let returning = Query::returning().exprs(
+                <A::Entity as EntityTrait>::Column::iter()
+                    .map(|c| c.select_as(c.into_returning_expr(db_backend))),
+            );
+            insert_statement.returning(returning);
+            let insert_statement = db_backend.build(&insert_statement);
+            let models: Vec<<A::Entity as EntityTrait>::Model> =
+                SelectorRaw::<SelectModel<<A::Entity as EntityTrait>::Model>>::from_statement(
+                    insert_statement,
+                )
+                .all(db)
+                .await?;
+            Ok(models)
+        }
+        false => unimplemented!("Database backend doesn't support RETURNING"),
+    }
+}