DerivePrimaryKey with custom primary key column name

SeaQL · May 3, 2022 · c837d12 · c837d12
1 parent 86e7e80
commit c837d12
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 11 deletions.
diff --git a/sea-orm-macros/src/derives/entity_model.rs b/sea-orm-macros/src/derives/entity_model.rs
@@ -195,23 +195,27 @@ 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
                         });
                     }
 
                     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
@@ -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
@@ -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);