Cast primary key column inside returning expression of insert stateme…

…nt (#1427)
SeaQL · Jan 29, 2023 · 51e6ad7 · 51e6ad7
1 parent 5ffd802
commit 51e6ad7
Show file tree

Hide file tree

Showing 3 changed files with 140 additions and 6 deletions.
diff --git a/sea-orm-macros/src/derives/entity_model.rs b/sea-orm-macros/src/derives/entity_model.rs
@@ -251,15 +251,15 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                             Self::#field_name => sea_orm::sea_query::SimpleExpr::cast_as(
                                 Into::<sea_orm::sea_query::SimpleExpr>::into(expr),
                                 sea_orm::sea_query::Alias::new(&#select_as),
-                            ),
+                            )
                         });
                     }
                     if let Some(save_as) = save_as {
                         columns_save_as.push(quote! {
                             Self::#field_name => sea_orm::sea_query::SimpleExpr::cast_as(
                                 Into::<sea_orm::sea_query::SimpleExpr>::into(val),
                                 sea_orm::sea_query::Alias::new(&#save_as),
-                            ),
+                            )
                         });
                     }
 
@@ -349,6 +349,14 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
         }
     }
 
+    // Add tailing comma
+    if !columns_select_as.is_empty() {
+        columns_select_as.push_punct(Comma::default());
+    }
+    if !columns_save_as.is_empty() {
+        columns_save_as.push_punct(Comma::default());
+    }
+
     let primary_key = {
         let auto_increment = auto_increment && primary_keys.len() == 1;
         let primary_key_types = if primary_key_types.len() == 1 {

diff --git a/src/executor/insert.rs b/src/executor/insert.rs
@@ -1,6 +1,6 @@
 use crate::{
     error::*, ActiveModelTrait, ColumnTrait, ConnectionTrait, EntityTrait, Insert, IntoActiveModel,
-    Iterable, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64,
+    Iterable, PrimaryKeyToColumn, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64,
 };
 use sea_query::{Expr, FromValueTuple, Iden, InsertStatement, IntoColumnRef, Query, ValueTuple};
 use std::{future::Future, marker::PhantomData};
@@ -40,8 +40,9 @@ where
         // so that self is dropped before entering await
         let mut query = self.query;
         if db.support_returning() && <A::Entity as EntityTrait>::PrimaryKey::iter().count() > 0 {
-            let returning = Query::returning().columns(
-                <A::Entity as EntityTrait>::PrimaryKey::iter().map(|c| c.into_column_ref()),
+            let returning = Query::returning().exprs(
+                <A::Entity as EntityTrait>::PrimaryKey::iter()
+                    .map(|c| c.into_column().select_as(Expr::col(c.into_column_ref()))),
             );
             query.returning(returning);
         }

diff --git a/src/query/insert.rs b/src/query/insert.rs
@@ -227,7 +227,7 @@ mod tests {
     use sea_query::OnConflict;
 
     use crate::tests_cfg::cake;
-    use crate::{ActiveValue, DbBackend, EntityTrait, Insert, QueryTrait};
+    use crate::{ActiveValue, DbBackend, DbErr, EntityTrait, Insert, IntoActiveModel, QueryTrait};
 
     #[test]
     fn insert_1() {
@@ -350,4 +350,129 @@ mod tests {
             r#"INSERT INTO "cake" ("id", "name") VALUES (2, 'Orange') ON CONFLICT ("name") DO UPDATE SET "name" = "excluded"."name""#,
         );
     }
+
+    #[smol_potat::test]
+    async fn insert_8() -> Result<(), DbErr> {
+        use crate::{ActiveModelTrait, DbBackend, MockDatabase, Statement, Transaction};
+
+        mod post {
+            use crate as sea_orm;
+            use crate::entity::prelude::*;
+
+            #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+            #[sea_orm(table_name = "posts")]
+            pub struct Model {
+                #[sea_orm(primary_key, select_as = "INTEGER", save_as = "TEXT")]
+                pub id: i32,
+                pub title: String,
+                pub text: String,
+            }
+
+            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+            pub enum Relation {}
+
+            impl ActiveModelBehavior for ActiveModel {}
+        }
+
+        let model = post::Model {
+            id: 1,
+            title: "News wrap up 2022".into(),
+            text: "brbrbrrrbrbrbrr...".into(),
+        };
+
+        let db = MockDatabase::new(DbBackend::Postgres)
+            .append_query_results([[model.clone()]])
+            .into_connection();
+
+        post::Entity::insert(model.into_active_model())
+            .exec(&db)
+            .await?;
+
+        assert_eq!(
+            db.into_transaction_log(),
+            [Transaction::many([Statement::from_sql_and_values(
+                DbBackend::Postgres,
+                r#"INSERT INTO "posts" ("id", "title", "text") VALUES (CAST($1 AS TEXT), $2, $3) RETURNING CAST("id" AS INTEGER)"#,
+                [
+                    1.into(),
+                    "News wrap up 2022".into(),
+                    "brbrbrrrbrbrbrr...".into(),
+                ]
+            )])]
+        );
+
+        Ok(())
+    }
+
+    #[smol_potat::test]
+    async fn insert_9() -> Result<(), DbErr> {
+        use crate::{
+            ActiveModelTrait, DbBackend, MockDatabase, MockExecResult, Statement, Transaction,
+        };
+
+        mod post {
+            use crate as sea_orm;
+            use crate::entity::prelude::*;
+
+            #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+            #[sea_orm(table_name = "posts")]
+            pub struct Model {
+                #[sea_orm(
+                    primary_key,
+                    auto_increment = false,
+                    select_as = "INTEGER",
+                    save_as = "TEXT"
+                )]
+                pub id_primary: i32,
+                #[sea_orm(
+                    primary_key,
+                    auto_increment = false,
+                    select_as = "INTEGER",
+                    save_as = "TEXT"
+                )]
+                pub id_secondary: i32,
+                pub title: String,
+                pub text: String,
+            }
+
+            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+            pub enum Relation {}
+
+            impl ActiveModelBehavior for ActiveModel {}
+        }
+
+        let model = post::Model {
+            id_primary: 1,
+            id_secondary: 1001,
+            title: "News wrap up 2022".into(),
+            text: "brbrbrrrbrbrbrr...".into(),
+        };
+
+        let db = MockDatabase::new(DbBackend::Postgres)
+            .append_exec_results([MockExecResult {
+                last_insert_id: 1,
+                rows_affected: 1,
+            }])
+            .into_connection();
+
+        post::Entity::insert(model.into_active_model())
+            .exec(&db)
+            .await?;
+
+        assert_eq!(
+            db.into_transaction_log(),
+            [Transaction::many([Statement::from_sql_and_values(
+                DbBackend::Postgres,
+                r#"INSERT INTO "posts" ("id_primary", "id_secondary", "title", "text") VALUES (CAST($1 AS TEXT), CAST($2 AS TEXT), $3, $4) RETURNING CAST("id_primary" AS INTEGER), CAST("id_secondary" AS INTEGER)"#,
+                [
+                    1.into(),
+                    1001.into(),
+                    "News wrap up 2022".into(),
+                    "brbrbrrrbrbrbrr...".into(),
+                ]
+            )])]
+        );
+
+        Ok(())
+    }
 }