Merge pull request #210 from SeaQL/active-model-behavior

Update `ActiveModelBehavior` API
SeaQL · Oct 13, 2021 · 6defdaf · 6defdaf
2 parents 608a492 + 6ecc313
commit 6defdaf
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 26 deletions.
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
@@ -10,8 +10,7 @@ pub struct ActiveValue<V>
 where
     V: Into<Value>,
 {
-    // Don't want to call ActiveValue::unwrap() and cause panic
-    pub(self) value: Option<V>,
+    value: Option<V>,
     state: ActiveValueState,
 }
 
@@ -74,7 +73,7 @@ pub trait ActiveModelTrait: Clone + Debug {
         macro_rules! next {
             () => {
                 if let Some(col) = cols.next() {
-                    if let Some(val) = self.get(col.into_column()).value {
+                    if let Some(val) = self.get(col.into_column()).into_value() {
                         val
                     } else {
                         return None;
@@ -107,12 +106,11 @@ pub trait ActiveModelTrait: Clone + Debug {
     async fn insert<'a, C>(self, db: &'a C) -> Result<Self, DbErr>
     where
         <Self::Entity as EntityTrait>::Model: IntoActiveModel<Self>,
+        Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait<'a>,
-        Self: 'a,
     {
-        let am = self;
-        let exec = <Self::Entity as EntityTrait>::insert(am).exec(db);
-        let res = exec.await?;
+        let am = ActiveModelBehavior::before_save(self, true)?;
+        let res = <Self::Entity as EntityTrait>::insert(am).exec(db).await?;
         let found = <Self::Entity as EntityTrait>::find_by_id(res.last_insert_id)
             .one(db)
             .await?;
@@ -124,23 +122,23 @@ pub trait ActiveModelTrait: Clone + Debug {
 
     async fn update<'a, C>(self, db: &'a C) -> Result<Self, DbErr>
     where
+        Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait<'a>,
-        Self: 'a,
     {
-        let exec = Self::Entity::update(self).exec(db);
-        exec.await
+        let am = ActiveModelBehavior::before_save(self, false)?;
+        let am = Self::Entity::update(am).exec(db).await?;
+        ActiveModelBehavior::after_save(am, false)
     }
 
     /// Insert the model if primary key is unset, update otherwise.
     /// Only works if the entity has auto increment primary key.
     async fn save<'a, C>(self, db: &'a C) -> Result<Self, DbErr>
     where
-        Self: ActiveModelBehavior + 'a,
         <Self::Entity as EntityTrait>::Model: IntoActiveModel<Self>,
+        Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait<'a>,
     {
         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();
@@ -154,7 +152,6 @@ pub trait ActiveModelTrait: Clone + Debug {
         } else {
             am = am.update(db).await?;
         }
-        am = ActiveModelBehavior::after_save(am);
         Ok(am)
     }
 
@@ -164,33 +161,40 @@ pub trait ActiveModelTrait: Clone + Debug {
         Self: ActiveModelBehavior + 'a,
         C: ConnectionTrait<'a>,
     {
-        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)?;
+        let am_clone = am.clone();
+        let delete_res = Self::Entity::delete(am).exec(db).await?;
+        ActiveModelBehavior::after_delete(am_clone)?;
+        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) -> Result<Self, DbErr> {
+        Ok(self)
     }
 
     /// Will be called after saving
-    fn after_save(self) -> Self {
-        self
+    fn after_save(self, insert: bool) -> Result<Self, DbErr> {
+        Ok(self)
     }
 
     /// Will be called before deleting
-    fn before_delete(self) -> Self {
-        self
+    fn before_delete(self) -> Result<Self, DbErr> {
+        Ok(self)
+    }
+
+    /// Will be called after deleting
+    fn after_delete(self) -> Result<Self, DbErr> {
+        Ok(self)
     }
 }
 

diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs
@@ -1,8 +1,8 @@
 pub use crate::{
     error::*, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType,
-    EntityName, EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, ModelTrait,
-    PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, Related, RelationDef,
-    RelationTrait, Select, Value,
+    DatabaseConnection, DbConn, EntityName, EntityTrait, EnumIter, ForeignKeyAction, Iden,
+    IdenStatic, Linked, ModelTrait, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult,
+    Related, RelationDef, RelationTrait, Select, Value,
 };
 
 #[cfg(feature = "macros")]

diff --git a/src/error.rs b/src/error.rs
@@ -4,6 +4,7 @@ pub enum DbErr {
     Exec(String),
     Query(String),
     RecordNotFound(String),
+    Custom(String),
 }
 
 impl std::error::Error for DbErr {}
@@ -15,6 +16,7 @@ impl std::fmt::Display for DbErr {
             Self::Exec(s) => write!(f, "Execution Error: {}", s),
             Self::Query(s) => write!(f, "Query Error: {}", s),
             Self::RecordNotFound(s) => write!(f, "RecordNotFound Error: {}", s),
+            Self::Custom(s) => write!(f, "Custom Error: {}", s),
         }
     }
 }

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
@@ -306,6 +306,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(),
     };