From c837d12c93e91c3bdff97307756b470f0b5a13c8 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Tue, 3 May 2022 16:29:38 +0800
Subject: [PATCH] `DerivePrimaryKey` with custom primary key column name

---
 sea-orm-macros/src/derives/entity_model.rs | 18 ++++++++------
 sea-orm-macros/src/derives/primary_key.rs  | 29 +++++++++++++++++++---
 sea-orm-macros/src/lib.rs                  |  2 +-
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/sea-orm-macros/src/derives/entity_model.rs b/sea-orm-macros/src/derives/entity_model.rs
index 0c7a0148d..bdb9c05a6 100644
--- a/sea-orm-macros/src/derives/entity_model.rs
+++ b/sea-orm-macros/src/derives/entity_model.rs
@@ -195,15 +195,16 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
 
                     field_name = Ident::new(&escape_rust_keyword(field_name), Span::call_site());
 
+                    let variant_attrs = match &column_name {
+                        Some(column_name) => quote! {
+                            #[sea_orm(column_name = #column_name)]
+                        },
+                        None => quote! {},
+                    };
+
                     if ignore {
                         continue;
                     } else {
-                        let variant_attrs = match &column_name {
-                            Some(column_name) => quote! {
-                                #[sea_orm(column_name = #column_name)]
-                            },
-                            None => quote! {},
-                        };
                         columns_enum.push(quote! {
                             #variant_attrs
                             #field_name
@@ -211,7 +212,10 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                     }
 
                     if is_primary_key {
-                        primary_keys.push(quote! { #field_name });
+                        primary_keys.push(quote! {
+                            #variant_attrs
+                            #field_name
+                        });
                     }
 
                     let col_type = match sql_type {
diff --git a/sea-orm-macros/src/derives/primary_key.rs b/sea-orm-macros/src/derives/primary_key.rs
index 13c7b5aba..3a05617bf 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 syn::{punctuated::Punctuated, token::Comma, Data, DataEnum, Fields, Lit, Meta, Variant};
 
 /// Method to derive a Primary Key for a Model using the [PrimaryKeyTrait](sea_orm::PrimaryKeyTrait)
 pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenStream> {
@@ -26,8 +26,31 @@ pub fn expand_derive_primary_key(ident: Ident, data: Data) -> syn::Result<TokenS
     let name: Vec<TokenStream> = variants
         .iter()
         .map(|v| {
-            let ident = v.ident.to_string().to_snake_case();
-            quote! { #ident }
+            let mut column_name = v.ident.to_string().to_snake_case();
+            for attr in v.attrs.iter() {
+                if let Some(ident) = attr.path.get_ident() {
+                    if ident != "sea_orm" {
+                        continue;
+                    }
+                } else {
+                    continue;
+                }
+                if let Ok(list) = attr.parse_args_with(Punctuated::<Meta, Comma>::parse_terminated)
+                {
+                    for meta in list.iter() {
+                        if let Meta::NameValue(nv) = meta {
+                            if let Some(name) = nv.path.get_ident() {
+                                if name == "column_name" {
+                                    if let Lit::Str(litstr) = &nv.lit {
+                                        column_name = litstr.value();
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+            quote! { #column_name }
         })
         .collect();
 
diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs
index 90cfee60e..3201907ea 100644
--- a/sea-orm-macros/src/lib.rs
+++ b/sea-orm-macros/src/lib.rs
@@ -195,7 +195,7 @@ pub fn derive_entity_model(input: TokenStream) -> TokenStream {
 /// #
 /// # impl ActiveModelBehavior for ActiveModel {}
 /// ```
-#[proc_macro_derive(DerivePrimaryKey)]
+#[proc_macro_derive(DerivePrimaryKey, attributes(sea_orm))]
 pub fn derive_primary_key(input: TokenStream) -> TokenStream {
     let DeriveInput { ident, data, .. } = parse_macro_input!(input);