From 163835c0fcabe9b8f976dae3c26855747fe230ff Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Thu, 12 Jan 2023 15:42:13 +0800
Subject: [PATCH] Unchanged ActiveValue as Set (#1177)

* Unchanged ActiveValue as Set

* Renaming

* Rename

* Rename methods
---
 sea-orm-macros/src/derives/active_model.rs |   7 +
 src/entity/active_model.rs                 | 150 +++++++++++++++++++++
 2 files changed, 157 insertions(+)

diff --git a/sea-orm-macros/src/derives/active_model.rs b/sea-orm-macros/src/derives/active_model.rs
index 43efe6b6c2..5dbba6692c 100644
--- a/sea-orm-macros/src/derives/active_model.rs
+++ b/sea-orm-macros/src/derives/active_model.rs
@@ -154,6 +154,13 @@ fn derive_active_model(all_fields: IntoIter<Field>) -> syn::Result<TokenStream>
                     #(#field: sea_orm::ActiveValue::not_set()),*
                 }
             }
+
+            fn reset(&mut self, c: <Self::Entity as EntityTrait>::Column) {
+                match c {
+                    #(<Self::Entity as EntityTrait>::Column::#name => self.#field.reset(),)*
+                    _ => panic!("This ActiveModel does not have this field"),
+                }
+            }
         }
     ))
 }
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
index 1971a876c4..b5c078f6d9 100644
--- a/src/entity/active_model.rs
+++ b/src/entity/active_model.rs
@@ -105,6 +105,19 @@ pub trait ActiveModelTrait: Clone + Debug {
     /// The default implementation of the ActiveModel
     fn default() -> Self;
 
+    /// Reset the value from [ActiveValue::Unchanged] to [ActiveValue::Set],
+    /// leaving [ActiveValue::NotSet] untouched.
+    fn reset(&mut self, c: <Self::Entity as EntityTrait>::Column);
+
+    /// Reset all values from [ActiveValue::Unchanged] to [ActiveValue::Set],
+    /// leaving [ActiveValue::NotSet] untouched.
+    fn reset_all(mut self) -> Self {
+        for col in <Self::Entity as EntityTrait>::Column::iter() {
+            self.reset(col);
+        }
+        self
+    }
+
     /// Get the primary key of the ActiveModel
     #[allow(clippy::question_mark)]
     fn get_primary_key_value(&self) -> Option<ValueTuple> {
@@ -819,6 +832,15 @@ where
             Self::NotSet => ActiveValue::not_set(),
         }
     }
+
+    /// Reset the value from [ActiveValue::Unchanged] to [ActiveValue::Set],
+    /// leaving [ActiveValue::NotSet] untouched.
+    pub fn reset(&mut self) {
+        *self = match self.take() {
+            Some(value) => ActiveValue::Set(value),
+            None => ActiveValue::NotSet,
+        };
+    }
 }
 
 impl<V> std::convert::AsRef<V> for ActiveValue<V>
@@ -1272,4 +1294,132 @@ mod tests {
         fruit.set(fruit::Column::Name, "apple".into());
         assert!(fruit.is_changed());
     }
+
+    #[test]
+    fn test_reset_1() {
+        assert_eq!(
+            fruit::Model {
+                id: 1,
+                name: "Apple".into(),
+                cake_id: None,
+            }
+            .into_active_model(),
+            fruit::ActiveModel {
+                id: Unchanged(1),
+                name: Unchanged("Apple".into()),
+                cake_id: Unchanged(None)
+            },
+        );
+
+        assert_eq!(
+            fruit::Model {
+                id: 1,
+                name: "Apple".into(),
+                cake_id: None,
+            }
+            .into_active_model()
+            .reset_all(),
+            fruit::ActiveModel {
+                id: Set(1),
+                name: Set("Apple".into()),
+                cake_id: Set(None)
+            },
+        );
+
+        assert_eq!(
+            fruit::Model {
+                id: 1,
+                name: "Apple".into(),
+                cake_id: Some(2),
+            }
+            .into_active_model(),
+            fruit::ActiveModel {
+                id: Unchanged(1),
+                name: Unchanged("Apple".into()),
+                cake_id: Unchanged(Some(2)),
+            },
+        );
+
+        assert_eq!(
+            fruit::Model {
+                id: 1,
+                name: "Apple".into(),
+                cake_id: Some(2),
+            }
+            .into_active_model()
+            .reset_all(),
+            fruit::ActiveModel {
+                id: Set(1),
+                name: Set("Apple".into()),
+                cake_id: Set(Some(2)),
+            },
+        );
+    }
+
+    #[smol_potat::test]
+    async fn test_reset_2() -> Result<(), DbErr> {
+        use crate::*;
+
+        let db = MockDatabase::new(DbBackend::Postgres)
+            .append_exec_results(vec![
+                MockExecResult {
+                    last_insert_id: 1,
+                    rows_affected: 1,
+                },
+                MockExecResult {
+                    last_insert_id: 1,
+                    rows_affected: 1,
+                },
+            ])
+            .append_query_results(vec![
+                vec![fruit::Model {
+                    id: 1,
+                    name: "Apple".to_owned(),
+                    cake_id: None,
+                }],
+                vec![fruit::Model {
+                    id: 1,
+                    name: "Apple".to_owned(),
+                    cake_id: None,
+                }],
+            ])
+            .into_connection();
+
+        fruit::Model {
+            id: 1,
+            name: "Apple".into(),
+            cake_id: None,
+        }
+        .into_active_model()
+        .update(&db)
+        .await?;
+
+        fruit::Model {
+            id: 1,
+            name: "Apple".into(),
+            cake_id: None,
+        }
+        .into_active_model()
+        .reset_all()
+        .update(&db)
+        .await?;
+
+        assert_eq!(
+            db.into_transaction_log(),
+            vec![
+                Transaction::from_sql_and_values(
+                    DbBackend::Postgres,
+                    r#"UPDATE "fruit" SET  WHERE "fruit"."id" = $1 RETURNING "id", "name", "cake_id""#,
+                    vec![1i32.into()],
+                ),
+                Transaction::from_sql_and_values(
+                    DbBackend::Postgres,
+                    r#"UPDATE "fruit" SET "name" = $1, "cake_id" = $2 WHERE "fruit"."id" = $3 RETURNING "id", "name", "cake_id""#,
+                    vec!["Apple".into(), Option::<i32>::None.into(), 1i32.into()],
+                ),
+            ]
+        );
+
+        Ok(())
+    }
 }