feat: expose database connection to ActiveModelBehaviour's methods (#…

…1145)
SeaQL · Dec 22, 2022 · da14510 · da14510
1 parent 83d46d9
commit da14510
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
@@ -278,11 +278,11 @@ pub trait ActiveModelTrait: Clone + Debug {
         Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait,
     {
-        let am = ActiveModelBehavior::before_save(self, true)?;
+        let am = ActiveModelBehavior::before_save(self, db, true)?;
         let model = <Self::Entity as EntityTrait>::insert(am)
             .exec_with_returning(db)
             .await?;
-        Self::after_save(model, true)
+        Self::after_save(model, db, true)
     }
 
     /// Perform the `UPDATE` operation on an ActiveModel
@@ -400,9 +400,9 @@ pub trait ActiveModelTrait: Clone + Debug {
         Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait,
     {
-        let am = ActiveModelBehavior::before_save(self, false)?;
+        let am = ActiveModelBehavior::before_save(self, db, false)?;
         let model: <Self::Entity as EntityTrait>::Model = Self::Entity::update(am).exec(db).await?;
-        Self::after_save(model, false)
+        Self::after_save(model, db, false)
     }
 
     /// Insert the model if primary key is `NotSet`, update otherwise.
@@ -477,10 +477,10 @@ pub trait ActiveModelTrait: Clone + Debug {
         Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait,
     {
-        let am = ActiveModelBehavior::before_delete(self)?;
+        let am = ActiveModelBehavior::before_delete(self, db)?;
         let am_clone = am.clone();
         let delete_res = Self::Entity::delete(am).exec(db).await?;
-        ActiveModelBehavior::after_delete(am_clone)?;
+        ActiveModelBehavior::after_delete(am_clone, db)?;
         Ok(delete_res)
     }
 
@@ -591,25 +591,26 @@ pub trait ActiveModelBehavior: ActiveModelTrait {
     }
 
     /// Will be called before saving
-    fn before_save(self, insert: bool) -> Result<Self, DbErr> {
+    fn before_save(self, db: &impl ConnectionTrait, insert: bool) -> Result<Self, DbErr> {
         Ok(self)
     }
 
     /// Will be called after saving
     fn after_save(
         model: <Self::Entity as EntityTrait>::Model,
+        db: &impl ConnectionTrait,
         insert: bool,
     ) -> Result<<Self::Entity as EntityTrait>::Model, DbErr> {
         Ok(model)
     }
 
     /// Will be called before deleting
-    fn before_delete(self) -> Result<Self, DbErr> {
+    fn before_delete(self, db: &impl ConnectionTrait) -> Result<Self, DbErr> {
         Ok(self)
     }
 
     /// Will be called after deleting
-    fn after_delete(self) -> Result<Self, DbErr> {
+    fn after_delete(self, db: &impl ConnectionTrait) -> Result<Self, DbErr> {
         Ok(self)
     }
 }

diff --git a/tests/common/bakery_chain/cake.rs b/tests/common/bakery_chain/cake.rs
@@ -1,4 +1,4 @@
-use sea_orm::entity::prelude::*;
+use sea_orm::{entity::prelude::*, ConnectionTrait};
 
 #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
 #[sea_orm(table_name = "cake")]
@@ -58,7 +58,7 @@ impl ActiveModelBehavior for ActiveModel {
         }
     }
 
-    fn before_save(self, insert: bool) -> Result<Self, DbErr> {
+    fn before_save(self, _db: &impl ConnectionTrait, insert: bool) -> Result<Self, DbErr> {
         use rust_decimal_macros::dec;
         if self.price.as_ref() == &dec!(0) {
             Err(DbErr::Custom(format!(
@@ -70,7 +70,7 @@ impl ActiveModelBehavior for ActiveModel {
         }
     }
 
-    fn after_save(model: Model, insert: bool) -> Result<Model, DbErr> {
+    fn after_save(model: Model, _db: &impl ConnectionTrait, insert: bool) -> Result<Model, DbErr> {
         use rust_decimal_macros::dec;
         if model.price < dec!(0) {
             Err(DbErr::Custom(format!(
@@ -82,7 +82,7 @@ impl ActiveModelBehavior for ActiveModel {
         }
     }
 
-    fn before_delete(self) -> Result<Self, DbErr> {
+    fn before_delete(self, _db: &impl ConnectionTrait) -> Result<Self, DbErr> {
         if self.name.as_ref().contains("(err_on_before_delete)") {
             Err(DbErr::Custom(
                 "[before_delete] Cannot be deleted".to_owned(),
@@ -92,7 +92,7 @@ impl ActiveModelBehavior for ActiveModel {
         }
     }
 
-    fn after_delete(self) -> Result<Self, DbErr> {
+    fn after_delete(self, _db: &impl ConnectionTrait) -> Result<Self, DbErr> {
         if self.name.as_ref().contains("(err_on_after_delete)") {
             Err(DbErr::Custom("[after_delete] Cannot be deleted".to_owned()))
         } else {