Insert many with returning many

Co-authored-by: Ivan Porto Wigner <[email protected]>

src/executor/insert.rs

@@ -102,6 +102,28 @@ where
             Err(err) => Err(err),
         }
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if database supported)
+    pub async fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> Result<TryInsertResult<Vec<<A::Entity as EntityTrait>::Model>>, DbErr>
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        if self.insert_struct.columns.is_empty() {
+            return Ok(TryInsertResult::Empty);
+        }
+
+        let res = self.insert_struct.exec_with_returning_many(db).await;
+        match res {
+            Ok(res) => Ok(TryInsertResult::Inserted(res)),
+            Err(DbErr::RecordNotInserted) => Ok(TryInsertResult::Conflicted),
+            Err(err) => Err(err),
+        }
+    }
 }
 
 impl<A> Insert<A>
@@ -155,6 +177,25 @@ where
     {
         Inserter::<A>::new(self.primary_key, self.query).exec_with_returning(db)
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if supported)
+    pub async fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        if self.columns.is_empty() {
+            Ok(vec![])
+        } else {
+            Inserter::<A>::new(self.primary_key, self.query)
+                .exec_with_returning_many(db)
+                .await
+        }
+    }
 }
 
 impl<A> Inserter<A>
@@ -203,6 +244,19 @@ where
     {
         exec_insert_with_returning::<A, _>(self.primary_key, self.query, db)
     }
+
+    /// Execute an insert many operation and return the inserted model (use `RETURNING` syntax if supported)
+    pub fn exec_with_returning_many<'a, C>(
+        self,
+        db: &'a C,
+    ) -> impl Future<Output = Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>> + '_
+    where
+        <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+        C: ConnectionTrait,
+        A: 'a,
+    {
+        exec_insert_many_with_returning::<A, _>(self.query, db)
+    }
 }
 
 async fn exec_insert<A, C>(
@@ -313,3 +367,33 @@ where
         )),
     }
 }
+
+async fn exec_insert_many_with_returning<A, C>(
+    mut insert_statement: InsertStatement,
+    db: &C,
+) -> Result<Vec<<A::Entity as EntityTrait>::Model>, DbErr>
+where
+    <A::Entity as EntityTrait>::Model: IntoActiveModel<A>,
+    C: ConnectionTrait,
+    A: ActiveModelTrait,
+{
+    match db.support_returning() {
+        true => {
+            let db_backend = db.get_database_backend();
+            let returning = Query::returning().exprs(
+                <A::Entity as EntityTrait>::Column::iter()
+                    .map(|c| c.select_as(c.into_returning_expr(db_backend))),
+            );
+            insert_statement.returning(returning);
+            let insert_statement = db_backend.build(&insert_statement);
+            let models: Vec<<A::Entity as EntityTrait>::Model> =
+                SelectorRaw::<SelectModel<<A::Entity as EntityTrait>::Model>>::from_statement(
+                    insert_statement,
+                )
+                .all(db)
+                .await?;
+            Ok(models)
+        }
+        false => unimplemented!("Database backend doesn't support RETURNING"),
+    }
+}

tests/returning_tests.rs

@@ -3,7 +3,7 @@
 pub mod common;
 
 pub use common::{bakery_chain::*, setup::*, TestContext};
-use sea_orm::{entity::prelude::*, IntoActiveModel, Set};
+use sea_orm::{entity::prelude::*, IntoActiveModel, NotSet, Set};
 pub use sea_query::{Expr, Query};
 use serde_json::json;
 
@@ -67,6 +67,79 @@ async fn main() -> Result<(), DbErr> {
     Ok(())
 }
 
+#[sea_orm_macros::test]
+#[cfg_attr(
+    any(
+        feature = "sqlx-mysql",
+        all(
+            feature = "sqlx-sqlite",
+            not(feature = "sqlite-use-returning-for-3_35")
+        )
+    ),
+    should_panic(expected = "Database backend doesn't support RETURNING")
+)]
+async fn insert_many() {
+    pub use common::{features::*, TestContext};
+    use edit_log::*;
+
+    let run = || async {
+        let ctx = TestContext::new("returning_tests_insert_many").await;
+        let db = &ctx.db;
+
+        create_tables(db).await?;
+
+        // Insert many with returning
+        assert_eq!(
+            Entity::insert_many(vec![
+                ActiveModel {
+                    id: NotSet,
+                    action: Set("before_save".to_string()),
+                    values: Set(json!({ "id": "unique-id-001" })),
+                },
+                ActiveModel {
+                    id: NotSet,
+                    action: Set("before_save".to_string()),
+                    values: Set(json!({ "id": "unique-id-002" })),
+                },
+            ])
+            .exec_with_returning_many(db)
+            .await?,
+            [
+                Model {
+                    id: 1,
+                    action: "before_save".to_string(),
+                    values: json!({ "id": "unique-id-001" }),
+                },
+                Model {
+                    id: 2,
+                    action: "before_save".to_string(),
+                    values: json!({ "id": "unique-id-002" }),
+                },
+            ]
+        );
+
+        // No-op
+        assert_eq!(
+            Entity::insert_many(
+                vec![ActiveModel {
+                    id: NotSet,
+                    action: Set("before_save".to_string()),
+                    values: Set(json!({ "id": "unique-id-001" })),
+                }]
+                .into_iter()
+                .filter(|_| false)
+            )
+            .exec_with_returning_many(db)
+            .await?,
+            []
+        );
+
+        Result::<(), DbErr>::Ok(())
+    };
+
+    run().await.unwrap();
+}
+
 #[sea_orm_macros::test]
 #[cfg_attr(
     any(