From 71d0cde77181bd1c66ce6f8ee6d2626940e5eed5 Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Sat, 12 Jun 2021 19:03:29 +0800
Subject: [PATCH] PrimaryKey auto_increment

---
 examples/sqlx-mysql/src/example_cake.rs         | 6 ++++++
 examples/sqlx-mysql/src/example_cake_filling.rs | 6 ++++++
 examples/sqlx-mysql/src/example_filling.rs      | 6 ++++++
 examples/sqlx-mysql/src/example_fruit.rs        | 6 ++++++
 sea-orm-macros/src/derives/primary_key.rs       | 2 --
 src/entity/active_model.rs                      | 7 +++++--
 src/entity/primary_key.rs                       | 4 +++-
 src/tests_cfg/cake.rs                           | 6 ++++++
 src/tests_cfg/cake_filling.rs                   | 6 ++++++
 src/tests_cfg/filling.rs                        | 6 ++++++
 src/tests_cfg/fruit.rs                          | 6 ++++++
 11 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/examples/sqlx-mysql/src/example_cake.rs b/examples/sqlx-mysql/src/example_cake.rs
index c420ab34..63126cfd 100644
--- a/examples/sqlx-mysql/src/example_cake.rs
+++ b/examples/sqlx-mysql/src/example_cake.rs
@@ -26,6 +26,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {
     Fruit,
diff --git a/examples/sqlx-mysql/src/example_cake_filling.rs b/examples/sqlx-mysql/src/example_cake_filling.rs
index f23d9d9d..bae780af 100644
--- a/examples/sqlx-mysql/src/example_cake_filling.rs
+++ b/examples/sqlx-mysql/src/example_cake_filling.rs
@@ -27,6 +27,12 @@ pub enum PrimaryKey {
     FillingId,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        false
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {
     Cake,
diff --git a/examples/sqlx-mysql/src/example_filling.rs b/examples/sqlx-mysql/src/example_filling.rs
index 442a0701..c047c8f5 100644
--- a/examples/sqlx-mysql/src/example_filling.rs
+++ b/examples/sqlx-mysql/src/example_filling.rs
@@ -26,6 +26,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {}
 
diff --git a/examples/sqlx-mysql/src/example_fruit.rs b/examples/sqlx-mysql/src/example_fruit.rs
index 9d35b709..15e14768 100644
--- a/examples/sqlx-mysql/src/example_fruit.rs
+++ b/examples/sqlx-mysql/src/example_fruit.rs
@@ -28,6 +28,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {}
 
diff --git a/sea-orm-macros/src/derives/primary_key.rs b/sea-orm-macros/src/derives/primary_key.rs
index 2dfbcbe4..e677e2e6 100644
--- a/sea-orm-macros/src/derives/primary_key.rs
+++ b/sea-orm-macros/src/derives/primary_key.rs
@@ -45,8 +45,6 @@ pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenS
             }
         }
 
-        impl sea_orm::PrimaryKeyTrait for #ident {}
-
         impl sea_orm::PrimaryKeyToColumn for #ident {
             type Column = Column;
 
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
index bc1a478e..b70b3df8 100644
--- a/src/entity/active_model.rs
+++ b/src/entity/active_model.rs
@@ -1,4 +1,7 @@
-use crate::{Database, DeleteResult, EntityTrait, ExecErr, Iterable, PrimaryKeyToColumn, Value};
+use crate::{
+    Database, DeleteResult, EntityTrait, ExecErr, Iterable, PrimaryKeyToColumn, PrimaryKeyTrait,
+    Value,
+};
 use std::fmt::Debug;
 
 #[derive(Clone, Debug, Default)]
@@ -258,7 +261,7 @@ where
     let exec = E::insert(am).exec(db);
     let res = exec.await?;
     // TODO: if the entity does not have auto increment primary key, then last_insert_id is a wrong value
-    if res.last_insert_id != 0 {
+    if <E::PrimaryKey as PrimaryKeyTrait>::auto_increment() && res.last_insert_id != 0 {
         let find = E::find_by(res.last_insert_id).one(db);
         let res = find.await;
         let model: Option<E::Model> = res.map_err(|_| ExecErr)?;
diff --git a/src/entity/primary_key.rs b/src/entity/primary_key.rs
index 5fa3658d..8023c5ee 100644
--- a/src/entity/primary_key.rs
+++ b/src/entity/primary_key.rs
@@ -1,6 +1,8 @@
 use super::{ColumnTrait, IdenStatic, Iterable};
 
-pub trait PrimaryKeyTrait: IdenStatic + Iterable {}
+pub trait PrimaryKeyTrait: IdenStatic + Iterable {
+    fn auto_increment() -> bool;
+}
 
 pub trait PrimaryKeyToColumn {
     type Column: ColumnTrait;
diff --git a/src/tests_cfg/cake.rs b/src/tests_cfg/cake.rs
index bb7c2130..1ad182f8 100644
--- a/src/tests_cfg/cake.rs
+++ b/src/tests_cfg/cake.rs
@@ -27,6 +27,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {
     Fruit,
diff --git a/src/tests_cfg/cake_filling.rs b/src/tests_cfg/cake_filling.rs
index 7534634f..fa5e3a74 100644
--- a/src/tests_cfg/cake_filling.rs
+++ b/src/tests_cfg/cake_filling.rs
@@ -28,6 +28,12 @@ pub enum PrimaryKey {
     FillingId,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        false
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {
     Cake,
diff --git a/src/tests_cfg/filling.rs b/src/tests_cfg/filling.rs
index 34db602b..34df9cba 100644
--- a/src/tests_cfg/filling.rs
+++ b/src/tests_cfg/filling.rs
@@ -27,6 +27,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {}
 
diff --git a/src/tests_cfg/fruit.rs b/src/tests_cfg/fruit.rs
index 4475945c..32a0e553 100644
--- a/src/tests_cfg/fruit.rs
+++ b/src/tests_cfg/fruit.rs
@@ -29,6 +29,12 @@ pub enum PrimaryKey {
     Id,
 }
 
+impl PrimaryKeyTrait for PrimaryKey {
+    fn auto_increment() -> bool {
+        true
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {}