diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs index d46c6dbf..ea8d47a2 100644 --- a/src/entity/active_model.rs +++ b/src/entity/active_model.rs @@ -147,14 +147,9 @@ pub trait ActiveModelTrait: Clone + Debug { C: ConnectionTrait<'a>, { let am = ActiveModelBehavior::before_save(self, true)?; - let res = ::insert(am).exec(db).await?; - let found = ::find_by_id(res.last_insert_id) - .one(db) + let am = ::insert(am) + .exec_with_returning(db) .await?; - let am = match found { - Some(model) => model.into_active_model(), - None => return Err(DbErr::Exec("Failed to find inserted item".to_owned())), - }; ActiveModelBehavior::after_save(am, true) } diff --git a/src/executor/insert.rs b/src/executor/insert.rs index 8e97ccbc..363dccd9 100644 --- a/src/executor/insert.rs +++ b/src/executor/insert.rs @@ -1,6 +1,6 @@ use crate::{ - error::*, ActiveModelTrait, ConnectionTrait, DbBackend, EntityTrait, Insert, Iterable, - PrimaryKeyTrait, Statement, TryFromU64, + error::*, ActiveModelTrait, ConnectionTrait, DbBackend, EntityTrait, Insert, IntoActiveModel, + Iterable, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64, }; use sea_query::{FromValueTuple, InsertStatement, IntoColumnRef, Returning, ValueTuple}; use std::{future::Future, marker::PhantomData}; @@ -50,6 +50,19 @@ where } Inserter::::new(self.primary_key, query).exec(db) } + + /// Execute an insert operation and return inserted row + pub fn exec_with_returning<'a, C>( + self, + db: &'a C, + ) -> impl Future> + '_ + where + ::Model: IntoActiveModel, + C: ConnectionTrait<'a>, + A: 'a, + { + Inserter::::new(self.primary_key, self.query).exec_with_returning(db) + } } impl Inserter @@ -74,6 +87,19 @@ where let builder = db.get_database_backend(); exec_insert(self.primary_key, builder.build(&self.query), db) } + + /// Execute an insert operation and return inserted row + pub fn exec_with_returning<'a, C>( + self, + db: &'a C, + ) -> impl Future> + '_ + where + ::Model: IntoActiveModel, + C: ConnectionTrait<'a>, + A: 'a, + { + exec_insert_with_returning(self.primary_key, self.query, db) + } } async fn exec_insert<'a, A, C>( @@ -89,7 +115,7 @@ where type ValueTypeOf = as PrimaryKeyTrait>::ValueType; let last_insert_id_opt = match db.get_database_backend() { DbBackend::Postgres => { - use crate::{sea_query::Iden, Iterable}; + use crate::sea_query::Iden; let cols = PrimaryKey::::iter() .map(|col| col.to_string()) .collect::>(); @@ -110,3 +136,41 @@ where }; Ok(InsertResult { last_insert_id }) } + +async fn exec_insert_with_returning<'a, A, C>( + primary_key: Option, + mut insert_statement: InsertStatement, + db: &'a C, +) -> Result +where + ::Model: IntoActiveModel, + C: ConnectionTrait<'a>, + A: ActiveModelTrait, +{ + let db_backend = db.get_database_backend(); + let found = match db_backend { + DbBackend::Postgres => { + insert_statement.returning(Returning::Columns( + ::Column::iter() + .map(|c| c.into_column_ref()) + .collect(), + )); + SelectorRaw::::Model>>::from_statement( + db_backend.build(&insert_statement), + ) + .one(db) + .await? + } + _ => { + let insert_res = + exec_insert::(primary_key, db_backend.build(&insert_statement), db).await?; + ::find_by_id(insert_res.last_insert_id) + .one(db) + .await? + } + }; + match found { + Some(model) => Ok(model.into_active_model()), + None => Err(DbErr::Exec("Failed to find inserted item".to_owned())), + } +}