From f56ac7b7f6bb542956803ebc08d8bd3b088881b8 Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Fri, 10 Sep 2021 15:21:06 +0800
Subject: [PATCH] #129 Add `set` method to `UpdateMany`

---
 src/query/update.rs | 48 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/query/update.rs b/src/query/update.rs
index 2f805511..fec7757c 100644
--- a/src/query/update.rs
+++ b/src/query/update.rs
@@ -86,7 +86,7 @@ impl<A> UpdateOne<A>
 where
     A: ActiveModelTrait,
 {
-    pub(crate) fn prepare_filters(mut self) -> Self {
+    fn prepare_filters(mut self) -> Self {
         for key in <A::Entity as EntityTrait>::PrimaryKey::iter() {
             let col = key.into_column();
             let av = self.model.get(col);
@@ -99,7 +99,7 @@ where
         self
     }
 
-    pub(crate) fn prepare_values(mut self) -> Self {
+    fn prepare_values(mut self) -> Self {
         for col in <A::Entity as EntityTrait>::Column::iter() {
             if <A::Entity as EntityTrait>::PrimaryKey::from_column(col).is_some() {
                 continue;
@@ -177,6 +177,19 @@ impl<E> UpdateMany<E>
 where
     E: EntityTrait,
 {
+    pub fn set<A>(mut self, model: A) -> Self
+    where
+        A: ActiveModelTrait<Entity = E>,
+    {
+        for col in E::Column::iter() {
+            let av = model.get(col);
+            if av.is_set() {
+                self.query.value(col, av.unwrap());
+            }
+        }
+        self
+    }
+
     pub fn col_expr<T>(mut self, col: T, expr: SimpleExpr) -> Self
     where
         T: IntoIden,
@@ -244,4 +257,35 @@ mod tests {
             r#"UPDATE "fruit" SET "cake_id" = NULL WHERE "fruit"."id" = 2"#,
         );
     }
+
+    #[test]
+    fn update_5() {
+        assert_eq!(
+            Update::many(fruit::Entity)
+                .set(fruit::ActiveModel {
+                    name: ActiveValue::set("Apple".to_owned()),
+                    cake_id: ActiveValue::set(Some(3)),
+                    ..Default::default()
+                })
+                .filter(fruit::Column::Id.eq(2))
+                .build(DbBackend::Postgres)
+                .to_string(),
+            r#"UPDATE "fruit" SET "name" = 'Apple', "cake_id" = 3 WHERE "fruit"."id" = 2"#,
+        );
+    }
+
+    #[test]
+    fn update_6() {
+        assert_eq!(
+            Update::many(fruit::Entity)
+                .set(fruit::ActiveModel {
+                    id: ActiveValue::set(3),
+                    ..Default::default()
+                })
+                .filter(fruit::Column::Id.eq(2))
+                .build(DbBackend::Postgres)
+                .to_string(),
+            r#"UPDATE "fruit" SET "id" = 3 WHERE "fruit"."id" = 2"#,
+        );
+    }
 }