use crate::{ error::*, ActiveModelTrait, ColumnTrait, ConnectionTrait, EntityTrait, Insert, IntoActiveModel, Iterable, PrimaryKeyToColumn, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64, TryInsert, }; use sea_query::{Expr, FromValueTuple, Iden, InsertStatement, IntoColumnRef, Query, ValueTuple}; use std::{future::Future, marker::PhantomData}; /// Defines a structure to perform INSERT operations in an ActiveModel #[derive(Debug)] pub struct Inserter where A: ActiveModelTrait, { primary_key: Option, query: InsertStatement, model: PhantomData, } /// The result of an INSERT operation on an ActiveModel #[derive(Debug)] pub struct InsertResult where A: ActiveModelTrait, { /// The id performed when AUTOINCREMENT was performed on the PrimaryKey pub last_insert_id: <<::Entity as EntityTrait>::PrimaryKey as PrimaryKeyTrait>::ValueType, } /// The types of results for an INSERT operation #[derive(Debug)] pub enum TryInsertResult { /// The INSERT statement did not have any value to insert Empty, /// The INSERT operation did not insert any valid value Conflicted, /// Successfully inserted Inserted(T), } impl TryInsert where A: ActiveModelTrait, { /// Execute an insert operation #[allow(unused_mut)] pub async fn exec<'a, C>(self, db: &'a C) -> Result>, DbErr> where C: ConnectionTrait, A: 'a, { if self.insert_struct.columns.is_empty() { return Ok(TryInsertResult::Empty); } let res = self.insert_struct.exec(db).await; match res { Ok(res) => Ok(TryInsertResult::Inserted(res)), Err(DbErr::RecordNotInserted) => Ok(TryInsertResult::Conflicted), Err(err) => Err(err), } } /// Execute an insert operation without returning (don't use `RETURNING` syntax) /// Number of rows affected is returned pub async fn exec_without_returning<'a, C>( self, db: &'a C, ) -> Result, DbErr> where ::Model: IntoActiveModel, C: ConnectionTrait, A: 'a, { if self.insert_struct.columns.is_empty() { return Ok(TryInsertResult::Empty); } let res = self.insert_struct.exec_without_returning(db).await; match res { Ok(res) => Ok(TryInsertResult::Inserted(res)), Err(DbErr::RecordNotInserted) => Ok(TryInsertResult::Conflicted), Err(err) => Err(err), } } /// Execute an insert operation and return the inserted model (use `RETURNING` syntax if database supported) pub async fn exec_with_returning<'a, C>( self, db: &'a C, ) -> Result::Model>, DbErr> where ::Model: IntoActiveModel, C: ConnectionTrait, A: 'a, { if self.insert_struct.columns.is_empty() { return Ok(TryInsertResult::Empty); } let res = self.insert_struct.exec_with_returning(db).await; match res { Ok(res) => Ok(TryInsertResult::Inserted(res)), Err(DbErr::RecordNotInserted) => Ok(TryInsertResult::Conflicted), Err(err) => Err(err), } } } impl Insert where A: ActiveModelTrait, { /// Execute an insert operation #[allow(unused_mut)] pub fn exec<'a, C>(self, db: &'a C) -> impl Future, DbErr>> + '_ where C: ConnectionTrait, A: 'a, { // so that self is dropped before entering await let mut query = self.query; if db.support_returning() && ::PrimaryKey::iter().count() > 0 { let returning = Query::returning().exprs( ::PrimaryKey::iter() .map(|c| c.into_column().select_as(Expr::col(c.into_column_ref()))), ); query.returning(returning); } Inserter::::new(self.primary_key, query).exec(db) } /// Execute an insert operation without returning (don't use `RETURNING` syntax) /// Number of rows affected is returned pub fn exec_without_returning<'a, C>( self, db: &'a C, ) -> impl Future> + '_ where ::Model: IntoActiveModel, C: ConnectionTrait, A: 'a, { Inserter::::new(self.primary_key, self.query).exec_without_returning(db) } /// Execute an insert operation and return the inserted model (use `RETURNING` syntax if database supported) pub fn exec_with_returning<'a, C>( self, db: &'a C, ) -> impl Future::Model, DbErr>> + '_ where ::Model: IntoActiveModel, C: ConnectionTrait, A: 'a, { Inserter::::new(self.primary_key, self.query).exec_with_returning(db) } } impl Inserter where A: ActiveModelTrait, { /// Instantiate a new insert operation pub fn new(primary_key: Option, query: InsertStatement) -> Self { Self { primary_key, query, model: PhantomData, } } /// Execute an insert operation, returning the last inserted id pub fn exec<'a, C>(self, db: &'a C) -> impl Future, DbErr>> + '_ where C: ConnectionTrait, A: 'a, { let builder = db.get_database_backend(); exec_insert(self.primary_key, builder.build(&self.query), db) } /// Execute an insert operation pub fn exec_without_returning<'a, C>( self, db: &'a C, ) -> impl Future> + '_ where C: ConnectionTrait, A: 'a, { exec_insert_without_returning(self.query, db) } /// Execute an insert operation and return the inserted model (use `RETURNING` syntax if database supported) pub fn exec_with_returning<'a, C>( self, db: &'a C, ) -> impl Future::Model, DbErr>> + '_ where ::Model: IntoActiveModel, C: ConnectionTrait, A: 'a, { exec_insert_with_returning::(self.primary_key, self.query, db) } } async fn exec_insert( primary_key: Option, statement: Statement, db: &C, ) -> Result, DbErr> where C: ConnectionTrait, A: ActiveModelTrait, { type PrimaryKey = <::Entity as EntityTrait>::PrimaryKey; type ValueTypeOf = as PrimaryKeyTrait>::ValueType; let last_insert_id = match (primary_key, db.support_returning()) { (Some(value_tuple), _) => { let res = db.execute(statement).await?; if res.rows_affected() == 0 { return Err(DbErr::RecordNotInserted); } FromValueTuple::from_value_tuple(value_tuple) } (None, true) => { let mut rows = db.query_all(statement).await?; let row = match rows.pop() { Some(row) => row, None => return Err(DbErr::RecordNotInserted), }; let cols = PrimaryKey::::iter() .map(|col| col.to_string()) .collect::>(); row.try_get_many("", cols.as_ref()) .map_err(|_| DbErr::UnpackInsertId)? } (None, false) => { let res = db.execute(statement).await?; if res.rows_affected() == 0 { return Err(DbErr::RecordNotInserted); } let last_insert_id = res.last_insert_id(); ValueTypeOf::::try_from_u64(last_insert_id).map_err(|_| DbErr::UnpackInsertId)? } }; Ok(InsertResult { last_insert_id }) } async fn exec_insert_without_returning( insert_statement: InsertStatement, db: &C, ) -> Result where C: ConnectionTrait, { let db_backend = db.get_database_backend(); let exec_result = db.execute(db_backend.build(&insert_statement)).await?; Ok(exec_result.rows_affected()) } async fn exec_insert_with_returning( primary_key: Option, mut insert_statement: InsertStatement, db: &C, ) -> Result<::Model, DbErr> where ::Model: IntoActiveModel, C: ConnectionTrait, A: ActiveModelTrait, { let db_backend = db.get_database_backend(); let found = match db.support_returning() { true => { let returning = Query::returning().exprs( ::Column::iter().map(|c| c.select_as(Expr::col(c))), ); insert_statement.returning(returning); SelectorRaw::::Model>>::from_statement( db_backend.build(&insert_statement), ) .one(db) .await? } false => { 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), None => Err(DbErr::RecordNotFound( "Failed to find inserted item".to_owned(), )), } }