Update ActiveModelBehavior API

SeaQL · Sep 30, 2021 · ef1eb57 · ef1eb57
1 parent beb3ec6
commit ef1eb57
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 26 deletions.
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
@@ -69,30 +69,35 @@ pub trait ActiveModelTrait: Clone + Debug {
 
     async fn insert(self, db: &DatabaseConnection) -> Result<Self, DbErr>
     where
+        Self: ActiveModelBehavior,
         <Self::Entity as EntityTrait>::Model: IntoActiveModel<Self>,
     {
-        let am = self;
-        let exec = <Self::Entity as EntityTrait>::insert(am).exec(db);
-        let res = exec.await?;
+        let am = ActiveModelBehavior::before_save(self, true, db)?;
+        let res = <Self::Entity as EntityTrait>::insert(am).exec(db).await?;
         // Assume valid last_insert_id is not equals to Default::default()
         if res.last_insert_id
             != <<Self::Entity as EntityTrait>::PrimaryKey as PrimaryKeyTrait>::ValueType::default()
         {
             let found = <Self::Entity as EntityTrait>::find_by_id(res.last_insert_id)
                 .one(db)
                 .await?;
-            match found {
+            let am = match found {
                 Some(model) => Ok(model.into_active_model()),
                 None => Err(DbErr::Exec("Failed to find inserted item".to_owned())),
-            }
+            }?;
+            ActiveModelBehavior::after_save(am, true, db)
         } else {
             Ok(Self::default())
         }
     }
 
-    async fn update(self, db: &DatabaseConnection) -> Result<Self, DbErr> {
-        let exec = Self::Entity::update(self).exec(db);
-        exec.await
+    async fn update(self, db: &DatabaseConnection) -> Result<Self, DbErr>
+    where
+        Self: ActiveModelBehavior,
+    {
+        let am = ActiveModelBehavior::before_save(self, false, db)?;
+        let am = Self::Entity::update(am).exec(db).await?;
+        ActiveModelBehavior::after_save(am, false, db)
     }
 
     /// Insert the model if primary key is unset, update otherwise.
@@ -103,7 +108,6 @@ pub trait ActiveModelTrait: Clone + Debug {
         <Self::Entity as EntityTrait>::Model: IntoActiveModel<Self>,
     {
         let mut am = self;
-        am = ActiveModelBehavior::before_save(am);
         let mut is_update = true;
         for key in <Self::Entity as EntityTrait>::PrimaryKey::iter() {
             let col = key.into_column();
@@ -117,7 +121,6 @@ pub trait ActiveModelTrait: Clone + Debug {
         } else {
             am = am.update(db).await?;
         }
-        am = ActiveModelBehavior::after_save(am);
         Ok(am)
     }
 
@@ -126,33 +129,40 @@ pub trait ActiveModelTrait: Clone + Debug {
     where
         Self: ActiveModelBehavior,
     {
-        let mut am = self;
-        am = ActiveModelBehavior::before_delete(am);
-        let exec = Self::Entity::delete(am).exec(db);
-        exec.await
+        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, db)?;
+        Ok(delete_res)
     }
 }
 
 /// Behaviors for users to override
+#[allow(unused_variables)]
 pub trait ActiveModelBehavior: ActiveModelTrait {
     /// Create a new ActiveModel with default values. Also used by `Default::default()`.
     fn new() -> Self {
         <Self as ActiveModelTrait>::default()
     }
 
     /// Will be called before saving
-    fn before_save(self) -> Self {
-        self
+    fn before_save(self, insert: bool, db: &DatabaseConnection) -> Result<Self, DbErr> {
+        Ok(self)
     }
 
     /// Will be called after saving
-    fn after_save(self) -> Self {
-        self
+    fn after_save(self, insert: bool, db: &DatabaseConnection) -> Result<Self, DbErr> {
+        Ok(self)
     }
 
     /// Will be called before deleting
-    fn before_delete(self) -> Self {
-        self
+    fn before_delete(self, db: &DatabaseConnection) -> Result<Self, DbErr> {
+        Ok(self)
+    }
+
+    /// Will be called after deleting
+    fn after_delete(self, db: &DatabaseConnection) -> Result<Self, DbErr> {
+        Ok(self)
     }
 }
 

diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs
@@ -1,9 +1,10 @@
 pub use crate::{
     error::*, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType,
-    DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn, DeriveCustomColumn, DeriveEntity,
-    DeriveEntityModel, DeriveModel, DerivePrimaryKey, DeriveRelation, EntityName, EntityTrait,
-    EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, ModelTrait, PrimaryKeyToColumn,
-    PrimaryKeyTrait, QueryFilter, QueryResult, Related, RelationDef, RelationTrait, Select, Value,
+    DatabaseConnection, DbConn, DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn,
+    DeriveCustomColumn, DeriveEntity, DeriveEntityModel, DeriveModel, DerivePrimaryKey,
+    DeriveRelation, EntityName, EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked,
+    ModelTrait, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, Related,
+    RelationDef, RelationTrait, Select, Value,
 };
 
 #[cfg(feature = "with-json")]

diff --git a/src/executor/insert.rs b/src/executor/insert.rs
@@ -1,6 +1,6 @@
 use crate::{
-    error::*, ActiveModelTrait, DatabaseConnection, DbBackend, EntityTrait, Insert, PrimaryKeyTrait,
-    Statement, TryFromU64,
+    error::*, ActiveModelTrait, DatabaseConnection, DbBackend, EntityTrait, Insert,
+    PrimaryKeyTrait, Statement, TryFromU64,
 };
 use sea_query::InsertStatement;
 use std::{future::Future, marker::PhantomData};

diff --git a/tests/common/bakery_chain/applog.rs b/tests/common/bakery_chain/applog.rs
@@ -5,6 +5,7 @@ use sea_orm::entity::prelude::*;
 pub struct Model {
     #[sea_orm(primary_key)]
     pub id: i32,
+    pub action: String,
     pub json: Json,
     pub created_at: DateTimeWithTimeZone,
 }

diff --git a/tests/common/setup/schema.rs b/tests/common/setup/schema.rs
@@ -301,6 +301,7 @@ pub async fn create_log_table(db: &DbConn) -> Result<ExecResult, DbErr> {
                 .auto_increment()
                 .primary_key(),
         )
+        .col(ColumnDef::new(applog::Column::Action).string().not_null())
         .col(ColumnDef::new(applog::Column::Json).json().not_null())
         .col(
             ColumnDef::new(applog::Column::CreatedAt)

diff --git a/tests/timestamp_tests.rs b/tests/timestamp_tests.rs
@@ -16,6 +16,7 @@ async fn main() -> Result<(), DbErr> {
 pub async fn create_applog(db: &DatabaseConnection) -> Result<(), DbErr> {
     let log = applog::Model {
         id: 1,
+        action: "Testing".to_owned(),
         json: Json::String("HI".to_owned()),
         created_at: "2021-09-17T17:50:20+08:00".parse().unwrap(),
     };