Add exec_with_returning to delete SeaQL#2432

Author: tyt2y3

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)
+}

tests/returning_tests.rs

@@ -351,3 +351,97 @@ async fn update_many() {
 
     run().await.unwrap();
 }
+
+#[sea_orm_macros::test]
+#[cfg_attr(
+    any(
+        feature = "sqlx-mysql",
+        all(
+            feature = "sqlx-sqlite",
+            not(feature = "sqlite-use-returning-for-3_35")
+        )
+    ),
+    should_panic(expected = "Database backend doesn't support RETURNING")
+)]
+async fn delete_many() {
+    pub use common::{features::*, TestContext};
+    use edit_log::*;
+
+    let run = || async {
+        let ctx = TestContext::new("returning_tests_delete_many").await;
+        let db = &ctx.db;
+
+        create_tables(db).await?;
+
+        let inserted_models = [
+            Model {
+                id: 1,
+                action: "before_save".to_string(),
+                values: json!({ "id": "unique-id-001" }),
+            },
+            Model {
+                id: 2,
+                action: "before_save".to_string(),
+                values: json!({ "id": "unique-id-002" }),
+            },
+        ];
+        // Delete many with returning
+        assert_eq!(
+            Entity::insert_many(vec![
+                ActiveModel {
+                    id: NotSet,
+                    action: Set("before_save".to_string()),
+                    values: Set(json!({ "id": "unique-id-001" })),
+                },
+                ActiveModel {
+                    id: NotSet,
+                    action: Set("before_save".to_string()),
+                    values: Set(json!({ "id": "unique-id-002" })),
+                },
+            ])
+            .exec_with_returning_many(db)
+            .await?,
+            inserted_models
+        );
+
+        assert_eq!(
+            Entity::delete_many()
+                .filter(Column::Action.eq("before_save"))
+                .exec_with_returning(db)
+                .await?,
+            inserted_models
+        );
+
+        let inserted_model_3 = Model {
+            id: 3,
+            action: "before_save".to_string(),
+            values: json!({ "id": "unique-id-003" }),
+        };
+
+        Entity::insert(ActiveModel {
+            id: NotSet,
+            action: Set("before_save".to_string()),
+            values: Set(json!({ "id": "unique-id-003" })),
+        })
+        .exec(db)
+        .await?;
+
+        // One
+        assert_eq!(
+            Entity::delete(ActiveModel {
+                id: Set(3),
+                ..Default::default()
+            })
+            .exec_with_returning(db)
+            .await?,
+            Some(inserted_model_3)
+        );
+
+        // No-op
+        assert_eq!(Entity::delete_many().exec_with_returning(db).await?, []);
+
+        Result::<(), DbErr>::Ok(())
+    };
+
+    run().await.unwrap();
+}