From 4f090d192bf7a26024c60792cd0b9bfdddbda97b Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Mon, 27 Sep 2021 18:02:12 +0800
Subject: [PATCH] Update `DerivePrimaryKey`

---
 sea-orm-macros/src/derives/primary_key.rs | 30 +++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/sea-orm-macros/src/derives/primary_key.rs b/sea-orm-macros/src/derives/primary_key.rs
index e677e2e6..ce7ca840 100644
--- a/sea-orm-macros/src/derives/primary_key.rs
+++ b/sea-orm-macros/src/derives/primary_key.rs
@@ -1,7 +1,7 @@
 use heck::SnakeCase;
 use proc_macro2::{Ident, TokenStream};
-use quote::{quote, quote_spanned};
-use syn::{Data, DataEnum, Fields, Variant};
+use quote::{quote, quote_spanned, ToTokens};
+use syn::{punctuated::Punctuated, token::Comma, Data, DataEnum, Fields, Variant};
 
 pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenStream> {
     let variants = match data {
@@ -30,6 +30,21 @@ pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenS
         })
         .collect();
 
+    let primary_key_value: Punctuated<_, Comma> =
+        variants.iter().fold(Punctuated::new(), |mut acc, v| {
+            let variant = &v.ident;
+            acc.push(
+                quote! { active_model.take(#ident::#variant.into_column()).unwrap().unwrap() },
+            );
+            acc
+        });
+    let mut primary_key_value = primary_key_value.to_token_stream();
+    if variants.len() > 1 {
+        primary_key_value = quote! {
+            (#primary_key_value)
+        };
+    }
+
     Ok(quote!(
         impl sea_orm::Iden for #ident {
             fn unquoted(&self, s: &mut dyn std::fmt::Write) {
@@ -61,5 +76,16 @@ pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenS
                 }
             }
         }
+
+        impl PrimaryKeyValue<Entity> for #ident {
+            fn get_primary_key_value<A>(
+                mut active_model: A,
+            ) -> <<Entity as EntityTrait>::PrimaryKey as PrimaryKeyTrait>::ValueType
+            where
+                A: ActiveModelTrait<Entity = Entity>,
+            {
+                #primary_key_value
+            }
+        }
     ))
 }